"""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", ]