from __future__ import annotations import json import uuid from dataclasses import dataclass from datetime import datetime from pathlib import Path from typing import Any @dataclass(frozen=True) class RetrievalTrace: trace_id: str query: str subject_scope: str intent: str decision: str reason: str resolved_drug_id: str | None citations: tuple[dict[str, Any], ...] correlation_id: str | None = None otel_trace_id: str | None = None conversation_id: str | None = None created_at: datetime | None = None # The full assistant turn (answer text, blocks, plan, ...) — see # `PostgresTraceRepository.save`'s `response_payload` param. None for # rows persisted before that column existed, or when persistence raced # a legacy caller that never passed it. response_payload: dict[str, Any] | None = None class FeedbackTraceNotFound(LookupError): """The answer trace was never persisted, so feedback cannot be linked.""" class PostgresTraceRepository: """Opens a new connection per call — no pooling (F-09: a real pool, with startup-time lifecycle, is a further improvement not made here). `connect_timeout` matters more than it looks: found live 2026-08-06 that an unreachable Postgres (packets dropped, not actively refused) makes a bare `psycopg.connect()` hang on the OS-level TCP timeout — tens of seconds, not immediate — which defeats a caller's try/except fail-open around `save()` just as effectively as no try/except at all, since the exception it's waiting for never arrives in time. `routers/rag.py` wraps `save()` to keep a trace outage from failing an already-computed answer; this bounds how long that protection can take to kick in. """ def __init__(self, dsn: str) -> None: self._dsn = dsn def migrate(self, migration_path: Path) -> None: import psycopg statement = migration_path.read_text(encoding="utf-8") with psycopg.connect(self._dsn, connect_timeout=5) as connection: connection.execute(statement) def save( self, *, query: str, subject_scope: str, intent: str, decision: str, reason: str, resolved_drug_id: str | None, citations: tuple[dict[str, Any], ...], correlation_id: str | None = None, otel_trace_id: str | None = None, conversation_id: str | None = None, response_payload: dict[str, Any] | None = None, ) -> str: import psycopg trace_id = str(uuid.uuid4()) with psycopg.connect(self._dsn, connect_timeout=5) as connection: connection.execute( """ INSERT INTO rag_retrieval_trace ( trace_id, query_text, subject_scope, query_intent, decision, reason, resolved_drug_id, citations, correlation_id, otel_trace_id, conversation_id, response_payload ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s::jsonb, %s, %s, %s, %s::jsonb) """, ( trace_id, query, subject_scope, intent, decision, reason, resolved_drug_id, json.dumps(citations, ensure_ascii=False), correlation_id, otel_trace_id, conversation_id, json.dumps(response_payload, ensure_ascii=False) if response_payload is not None else None, ), ) return trace_id def get(self, trace_id: str) -> RetrievalTrace | None: import psycopg with psycopg.connect(self._dsn, connect_timeout=5) as connection: row = connection.execute( """ SELECT trace_id::text, query_text, subject_scope, query_intent, decision, reason, resolved_drug_id, citations, correlation_id, otel_trace_id, conversation_id, created_at, response_payload FROM rag_retrieval_trace WHERE trace_id = %s """, (trace_id,), ).fetchone() if row is None: return None return RetrievalTrace( trace_id=row[0], query=row[1], subject_scope=row[2], intent=row[3], decision=row[4], reason=row[5], resolved_drug_id=row[6], citations=tuple(row[7]), correlation_id=row[8], otel_trace_id=row[9], conversation_id=row[10], created_at=row[11], response_payload=row[12], ) def list_by_conversation( self, conversation_id: str, limit: int ) -> list[RetrievalTrace]: """Past queries for one session — Feature-List #25 (lịch sử tra cứu), most recent first. Scoped to `conversation_id` on purpose: this system has no auth anywhere (`apps/api-gateway`/`auth-service` are unbuilt — see README), so an unscoped listing would mix every browser's/user's queries together. Each row's full replayable answer, when one was persisted, rides along in `response_payload` — see `routers/rag.py`'s `/transcript` endpoint, which is what turns this into an actual resumable conversation rather than just a re-run-the-query shortcut. """ import psycopg with psycopg.connect(self._dsn, connect_timeout=5) as connection: rows = connection.execute( """ SELECT trace_id::text, query_text, subject_scope, query_intent, decision, reason, resolved_drug_id, citations, correlation_id, otel_trace_id, conversation_id, created_at, response_payload FROM rag_retrieval_trace WHERE conversation_id = %s ORDER BY created_at DESC LIMIT %s """, (conversation_id, limit), ).fetchall() return [ RetrievalTrace( trace_id=row[0], query=row[1], subject_scope=row[2], intent=row[3], decision=row[4], reason=row[5], resolved_drug_id=row[6], citations=tuple(row[7]), correlation_id=row[8], otel_trace_id=row[9], conversation_id=row[10], created_at=row[11], response_payload=row[12], ) for row in rows ] def save_feedback( self, *, trace_id: str, rating: str, comment: str | None, conversation_id: str | None, ) -> str: """Create or replace one user's verdict for one persisted answer.""" import psycopg feedback_id = str(uuid.uuid4()) with psycopg.connect(self._dsn, connect_timeout=5) as connection: row = connection.execute( """ INSERT INTO rag_answer_feedback ( feedback_id, trace_id, conversation_id, rating, comment ) SELECT %s, trace_id, %s, %s, %s FROM rag_retrieval_trace WHERE trace_id = %s ON CONFLICT (trace_id) DO UPDATE SET conversation_id = EXCLUDED.conversation_id, rating = EXCLUDED.rating, comment = EXCLUDED.comment, updated_at = now() RETURNING feedback_id::text """, (feedback_id, conversation_id, rating, comment, trace_id), ).fetchone() if row is None: raise FeedbackTraceNotFound(trace_id) return str(row[0]) class PostgresConversationStore: """Durable, cross-worker replacement for `RagAgent`'s in-process `dict[str, list[str]]` history — the F-08 gap named in ADR 0008: history was lost on restart and not shared across workers. Append-only, one connection per call — the same tradeoffs as `PostgresTraceRepository` (F-09: real pooling is a further improvement not made here) and the same `connect_timeout=5` for the same reason (an unreachable-but-not-refusing host hangs on the OS TCP timeout otherwise, defeating a caller's fail-open try/except just as completely as no try/except at all). Raises on any error rather than swallowing it — `RagAgent` is the caller that decides fail-open (lose this turn's memory, not the response), matching how `routers/rag.py` already wraps `PostgresTraceRepository.save()` rather than the repository hiding its own failures. """ def __init__(self, dsn: str) -> None: self._dsn = dsn def migrate(self, migration_path: Path) -> None: import psycopg statement = migration_path.read_text(encoding="utf-8") with psycopg.connect(self._dsn, connect_timeout=5) as connection: connection.execute(statement) def recent(self, conversation_id: str, limit: int) -> list[str]: """The last `limit` lines, oldest first — matches the ordering the understanding LLM prompt already expects ("LỊCH SỬ HỘI THOẠI (cũ -> mới)").""" import psycopg with psycopg.connect(self._dsn, connect_timeout=5) as connection: rows = connection.execute( """ SELECT line FROM rag_conversation_turn WHERE conversation_id = %s ORDER BY id DESC LIMIT %s """, (conversation_id, limit), ).fetchall() return [row[0] for row in reversed(rows)] def append(self, conversation_id: str, line: str) -> None: import psycopg with psycopg.connect(self._dsn, connect_timeout=5) as connection: connection.execute( "INSERT INTO rag_conversation_turn (conversation_id, line) " "VALUES (%s, %s)", (conversation_id, line), )