File size: 3,680 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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | 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)
|