96 lines
3.2 KiB
Python
96 lines
3.2 KiB
Python
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], ...]
|
|
created_at: datetime | None = None
|
|
|
|
|
|
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], ...],
|
|
) -> 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
|
|
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s::jsonb)
|
|
""",
|
|
(
|
|
trace_id, query, subject_scope, intent, decision, reason,
|
|
resolved_drug_id, json.dumps(citations, ensure_ascii=False),
|
|
),
|
|
)
|
|
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, created_at
|
|
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]), created_at=row[8],
|
|
)
|