query_rewriter.py

from typing import Optional
import config as cfg
from llm_handler import LLMHandler


_REFORMULATION_PROMPT = """Given a user question about a technical document, rewrite it into a clear, search-optimized query that will retrieve the most relevant passages.

Original question: {question}

Return ONLY the rewritten query, nothing else. Make it concise and focused on key technical terms."""


class QueryRewriter:
    def __init__(self, llm: Optional[LLMHandler] = None):
        self.llm = llm

    def rewrite(self, question: str) -> str:
        if not self.llm:
            return question
        try:
            prompt = _REFORMULATION_PROMPT.format(question=question)
            result = self.llm.generate(prompt)
            if result and result.strip():
                return result.strip()
        except Exception:
            pass
        return question