226 lines
10 KiB
Python
226 lines
10 KiB
Python
"""Durable sanitized context assembly traces for PMFlow AI."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
from datetime import datetime
|
|
from typing import Any
|
|
|
|
from sqlalchemy import and_, desc, select, update
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from ..models import PMFlowAIContextAssembly, PMFlowAIThread
|
|
from .pmflow_ai_context_pack_service import (
|
|
context_pack_observability_packet,
|
|
context_pack_used_context_packet,
|
|
evaluate_context_replay_case,
|
|
)
|
|
|
|
|
|
class PMFlowAIContextAssemblyService:
|
|
"""Persist and read customer-safe context assembly observability."""
|
|
|
|
def __init__(self, db: AsyncSession, *, tenant_id: int) -> None:
|
|
self.db = db
|
|
self.tenant_id = int(tenant_id or 1)
|
|
|
|
async def persist_for_run(
|
|
self,
|
|
*,
|
|
thread: PMFlowAIThread | None,
|
|
current_user: Any,
|
|
context_pack: dict[str, Any],
|
|
user_question: str,
|
|
run_public_id: str | None,
|
|
grounding_packet: dict[str, Any] | None = None,
|
|
) -> dict[str, Any]:
|
|
observability = context_pack_observability_packet(context_pack)
|
|
used_context = context_pack_used_context_packet(context_pack)
|
|
retrieval_critic = context_pack.get("retrieval_critic_packet") if isinstance(context_pack, dict) else {}
|
|
retrieval_critic = retrieval_critic if isinstance(retrieval_critic, dict) else {}
|
|
replay = evaluate_context_replay_case(
|
|
context_pack={
|
|
"packet_version": "pmflow_ai_context_assembly_replay_source.v1",
|
|
"used_context": used_context,
|
|
"observability": observability,
|
|
"retrieval_critic": retrieval_critic,
|
|
},
|
|
expected_evidence=[],
|
|
)
|
|
row = PMFlowAIContextAssembly(
|
|
tenant_id=self.tenant_id,
|
|
public_assembly_id=f"ctxasm_{uuid.uuid4().hex[:16]}",
|
|
thread_id=int(thread.id) if thread is not None and getattr(thread, "id", None) else None,
|
|
user_id=_safe_int(getattr(current_user, "id", None)),
|
|
qa_session_id=_safe_int(getattr(thread, "qa_session_id", None)) if thread is not None else None,
|
|
run_public_id=str(run_public_id or "")[:80] or None,
|
|
context_packet_version=str(context_pack.get("packet_version") or "")[:80] if isinstance(context_pack, dict) else "",
|
|
user_question_preview=_clean_text(user_question, limit=800),
|
|
assembly_payload=sanitized_context_assembly_payload(context_pack),
|
|
observability_payload=observability,
|
|
used_context_payload=used_context,
|
|
retrieval_critic_payload=retrieval_critic,
|
|
grounding_payload=grounding_packet if isinstance(grounding_packet, dict) else {},
|
|
replay_payload=replay,
|
|
assembled_at=datetime.utcnow(),
|
|
)
|
|
self.db.add(row)
|
|
await self.db.flush()
|
|
return context_assembly_packet(row)
|
|
|
|
async def attach_thread_to_run(self, *, run_public_id: str | None, thread: PMFlowAIThread) -> int:
|
|
if not run_public_id or not getattr(thread, "id", None):
|
|
return 0
|
|
result = await self.db.execute(
|
|
update(PMFlowAIContextAssembly)
|
|
.where(
|
|
and_(
|
|
PMFlowAIContextAssembly.tenant_id == self.tenant_id,
|
|
PMFlowAIContextAssembly.run_public_id == str(run_public_id)[:80],
|
|
PMFlowAIContextAssembly.thread_id.is_(None),
|
|
)
|
|
)
|
|
.values(thread_id=int(thread.id), qa_session_id=_safe_int(getattr(thread, "qa_session_id", None)))
|
|
)
|
|
return int(result.rowcount or 0)
|
|
|
|
async def latest_for_thread(self, *, thread_id: int, limit: int = 8) -> list[dict[str, Any]]:
|
|
rows = await self._assembly_rows(
|
|
filters=[PMFlowAIContextAssembly.thread_id == int(thread_id)],
|
|
limit=limit,
|
|
)
|
|
return [context_assembly_packet(row) for row in rows]
|
|
|
|
async def latest_for_run(self, *, run_public_id: str) -> dict[str, Any] | None:
|
|
rows = await self._assembly_rows(
|
|
filters=[PMFlowAIContextAssembly.run_public_id == str(run_public_id)[:80]],
|
|
limit=1,
|
|
)
|
|
return context_assembly_packet(rows[0]) if rows else None
|
|
|
|
async def _assembly_rows(self, *, filters: list[Any], limit: int) -> list[PMFlowAIContextAssembly]:
|
|
result = await self.db.execute(
|
|
select(PMFlowAIContextAssembly)
|
|
.where(and_(PMFlowAIContextAssembly.tenant_id == self.tenant_id, *filters))
|
|
.order_by(desc(PMFlowAIContextAssembly.assembled_at), desc(PMFlowAIContextAssembly.id))
|
|
.limit(max(1, min(50, int(limit or 8))))
|
|
)
|
|
return list(result.scalars().all())
|
|
|
|
|
|
def sanitized_context_assembly_payload(context_pack: Any) -> dict[str, Any]:
|
|
if not isinstance(context_pack, dict):
|
|
return {}
|
|
thread_context = context_pack.get("thread_context") if isinstance(context_pack.get("thread_context"), dict) else {}
|
|
memory_context = context_pack.get("memory_context") if isinstance(context_pack.get("memory_context"), dict) else {}
|
|
topic_context = context_pack.get("topic_context") if isinstance(context_pack.get("topic_context"), dict) else {}
|
|
evidence = context_pack.get("context_evidence_packet") if isinstance(context_pack.get("context_evidence_packet"), dict) else {}
|
|
budget = context_pack.get("budget_report") if isinstance(context_pack.get("budget_report"), dict) else {}
|
|
return {
|
|
"packet_version": "pmflow_ai_context_assembly.v1",
|
|
"context_packet_version": str(context_pack.get("packet_version") or "")[:80],
|
|
"thread_refs": [
|
|
{
|
|
"message_id": str(item.get("message_id") or "")[:80],
|
|
"role": str(item.get("role") or "")[:32],
|
|
"message_index": item.get("message_index"),
|
|
"preview": _clean_text(item.get("content") or item.get("summary") or "", limit=260),
|
|
}
|
|
for item in [*list(thread_context.get("relevant_messages") or []), *list(thread_context.get("recent_messages") or [])][-18:]
|
|
if isinstance(item, dict)
|
|
],
|
|
"memory_refs": [
|
|
{
|
|
"memory_id": str(item.get("memory_id") or "")[:80],
|
|
"memory_scope": str(item.get("memory_scope") or "")[:40],
|
|
"confidence_score": item.get("confidence_score"),
|
|
"summary_preview": _clean_text(item.get("summary") or "", limit=260),
|
|
}
|
|
for item in list(memory_context.get("memories") or [])[:18]
|
|
if isinstance(item, dict)
|
|
],
|
|
"topic_refs": [
|
|
{
|
|
"topic_id": str(item.get("topic_id") or "")[:80],
|
|
"topic_signature": str(item.get("topic_signature") or "")[:80],
|
|
"summary_preview": _clean_text(item.get("topic_summary") or item.get("current_objective") or "", limit=260),
|
|
"confidence_score": item.get("confidence_score"),
|
|
}
|
|
for item in list(topic_context.get("topic_workspaces") or [])[:12]
|
|
if isinstance(item, dict)
|
|
],
|
|
"context_evidence": {
|
|
"claim_count": len(list(evidence.get("claims") or [])),
|
|
"obligation_count": len(list(evidence.get("continuation_obligations") or [])),
|
|
"required_claims": [
|
|
{
|
|
"claim_id": str(item.get("claim_id") or "")[:80],
|
|
"semantic_role": str(item.get("semantic_role") or "")[:80],
|
|
"source_kind": str(item.get("source_kind") or "")[:80],
|
|
"preview": _clean_text(item.get("claim_text"), limit=220),
|
|
}
|
|
for item in list(evidence.get("claims") or [])[:12]
|
|
if isinstance(item, dict) and bool(item.get("required_for_answer"))
|
|
],
|
|
},
|
|
"budget": {
|
|
"max_chars": budget.get("max_chars"),
|
|
"used_chars": budget.get("used_chars"),
|
|
"sections": [
|
|
{
|
|
"section": str(item.get("section") or "")[:80],
|
|
"status": str(item.get("status") or "")[:40],
|
|
"priority": item.get("priority"),
|
|
"included_chars": item.get("included_chars"),
|
|
}
|
|
for item in list(budget.get("sections") or [])
|
|
if isinstance(item, dict)
|
|
],
|
|
},
|
|
}
|
|
|
|
|
|
def context_assembly_packet(row: PMFlowAIContextAssembly) -> dict[str, Any]:
|
|
return {
|
|
"packet_version": "pmflow_ai_context_assembly_record.v1",
|
|
"assembly_id": str(row.public_assembly_id or "")[:80],
|
|
"run_id": str(row.run_public_id or "")[:80],
|
|
"context_packet_version": str(row.context_packet_version or "")[:80],
|
|
"user_question_preview": str(row.user_question_preview or "")[:800],
|
|
"assembly": row.assembly_payload if isinstance(row.assembly_payload, dict) else {},
|
|
"observability": row.observability_payload if isinstance(row.observability_payload, dict) else {},
|
|
"used_context": row.used_context_payload if isinstance(row.used_context_payload, dict) else {},
|
|
"retrieval_critic": row.retrieval_critic_payload if isinstance(row.retrieval_critic_payload, dict) else {},
|
|
"grounding": row.grounding_payload if isinstance(row.grounding_payload, dict) else {},
|
|
"replay": row.replay_payload if isinstance(row.replay_payload, dict) else {},
|
|
"assembled_at": _iso(row.assembled_at),
|
|
}
|
|
|
|
|
|
def _clean_text(value: Any, *, limit: int) -> str:
|
|
return " ".join(str(value or "").replace("\x00", " ").split())[: max(0, int(limit or 0))]
|
|
|
|
|
|
def _safe_int(value: Any) -> int | None:
|
|
try:
|
|
if value in (None, ""):
|
|
return None
|
|
return int(value)
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|
|
|
|
def _iso(value: Any) -> str | None:
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, datetime):
|
|
return value.isoformat()
|
|
return str(value)
|
|
|
|
|
|
__all__ = [
|
|
"PMFlowAIContextAssemblyService",
|
|
"context_assembly_packet",
|
|
"sanitized_context_assembly_payload",
|
|
]
|