| """TraceabilityStore — the seam the chat pipeline persists provenance through (KM-691). |
| |
| `ChatHandler` (and the v2 chat endpoint's greeting/cache branches) flush one |
| `TraceabilityPayload` per assistant turn through this seam, right before the `done` |
| SSE event; `GET /api/v1/traceability` reads it back by (analysis_id, message_id). |
| |
| - `NullTraceabilityStore` logs and stores nothing (tests / disabled persistence). |
| - `PostgresTraceabilityStore` writes one `message_traceability` row per turn |
| (dedorch, `AsyncSessionLocal`), mirroring `PostgresReportInputStore`. |
| |
| `save` must NEVER raise on the caller's path — a persistence failure must not break |
| the user's answer. `get` is the endpoint read and returns `None` on a miss (→ 404). |
| """ |
|
|
| from __future__ import annotations |
|
|
| from typing import Protocol, runtime_checkable |
|
|
| from sqlalchemy import select |
| from sqlalchemy.dialects.postgresql import insert |
|
|
| from src.db.postgres.connection import AsyncSessionLocal |
| from src.db.postgres.models import MessageTraceabilityRow |
| from src.middlewares.logging import get_logger |
|
|
| from .schemas import TraceabilityPayload |
|
|
| logger = get_logger("traceability_store") |
|
|
|
|
| @runtime_checkable |
| class TraceabilityStore(Protocol): |
| """Persist + read one provenance record per assistant `message_id`. |
| |
| `save` must never raise on the caller's path. `get` returns the payload for one |
| turn, or `None` if none exists yet (the turn is still running or the id is unknown). |
| """ |
|
|
| async def save(self, payload: TraceabilityPayload) -> None: ... |
|
|
| async def get( |
| self, analysis_id: str, message_id: str |
| ) -> TraceabilityPayload | None: ... |
|
|
|
|
| class NullTraceabilityStore: |
| """No-op store: logs the payload, persists nothing. Reads return `None`.""" |
|
|
| async def save(self, payload: TraceabilityPayload) -> None: |
| logger.info( |
| "traceability produced (not persisted — NullTraceabilityStore)", |
| message_id=payload.message_id, |
| intent=payload.intent, |
| n_tool_calls=len(payload.tool_calls), |
| ) |
|
|
| async def get( |
| self, analysis_id: str, message_id: str |
| ) -> TraceabilityPayload | None: |
| return None |
|
|
|
|
| class PostgresTraceabilityStore: |
| """Writes/reads `message_traceability` jsonb rows. Upsert on `message_id`.""" |
|
|
| async def save(self, payload: TraceabilityPayload) -> None: |
| try: |
| data = payload.model_dump(mode="json", by_alias=True) |
| async with AsyncSessionLocal() as session: |
| stmt = insert(MessageTraceabilityRow).values( |
| message_id=payload.message_id, |
| analysis_id=payload.analysis_id, |
| user_id=payload.user_id, |
| intent=payload.intent, |
| data=data, |
| ) |
| |
| stmt = stmt.on_conflict_do_update( |
| index_elements=[MessageTraceabilityRow.message_id], |
| set_={"data": stmt.excluded.data, "intent": stmt.excluded.intent}, |
| ) |
| await session.execute(stmt) |
| await session.commit() |
| logger.info( |
| "traceability persisted", |
| message_id=payload.message_id, |
| analysis_id=payload.analysis_id, |
| intent=payload.intent, |
| ) |
| except Exception as exc: |
| logger.error( |
| "traceability persist failed", |
| message_id=payload.message_id, |
| error=str(exc), |
| ) |
|
|
| async def get( |
| self, analysis_id: str, message_id: str |
| ) -> TraceabilityPayload | None: |
| async with AsyncSessionLocal() as session: |
| result = await session.execute( |
| select(MessageTraceabilityRow.data).where( |
| MessageTraceabilityRow.message_id == message_id, |
| MessageTraceabilityRow.analysis_id == analysis_id, |
| ) |
| ) |
| row = result.scalar_one_or_none() |
| if row is None: |
| return None |
| return TraceabilityPayload.model_validate(row) |
|
|
