230 lines
9.0 KiB
Python
230 lines
9.0 KiB
Python
"""Adaptive RAG mixin module."""
|
|
from __future__ import annotations
|
|
|
|
from .shared import *
|
|
|
|
|
|
class AdaptiveRAGMultiStepFlowMixin:
|
|
async def _multi_step_search(
|
|
self,
|
|
*,
|
|
question: str,
|
|
domains: list[str],
|
|
project_id: Optional[int],
|
|
source_chat_id: Optional[str],
|
|
session_state: Optional[QASessionState],
|
|
max_steps_override: Optional[int] = None,
|
|
vector_top_k: Optional[int] = None,
|
|
) -> tuple[list[RetrievedItem], dict[str, Any]]:
|
|
resolved = await self.entity_resolution(question=question, project_id=project_id, limit=6)
|
|
manifest = await self.get_project_manifest(project_id=project_id)
|
|
system_snapshot = await self.context_cache.get_system_snapshot(
|
|
db=self.db,
|
|
tenant_id=self.tenant_id,
|
|
)
|
|
symbol_hits: dict[str, Any] = {"count": 0, "items": []}
|
|
try:
|
|
symbol_service = SymbolIndexService(self.db, tenant_id=self.tenant_id)
|
|
symbol_hits = await symbol_service.search_symbols(query=question, limit=6)
|
|
except Exception:
|
|
symbol_hits = {"count": 0, "items": []}
|
|
feature_hits = self._match_feature_catalog_entries(
|
|
question=question,
|
|
system_snapshot=system_snapshot,
|
|
limit=6,
|
|
)
|
|
max_steps = max(1, min(int(max_steps_override or AGENT_MAX_RETRIEVAL_STEPS), 8))
|
|
tool_trace: list[dict[str, Any]] = [
|
|
{
|
|
"step": 0,
|
|
"tool": "entity_resolution",
|
|
"matched_entities": len(resolved.get("matched_entities") or []),
|
|
"canonical_terms": resolved.get("canonical_terms") if isinstance(resolved.get("canonical_terms"), list) else [],
|
|
},
|
|
{
|
|
"step": 0,
|
|
"tool": "get_project_manifest",
|
|
"symbol_total": int(
|
|
(
|
|
manifest.get("symbol_summary")
|
|
if isinstance(manifest.get("symbol_summary"), dict)
|
|
else {}
|
|
).get("symbol_total", 0)
|
|
),
|
|
"kg_entity_total": int(
|
|
(
|
|
manifest.get("kg_overview")
|
|
if isinstance(manifest.get("kg_overview"), dict)
|
|
else {}
|
|
).get("entity_total", 0)
|
|
),
|
|
},
|
|
{
|
|
"step": 0,
|
|
"tool": "search_symbols",
|
|
"hit_count": int(symbol_hits.get("count") or 0),
|
|
},
|
|
{
|
|
"step": 0,
|
|
"tool": "feature_catalog_lookup",
|
|
"hit_count": len(feature_hits),
|
|
"feature_names": [str(x.get("name") or "") for x in feature_hits[:4]],
|
|
},
|
|
]
|
|
if session_state is not None and isinstance(session_state.last_entity_ids, list):
|
|
anchor_ids = [int(x) for x in session_state.last_entity_ids if str(x).isdigit()][:6]
|
|
else:
|
|
anchor_ids = []
|
|
current_question = str(question or "").strip()[:1200]
|
|
current_domains: list[str] = []
|
|
for item in domains:
|
|
if item not in current_domains:
|
|
current_domains.append(item)
|
|
merged_rows: list[RetrievedItem] = []
|
|
merged_metrics: dict[str, Any] = {
|
|
"db_query_ms": 0,
|
|
"vector_search_ms": 0,
|
|
"graph_query_ms": 0,
|
|
"db_timeout": False,
|
|
"vector_timeout": False,
|
|
"graph_timeout": False,
|
|
}
|
|
merge_limit = max(DEFAULT_TOP_K * 2, 8)
|
|
if vector_top_k is not None:
|
|
merge_limit = max(merge_limit, max(2, int(vector_top_k)) * 3)
|
|
|
|
for step in range(1, max_steps + 1):
|
|
resolved_payload = {
|
|
**resolved,
|
|
"anchor_entity_ids": anchor_ids,
|
|
}
|
|
step_rows, step_metrics = await self._run_single_retrieval_pass(
|
|
question=current_question,
|
|
domains=current_domains,
|
|
project_id=project_id,
|
|
source_chat_id=source_chat_id,
|
|
resolved_entities=resolved_payload,
|
|
vector_top_k=vector_top_k,
|
|
)
|
|
feature_rows = self._retrieve_feature_catalog_context(
|
|
question=current_question,
|
|
system_snapshot=system_snapshot,
|
|
max_items=4,
|
|
)
|
|
if feature_rows:
|
|
step_rows = self._merge_retrieved_items(feature_rows, step_rows, limit=merge_limit)
|
|
merged_rows = self._merge_retrieved_items(merged_rows, step_rows, limit=merge_limit)
|
|
merged_metrics = self._merge_search_metrics(merged_metrics, step_metrics)
|
|
top_score = round(float(step_rows[0].score), 4) if step_rows else 0.0
|
|
tool_trace.append(
|
|
{
|
|
"step": step,
|
|
"tool": "parallel_retrieval",
|
|
"query": current_question[:240],
|
|
"domains": current_domains,
|
|
"items": len(step_rows),
|
|
"feature_items": len(feature_rows),
|
|
"top_score": top_score,
|
|
"merged_items": len(merged_rows),
|
|
}
|
|
)
|
|
|
|
if step >= max_steps:
|
|
tool_trace.append(
|
|
{
|
|
"step": step,
|
|
"tool": "planner_stop",
|
|
"reason": "max_steps_reached",
|
|
"max_steps": max_steps,
|
|
}
|
|
)
|
|
break
|
|
if not self._should_second_pass(merged_rows):
|
|
tool_trace.append(
|
|
{
|
|
"step": step,
|
|
"tool": "planner_stop",
|
|
"reason": "enough_evidence",
|
|
"top_score": round(float(merged_rows[0].score), 4) if merged_rows else 0.0,
|
|
}
|
|
)
|
|
break
|
|
|
|
next_question, rewrite_meta = await self._build_second_pass_query(
|
|
question=current_question,
|
|
project_id=project_id,
|
|
)
|
|
row_hint = self._summarize_retrieval_hints(step_rows)
|
|
if row_hint and row_hint not in next_question:
|
|
next_question = f"{next_question}。已命中线索：{row_hint}"
|
|
next_question = next_question[:1200]
|
|
if next_question.strip() == current_question.strip():
|
|
tool_trace.append(
|
|
{
|
|
"step": step,
|
|
"tool": "planner_stop",
|
|
"reason": "query_not_changed",
|
|
"rewrite_meta": rewrite_meta,
|
|
}
|
|
)
|
|
break
|
|
|
|
next_domains: list[str] = []
|
|
for item in [*current_domains, "graph", "db", "doc", "pmflow", "feishu"]:
|
|
if item not in next_domains:
|
|
next_domains.append(item)
|
|
tool_trace.append(
|
|
{
|
|
"step": step + 1,
|
|
"tool": "query_refine",
|
|
"rewrite_meta": rewrite_meta,
|
|
"next_query": next_question[:240],
|
|
"next_domains": next_domains,
|
|
"anchor_entity_ids": anchor_ids,
|
|
}
|
|
)
|
|
current_question = next_question
|
|
current_domains = next_domains
|
|
|
|
result_cap = max(DEFAULT_TOP_K * 2, 8)
|
|
if vector_top_k is not None:
|
|
result_cap = max(result_cap, max(2, int(vector_top_k)) * 2)
|
|
|
|
return merged_rows[:result_cap], {
|
|
**merged_metrics,
|
|
"domains_used": current_domains,
|
|
"tool_trace": tool_trace,
|
|
}
|
|
|
|
async def multi_step_search(
|
|
self,
|
|
*,
|
|
question: str,
|
|
project_id: Optional[int] = None,
|
|
source_chat_id: Optional[str] = None,
|
|
session_state: Optional[QASessionState] = None,
|
|
) -> dict[str, Any]:
|
|
domains = self._route_domains(question)
|
|
rows, metrics = await self._multi_step_search(
|
|
question=question,
|
|
domains=domains,
|
|
project_id=project_id,
|
|
source_chat_id=source_chat_id,
|
|
session_state=session_state,
|
|
)
|
|
return {
|
|
"question": question,
|
|
"domains": metrics.get("domains_used") or domains,
|
|
"items": self._dedupe_sources(rows),
|
|
"retrieval_metrics": {
|
|
"db_query_ms": int(metrics.get("db_query_ms") or 0),
|
|
"vector_search_ms": int(metrics.get("vector_search_ms") or 0),
|
|
"graph_query_ms": int(metrics.get("graph_query_ms") or 0),
|
|
"db_timeout": bool(metrics.get("db_timeout")),
|
|
"vector_timeout": bool(metrics.get("vector_timeout")),
|
|
"graph_timeout": bool(metrics.get("graph_timeout")),
|
|
},
|
|
"tool_trace": metrics.get("tool_trace") if isinstance(metrics.get("tool_trace"), list) else [],
|
|
}
|
|
|