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")) ), )