query_processing.py

"""
query_processing.py
--------------------
Two techniques applied to the raw user question before retrieval:

1. Query Rewriting — resolves conversational references using chat
   history. "What about the second one?" becomes a self-contained
   question like "What is the pricing for Plan B?" This matters because
   retrieval runs BEFORE the LLM sees history — an ambiguous query
   retrieves the wrong chunks no matter how good the LLM's memory is.

2. Query Expansion — generates 2-3 alternate phrasings of the same
   question (synonyms, rephrasing) and retrieves for all of them, then
   merges results. This catches cases where the document uses different
   terminology than the user's question (e.g. user asks "cost", document
   says "pricing" or "fees").

Both use a cheap, fast LLM call — the cost is worth it because a bad
retrieval can't be fixed by a smarter generation step downstream.
"""

from typing import List, Tuple
from backend.logging_config import get_logger

logger = get_logger(__name__)

REWRITE_PROMPT = """Rewrite the user's latest question into a fully self-contained \
question that does not depend on the conversation history. If it is already \
self-contained, return it unchanged. Return ONLY the rewritten question, \
nothing else.

Conversation history:
{history}

Latest question: {question}

Rewritten question:"""

EXPANSION_PROMPT = """Generate {n} alternative phrasings of the following question. \
The alternatives should use different wording or synonyms but ask for the same \
information. Return ONLY the alternatives, one per line, no numbering, no extra text.

Question: {question}"""


def rewrite_query(question: str, history: List[Tuple[str, str]], provider) -> str:
    """Resolve pronouns / follow-up references using recent chat history."""
    if not history:
        return question

    recent = history[-3:]
    history_text = "\n".join(f"User: {u}\nAssistant: {a}" for u, a in recent)

    try:
        rewritten = provider.generate(
            messages=[{"role": "user", "content": REWRITE_PROMPT.format(
                history=history_text, question=question)}],
            temperature=0.0,
            max_tokens=120,
        )
        return rewritten.strip().strip('"')
    except Exception as e:
        logger.warning(f"Query rewrite failed, falling back to original question: {e}")
        return question


def expand_query(question: str, provider, n: int = 2) -> List[str]:
    """Generate a small number of alternate phrasings to widen retrieval recall."""
    try:
        raw = provider.generate(
            messages=[{"role": "user", "content": EXPANSION_PROMPT.format(n=n, question=question)}],
            temperature=0.3,
            max_tokens=150,
        )
        variants = [line.strip("- ").strip() for line in raw.strip().split("\n") if line.strip()]
        return [question] + variants[:n]
    except Exception as e:
        logger.warning(f"Query expansion failed, using original question only: {e}")
        return [question]