Retrieval

Vector stores need the retrieval extra. Stores hold vectors and metadata, not documents: keep each chunk's text in the metadata under "text".

Vector stores

openbtk.retrieval.faiss.FAISSVectorStore

FAISSVectorStore(*, dimension: int, metric: Metric = 'l2')

Bases: BaseVectorStore

Store and query vectors with a local FAISS index.

Parameters:

Name	Type	Description	Default
dimension	int	Vector width -- fixed for the life of the index, like every FAISS index.	required
metric	Metric	"l2" (Euclidean distance, smaller is closer) or "ip" (inner product, larger is closer -- the right choice for pre-normalised embeddings, where inner product is cosine similarity). SearchResult.score is always reported so that higher means more relevant regardless of metric: "l2" distances are negated, "ip" scores are used as-is.	'l2'

No I/O happens in __init__ (docs/09_CODING_STANDARDS.md rule 11): the faiss import (and the index it builds) is deferred to first real use, so constructing this class -- e.g. while validating a pipeline config -- never requires faiss-cpu to be installed unless a step actually runs.

Source code in src/openbtk/retrieval/faiss.py

def __init__(self, *, dimension: int, metric: Metric = "l2") -> None:
    self._dimension = dimension
    self._metric = metric
    self._index: Any = None
    self._id_to_int: dict[str, int] = {}
    self._int_to_id: dict[int, str] = {}
    self._metadata: dict[str, dict[str, Any]] = {}
    self._next_int_id = 0

openbtk.retrieval.chroma.ChromaVectorStore

ChromaVectorStore(*, path: str | None = None, collection_name: str = _DEFAULT_COLLECTION)

Bases: BaseVectorStore

Store and query vectors with an embedded Chroma collection.

Parameters:

Name	Type	Description	Default
path	str | None	Directory for on-disk persistence via chromadb.PersistentClient. None (the default) uses an ephemeral in-memory client -- gone once the process exits, same trade-off as :class:~openbtk.retrieval.faiss.FAISSVectorStore with no persist() call.	None
collection_name	str	Chroma requires 3-512 characters from [a-zA-Z0-9._-]; the default satisfies that on its own.	_DEFAULT_COLLECTION

No I/O happens in __init__ (docs/09_CODING_STANDARDS.md rule 11): the chromadb import and the client/collection it builds are deferred to first real use.

Source code in src/openbtk/retrieval/chroma.py

def __init__(
    self, *, path: str | None = None, collection_name: str = _DEFAULT_COLLECTION
) -> None:
    self._path = path
    self._collection_name = collection_name
    self._collection: Any = None

openbtk.retrieval.qdrant.QdrantVectorStore

QdrantVectorStore(*, dimension: int, metric: Metric = 'cosine', path: str | None = None, collection_name: str = _DEFAULT_COLLECTION)

Bases: BaseVectorStore

Store and query vectors with an embedded Qdrant collection.

Parameters:

Name	Type	Description	Default
dimension	int	Vector width, required up front by Qdrant's own create_collection.	required
metric	Metric	"cosine" (the default) or "dot" report higher-is-better similarity natively; "euclid" reports a raw distance, negated here so SearchResult.score means the same thing (higher is more relevant) regardless of metric -- verified directly against the real client, not assumed from Qdrant's docs alone.	'cosine'
path	str | None	Directory for on-disk persistence. None (the default) uses ":memory:".	None
collection_name	str	The Qdrant collection to use.	_DEFAULT_COLLECTION

No I/O happens in __init__ (docs/09_CODING_STANDARDS.md rule 11): the qdrant_client import and the client/collection it builds are deferred to first real use.

Source code in src/openbtk/retrieval/qdrant.py

def __init__(
    self,
    *,
    dimension: int,
    metric: Metric = "cosine",
    path: str | None = None,
    collection_name: str = _DEFAULT_COLLECTION,
) -> None:
    self._dimension = dimension
    self._metric = metric
    self._path = path
    self._collection_name = collection_name
    self._client: Any = None

Reranking

openbtk.retrieval.reranker.ConceptOverlapReranker

ConceptOverlapReranker(*, extract_concepts: Callable[[str], Iterable[str]], cuis_metadata_key: str = _DEFAULT_CUIS_KEY)

