106 lines
4.3 KiB
Python
106 lines
4.3 KiB
Python
|
"""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.5-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
|