333 lines
13 KiB
Python
333 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import uuid
|
|
from contextlib import suppress
|
|
from functools import lru_cache
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
pytestmark = pytest.mark.skipif(
|
|
os.getenv("RUN_INTEGRATION") != "1",
|
|
reason="set RUN_INTEGRATION=1 with local PostgreSQL and Qdrant running",
|
|
)
|
|
|
|
ROOT = Path(__file__).resolve().parents[3]
|
|
CHUNKS = ROOT / "ingestion/data/processed/chunks.jsonl"
|
|
PDF = ROOT / "ingestion/data/raw/duoc-thu-quoc-gia-viet-nam-2018.pdf"
|
|
MIGRATION = Path(__file__).resolve().parents[1] / "migrations/001_rag_retrieval_trace.sql"
|
|
|
|
|
|
class _PlumbingEmbedder:
|
|
"""Deterministic local vectors for the Qdrant round-trip plumbing tests.
|
|
|
|
Not a semantic model — it exists only so an integration test can upsert and
|
|
query real chunks without a cloud call. Production has exactly one query
|
|
embedder (`BedrockCohereQueryEmbedder`); the old local/section-only stubs
|
|
were removed, so this lives with the test that needs it.
|
|
"""
|
|
|
|
def __init__(self, dimensions: int) -> None:
|
|
self._dimensions = dimensions
|
|
|
|
@property
|
|
def dimensions(self) -> int:
|
|
return self._dimensions
|
|
|
|
def embed_query(self, text: str) -> list[float]:
|
|
import hashlib
|
|
import math
|
|
|
|
vector = [0.0] * self._dimensions
|
|
for token in text.casefold().split():
|
|
digest = hashlib.sha256(token.encode("utf-8")).digest()
|
|
index = int.from_bytes(digest[:4], "big") % self._dimensions
|
|
vector[index] += 1.0 if digest[4] & 1 else -1.0
|
|
norm = math.sqrt(sum(value * value for value in vector))
|
|
return [value / norm for value in vector] if norm else vector
|
|
|
|
|
|
@lru_cache
|
|
def _first_real_chunk() -> dict:
|
|
with CHUNKS.open(encoding="utf-8") as handle:
|
|
record = json.loads(next(handle))
|
|
import fitz
|
|
|
|
from ingestion.extract.page_map import build_page_map
|
|
|
|
with fitz.open(PDF) as document:
|
|
page_map = build_page_map(document)
|
|
physical_start, physical_end = record["source_page_range"]
|
|
printed_start = page_map[physical_start]
|
|
printed_end = page_map[physical_end]
|
|
assert printed_start is not None and printed_end is not None
|
|
record["printed_page_range"] = [printed_start, printed_end]
|
|
return record
|
|
|
|
|
|
def test_real_qdrant_round_trip_uses_real_chunk_and_printed_folio():
|
|
from qdrant_client import QdrantClient
|
|
from qdrant_client.models import Distance, PointStruct, VectorParams
|
|
|
|
from adapters.qdrant import QdrantRetriever
|
|
|
|
client = QdrantClient(url="http://localhost:6333")
|
|
collection = f"integration_{uuid.uuid4().hex}"
|
|
embedder = _PlumbingEmbedder(32)
|
|
record = _first_real_chunk()
|
|
try:
|
|
client.create_collection(
|
|
collection_name=collection,
|
|
vectors_config=VectorParams(size=32, distance=Distance.COSINE),
|
|
)
|
|
client.upsert(
|
|
collection_name=collection,
|
|
points=[PointStruct(
|
|
id=str(uuid.uuid4()),
|
|
vector=embedder.embed_query(record["text"]),
|
|
payload=record,
|
|
)],
|
|
wait=True,
|
|
)
|
|
hits = QdrantRetriever(client, collection, embedder).search(
|
|
record["text"], record["drug_id"], 3,
|
|
)
|
|
assert [hit.document.doc_id for hit in hits] == [record["chunk_id"]]
|
|
assert hits[0].document.text == record["text"]
|
|
assert hits[0].document.source_refs[0].printed_page_range == tuple(
|
|
record["printed_page_range"]
|
|
)
|
|
finally:
|
|
with suppress(Exception):
|
|
client.delete_collection(collection)
|
|
|
|
|
|
def test_real_postgres_migration_insert_and_read_back():
|
|
from adapters.postgres import PostgresTraceRepository
|
|
|
|
repository = PostgresTraceRepository(
|
|
"postgresql://duoc_thu:duoc_thu@localhost:5432/duoc_thu"
|
|
)
|
|
repository.migrate(MIGRATION)
|
|
trace_id = repository.save(
|
|
query="Liều abacavir?",
|
|
subject_scope="human",
|
|
intent="fact_lookup",
|
|
decision="answerable",
|
|
reason="grounded_evidence_available",
|
|
resolved_drug_id="abacavir",
|
|
citations=({
|
|
"chunk_id": "abacavir__ten_chung_quoc_te__0",
|
|
"printed_page_start": 101,
|
|
"printed_page_end": 103,
|
|
},),
|
|
)
|
|
stored = repository.get(trace_id)
|
|
assert stored is not None
|
|
assert stored.query == "Liều abacavir?"
|
|
assert stored.resolved_drug_id == "abacavir"
|
|
assert stored.citations[0]["printed_page_start"] == 101
|
|
|
|
|
|
class _FakeJsonLlm:
|
|
"""Deterministic stand-in for the Bedrock Converse generator. Satisfies
|
|
both `JsonLlm` (query understanding) and `AnswerGenerator` (answer +
|
|
entailment) — both ports are just `generate(system, user, schema) ->
|
|
str` — and tells the three call shapes apart the same way
|
|
`tests/test_grounded_generation.py`'s fake does: by schema shape.
|
|
|
|
Deliberately not a real Bedrock call: this suite asserts exact
|
|
drug id / section / citation / grounding outcomes, and this session's
|
|
own live probing (`docs/progress-log.md`, F-01/F-03 entries) found real
|
|
generation and entailment calls genuinely non-deterministic — the wrong
|
|
foundation for a regression assertion. The wiring under test — real
|
|
`RagAgent`, real `RetrievalService`/`QdrantRetriever` against a real
|
|
(temporary) Qdrant collection, real `GroundedAnswerService` — is
|
|
identical to production; only the cloud model call is faked.
|
|
"""
|
|
|
|
def __init__(self, frame_payload: dict, answer_payload: dict) -> None:
|
|
self._frame_payload = frame_payload
|
|
self._answer_payload = answer_payload
|
|
|
|
def generate(self, system: str, user: str, schema: dict) -> str:
|
|
if "turn_type" in schema: # FRAME_SCHEMA: flat, not JSON-Schema-shaped
|
|
return json.dumps(self._frame_payload, ensure_ascii=False)
|
|
if "entailed" in schema.get("properties", {}):
|
|
return json.dumps({"entailed": True, "unsupported": []})
|
|
return json.dumps(self._answer_payload, ensure_ascii=False)
|
|
|
|
|
|
def test_real_rag_agent_end_to_end_through_the_http_api():
|
|
"""F-10: the *production* orchestrator (`RagAgent`), driven through the
|
|
real `/v1/rag/query` HTTP endpoint against a real (temporary) Qdrant
|
|
collection seeded with one real corpus chunk, with a real Postgres trace
|
|
persisted and read back. Reproduces Codex's exact 2026-08-06 finding —
|
|
"no current test imports RagAgent, LlmQueryUnderstander, QueryFrame, or
|
|
retrieve_framed" and "the evaluation runner constructs an in-memory
|
|
lexical retriever and the old resolver rather than executing the same
|
|
dependency graph as the live HTTP service" — both false as of this test.
|
|
"""
|
|
from fastapi.testclient import TestClient
|
|
from qdrant_client import QdrantClient
|
|
from qdrant_client.models import Distance, PointStruct, VectorParams
|
|
|
|
from adapters.postgres import PostgresTraceRepository
|
|
from adapters.qdrant import QdrantParentStore, QdrantRetriever
|
|
from config import Settings
|
|
from main import create_app
|
|
from rag.agent import RagAgent
|
|
from rag.answer import GroundedAnswerService
|
|
from rag.routing import CatalogDrugResolver, QueryRoutingService
|
|
from rag.service import EvidencePolicy, RetrievalService
|
|
from rag.understanding import LlmQueryUnderstander
|
|
|
|
qdrant = QdrantClient(url="http://localhost:6333")
|
|
collection = f"integration_{uuid.uuid4().hex}"
|
|
embedder = _PlumbingEmbedder(32)
|
|
record = dict(_first_real_chunk())
|
|
traces = PostgresTraceRepository(
|
|
"postgresql://duoc_thu:duoc_thu@localhost:5432/duoc_thu"
|
|
)
|
|
traces.migrate(MIGRATION)
|
|
try:
|
|
qdrant.create_collection(
|
|
collection_name=collection,
|
|
vectors_config=VectorParams(size=32, distance=Distance.COSINE),
|
|
)
|
|
qdrant.upsert(
|
|
collection_name=collection,
|
|
points=[PointStruct(
|
|
id=str(uuid.uuid4()),
|
|
vector=embedder.embed_query(record["text"]),
|
|
payload=record,
|
|
)],
|
|
wait=True,
|
|
)
|
|
retrieval = RetrievalService(
|
|
QdrantRetriever(qdrant, collection, embedder),
|
|
QdrantParentStore(qdrant, collection),
|
|
EvidencePolicy(minimum_score=0.01),
|
|
)
|
|
# A minimal one-drug catalog, real drug_id — F-04's candidate
|
|
# bounding runs for real here (`RagAgent`/`LlmQueryUnderstander`
|
|
# are not mocked), so the query must literally name the drug for
|
|
# the deterministic resolver to find it as a candidate.
|
|
resolver = CatalogDrugResolver({record["drug_id"]: {record["drug_name"]}})
|
|
llm = _FakeJsonLlm(
|
|
frame_payload={
|
|
"turn_type": "drug_attribute", "drugs": [record["drug_id"]],
|
|
"unknown_drugs": [], "attribute": None, "population": None,
|
|
"weight_kg": None, "age_text": None, "indication": None,
|
|
"needs_clarify": False, "clarify_reason": None,
|
|
},
|
|
answer_payload={
|
|
"answer": f"{record['text']} [1].", "evidence_sufficient": True,
|
|
},
|
|
)
|
|
understander = LlmQueryUnderstander(
|
|
llm, {record["drug_id"]: record["drug_name"]}, resolver,
|
|
)
|
|
answers = GroundedAnswerService(
|
|
QueryRoutingService(retrieval, resolver), generator=llm,
|
|
)
|
|
agent = RagAgent(understander, retrieval, answers)
|
|
|
|
app = create_app(
|
|
settings=Settings(), answer_service=answers,
|
|
conversational=agent, trace_writer=traces,
|
|
)
|
|
response = TestClient(app).post("/v1/rag/query", json={
|
|
"query": record["drug_name"],
|
|
"subject_scope": "human", "intent": "fact_lookup",
|
|
})
|
|
|
|
assert response.status_code == 200
|
|
body = response.json()
|
|
assert body["decision"] == "answerable"
|
|
assert body["resolved_drug_id"] == record["drug_id"]
|
|
assert body["citations"][0]["chunk_id"] == record["chunk_id"]
|
|
assert body["citations"][0]["printed_page_start"] == (
|
|
record["printed_page_range"][0]
|
|
)
|
|
assert record["text"] in body["answer"]
|
|
|
|
stored = traces.get(body["trace_id"])
|
|
assert stored is not None
|
|
assert stored.decision == "answerable"
|
|
assert stored.resolved_drug_id == record["drug_id"]
|
|
assert stored.citations[0]["chunk_id"] == record["chunk_id"]
|
|
finally:
|
|
with suppress(Exception):
|
|
qdrant.delete_collection(collection)
|
|
|
|
|
|
def test_api_round_trip_uses_qdrant_and_persists_postgres_trace():
|
|
from fastapi.testclient import TestClient
|
|
from qdrant_client import QdrantClient
|
|
from qdrant_client.models import Distance, PointStruct, VectorParams
|
|
|
|
from adapters.postgres import PostgresTraceRepository
|
|
from adapters.qdrant import QdrantParentStore, QdrantRetriever
|
|
from config import Settings
|
|
from main import create_app
|
|
from rag.answer import GroundedAnswerService
|
|
from rag.routing import CatalogDrugResolver, QueryRoutingService
|
|
from rag.service import EvidencePolicy, RetrievalService
|
|
|
|
qdrant = QdrantClient(url="http://localhost:6333")
|
|
collection = f"integration_{uuid.uuid4().hex}"
|
|
embedder = _PlumbingEmbedder(32)
|
|
record = dict(_first_real_chunk())
|
|
traces = PostgresTraceRepository(
|
|
"postgresql://duoc_thu:duoc_thu@localhost:5432/duoc_thu"
|
|
)
|
|
traces.migrate(MIGRATION)
|
|
try:
|
|
qdrant.create_collection(
|
|
collection_name=collection,
|
|
vectors_config=VectorParams(size=32, distance=Distance.COSINE),
|
|
)
|
|
qdrant.upsert(
|
|
collection_name=collection,
|
|
points=[PointStruct(
|
|
id=str(uuid.uuid4()),
|
|
vector=embedder.embed_query(record["text"]),
|
|
payload=record,
|
|
)],
|
|
wait=True,
|
|
)
|
|
retrieval = RetrievalService(
|
|
QdrantRetriever(qdrant, collection, embedder),
|
|
QdrantParentStore(qdrant, collection),
|
|
EvidencePolicy(minimum_score=0.01),
|
|
)
|
|
answers = GroundedAnswerService(QueryRoutingService(
|
|
retrieval,
|
|
CatalogDrugResolver({record["drug_id"]: {record["drug_name"]}}),
|
|
))
|
|
app = create_app(
|
|
settings=Settings(), answer_service=answers, trace_writer=traces,
|
|
)
|
|
response = TestClient(app).post("/v1/rag/query", json={
|
|
"query": record["text"],
|
|
"subject_scope": "human",
|
|
"intent": "fact_lookup",
|
|
})
|
|
assert response.status_code == 200
|
|
body = response.json()
|
|
assert body["decision"] == "answerable"
|
|
assert body["citations"][0]["chunk_id"] == record["chunk_id"]
|
|
assert body["citations"][0]["printed_page_start"] == (
|
|
record["printed_page_range"][0]
|
|
)
|
|
stored = traces.get(body["trace_id"])
|
|
assert stored is not None
|
|
assert stored.decision == "answerable"
|
|
assert stored.citations[0]["chunk_id"] == record["chunk_id"]
|
|
finally:
|
|
with suppress(Exception):
|
|
qdrant.delete_collection(collection)
|