Bases: BaseReranker

Rerank results by the count of UMLS CUIs they share with the query.

Parameters:

Name	Type	Description	Default
extract_concepts	Callable[[str], Iterable[str]]	Maps a text string to the CUIs it mentions. Called once, on the query text, per :meth:rerank call. See this module's own docstring for why this is injected rather than built in.	required
cuis_metadata_key	str	The SearchResult.metadata key each result's own pre-computed CUI list is read from. A result with no such key, or an empty list there, contributes zero overlap -- it is never excluded, just never promoted by this criterion.	_DEFAULT_CUIS_KEY

Sorts by (overlap_count, original_score), both descending: shared concepts are the primary criterion (the entire point of this reranker), the original retrieval score the tiebreaker among results that share the same number of concepts with the query -- including "none at all", so a result set with no concept metadata populated anywhere degrades gracefully to the original score order rather than an arbitrary one.

No I/O happens in __init__ (docs/09_CODING_STANDARDS.md rule 11): extract_concepts is only ever called from :meth:rerank.

Source code in src/openbtk/retrieval/reranker.py

def __init__(
    self,
    *,
    extract_concepts: Callable[[str], Iterable[str]],
    cuis_metadata_key: str = _DEFAULT_CUIS_KEY,
) -> None:
    self._extract_concepts = extract_concepts
    self._cuis_metadata_key = cuis_metadata_key

openbtk.retrieval.cross_encoder.CrossEncoderReranker

CrossEncoderReranker(*, model: str = DEFAULT_MODEL, revision: str | None = None, text_metadata_key: str = _DEFAULT_TEXT_KEY, batch_size: int = 16, max_length: int = 512, device: str = 'cpu')

Bases: BaseReranker

Rerank search results with a local cross-encoder model.

Parameters:

Name	Type	Description	Default
model	str	A Hugging Face Hub model id (or local path) of a sequence-classification model with a single relevance logit. Defaults to MedCPT's cross-encoder.	DEFAULT_MODEL
revision	str | None	The pinned commit of model. Optional only for the default model, whose verified commit is built in; for any other model it is required, so a run never silently follows a moving branch.	None
text_metadata_key	str	The SearchResult.metadata key holding each chunk's text.	_DEFAULT_TEXT_KEY
batch_size	int	Query-chunk pairs scored per model pass.	16
max_length	int	Tokens kept per pair (longer pairs are truncated).	512
device	str	A torch device string, "cpu" by default.	'cpu'

Example

reranker = CrossEncoderReranker() # nothing is loaded yet reranker.provenance().model_identity.name 'ncbi/MedCPT-Cross-Encoder'

Source code in src/openbtk/retrieval/cross_encoder.py

def __init__(
    self,
    *,
    model: str = DEFAULT_MODEL,
    revision: str | None = None,
    text_metadata_key: str = _DEFAULT_TEXT_KEY,
    batch_size: int = 16,
    max_length: int = 512,
    device: str = "cpu",
) -> None:
    if revision is None:
        if model != DEFAULT_MODEL:
            raise ConfigError(
                f"A revision is required for model {model!r}: pin a commit "
                "so a run never follows a moving branch.",
                context={"component": "reranker.general.cross_encoder"},
            )
        revision = DEFAULT_REVISION
    if batch_size < 1 or max_length < 2:
        raise ConfigError(
            "batch_size must be at least 1 and max_length at least 2.",
            context={"component": "reranker.general.cross_encoder"},
        )
    self._model_name = model
    self._revision = revision
    self._text_key = text_metadata_key
    self._batch_size = batch_size
    self._max_length = max_length
    self._device = device
    self._loaded: tuple[Any, Any, Any] | None = None

Hybrid (dense + BM25)

See the retrieval guide for how the pieces fit and what BM25 will and will not find.

openbtk.retrieval.hybrid.BM25Index

BM25Index(*, k1: float = 1.5, b: float = 0.75, tokenizer: Callable[[str], list[str]] = default_tokenizer, text_metadata_key: str = 'text')

An in-memory Okapi BM25 index. See the module docstring for its limits.

Parameters:

