Wire the guarded conversational RAG answer layer end-to-end
This commit is contained in:
@@ -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