236 lines
9.0 KiB
Python
236 lines
9.0 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], ...]
|
|
correlation_id: str | None = None
|
|
otel_trace_id: str | None = None
|
|
conversation_id: str | None = None
|
|
created_at: datetime | 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,
|
|
) -> 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
|
|
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s::jsonb, %s, %s, %s)
|
|
""",
|
|
(
|
|
trace_id, query, subject_scope, intent, decision, reason,
|
|
resolved_drug_id, json.dumps(citations, ensure_ascii=False),
|
|
correlation_id, otel_trace_id, conversation_id,
|
|
),
|
|
)
|
|
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
|
|
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],
|
|
)
|
|
|
|
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. Citations/answer text are NOT
|
|
persisted here (only decision/reason/resolved_drug_id) — a history
|
|
entry is for re-running the same query, not replaying its old
|
|
answer verbatim.
|
|
"""
|
|
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
|
|
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],
|
|
)
|
|
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),
|
|
)
|