core/query_rewriter.py

# core/query_rewriter.py
# =======================
# Capability 2 — Query rewriting for self-healing RAG
#
# Provides two strategies:
#
#   rewrite_query()
#       Classic paraphrase rewrite: the LLM is asked to rephrase the user's
#       question in vocabulary more likely to match document passages. This is
#       the default first-retry strategy.
#
#   hyde_query()
#       Hypothetical Document Embedding (HyDE): the LLM generates a short
#       hypothetical passage that *would* answer the question, then that
#       passage — not the original question — is used as the embedding query.
#       Because the hypothetical passage looks like a real document chunk, it
#       maps to a closer point in the embedding space and significantly improves
#       recall for abstract or jargon-heavy questions.
#       Used as the second-retry strategy when ENABLE_HYDE=true in config.
#
# Both functions take the ChatGenerator instance so no new API client is needed.
# They make a single, cheap LLM call each (max_tokens capped at 256).
#
# Usage (in app.py):
#   from core.query_rewriter import rewrite_query, hyde_query
#   rewritten = rewrite_query(original_question, generator)
#   hypothetical = hyde_query(original_question, generator)

from __future__ import annotations

import logging
import re
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from core.generator import ChatGenerator

logger = logging.getLogger(__name__)


# ── Internal LLM caller ───────────────────────────────────────────────────────

def _call_small(system: str, user: str, generator: "ChatGenerator") -> str:
    """
    Make a lightweight LLM call via the generator's existing client.
    Caps at 256 tokens — rewrites and hypothetical passages are short.
    Strips <think>…</think> tags that reasoning models emit.
    """
    from config import IS_HF_SPACES
    try:
        if IS_HF_SPACES:
            # Temporarily override max_tokens for this lightweight call
            original_create = generator._client.chat.completions.create

            def _create_short(**kwargs):
                kwargs["max_tokens"] = 256
                return original_create(**kwargs)

            generator._client.chat.completions.create = _create_short
            result = generator._call_hf(system=system, user=user)
            generator._client.chat.completions.create = original_create
        else:
            result = generator._call_ollama(system=system, user=user)

        # Strip reasoning tags if present
        result = re.sub(r"<think>.*?</think>", "", result, flags=re.DOTALL).strip()
        return result
    except Exception as exc:
        logger.warning("QueryRewriter LLM call failed: %s", exc)
        return ""


# ── Public API ────────────────────────────────────────────────────────────────

def rewrite_query(original: str, generator: "ChatGenerator") -> str:
    """
    Rephrase the user's question to improve document retrieval.

    Strategy: ask the LLM to restate the question using different vocabulary —
    particularly substituting domain terms, expanding acronyms, and adding
    synonyms that are more likely to appear in a formal document.

    Returns the rewritten query, or the original if the rewrite fails or is
    suspiciously short / identical.

    Parameters
    ----------
    original  : str — the user's raw question
    generator : ChatGenerator — existing generator instance (no new client)
    """
    system = (
        "You are a search-query optimizer. "
        "Your only job is to rewrite a user's question so it matches the "
        "vocabulary used in formal documents, reports, and articles. "
        "Rules:\n"
        "1. Output ONLY the rewritten query — no explanation, no quotes, no preamble.\n"
        "2. Keep it under 30 words.\n"
        "3. Preserve the original intent exactly.\n"
        "4. Expand acronyms, replace informal phrases with formal equivalents, "
        "   and add one or two relevant synonyms if helpful.\n"
        "5. If the query is already clear and formal, return it unchanged."
    )
    user = f"Rewrite this search query:\n\n{original}"

    rewritten = _call_small(system=system, user=user, generator=generator)

    # Safety checks — fall back to original if rewrite looks bad
    if (
        not rewritten
        or len(rewritten) < 5
        or rewritten.lower() == original.lower()
        or len(rewritten) > len(original) * 4  # runaway output
    ):
        logger.debug("Query rewrite skipped (bad output); using original.")
        return original

    logger.info("Query rewritten: %r → %r", original, rewritten)
    return rewritten


def hyde_query(original: str, generator: "ChatGenerator") -> str:
    """
    Generate a Hypothetical Document Embedding (HyDE) passage.

    The LLM writes a short (3-5 sentence) hypothetical document passage that
    would *answer* the user's question. That passage — which looks like real
    document text — is then used as the embedding query instead of the question.
    This dramatically narrows the distance in embedding space between the query
    and relevant chunks.

    Returns the hypothetical passage, or the original question if generation
    fails.

    Parameters
    ----------
    original  : str — the user's raw question
    generator : ChatGenerator — existing generator instance (no new client)

    References
    ----------
    Gao et al., 2022 — "Precise Zero-Shot Dense Retrieval without Relevance Labels"
    https://arxiv.org/abs/2212.10496
    """
    system = (
        "You are an expert document writer. "
        "Given a question, write a short passage (3-5 sentences) from a "
        "formal document, report, or article that would DIRECTLY ANSWER the "
        "question. The passage should:\n"
        "1. Sound like it was extracted from a real document.\n"
        "2. Use formal, domain-appropriate vocabulary.\n"
        "3. Contain the key facts or explanations that answer the question.\n"
        "4. Be 50-120 words long.\n"
        "Output ONLY the passage — no heading, no preamble, no explanation."
    )
    user = f"Write a document passage that answers this question:\n\n{original}"

    passage = _call_small(system=system, user=user, generator=generator)

    if not passage or len(passage) < 20:
        logger.debug("HyDE generation failed; falling back to original query.")
        return original

    logger.info(
        "HyDE passage generated (%d chars) for query: %r",
        len(passage),
        original,
    )
    return passage