Rewriter prompt and LangGraph node

The rewriter is a single async function. It takes a query and an integer n, calls the LLM with a structured prompt, parses the output into a list of strings, and always prepends the original query so downstream retrieval never loses the user literal phrasing.

rewriter.py
python
REWRITE_PROMPT = """You are a query rewriting assistant for a retrieval system.
Given the user's question, produce {n} diverse reformulations that will help
retrieve relevant documents. Vary the phrasing:
- one keyword-heavy version
- one natural-language question
- others as paraphrases with different terminology

Return ONLY the rewritten queries, one per line, no numbering, no extra text.

User question: {query}

Rewrites:"""

The prompt is explicit about the number of rewrites, the diversity requirement, and the output format. Explicit output formatting makes parsing reliable across providers.

rewriter.py
python
async def rewrite_query(query: str, n: int = 3) -> list[str]:
    if n <= 0:
        return [query]

    provider = get_llm_provider()
    prompt = REWRITE_PROMPT.format(n=n, query=query)

    output = ""
    try:
        async for chunk in provider.generate_stream(prompt, temperature=0.4, max_tokens=400):
            output += chunk
    except Exception as e:
        logger.warning(f"Rewriter failed, returning original query only: {e}")
        return [query]

    lines = [l.strip() for l in output.splitlines() if l.strip()]
    cleaned = [re.sub(r"^[\-\*\d\.\)]+\s*", "", l).strip() for l in lines]
    cleaned = [c for c in cleaned if c and c.lower() != query.lower()]

    seen, deduped = set(), []
    for q in [query] + cleaned:
        key = q.lower()
        if key not in seen:
            seen.add(key); deduped.append(q)
        if len(deduped) >= n + 1:
            break
    return deduped[: n + 1]

A few defensive habits worth stealing: always include the original query, strip numbering or bullets the LLM adds anyway, and dedupe case-insensitively. Temperature 0.4 gives enough variety without drifting off topic.

rag_graph.py
python
async def _rewrite_node(state: GraphState) -> GraphState:
    """Generate N diverse rewrites of the (anonymized) question."""
    q = state.get("anonymized_question") or state["question"]
    n = state.get("rewrite_n", 3)
    rewrites = await rewrite_query(q, n=n)
    return {"rewrites": rewrites}

The LangGraph node is a thin wrapper: read state, call the rewriter, write the rewrites back to state. Keeping nodes this small is what makes the graph easy to test and reason about.

Fill in the blanks: Complete the rewriter fallback

Loading practice…