Name	Type	Description	Default
k1	float	Term-frequency saturation. Higher rewards repeated terms for longer.	1.5
b	float	Length normalisation, 0 (none) to 1 (full).	0.75
tokenizer	Callable[[str], list[str]]	Splits text into terms; used for documents and queries alike.	default_tokenizer
text_metadata_key	str	Where :meth:search results carry the indexed text, so they look like vector-store results (the key RAGPipeline prompts from).	'text'

Example

index = BM25Index() index.upsert( ... ["a", "b", "c"], ... ["metformin for diabetes", "aspirin for pain", "insulin for diabetes"], ... ) ids = [hit.id for hit in index.search("metformin")] ids ['a']

Source code in src/openbtk/retrieval/hybrid.py

def __init__(
    self,
    *,
    k1: float = 1.5,
    b: float = 0.75,
    tokenizer: Callable[[str], list[str]] = default_tokenizer,
    text_metadata_key: str = "text",
) -> None:
    if k1 < 0 or not 0 <= b <= 1:
        raise ValueError("k1 must be non-negative and b between 0 and 1")
    self._k1 = k1
    self._b = b
    self._tokenize = tokenizer
    self._text_key = text_metadata_key
    self._ids: dict[str, int] = {}  # id -> slot
    self._slots: list[str | None] = []  # slot -> id (None once deleted)
    self._lengths: list[int] = []
    self._terms: list[tuple[str, ...]] = []  # the distinct terms of each slot
    self._metadata: list[dict[str, JsonValue]] = []
    self._postings: dict[str, dict[int, int]] = {}  # term -> {slot: tf}
    self._total_length = 0
    self._live = 0

delete

delete(ids: Sequence[str]) -> None

Remove chunks by id. An id that is not indexed is ignored.

Source code in src/openbtk/retrieval/hybrid.py

def delete(self, ids: Sequence[str]) -> None:
    """Remove chunks by id. An id that is not indexed is ignored."""
    for chunk_id in ids:
        slot = self._ids.pop(chunk_id, None)
        if slot is None:
            continue
        for term in self._terms[slot]:
            posting = self._postings.get(term)
            if posting is not None:
                posting.pop(slot, None)
                if not posting:
                    del self._postings[term]
        self._total_length -= self._lengths[slot]
        self._slots[slot] = None
        # Nothing about a deleted chunk is kept, its text and metadata included.
        self._terms[slot] = ()
        self._metadata[slot] = {}
        self._live -= 1

search

search(query: str, top_k: int = 5) -> list[SearchResult]

The top_k chunks that best match query, best first.

Only chunks sharing at least one term with the query are returned, so the result can be shorter than top_k (or empty). Ties break on id, so the order is deterministic.

Source code in src/openbtk/retrieval/hybrid.py

def search(self, query: str, top_k: int = 5) -> list[SearchResult]:
    """The ``top_k`` chunks that best match ``query``, best first.

    Only chunks sharing at least one term with the query are returned, so the result
    can be shorter than ``top_k`` (or empty). Ties break on id, so the order is
    deterministic.
    """
    if top_k < 1 or self._live == 0:
        return []
    average = self._total_length / self._live if self._live else 0.0
    scores: dict[int, float] = {}
    for term in set(self._tokenize(query)):
        posting = self._postings.get(term)
        if not posting:
            continue
        df = len(posting)
        idf = math.log(1 + (self._live - df + 0.5) / (df + 0.5))
        for slot, tf in posting.items():
            norm = (
                1
                - self._b
                + (self._b * self._lengths[slot] / average if average else 0)
            )
            scores[slot] = scores.get(slot, 0.0) + idf * tf * (self._k1 + 1) / (
                tf + self._k1 * norm
            )
    ranked = sorted(
        scores.items(), key=lambda kv: (-kv[1], self._slots[kv[0]] or "")
    )
    return [
        SearchResult(
            id=self._slots[slot] or "",
            score=score,
            metadata=dict(self._metadata[slot]),
        )
        for slot, score in ranked[:top_k]
    ]

upsert

upsert(ids: Sequence[str], texts: Sequence[str], metadata: Sequence[dict[str, JsonValue]] | None = None) -> None

Index texts under ids, replacing any chunk that already has an id.

