File size: 3,610 Bytes
198bd6d | 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 86 87 88 89 90 91 92 93 94 95 96 97 | from __future__ import annotations
from typing import Any, Literal, TypedDict
from .core import BayesianIndex, ExtractiveGroundedGenerator, RetrievalHit
class RAGState(TypedDict, total=False):
query: str
hits: list[RetrievalHit]
decision: Literal["generate", "abstain"]
answer: str
citations: list[str]
abstained: bool
class BayesianRAGGraph:
"""Actual LangGraph orchestration around the Bayesian retriever.
Graph topology:
START -> retrieve -> assess -> {generate | abstain} -> END
"""
def __init__(self, index: BayesianIndex, generator: Any | None = None, k: int = 4):
self.index = index
self.generator = generator or ExtractiveGroundedGenerator()
self.k = k
self.graph = self._build_graph()
def _retrieve(self, state: RAGState) -> dict[str, Any]:
return {"hits": self.index.search(state["query"], k=self.k)}
def _assess(self, state: RAGState) -> dict[str, Any]:
decision: Literal["generate", "abstain"] = "abstain" if self.index.should_abstain(state.get("hits", [])) else "generate"
return {"decision": decision}
@staticmethod
def _route(state: RAGState) -> str:
return state.get("decision", "abstain")
def _generate(self, state: RAGState) -> dict[str, Any]:
hits = state.get("hits", [])
if hasattr(self.generator, "generate"):
answer, citations = self.generator.generate(state["query"], 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: {state['query']}"
)
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 {"answer": answer, "citations": citations, "abstained": False}
@staticmethod
def _abstain(state: RAGState) -> dict[str, Any]:
return {
"answer": ExtractiveGroundedGenerator.abstain_text,
"citations": [],
"abstained": True,
}
def _build_graph(self):
try:
from langgraph.graph import StateGraph, START, END
except ImportError as exc:
raise ImportError("bayesian-langgraph-rag requires langgraph>=1.2,<2.0") from exc
builder = StateGraph(RAGState)
builder.add_node("retrieve", self._retrieve)
builder.add_node("assess", self._assess)
builder.add_node("generate", self._generate)
builder.add_node("abstain", self._abstain)
builder.add_edge(START, "retrieve")
builder.add_edge("retrieve", "assess")
builder.add_conditional_edges(
"assess",
self._route,
{"generate": "generate", "abstain": "abstain"},
)
builder.add_edge("generate", END)
builder.add_edge("abstain", END)
return builder.compile()
def invoke(self, query: str) -> dict[str, Any]:
state = self.graph.invoke({"query": query})
hits = state.get("hits", [])
return {
"question": query,
"answer": state.get("answer", ExtractiveGroundedGenerator.abstain_text),
"citations": state.get("citations", []),
"abstained": state.get("abstained", True),
"hits": [h.as_dict() for h in hits],
}
|