psq_rag/llm/rewrite.py

from .openrouter_client import openrouter_chat


REWRITE_SYSTEM = """Rewrite the input into a concise, comma-separated list of short phrases
that resemble image tags.

Use short, literal phrases that reflect how visual concepts are commonly
written in image tag vocabularies.

Multi-word phrases are appropriate when they represent one coherent
visual idea.

Examples of tag-shaped phrases:
- wolf, angry
- blue jacket, striped tail
- long hair, raised ears
- holding object, hand on shoulder
- looking at viewer, looking down
- simple background, outdoor scene
- wooden table, plant
- running, sleeping
- smiling, angry expression
- bedroom, forest
- sonic the hedgehog, princess peach

Do not invent details or guess identities.
Do not infer demographic attributes (e.g., gender/age) unless explicitly stated.

Output ONLY the rewritten list.
"""


def llm_rewrite_prompt(prompt_in: str, log) -> str:
    messages = [
        {"role": "system", "content": REWRITE_SYSTEM},
        {"role": "user", "content": prompt_in},
    ]

    raw, _parsed_unused, err = openrouter_chat(
        messages,
        response_format=None,
        temperature=0.0,
        max_tokens=256,
    )

    if err:
        log(f"LLM rewrite: fallback (error: {err})")
        # NEW: if we got a refusal-like completion, log the refusal text for debugging
        if raw and err.lower().startswith("refusal-like"):
            log(f"LLM rewrite refusal text: {raw.strip()[:300]}")
        return ""

    out = (raw or "").strip()
    if not out:
        log("LLM rewrite: fallback (empty response)")
        return ""

    out = " ".join(out.split())
    if len(out) > 800:
        out = out[:800].rstrip()

    log("LLM rewrite: ok")
    return out


if __name__ == "__main__":
    print("rewrite.py imports ok")