Parameters:

Name	Type	Description	Default
ids	Sequence[str]	One unique id per text (the ids you give the vector store).	required
texts	Sequence[str]	The chunk texts.	required
metadata	Sequence[dict[str, JsonValue]] | None	Optional metadata per chunk, returned with each hit. The text itself is added under text_metadata_key if it is not already there.	None

Raises:

Type	Description
RetrievalError	If the lists differ in length or ids repeats an id.

Source code in src/openbtk/retrieval/hybrid.py

def upsert(
    self,
    ids: Sequence[str],
    texts: Sequence[str],
    metadata: Sequence[dict[str, JsonValue]] | None = None,
) -> None:
    """Index ``texts`` under ``ids``, replacing any chunk that already has an id.

    Args:
        ids: One unique id per text (the ids you give the vector store).
        texts: The chunk texts.
        metadata: Optional metadata per chunk, returned with each hit. The text
            itself is added under ``text_metadata_key`` if it is not already there.

    Raises:
        RetrievalError: If the lists differ in length or ``ids`` repeats an id.
    """
    if len(ids) != len(texts) or (
        metadata is not None and len(metadata) != len(ids)
    ):
        raise RetrievalError(
            "ids, texts and metadata must be the same length.",
            context={"component": "retrieval.bm25"},
        )
    if len(set(ids)) != len(ids):
        raise RetrievalError(
            "ids must be unique within one upsert.",
            context={"component": "retrieval.bm25"},
        )
    self.delete([i for i in ids if i in self._ids])
    for n, (chunk_id, text) in enumerate(zip(ids, texts, strict=True)):
        terms = self._tokenize(text)
        slot = len(self._slots)
        self._slots.append(chunk_id)
        self._ids[chunk_id] = slot
        counts = Counter(terms)
        self._lengths.append(len(terms))
        self._terms.append(tuple(counts))
        meta = dict(metadata[n]) if metadata is not None else {}
        meta.setdefault(self._text_key, text)
        self._metadata.append(meta)
        self._total_length += len(terms)
        self._live += 1
        for term, count in counts.items():
            self._postings.setdefault(term, {})[slot] = count

openbtk.retrieval.hybrid.reciprocal_rank_fusion

reciprocal_rank_fusion(rankings: Sequence[Sequence[SearchResult]], *, k: int = 60, weights: Sequence[float] | None = None, top_k: int | None = None) -> list[SearchResult]

Merge several rankings into one by reciprocal rank (Cormack et al., 2009).

Each ranking gives an id weight / (k + rank) (rank 1 is the best), the scores are summed per id, and ids are sorted by that sum. Only the order of each input matters, so rankings on different scales (cosine similarity, BM25) combine sensibly. A larger k flattens the advantage of a top rank.

The returned results carry the fused score and the metadata of the first ranking that contained the id, so put the ranking whose metadata you prefer first.

Parameters:

Name	Type	Description	Default
rankings	Sequence[Sequence[SearchResult]]	Best-first lists of results.	required
k	int	The smoothing constant; 60 is the value the original paper used.	60
weights	Sequence[float] | None	One weight per ranking (default all 1).	None
top_k	int | None	Keep at most this many.	None

Example

dense = [SearchResult(id="a", score=0.9), SearchResult(id="b", score=0.8)] sparse = [SearchResult(id="b", score=12.0), SearchResult(id="c", score=3.0)] ids = [hit.id for hit in reciprocal_rank_fusion([dense, sparse])] ids ['b', 'a', 'c']

Source code in src/openbtk/retrieval/hybrid.py

