| 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)] |
|
|
