File size: 4,213 Bytes
f873f92 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | """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,
)
# Idempotent: a re-flushed turn overwrites its own row.
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: # never break the user's answer
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)
|