240 lines
9.5 KiB
Python
240 lines
9.5 KiB
Python
"""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]
|
||
|
||
@staticmethod
|
||
def _extract_project_subject_from_answer(answer: str) -> str:
|
||
text = str(answer or "").strip()
|
||
if not text:
|
||
return ""
|
||
patterns = (
|
||
r"项目名(?:称)?[：: ]*[“\"]?([^”\"，。；、\n]{2,80})[”\"]?",
|
||
r"项目名[“\"]([^”\"\n]{2,80})[”\"]",
|
||
r"项目\s*#\d+[《“\"]([^》”\"\n]{2,80})[》”\"]",
|
||
r"#\d+\s+([^\n（(]{2,80})[（(]",
|
||
r"项目[《“\"]([^》”\"\n]{2,80})[》”\"]",
|
||
)
|
||
for pattern in patterns:
|
||
match = re.search(pattern, text)
|
||
if match:
|
||
return str(match.group(1)).strip()[:255]
|
||
return ""
|
||
|
||
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]
|
||
|
||
answer_subject = self._extract_project_subject_from_answer(answer)
|
||
if answer_subject:
|
||
return answer_subject[: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]
|
||
|
||
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,
|
||
}
|