212 lines
8.3 KiB
Python
212 lines
8.3 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"}
|
|
|
|
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()
|
|
keyword = self._extract_project_lookup_keyword(text)
|
|
if len(keyword) >= 2:
|
|
return keyword[: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,
|
|
}
|