| from __future__ import annotations |
|
|
| from typing import Any, Sequence |
| from langchain_core.documents import Document |
| from langchain_core.runnables import RunnableLambda |
|
|
| from .core import BayesianIndex, DocumentRecord, ExtractiveGroundedGenerator, RetrievalHit |
| from .retriever import BayesianRetriever |
|
|
|
|
| class BayesianRAG: |
| """End-to-end LangChain RAG with explicit posterior uncertainty. |
| |
| Retrieval is performed exactly once through ``BaseRetriever.invoke``. The |
| returned LangChain ``Document`` objects carry the posterior metadata and are |
| converted back to typed hits only for the uncertainty/grounding layer. |
| """ |
|
|
| def __init__(self, index: BayesianIndex, generator: Any | None = None, k: int = 4): |
| self.index = index |
| self.retriever = BayesianRetriever(index=index, k=k) |
| self.generator = generator or ExtractiveGroundedGenerator() |
| self.k = k |
|
|
| @staticmethod |
| def _documents_to_hits(documents: Sequence[Document]) -> list[RetrievalHit]: |
| hits: list[RetrievalHit] = [] |
| for fallback_rank, doc in enumerate(documents, start=1): |
| md = dict(doc.metadata) |
| rec = DocumentRecord( |
| id=str(md.pop("id", f"doc-{fallback_rank}")), |
| text=doc.page_content, |
| title=str(md.pop("title", "")), |
| source=str(md.pop("source", "unknown")), |
| metadata={ |
| k: v |
| for k, v in md.items() |
| if k not in { |
| "retrieval_score", "log_prior", "log_likelihood", |
| "log_posterior_unnormalized", "posterior", |
| "query_coverage", "rank", |
| } |
| }, |
| ) |
| hits.append(RetrievalHit( |
| document=rec, |
| retrieval_score=float(doc.metadata.get("retrieval_score", 0.0)), |
| log_prior=float(doc.metadata.get("log_prior", 0.0)), |
| log_likelihood=float(doc.metadata.get("log_likelihood", 0.0)), |
| log_posterior_unnormalized=float(doc.metadata.get("log_posterior_unnormalized", 0.0)), |
| posterior=float(doc.metadata.get("posterior", 0.0)), |
| query_coverage=float(doc.metadata.get("query_coverage", 0.0)), |
| rank=int(doc.metadata.get("rank", fallback_rank)), |
| )) |
| return hits |
|
|
| def invoke(self, question: str) -> dict[str, Any]: |
| documents = self.retriever.invoke(question) |
| hits = self._documents_to_hits(documents) |
| if self.index.should_abstain(hits): |
| answer = ExtractiveGroundedGenerator.abstain_text |
| citations: list[str] = [] |
| elif hasattr(self.generator, "generate"): |
| answer, citations = self.generator.generate(question, hits) |
| else: |
| context = "\n".join(f"[{h.document.id}] {h.document.text}" for h in hits) |
| prompt = ( |
| "Answer only from the supplied context. Cite document ids in square brackets. " |
| "If unsupported, say you do not know.\n\n" |
| f"Context:\n{context}\n\nQuestion: {question}" |
| ) |
| resp = self.generator.invoke(prompt) |
| answer = getattr(resp, "content", str(resp)) |
| citations = [h.document.id for h in hits if f"[{h.document.id}]" in answer] |
| return { |
| "question": question, |
| "answer": answer, |
| "citations": citations, |
| "abstained": not citations, |
| "hits": [h.as_dict() for h in hits], |
| } |
|
|
| def as_runnable(self) -> RunnableLambda: |
| return RunnableLambda(self.invoke) |
|
|
