Wire the guarded conversational RAG answer layer end-to-end
This commit is contained in:
@@ -0,0 +1,4 @@
|
||||
from .postgres import PostgresTraceRepository
|
||||
from .qdrant import QdrantParentStore, QdrantRetriever
|
||||
|
||||
__all__ = ["PostgresTraceRepository", "QdrantParentStore", "QdrantRetriever"]
|
||||
@@ -0,0 +1,130 @@
|
||||
"""Claude on Amazon Bedrock as the answer generator.
|
||||
|
||||
The only module that names the `anthropic` SDK, imported lazily — the same
|
||||
arrangement that confines `boto3` to `embedding.py` and `qdrant_client` to
|
||||
`qdrant.py`, so `rag/` imports and the whole suite runs with no SDK and no
|
||||
cloud account.
|
||||
|
||||
Two provider facts, taken from the Anthropic API reference rather than from
|
||||
memory: Bedrock model ids carry an `anthropic.` prefix (`anthropic.claude-opus-5`),
|
||||
and the Messages-API path on Bedrock is the Mantle client — not the legacy
|
||||
`bedrock-runtime` InvokeModel route the embedding adapter uses.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from rag.ports import AnswerGenerationUnavailable
|
||||
|
||||
BEDROCK_CLAUDE_OPUS_5 = "anthropic.claude-opus-5"
|
||||
|
||||
# Generation must not outrun the evidence it is rewriting. The section route
|
||||
# can return a long section, so this is sized for the rewrite, not the source.
|
||||
MAX_OUTPUT_TOKENS = 4096
|
||||
|
||||
|
||||
def _provider_error_types() -> tuple[type[BaseException], ...]:
|
||||
"""SDK and transport error classes, or none when the SDK is absent."""
|
||||
collected: list[type[BaseException]] = []
|
||||
try:
|
||||
import anthropic
|
||||
|
||||
collected.append(anthropic.APIError)
|
||||
except ImportError:
|
||||
pass
|
||||
try:
|
||||
from botocore.exceptions import BotoCoreError, ClientError
|
||||
|
||||
collected.extend((BotoCoreError, ClientError))
|
||||
except ImportError:
|
||||
pass
|
||||
return tuple(collected)
|
||||
|
||||
|
||||
class BedrockClaudeAnswerGenerator:
|
||||
"""Rewrites evidence into prose under a schema the API enforces.
|
||||
|
||||
The output shape is constrained by `output_config.format` rather than by
|
||||
asking for JSON in the prompt, so a malformed envelope is the provider's
|
||||
error rather than this code's parsing problem. What the schema cannot
|
||||
constrain is whether the *content* is faithful — that is
|
||||
`rag.grounding.verify`'s job, and it runs on every response this returns.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
region: str = "us-east-1",
|
||||
client: Any | None = None,
|
||||
model_id: str = BEDROCK_CLAUDE_OPUS_5,
|
||||
max_tokens: int = MAX_OUTPUT_TOKENS,
|
||||
) -> None:
|
||||
self._region = region
|
||||
self._client = client
|
||||
self._model_id = model_id
|
||||
self._max_tokens = max_tokens
|
||||
|
||||
@property
|
||||
def model_id(self) -> str:
|
||||
return self._model_id
|
||||
|
||||
def _runtime(self) -> Any:
|
||||
if self._client is None:
|
||||
from anthropic import AnthropicBedrockMantle
|
||||
|
||||
self._client = AnthropicBedrockMantle(aws_region=self._region)
|
||||
return self._client
|
||||
|
||||
def generate(self, system: str, user: str, schema: dict) -> str:
|
||||
try:
|
||||
response = self._runtime().messages.create(
|
||||
model=self._model_id,
|
||||
max_tokens=self._max_tokens,
|
||||
system=system,
|
||||
messages=[{"role": "user", "content": user}],
|
||||
output_config={"format": {"type": "json_schema", "schema": schema}},
|
||||
)
|
||||
except _provider_error_types() as error:
|
||||
raise AnswerGenerationUnavailable(
|
||||
f"{self._model_id} could not be invoked: {type(error).__name__}"
|
||||
) from error
|
||||
|
||||
# A refusal is a successful HTTP response with no usable content, not
|
||||
# an exception. Treating it as an outage routes it to the extractive
|
||||
# fallback instead of letting `content[0]` raise.
|
||||
if getattr(response, "stop_reason", None) == "refusal":
|
||||
raise AnswerGenerationUnavailable(
|
||||
f"{self._model_id} declined the request"
|
||||
)
|
||||
|
||||
text = "".join(
|
||||
block.text
|
||||
for block in response.content
|
||||
if getattr(block, "type", None) == "text"
|
||||
)
|
||||
if not text.strip():
|
||||
raise AnswerGenerationUnavailable(
|
||||
f"{self._model_id} returned no text content"
|
||||
)
|
||||
return text
|
||||
|
||||
|
||||
class StubAnswerGenerator:
|
||||
"""Returns a fixed payload; lets the whole answer path run with no cloud.
|
||||
|
||||
Not a fake for tests only — it is what `EMBEDDING_PROVIDER`-style local
|
||||
demos use to exercise prompt building, schema parsing, grounding
|
||||
verification and the fallback branch without spending anything.
|
||||
"""
|
||||
|
||||
def __init__(self, answer: str, evidence_sufficient: bool = True) -> None:
|
||||
self._payload = json.dumps(
|
||||
{"answer": answer, "evidence_sufficient": evidence_sufficient},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
self.calls: list[tuple[str, str]] = []
|
||||
|
||||
def generate(self, system: str, user: str, schema: dict) -> str: # noqa: ARG002
|
||||
# `schema` is unused here; the stub returns an already-valid payload.
|
||||
self.calls.append((system, user))
|
||||
return self._payload
|
||||
@@ -0,0 +1,178 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import math
|
||||
from typing import Any
|
||||
|
||||
from rag.ports import QueryEmbeddingUnavailable
|
||||
|
||||
BEDROCK_RUNTIME_SERVICE = "bedrock-runtime"
|
||||
COHERE_EMBED_V4 = "cohere.embed-v4:0"
|
||||
|
||||
# Cohere embeds corpus records and queries into different subspaces. Sending
|
||||
# the corpus value here raises no error — recall just drops silently — so this
|
||||
# constant exists to make the asymmetry visible rather than incidental.
|
||||
COHERE_QUERY_INPUT_TYPE = "search_query"
|
||||
|
||||
|
||||
def _provider_error_types() -> tuple[type[BaseException], ...]:
|
||||
"""botocore's error classes, or none when botocore is absent.
|
||||
|
||||
Resolved lazily and tolerantly so a test injecting a stub client — and a
|
||||
deployment that never enables a cloud provider — neither imports botocore
|
||||
nor depends on it being installed.
|
||||
"""
|
||||
try:
|
||||
from botocore.exceptions import BotoCoreError, ClientError
|
||||
except ImportError:
|
||||
return ()
|
||||
return (BotoCoreError, ClientError)
|
||||
|
||||
|
||||
class BedrockCohereQueryEmbedder:
|
||||
"""Embeds a query with the same model the collection was built from.
|
||||
|
||||
A collection loaded with Cohere vectors and queried with any other embedder
|
||||
still returns results and still raises nothing — the hits are simply
|
||||
meaningless. That failure is silent, which is why the corpus manifest
|
||||
records `model_id` and why this adapter names the model explicitly.
|
||||
|
||||
The request shape is duplicated from `ingestion/embed/bedrock_cohere.py`
|
||||
rather than imported: `ingestion` is a separate deployable and importing it
|
||||
here would couple the API to the batch pipeline. The duplication is one
|
||||
JSON body and is deliberate.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dimensions: int,
|
||||
region: str = "us-east-1",
|
||||
client: Any | None = None,
|
||||
model_id: str = COHERE_EMBED_V4,
|
||||
) -> None:
|
||||
if dimensions <= 0:
|
||||
raise ValueError("dimensions must be positive")
|
||||
self._dimensions = dimensions
|
||||
self._region = region
|
||||
self._client = client
|
||||
self._model_id = model_id
|
||||
|
||||
@property
|
||||
def dimensions(self) -> int:
|
||||
return self._dimensions
|
||||
|
||||
@property
|
||||
def model_id(self) -> str:
|
||||
return self._model_id
|
||||
|
||||
def _runtime(self) -> Any:
|
||||
if self._client is None:
|
||||
import boto3
|
||||
from botocore.config import Config
|
||||
|
||||
self._client = boto3.client(
|
||||
BEDROCK_RUNTIME_SERVICE,
|
||||
region_name=self._region,
|
||||
config=Config(
|
||||
connect_timeout=10,
|
||||
read_timeout=30,
|
||||
retries={"max_attempts": 3, "mode": "standard"},
|
||||
),
|
||||
)
|
||||
return self._client
|
||||
|
||||
def embed_query(self, text: str) -> list[float]:
|
||||
import json
|
||||
|
||||
try:
|
||||
response = self._runtime().invoke_model(
|
||||
modelId=self._model_id,
|
||||
body=json.dumps(
|
||||
{
|
||||
"texts": [text],
|
||||
"input_type": COHERE_QUERY_INPUT_TYPE,
|
||||
"embedding_types": ["float"],
|
||||
"output_dimension": self._dimensions,
|
||||
}
|
||||
),
|
||||
accept="*/*",
|
||||
contentType="application/json",
|
||||
)
|
||||
except _provider_error_types() as error:
|
||||
raise QueryEmbeddingUnavailable(
|
||||
f"{self._model_id} could not be invoked: {type(error).__name__}"
|
||||
) from error
|
||||
body = json.loads(response["body"].read())
|
||||
embeddings = body.get("embeddings")
|
||||
if isinstance(embeddings, dict):
|
||||
rows = embeddings.get("float")
|
||||
else:
|
||||
rows = embeddings
|
||||
if not rows:
|
||||
raise ValueError(
|
||||
f"{self._model_id} returned no float embedding; "
|
||||
f"response keys were {sorted(body)}"
|
||||
)
|
||||
values = list(rows[0])
|
||||
if len(values) != self._dimensions:
|
||||
raise ValueError(
|
||||
f"{self._model_id} returned {len(values)} dimensions, "
|
||||
f"expected {self._dimensions}"
|
||||
)
|
||||
return values
|
||||
|
||||
|
||||
class SectionOnlyQueryEmbedder:
|
||||
"""Refuses to embed, confining retrieval to the route that is verified.
|
||||
|
||||
The section route resolves the question's attribute to a `section_key` and
|
||||
filters on it; no vector is involved, and it measured 16/16 on the
|
||||
human-written golden questions on 2026-08-04. Similarity measured 0.544 and
|
||||
its provider is currently revoked.
|
||||
|
||||
Declining locally and immediately is better than the two alternatives it
|
||||
replaces: a `cohere-v4` round-trip spends the boto3 retry budget before
|
||||
failing, and `LocalHashQueryEmbedder` searches a SHA-256 vector against a
|
||||
Cohere collection, which returns confident and meaningless hits.
|
||||
"""
|
||||
|
||||
def __init__(self, dimensions: int) -> None:
|
||||
if dimensions <= 0:
|
||||
raise ValueError("dimensions must be positive")
|
||||
self._dimensions = dimensions
|
||||
|
||||
@property
|
||||
def dimensions(self) -> int:
|
||||
return self._dimensions
|
||||
|
||||
def embed_query(self, text: str) -> list[float]: # noqa: ARG002
|
||||
# The text is irrelevant: this embedder exists to refuse, not to embed.
|
||||
raise QueryEmbeddingUnavailable(
|
||||
"no query embedding provider is enabled; set EMBEDDING_PROVIDER to "
|
||||
"use the similarity fallback"
|
||||
)
|
||||
|
||||
|
||||
class LocalHashQueryEmbedder:
|
||||
"""Deterministic local plumbing probe; not a semantic retrieval model."""
|
||||
|
||||
def __init__(self, dimensions: int) -> None:
|
||||
if dimensions <= 0:
|
||||
raise ValueError("dimensions must be positive")
|
||||
self._dimensions = dimensions
|
||||
|
||||
@property
|
||||
def dimensions(self) -> int:
|
||||
return self._dimensions
|
||||
|
||||
def embed_query(self, text: str) -> list[float]:
|
||||
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
|
||||
sign = 1.0 if digest[4] & 1 else -1.0
|
||||
vector[index] += sign
|
||||
norm = math.sqrt(sum(value * value for value in vector))
|
||||
if norm == 0:
|
||||
return vector
|
||||
return [value / norm for value in vector]
|
||||
@@ -0,0 +1,82 @@
|
||||
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:
|
||||
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) 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) 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) 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],
|
||||
)
|
||||
@@ -0,0 +1,73 @@
|
||||
"""Exports the domain's counters; the only module that names prometheus_client.
|
||||
|
||||
`rag/metrics.py` defines what is counted and why. This decides how it leaves
|
||||
the process, and is imported lazily so the service runs — and the suite passes
|
||||
— with no metrics stack installed.
|
||||
|
||||
Counter names carry a `duocthu_` prefix and a `_total` suffix because that is
|
||||
what Prometheus expects of a counter; the dashboard queries them by name.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from rag.metrics import (
|
||||
ABSTENTION,
|
||||
ANSWER_EXTRACTIVE,
|
||||
GENERATION_REJECTED,
|
||||
GENERATION_SERVED,
|
||||
RETRIEVAL_ROUTE,
|
||||
)
|
||||
|
||||
_LABELS: dict[str, tuple[str, ...]] = {
|
||||
ABSTENTION: ("reason",),
|
||||
GENERATION_REJECTED: ("reason",),
|
||||
RETRIEVAL_ROUTE: ("route",),
|
||||
GENERATION_SERVED: (),
|
||||
ANSWER_EXTRACTIVE: (),
|
||||
}
|
||||
|
||||
_HELP = {
|
||||
ABSTENTION: "Answers refused, by the reason retrieval gave.",
|
||||
GENERATION_REJECTED: (
|
||||
"Generations discarded before reaching the caller. `reason=\"ungrounded_number\"` "
|
||||
"counts answers that stated a figure absent from the cited source."
|
||||
),
|
||||
GENERATION_SERVED: "Generations that passed grounding verification and were served.",
|
||||
ANSWER_EXTRACTIVE: "Answers served as verbatim source text.",
|
||||
RETRIEVAL_ROUTE: "Retrievals by route: section filter, or similarity fallback.",
|
||||
}
|
||||
|
||||
|
||||
class PrometheusMetrics:
|
||||
"""Domain `Metrics` backed by a Prometheus registry."""
|
||||
|
||||
def __init__(self, registry: Any | None = None) -> None:
|
||||
from prometheus_client import CollectorRegistry, Counter
|
||||
|
||||
self._registry = registry or CollectorRegistry()
|
||||
self._counters = {
|
||||
name: Counter(name, _HELP[name], labels, registry=self._registry)
|
||||
for name, labels in _LABELS.items()
|
||||
}
|
||||
|
||||
@property
|
||||
def registry(self) -> Any:
|
||||
return self._registry
|
||||
|
||||
def increment(self, name: str, **labels: str) -> None:
|
||||
counter = self._counters.get(name)
|
||||
if counter is None:
|
||||
return
|
||||
# An unexpected label would raise at scrape time, far from its cause.
|
||||
# Metrics must not be able to break a clinical answer, so a mismatch
|
||||
# drops the sample rather than the request.
|
||||
expected = set(_LABELS[name])
|
||||
if set(labels) != expected:
|
||||
return
|
||||
(counter.labels(**labels) if labels else counter).inc()
|
||||
|
||||
def render(self) -> tuple[bytes, str]:
|
||||
from prometheus_client import CONTENT_TYPE_LATEST, generate_latest
|
||||
|
||||
return generate_latest(self._registry), CONTENT_TYPE_LATEST
|
||||
@@ -0,0 +1,265 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Protocol, Sequence
|
||||
|
||||
from rag.models import ParentDocument, RetrievalDocument, SearchHit, SourceRef
|
||||
|
||||
|
||||
class QueryEmbedder(Protocol):
|
||||
@property
|
||||
def dimensions(self) -> int: ...
|
||||
|
||||
def embed_query(self, text: str) -> Sequence[float]: ...
|
||||
|
||||
|
||||
def _source_refs(payload: dict[str, Any]) -> tuple[SourceRef, ...]:
|
||||
explicit = payload.get("source_refs") or []
|
||||
if explicit:
|
||||
return tuple(
|
||||
SourceRef(
|
||||
physical_page=int(item["physical_page"]),
|
||||
precision=item.get("precision", "region"),
|
||||
block_id=item.get("block_id"),
|
||||
bbox=tuple(item["bbox"]) if item.get("bbox") else None,
|
||||
source_crop=item.get("source_crop"),
|
||||
page_range=(
|
||||
tuple(item["page_range"]) if item.get("page_range") else None
|
||||
),
|
||||
printed_page=item.get("printed_page"),
|
||||
printed_page_range=(
|
||||
tuple(item["printed_page_range"])
|
||||
if item.get("printed_page_range") else None
|
||||
),
|
||||
)
|
||||
for item in explicit
|
||||
)
|
||||
|
||||
physical_range = payload.get("source_page_range")
|
||||
printed_range = payload.get("printed_page_range")
|
||||
attachments = payload.get("attachments") or []
|
||||
|
||||
attachment_refs = tuple(
|
||||
SourceRef(
|
||||
physical_page=int(item["physical_page"]),
|
||||
precision="region",
|
||||
block_id=item.get("block_id"),
|
||||
bbox=tuple(item["bbox"]) if item.get("bbox") else None,
|
||||
source_crop=item.get("source_crop"),
|
||||
page_range=(int(item["physical_page"]), int(item["physical_page"])),
|
||||
printed_page=(
|
||||
int(item["printed_page"])
|
||||
if item.get("printed_page") is not None else None
|
||||
),
|
||||
printed_page_range=(
|
||||
(int(item["printed_page"]), int(item["printed_page"]))
|
||||
if item.get("printed_page") is not None else None
|
||||
),
|
||||
)
|
||||
for item in attachments
|
||||
if item.get("physical_page") is not None
|
||||
)
|
||||
if payload.get("chunk_kind") == "block_descriptor" and attachment_refs:
|
||||
return attachment_refs
|
||||
|
||||
physical_page = (
|
||||
physical_range[0]
|
||||
if physical_range else payload.get("heading_physical_page")
|
||||
)
|
||||
base_refs: tuple[SourceRef, ...] = ()
|
||||
if physical_page is not None:
|
||||
base_refs = (
|
||||
SourceRef(
|
||||
physical_page=int(physical_page),
|
||||
precision="chunk_page_range",
|
||||
page_range=tuple(physical_range) if physical_range else None,
|
||||
printed_page=(int(printed_range[0]) if printed_range else None),
|
||||
printed_page_range=tuple(printed_range) if printed_range else None,
|
||||
),
|
||||
)
|
||||
return base_refs + attachment_refs
|
||||
|
||||
|
||||
def _document(payload: dict[str, Any]) -> RetrievalDocument:
|
||||
return RetrievalDocument(
|
||||
doc_id=payload["chunk_id"],
|
||||
drug_id=payload["drug_id"],
|
||||
drug_name=payload.get("drug_name"),
|
||||
kind=payload.get("chunk_kind", "prose"),
|
||||
text=payload["text"],
|
||||
section_key=payload["section_key"],
|
||||
source_refs=_source_refs(payload),
|
||||
parent_id=payload.get("parent_id"),
|
||||
requires_visual_check=(
|
||||
bool(payload.get("requires_visual_check"))
|
||||
or bool(payload.get("has_quarantined_content"))
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class QdrantRetriever:
|
||||
def __init__(
|
||||
self,
|
||||
client: Any,
|
||||
collection_name: str,
|
||||
embedder: QueryEmbedder,
|
||||
) -> None:
|
||||
self._client = client
|
||||
self._collection_name = collection_name
|
||||
self._embedder = embedder
|
||||
|
||||
def search(self, query: str, drug_id: str, limit: int) -> list[SearchHit]:
|
||||
from qdrant_client.models import FieldCondition, Filter, MatchValue
|
||||
|
||||
vector = list(self._embedder.embed_query(query))
|
||||
if len(vector) != self._embedder.dimensions:
|
||||
raise ValueError(
|
||||
f"query vector has {len(vector)} dimensions; "
|
||||
f"expected {self._embedder.dimensions}"
|
||||
)
|
||||
points = self._client.search(
|
||||
collection_name=self._collection_name,
|
||||
query_vector=vector,
|
||||
query_filter=Filter(
|
||||
must=[FieldCondition(key="drug_id", match=MatchValue(value=drug_id))]
|
||||
),
|
||||
limit=limit,
|
||||
with_payload=True,
|
||||
)
|
||||
return [
|
||||
SearchHit(_document(dict(point.payload or {})), float(point.score))
|
||||
for point in points
|
||||
]
|
||||
|
||||
def find_by_section(self, drug_id: str, section_key: str) -> list[SearchHit]:
|
||||
"""Every chunk of one section, by payload filter — no vector involved.
|
||||
|
||||
A `scroll`, not a `search`: this must not be a top-k. Paging continues
|
||||
until the offset is exhausted, because Qdrant's default page is 256 and
|
||||
a long section silently truncated would read as a complete answer.
|
||||
|
||||
Score is 1.0 because the match is exact by construction. It is not a
|
||||
similarity and must not be compared against one.
|
||||
|
||||
Results are re-sorted by `part_index` before returning. Qdrant scrolls
|
||||
in point-id order, and point ids are `uuid5(chunk_id)`, so the natural
|
||||
order is effectively random: PARACETAMOL's dosing section came back
|
||||
3, 4, 1, 2, 0 — the answer opened mid-sentence on paediatric doses and
|
||||
buried "Liều lượng: Người lớn:" last. A section served out of order is
|
||||
a clinical hazard, not a formatting one: a reader who stops early
|
||||
stops in the middle of a different population's dose.
|
||||
"""
|
||||
from qdrant_client.models import FieldCondition, Filter, MatchValue
|
||||
|
||||
scroll_filter = Filter(
|
||||
must=[
|
||||
FieldCondition(key="drug_id", match=MatchValue(value=drug_id)),
|
||||
FieldCondition(key="section_key", match=MatchValue(value=section_key)),
|
||||
]
|
||||
)
|
||||
hits: list[tuple[dict, SearchHit]] = []
|
||||
offset = None
|
||||
while True:
|
||||
points, offset = self._client.scroll(
|
||||
collection_name=self._collection_name,
|
||||
scroll_filter=scroll_filter,
|
||||
limit=256,
|
||||
offset=offset,
|
||||
with_payload=True,
|
||||
)
|
||||
hits.extend(
|
||||
(dict(point.payload or {}), SearchHit(_document(dict(point.payload or {})), 1.0))
|
||||
for point in points
|
||||
)
|
||||
if offset is None:
|
||||
break
|
||||
|
||||
# `part_index` is the chunker's own position within the section. A
|
||||
# payload missing it sorts last rather than raising: an unordered
|
||||
# section is worse than a scrambled one only if it also disappears.
|
||||
hits.sort(key=lambda item: item[0].get("part_index", 1 << 30))
|
||||
return [hit for _, hit in hits]
|
||||
|
||||
|
||||
def find_by_drug(self, drug_id: str) -> list[SearchHit]:
|
||||
"""Every prose section of one drug, in book order — the monograph view.
|
||||
|
||||
For a query that names the drug but no attribute ("PARACETAMOL"), a drug
|
||||
reference shows the whole monograph, not a "specify an attribute" prompt.
|
||||
A `scroll` filtered on `drug_id`, prose only (block descriptors stay out
|
||||
of a text answer), ordered by the book's section sequence then
|
||||
`part_index`. Each section's first chunk gets a `【heading】` so the
|
||||
result reads as a monograph, not a wall of text.
|
||||
"""
|
||||
from qdrant_client.models import FieldCondition, Filter, MatchValue
|
||||
|
||||
from rag.sections import SECTION_ORDER
|
||||
|
||||
scroll_filter = Filter(
|
||||
must=[
|
||||
FieldCondition(key="drug_id", match=MatchValue(value=drug_id)),
|
||||
FieldCondition(key="chunk_kind", match=MatchValue(value="prose")),
|
||||
]
|
||||
)
|
||||
payloads: list[dict] = []
|
||||
offset = None
|
||||
while True:
|
||||
points, offset = self._client.scroll(
|
||||
collection_name=self._collection_name,
|
||||
scroll_filter=scroll_filter,
|
||||
limit=256,
|
||||
offset=offset,
|
||||
with_payload=True,
|
||||
)
|
||||
payloads.extend(dict(point.payload or {}) for point in points)
|
||||
if offset is None:
|
||||
break
|
||||
|
||||
order = {key: index for index, key in enumerate(SECTION_ORDER)}
|
||||
payloads.sort(
|
||||
key=lambda p: (
|
||||
order.get(p.get("section_key"), len(order)),
|
||||
p.get("part_index", 1 << 30),
|
||||
)
|
||||
)
|
||||
|
||||
hits: list[SearchHit] = []
|
||||
seen_sections: set[str] = set()
|
||||
for payload in payloads:
|
||||
section_key = payload.get("section_key")
|
||||
if section_key not in seen_sections:
|
||||
seen_sections.add(section_key)
|
||||
name = payload.get("section_display_name") or section_key or ""
|
||||
payload = {**payload, "text": f"【{name}】\n{payload.get('text', '')}"}
|
||||
hits.append(SearchHit(_document(payload), 1.0))
|
||||
return hits
|
||||
|
||||
|
||||
class QdrantParentStore:
|
||||
def __init__(self, client: Any, collection_name: str) -> None:
|
||||
self._client = client
|
||||
self._collection_name = collection_name
|
||||
|
||||
def get(self, parent_id: str) -> ParentDocument | None:
|
||||
from qdrant_client.models import FieldCondition, Filter, MatchValue
|
||||
|
||||
points, _ = self._client.scroll(
|
||||
collection_name=self._collection_name,
|
||||
scroll_filter=Filter(
|
||||
must=[FieldCondition(key="chunk_id", match=MatchValue(value=parent_id))]
|
||||
),
|
||||
limit=1,
|
||||
with_payload=True,
|
||||
)
|
||||
if not points:
|
||||
return None
|
||||
payload = dict(points[0].payload or {})
|
||||
return ParentDocument(
|
||||
parent_id=parent_id,
|
||||
kind=payload.get("chunk_kind", "parent"),
|
||||
text=payload["text"],
|
||||
source_refs=_source_refs(payload),
|
||||
requires_visual_check=(
|
||||
bool(payload.get("requires_visual_check"))
|
||||
or bool(payload.get("has_quarantined_content"))
|
||||
),
|
||||
)
|
||||
Reference in New Issue
Block a user