"""Adaptive RAG mixin module.""" from __future__ import annotations from .shared import * class AdaptiveRAGAnswerRetrievalMixin: async def _build_retrieval_response( self, *, clean_q: str, effective_q: str, asked_by: Optional[int], asked_via: str, project_id: Optional[int], source_chat_id: Optional[str], session_id: Optional[int], turn_index: int, context_session_id: Optional[int], conversation: list[dict[str, Any]], state_row: Optional[QASessionState], rewrite_meta: dict[str, Any], query_type: str, thinking_mode: str, active_thinking_mode: str, deep_thinking: bool, forced_deep_reason: str, deep_retrieval_steps: int, deep_vector_top_k: int, system_snapshot: dict[str, Any], user_snapshot: Optional[dict[str, Any]], started: float, message_intent: Optional[str] = None, mixed_signal_part: Optional[str] = None, ) -> dict[str, Any]: domains = self._route_domains(effective_q) if str(asked_via or "").strip().lower() == "feishu" and source_chat_id and "feishu" not in domains: domains.append("feishu") retrieved, search_meta = await self._multi_step_search( question=effective_q, domains=domains, project_id=project_id, source_chat_id=source_chat_id, session_state=state_row, max_steps_override=deep_retrieval_steps if deep_thinking else None, vector_top_k=deep_vector_top_k if deep_thinking else None, ) db_query_ms = int(search_meta.get("db_query_ms") or 0) vector_search_ms = int(search_meta.get("vector_search_ms") or 0) graph_query_ms = int(search_meta.get("graph_query_ms") or 0) db_timeout = bool(search_meta.get("db_timeout")) vector_timeout = bool(search_meta.get("vector_timeout")) graph_timeout = bool(search_meta.get("graph_timeout")) tool_trace = list(search_meta.get("tool_trace") or []) tool_trace.insert( 0, { "step": 0, "tool": "thinking_mode", "mode": active_thinking_mode, "fast_path_bypassed": bool(deep_thinking), "reason": forced_deep_reason or ( "manual_deep" if self._normalize_thinking_mode(thinking_mode) == THINKING_MODE_DEEP else "normal" ), }, ) final_domains = [str(x) for x in (search_meta.get("domains_used") or domains)] max_score = retrieved[0].score if retrieved else 0.0 answer_text = await self._generate_answer( effective_q, retrieved, conversation, project_id, system_snapshot=system_snapshot, user_snapshot=user_snapshot, relevance_max=max_score, thinking_mode=active_thinking_mode, context_block_limit=DEEP_CONTEXT_BLOCKS if deep_thinking else None, answer_max_tokens=DEEP_ANSWER_MAX_TOKENS if deep_thinking else None, asked_via=asked_via, message_intent=message_intent, mixed_signal_part=mixed_signal_part, ) used_sources = self._dedupe_sources(retrieved) if not used_sources: used_sources = [ { "source": "system", "label": "System Snapshot", "score": round(max(max_score, 0.6), 4), "metadata": {"mode": "snapshot_only", "relevance_max": round(max_score, 4)}, } ] if user_snapshot: used_sources.append( { "source": "system", "label": "User Snapshot", "score": round(max(max_score, 0.58), 4), "metadata": {"mode": "snapshot_only", "relevance_max": round(max_score, 4)}, } ) auto_retry_reason = "" if not deep_thinking: auto_retry_reason = self._should_auto_deep_retry( question=clean_q, effective_question=effective_q, sources=used_sources, relevance_max=max_score, ) if auto_retry_reason and not self._allow_auto_deep_retry_for_channel( asked_via=asked_via, question=clean_q, effective_question=effective_q, reason=auto_retry_reason, ): auto_retry_reason = "" if auto_retry_reason: deep_retrieved, deep_meta = await self._multi_step_search( question=effective_q, domains=domains, project_id=project_id, source_chat_id=source_chat_id, session_state=state_row, max_steps_override=deep_retrieval_steps, vector_top_k=deep_vector_top_k, ) deep_max_score = deep_retrieved[0].score if deep_retrieved else 0.0 deep_answer_text = await self._generate_answer( effective_q, deep_retrieved, conversation, project_id, system_snapshot=system_snapshot, user_snapshot=user_snapshot, relevance_max=deep_max_score, thinking_mode=THINKING_MODE_DEEP, context_block_limit=DEEP_CONTEXT_BLOCKS, answer_max_tokens=DEEP_ANSWER_MAX_TOKENS, asked_via=asked_via, message_intent=message_intent, mixed_signal_part=mixed_signal_part, ) deep_sources = self._dedupe_sources(deep_retrieved) if not deep_sources: deep_sources = [ { "source": "system", "label": "System Snapshot", "score": round(max(deep_max_score, 0.6), 4), "metadata": {"mode": "snapshot_only", "relevance_max": round(deep_max_score, 4)}, } ] if user_snapshot: deep_sources.append( { "source": "system", "label": "User Snapshot", "score": round(max(deep_max_score, 0.58), 4), "metadata": {"mode": "snapshot_only", "relevance_max": round(deep_max_score, 4)}, } ) merged_metrics = self._merge_search_metrics( { "db_query_ms": db_query_ms, "vector_search_ms": vector_search_ms, "graph_query_ms": graph_query_ms, "db_timeout": db_timeout, "vector_timeout": vector_timeout, "graph_timeout": graph_timeout, }, deep_meta, ) db_query_ms = int(merged_metrics.get("db_query_ms") or 0) vector_search_ms = int(merged_metrics.get("vector_search_ms") or 0) graph_query_ms = int(merged_metrics.get("graph_query_ms") or 0) db_timeout = bool(merged_metrics.get("db_timeout")) vector_timeout = bool(merged_metrics.get("vector_timeout")) graph_timeout = bool(merged_metrics.get("graph_timeout")) deep_trace = list(deep_meta.get("tool_trace") or []) tool_trace.append( { "step": 98, "tool": "auto_deep_retry", "reason": auto_retry_reason, "prev_relevance_max": round(float(max_score or 0.0), 4), "prev_snapshot_only": self._is_snapshot_only_sources(used_sources), } ) for row in deep_trace: tool_trace.append(row) final_domains = [str(x) for x in (deep_meta.get("domains_used") or final_domains or domains)] max_score = deep_max_score answer_text = deep_answer_text used_sources = deep_sources deep_thinking = True active_thinking_mode = THINKING_MODE_DEEP evidence_guard_triggered = self._should_block_answer_with_guard( question=clean_q, effective_question=effective_q, sources=used_sources, relevance_max=max_score, ) if evidence_guard_triggered: answer_text = self._build_evidence_guard_answer(clean_q) used_sources = [ { "source": "system", "label": "Evidence Guard", "score": 0.99, "metadata": {"mode": "guard", "reason": "snapshot_only_process_query"}, } ] max_score = min(float(max_score or 0.0), 0.35) if "system_context" not in final_domains: final_domains.append("system_context") tool_trace.append( { "step": 99, "tool": "answer_guard", "reason": "insufficient_call_chain_evidence", } ) elapsed = int((time.perf_counter() - started) * 1000) qa = QARecord( tenant_id=self.tenant_id, question=clean_q, answer=answer_text, sources=used_sources, intent="question", domains_used=final_domains, response_time_ms=elapsed, asked_by=asked_by, asked_via=(asked_via or "pmflow")[:20], project_id=project_id, session_id=session_id, turn_index=turn_index, ) self.db.add(qa) confidence = round(max(0.2, min(1.0, max_score if max_score > 0 else 0.55)), 4) if evidence_guard_triggered: confidence = min(confidence, 0.35) missing_fields = [] if max_score < 0.45: missing_fields = ["high_relevance_evidence"] if evidence_guard_triggered: missing_fields = sorted(set([*missing_fields, "call_chain_evidence", "target_scope"])) await self._update_session_state_after_turn( session_state=state_row, question=clean_q, effective_question=effective_q, answer=answer_text, query_type=query_type if query_type != QUERY_TYPE_UNKNOWN else QUERY_TYPE_EXPLAIN, ambiguity="clear" if not missing_fields else "low_evidence", rewrite_meta=rewrite_meta, sources=used_sources, ) await self.db.commit() await self.db.refresh(qa) await self.context_cache.append_session_messages( tenant_id=self.tenant_id, session_id=context_session_id, question=clean_q, answer=answer_text, ) return self._build_response_payload( answer=answer_text, sources=used_sources, domains_used=final_domains, response_time_ms=elapsed, db_query_ms=db_query_ms, vector_search_ms=vector_search_ms, graph_query_ms=graph_query_ms, db_timeout=db_timeout, vector_timeout=vector_timeout, graph_timeout=graph_timeout, tool_trace=tool_trace, qa_id=int(qa.id), session_id=int(session_id) if session_id is not None else None, turn_index=int(turn_index), relevance_max=round(max_score, 4), query_type=query_type if query_type != QUERY_TYPE_UNKNOWN else QUERY_TYPE_EXPLAIN, confidence=confidence, missing_fields=missing_fields, thinking_mode=active_thinking_mode, )