AIRegulation-DocAnalysis/backend/app/application/agent/hyde_expander.py

"""Implement HyDE (Hypothetical Document Embeddings) query expansion.

HyDE improves dense retrieval by addressing the vocabulary gap between
short user queries and longer document passages:

  User query  →  [LLM generates hypothetical answer]
                         ↓
             embed hypothetical answer (not original query)
                         ↓
             retrieve similar real passages from Milvus

The hypothetical answer uses the same vocabulary and phrasing as documents,
so its embedding is much closer to relevant chunks than a terse query embedding.

Usage:
    expander = HyDEExpander()
    retrieval_query = expander.expand(query, provider=..., model=...)
    chunks = retrieval_service.retrieve(query=retrieval_query, ...)

When the LLM call fails, expand() falls back to the original query so the
retrieval pipeline degrades gracefully.

References:
    Gao et al. (2022), "Precise Zero-Shot Dense Retrieval without Relevance Labels"
    https://arxiv.org/abs/2212.10496
"""

from __future__ import annotations

from loguru import logger

from app.config.settings import settings
from app.services.llm.llm_factory import get_llm_client

# Maximum chars to trim from the hypothetical answer to avoid token overrun.
_MAX_HYPOTHESIS_CHARS = 600

# System prompt that instructs the LLM to write a passage *as if* it were
# from a regulatory document, not a conversation answer.
_HYDE_SYSTEM = (
    "你是一位法规知识库专家。用户提出了一个问题，"
    "请用50-120字写一段话，模拟如果相关法规文档中存在完美答案，"
    "该段落会是什么内容。\n\n"
    "要求：\n"
    "- 使用与法规文档相同的正式书面语气\n"
    "- 包含可能的条款编号、标准名称等关键术语\n"
    "- 不要解释你在做什么，直接输出假设性段落\n"
    "- 如问题过于模糊，写一段合理的通用法规说明"
)


class HyDEExpander:
    """Generate a hypothetical document passage to improve dense retrieval.

    The expander is stateless — instantiate once and call expand() per query.
    It requires no external dependencies beyond the project's existing LLM
    client infrastructure.
    """

    def expand(self, query: str) -> str:
        """Return a combined retrieval query: original query + hypothetical passage.

        The combination ensures:
        - Dense retrieval uses the enriched hypothetical text (semantic match).
        - BM25 retrieval still benefits from the original query keywords.

        The model used is ``settings.hyde_llm_model`` (dedicated lightweight model)
        falling back to the main ``settings.llm_model`` when not configured.

        If the LLM call fails for any reason, returns the original query unchanged.
        """
        if not settings.hyde_enabled:
            return query

        # Use the dedicated HyDE model when configured; fall back to main LLM.
        # A lightweight model (e.g. qwen3.6-flash) is sufficient for generating
        # a short hypothetical passage and significantly reduces cost + latency.
        provider = settings.hyde_llm_provider or settings.llm_provider
        model = settings.hyde_llm_model or settings.llm_model

        try:
            client = get_llm_client(provider=provider, model=model)
            resp = client.chat(
                messages=[
                    {"role": "system", "content": _HYDE_SYSTEM},
                    {"role": "user", "content": f"问题：{query}"},
                ],
                max_tokens=settings.hyde_max_tokens,
                # Low temperature: we want a plausible, deterministic passage.
                temperature=0.3,
            )
            if not resp.is_success or not resp.content:
                logger.debug("HyDE LLM call failed or empty — using original query")
                return query

            hypothesis = resp.content.strip()[:_MAX_HYPOTHESIS_CHARS]
            logger.debug("HyDE expanded query ({}→{} chars)", len(query), len(hypothesis))

            # Concatenate: the embedding model will see the full combined text,
            # so the resulting vector leans toward the hypothetical document style.
            return f"{query}\n\n{hypothesis}"

        except Exception as exc:  # noqa: BLE001 — intentional broad catch for graceful fallback
            logger.warning("HyDE expansion failed: {} — using original query", exc)
            return query