def reciprocal_rank_fusion(
    rankings: Sequence[Sequence[SearchResult]],
    *,
    k: int = 60,
    weights: Sequence[float] | None = None,
    top_k: int | None = None,
) -> list[SearchResult]:
    """Merge several rankings into one by reciprocal rank (Cormack et al., 2009).

    Each ranking gives an id ``weight / (k + rank)`` (rank 1 is the best), the scores
    are summed per id, and ids are sorted by that sum. Only the *order* of each input
    matters, so rankings on different scales (cosine similarity, BM25) combine
    sensibly. A larger ``k`` flattens the advantage of a top rank.

    The returned results carry the fused score and the metadata of the **first ranking
    that contained the id**, so put the ranking whose metadata you prefer first.

    Args:
        rankings: Best-first lists of results.
        k: The smoothing constant; ``60`` is the value the original paper used.
        weights: One weight per ranking (default all ``1``).
        top_k: Keep at most this many.

    Example:
        >>> dense = [SearchResult(id="a", score=0.9), SearchResult(id="b", score=0.8)]
        >>> sparse = [SearchResult(id="b", score=12.0), SearchResult(id="c", score=3.0)]
        >>> ids = [hit.id for hit in reciprocal_rank_fusion([dense, sparse])]
        >>> ids
        ['b', 'a', 'c']
    """
    if k < 0:
        raise ValueError("k must not be negative")
    if weights is not None and len(weights) != len(rankings):
        raise ValueError("give one weight per ranking")
    fused: dict[str, float] = {}
    first: dict[str, SearchResult] = {}
    for n, ranking in enumerate(rankings):
        weight = 1.0 if weights is None else weights[n]
        seen: set[str] = set()
        for rank, result in enumerate(ranking, start=1):
            if result.id in seen:
                continue  # a repeated id in one ranking counts once, at its best rank
            seen.add(result.id)
            fused[result.id] = fused.get(result.id, 0.0) + weight / (k + rank)
            first.setdefault(result.id, result)
    ordered = sorted(fused, key=lambda i: (-fused[i], i))
    merged = [first[i].model_copy(update={"score": fused[i]}) for i in ordered]
    return merged if top_k is None else merged[:top_k]

openbtk.retrieval.hybrid.default_tokenizer

default_tokenizer(text: str) -> list[str]

Lowercase alphanumeric runs ("58410-2" -> ["58410", "2"]).

Example

default_tokenizer("Metformin 500 mg, BID; LOINC 58410-2") ['metformin', '500', 'mg', 'bid', 'loinc', '58410', '2']

Source code in src/openbtk/retrieval/hybrid.py

def default_tokenizer(text: str) -> list[str]:
    """Lowercase alphanumeric runs (``"58410-2"`` -> ``["58410", "2"]``).

    Example:
        >>> default_tokenizer("Metformin 500 mg, BID; LOINC 58410-2")
        ['metformin', '500', 'mg', 'bid', 'loinc', '58410', '2']
    """
    return _TOKEN.findall(text.casefold())

RAG

openbtk.pipelines.rag.RAGPipeline

RAGPipeline(*, embedding: BaseEmbeddingProvider, vectorstore: BaseVectorStore, llm: BaseLLMProvider, reranker: BaseReranker | None = None, bm25: BM25Index | None = None, fusion_k: int = 60, top_k: int = 5, candidate_pool_size: int | None = None, text_metadata_key: str = _DEFAULT_TEXT_KEY, record_id_metadata_key: str = _DEFAULT_RECORD_ID_KEY, chunk_id_metadata_key: str = _DEFAULT_CHUNK_ID_KEY)

Answer a question by retrieving grounded context and generating against it, real chunk-level provenance attached throughout.

Parameters:

Name	Type	Description	Default
embedding	BaseEmbeddingProvider	Embeds the question into the vector store's own space.	required
vectorstore	BaseVectorStore	Queried for the nearest indexed chunks.	required
llm	BaseLLMProvider	Generates the final answer from the retrieved context.	required
reranker	BaseReranker | None	Optional. When given, candidate_pool_size results are fetched and reranked down to top_k; without one, top_k results are used directly in the vector store's own order.	None
bm25	BM25Index | None	Optional. A :class:~openbtk.retrieval.hybrid.BM25Index over the same chunks (index both with the same ids). When given, the question is also run through BM25 and the two rankings are fused by reciprocal rank (FR-R-05), so an exact term the embedding model misses can still surface. See :mod:openbtk.retrieval.hybrid for its limits.	None
fusion_k	int	The reciprocal-rank-fusion smoothing constant (only used with bm25).	60
top_k	int	Number of chunks to ground the answer in.	5
candidate_pool_size	int | None	How many results to fetch (from each retriever) before fusing and reranking. Defaults to top_k * 4 when reranker or bm25 is given (both need a wider candidate pool to be worth running at all), or top_k itself otherwise.	None
text_metadata_key	str	Metadata key holding each chunk's own text -- see this module's own docstring for the indexing convention this whole class depends on.	_DEFAULT_TEXT_KEY
record_id_metadata_key	str	Metadata key holding the source record id.	_DEFAULT_RECORD_ID_KEY
chunk_id_metadata_key	str	Metadata key holding the chunk id; falls back to the vector store's own result id when absent.	_DEFAULT_CHUNK_ID_KEY

