"""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], query_semantics: Optional[dict[str, Any]] = None, session_state: Optional[QASessionState], max_steps_override: Optional[int] = None, vector_top_k: Optional[int] = None, ) -> tuple[list[RetrievedItem], dict[str, Any]]: semantics = query_semantics if isinstance(query_semantics, dict) else {} external_freshness_only = bool(semantics.get("external_freshness_only")) resolved: dict[str, Any] = {"matched_entities": [], "canonical_terms": []} manifest: dict[str, Any] = {} system_snapshot: dict[str, Any] = {} symbol_hits: dict[str, Any] = {"count": 0, "items": []} feature_hits: list[dict[str, Any]] = [] if not external_freshness_only: 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, ) 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 external_freshness_only: tool_trace.append( { "step": 0, "tool": "freshness_guard", "mode": "external_freshness_only", } ) 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, query_semantics=query_semantics, resolved_entities=resolved_payload, vector_top_k=vector_top_k, ) feature_rows = [] if not external_freshness_only: 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, query_semantics=query_semantics, ) 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] = [] if external_freshness_only: next_domains = list(current_domains) else: 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, query_semantics: Optional[dict[str, Any]] = None, session_state: Optional[QASessionState] = None, ) -> dict[str, Any]: semantics = query_semantics if isinstance(query_semantics, dict) else await self._resolve_query_semantics( message=question, source_system="adaptive_rag_multi_step_search", metadata={ "project_id": project_id, "source_chat_id": source_chat_id, }, ) domains = [str(item).strip().lower() for item in (semantics.get("retrieval_domains") or []) if str(item or "").strip()] if not domains: domains = ["db", "graph", "doc", "feishu", "pmflow", "url"] rows, metrics = await self._multi_step_search( question=question, domains=domains, project_id=project_id, source_chat_id=source_chat_id, query_semantics=semantics, 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 [], }