File size: 1,622 Bytes
ce3f6c1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | from __future__ import annotations
from typing import Any
from pydantic import ConfigDict
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever
from .core import BayesianIndex, RetrievalHit
class BayesianRetriever(BaseRetriever):
"""LangChain-native retriever backed by the generalized Bayesian index.
The LangChain contract returns ``Document`` objects. All posterior fields
required downstream are copied into document metadata so an end-to-end RAG
chain can use the *actual* retriever output without recomputing retrieval.
"""
model_config = ConfigDict(arbitrary_types_allowed=True)
index: Any
k: int = 4
def search_hits(self, query: str) -> list[RetrievalHit]:
return self.index.search(query, k=self.k)
@staticmethod
def _hit_to_document(h: RetrievalHit) -> Document:
md = dict(h.document.metadata)
md.update({
"id": h.document.id,
"title": h.document.title,
"source": h.document.source,
"retrieval_score": h.retrieval_score,
"log_prior": h.log_prior,
"log_likelihood": h.log_likelihood,
"log_posterior_unnormalized": h.log_posterior_unnormalized,
"posterior": h.posterior,
"query_coverage": h.query_coverage,
"rank": h.rank,
})
return Document(page_content=h.document.text, metadata=md)
def _get_relevant_documents(self, query: str, *, run_manager: Any = None) -> list[Document]:
return [self._hit_to_document(h) for h in self.search_hits(query)]
|