No I/O happens in __init__ (docs/09_CODING_STANDARDS.md rule 11): every real call happens inside :meth:ask.

Source code in src/openbtk/pipelines/rag.py

def __init__(
    self,
    *,
    embedding: BaseEmbeddingProvider,
    vectorstore: BaseVectorStore,
    llm: BaseLLMProvider,
    reranker: BaseReranker | None = None,
    bm25: BM25Index | None = None,
    fusion_k: int = 60,
    top_k: int = 5,
    candidate_pool_size: int | None = None,
    text_metadata_key: str = _DEFAULT_TEXT_KEY,
    record_id_metadata_key: str = _DEFAULT_RECORD_ID_KEY,
    chunk_id_metadata_key: str = _DEFAULT_CHUNK_ID_KEY,
) -> None:
    self._embedding = embedding
    self._vectorstore = vectorstore
    self._llm = llm
    self._reranker = reranker
    self._bm25 = bm25
    self._fusion_k = fusion_k
    self._top_k = top_k
    self._candidate_pool_size = candidate_pool_size or (
        top_k * 4 if reranker is not None or bm25 is not None else top_k
    )
    self._text_key = text_metadata_key
    self._record_id_key = record_id_metadata_key
    self._chunk_id_key = chunk_id_metadata_key

ask

ask(question: str, **generate_kwargs: Any) -> RAGAnswer

Retrieve context for question, then generate a grounded answer.

Parameters:

Name	Type	Description	Default
question	str	The user's question. Embedded, used as the vector store query, and (if a reranker is set) as the reranking query too.	required
**generate_kwargs	Any	Forwarded to the LLM provider's chat().	{}

Returns:

Type	Description
RAGAnswer	The generated text plus the exact chunks it was grounded in,
RAGAnswer	in the order given to the model.

Raises:

Type	Description
RetrievalError	If embedding or the vector store query fails.
ProviderError	If the LLM call fails.

Source code in src/openbtk/pipelines/rag.py

def ask(self, question: str, **generate_kwargs: Any) -> RAGAnswer:
    """Retrieve context for ``question``, then generate a grounded answer.

    Args:
        question: The user's question. Embedded, used as the vector
            store query, and (if a reranker is set) as the reranking
            query too.
        **generate_kwargs: Forwarded to the LLM provider's ``chat()``.

    Returns:
        The generated text plus the exact chunks it was grounded in,
        in the order given to the model.

    Raises:
        RetrievalError: If embedding or the vector store query fails.
        ProviderError: If the LLM call fails.
    """
    log.info(
        "rag.ask.start",
        top_k=self._top_k,
        candidate_pool_size=self._candidate_pool_size,
        reranking=self._reranker is not None,
        hybrid=self._bm25 is not None,
    )
    query_vector = self._embedding.embed_one(question)
    candidates = self._vectorstore.query(
        query_vector, top_k=self._candidate_pool_size
    )
    if self._bm25 is not None:
        sparse = self._bm25.search(question, top_k=self._candidate_pool_size)
        candidates = reciprocal_rank_fusion(
            [candidates, sparse],
            k=self._fusion_k,
            top_k=self._candidate_pool_size,
        )
    if self._reranker is not None:
        candidates = self._reranker.rerank(question, candidates, top_k=self._top_k)
    else:
        candidates = candidates[: self._top_k]

    sources = [self._source_ref(c) for c in candidates]
    prompt = self._build_prompt(question, candidates)
    response = self._llm.chat(
        [Message(role="user", content=prompt)], **generate_kwargs
    )
    log.info("rag.ask.complete", sources_count=len(sources))
    return RAGAnswer(text=response.text, sources=sources, usage=response.usage)