| 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], |
| } |
|
|
