eb3a394224
Co-authored-by: Junda (David) Su <90978028+Davids048@users.noreply.github.com> Co-authored-by: Matthew Noto <99706358+RandNMR73@users.noreply.github.com> Co-authored-by: XOR-op <17672363+XOR-op@users.noreply.github.com> Co-authored-by: Zhang Peiyuan <42993249+jzhang38@users.noreply.github.com>
83 lines
2.4 KiB
Python
83 lines
2.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Rewrite payload builder.
|
|
|
|
The UI's "rewrite seed prompts" flow asks the enhancer to produce a
|
|
batch of alternative prompts given one seed. This module packages the
|
|
seed + options into the payload the enhancer expects and unpacks the
|
|
response back into a typed :class:`RewriteResult`.
|
|
|
|
Separating this from :mod:`enhancer` keeps the enhancer provider-
|
|
agnostic; anything UI-specific (how many alternatives to request, how
|
|
to split the response, temperature) lives here.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
from dataclasses import dataclass
|
|
|
|
from fastvideo.entrypoints.streaming.prompt.enhancer import PromptEnhancer
|
|
|
|
_LEADING_MARKER_RE = re.compile(r"^(?:[-*•]\s*|\d+\s*[.)]\s*)+")
|
|
|
|
|
|
@dataclass
|
|
class RewriteOptions:
|
|
count: int = 3
|
|
"""Number of alternative prompts to request."""
|
|
temperature: float | None = None
|
|
|
|
|
|
@dataclass
|
|
class RewriteResult:
|
|
seed_prompt: str
|
|
alternatives: list[str]
|
|
provider: str
|
|
model: str
|
|
latency_ms: float
|
|
fallback_used: bool = False
|
|
|
|
|
|
async def build_rewrite(
|
|
enhancer: PromptEnhancer,
|
|
seed_prompt: str,
|
|
*,
|
|
options: RewriteOptions | None = None,
|
|
) -> RewriteResult:
|
|
"""Run a rewrite op through the enhancer and return a typed result."""
|
|
if not seed_prompt.strip():
|
|
raise ValueError("rewrite seed prompt must be non-empty")
|
|
options = options or RewriteOptions()
|
|
response = await enhancer.rewrite(seed_prompt)
|
|
alternatives = _split_response(response.content, limit=options.count)
|
|
return RewriteResult(
|
|
seed_prompt=seed_prompt,
|
|
alternatives=alternatives,
|
|
provider=response.provider,
|
|
model=response.model,
|
|
latency_ms=response.latency_ms,
|
|
fallback_used=response.fallback_used,
|
|
)
|
|
|
|
|
|
def _split_response(content: str, *, limit: int) -> list[str]:
|
|
"""Split the LLM response into discrete prompt candidates.
|
|
|
|
The shipped system prompt instructs the model to emit one prompt
|
|
per line; this function is forgiving about numbered lists or
|
|
leading bullets so user-supplied system prompts don't break it.
|
|
"""
|
|
lines = [line.strip() for line in content.splitlines() if line.strip()]
|
|
cleaned: list[str] = []
|
|
for line in lines:
|
|
stripped = _LEADING_MARKER_RE.sub("", line).strip()
|
|
if stripped:
|
|
cleaned.append(stripped)
|
|
return cleaned[:max(1, limit)]
|
|
|
|
|
|
__all__ = [
|
|
"RewriteOptions",
|
|
"RewriteResult",
|
|
"build_rewrite",
|
|
]
|