"""Adaptive RAG mixin module.""" from __future__ import annotations from .shared import * class AdaptiveRAGSessionStateMixin: async def _ensure_session_state(self, *, session_id: Optional[int]) -> Optional[QASessionState]: if session_id is None: return None row = ( await self.db.execute( select(QASessionState).where( and_( QASessionState.tenant_id == self.tenant_id, QASessionState.session_id == int(session_id), ) ) ) ).scalar_one_or_none() if row is None: row = QASessionState( tenant_id=self.tenant_id, session_id=int(session_id), last_subject=None, last_entity_ids=[], last_intent=None, ambiguity="unknown", rewrite_history=[], pending_clarification=None, pending_clarification_expires_at=None, ) self.db.add(row) await self.db.flush() return row def _rewrite_question_with_state( self, *, question: str, conversation: list[dict[str, str]], session_state: Optional[QASessionState], anchor_resolution: Optional[FollowupAnchorResolution] = None, ) -> tuple[str, dict[str, Any]]: _ = conversation current = str(question or "").strip() if not current: return current, {"mode": "empty"} resolution = anchor_resolution or FollowupAnchorResolution() if str(resolution.followup_mode or "").strip().lower() != "referential": return current, {"mode": "direct"} explicit_anchor_mode = str(resolution.anchor_mode or "").strip().lower() if explicit_anchor_mode == "explicit": anchored_question, anchor_meta = self._rewrite_question_with_followup_anchor( question=current, anchor_resolution=resolution, ) if str(anchor_meta.get("mode") or "").strip() not in {"", "direct"}: return anchored_question, anchor_meta state_subject = "" state_entity_ids: list[int] = [] if session_state is not None: state_subject = str(session_state.last_subject or "").strip() if isinstance(session_state.last_entity_ids, list): for value in session_state.last_entity_ids: try: state_entity_ids.append(int(value)) except Exception: continue if state_subject: rewritten = f"围绕“{state_subject}”回答：{current}" return rewritten[:1200], {"mode": "state_subject", "subject": state_subject} if state_entity_ids: anchor = ",".join(str(x) for x in state_entity_ids[:6]) rewritten = f"围绕会话实体ID[{anchor}]回答：{current}" return rewritten[:1200], {"mode": "state_entities", "entity_ids": state_entity_ids[:6]} anchored_question, anchor_meta = self._rewrite_question_with_followup_anchor( question=current, anchor_resolution=resolution, ) if str(anchor_meta.get("mode") or "").strip() not in {"", "direct"}: return anchored_question, anchor_meta return current, {"mode": "referential_without_subject"} async def _update_session_state_after_turn( self, *, session_state: Optional[QASessionState], question: str, effective_question: str, answer: str, query_type: str, ambiguity: str, rewrite_meta: dict[str, Any], sources: list[dict[str, Any]], ) -> None: if session_state is None: return session_state.last_intent = "question" session_state.ambiguity = str(ambiguity or "clear")[:20] if str(ambiguity or "").strip() != "needs_clarification": subject = self._extract_state_subject(question=question, answer=answer, sources=sources) if subject: session_state.last_subject = subject[:255] session_state.last_entity_ids = self._collect_entity_ids_from_sources(sources) self._clear_pending_clarification(session_state=session_state) history = session_state.rewrite_history if isinstance(session_state.rewrite_history, list) else [] history.append( { "at": now_cn_aware().isoformat(), "question": question[:500], "effective_question": effective_question[:500], "query_type": str(query_type or QUERY_TYPE_UNKNOWN), "ambiguity": str(ambiguity or "clear"), "rewrite_meta": rewrite_meta if isinstance(rewrite_meta, dict) else {}, } ) session_state.rewrite_history = history[-20:] await self.db.flush() def _collect_entity_ids_from_sources(self, sources: list[dict[str, Any]]) -> list[int]: ids: list[int] = [] for source in sources: if not isinstance(source, dict): continue meta = source.get("metadata") if isinstance(source.get("metadata"), dict) else {} for key in ( "project_id", "task_id", "signal_id", "user_id", "entity_id", "subject_entity_id", "object_entity_id", ): value = meta.get(key) try: if value is not None: ids.append(int(value)) except Exception: continue matched = meta.get("matched_project_ids") if isinstance(matched, list): for value in matched: try: ids.append(int(value)) except Exception: continue uniq: list[int] = [] for value in ids: if value not in uniq: uniq.append(value) return uniq[:20] def _extract_state_subject( self, *, question: str, answer: str, sources: list[dict[str, Any]], ) -> str: for source in sources: if not isinstance(source, dict): continue meta = source.get("metadata") if isinstance(source.get("metadata"), dict) else {} project_keyword = str(meta.get("query_keyword") or "").strip() if project_keyword: return project_keyword[:255] subject_name = str(meta.get("subject_name") or "").strip() object_name = str(meta.get("object_name") or "").strip() if subject_name and object_name: return f"{subject_name}->{object_name}"[:255] text = str(question or "").strip() normalizer = getattr(self, "_normalize_lookup_text", None) normalized_question = normalizer(text) if callable(normalizer) else text if 2 <= len(normalized_question) <= 64: return normalized_question[:255] answer_text = str(answer or "").strip() m = re.search(r"项目[「\"]?([^」\"\n]{2,40})", answer_text) if m: return str(m.group(1)).strip()[:255] return "" async def get_session_state_replay(self, *, session_id: int, user_id: int) -> dict[str, Any]: session = ( await self.db.execute( select(QASession).where( and_( QASession.tenant_id == self.tenant_id, QASession.id == int(session_id), QASession.user_id == int(user_id), ) ) ) ).scalar_one_or_none() if session is None: raise ValueError("qa session 不存在") state = await self._ensure_session_state(session_id=int(session_id)) history = state.rewrite_history if isinstance(state.rewrite_history, list) else [] return { "session_id": int(session_id), "last_subject": str(state.last_subject or "") if state else "", "last_entity_ids": state.last_entity_ids if state else [], "last_intent": str(state.last_intent or "") if state else "", "ambiguity": str(state.ambiguity or "") if state else "", "rewrite_history": history[-20:], "pending_clarification": state.pending_clarification if state else None, "pending_clarification_expires_at": ( state.pending_clarification_expires_at.isoformat() if state and state.pending_clarification_expires_at else None ), "updated_at": state.updated_at.isoformat() if state and state.updated_at else None, }