File size: 4,675 Bytes
7b4d34e | 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 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | from __future__ import annotations
from typing import Any, TypedDict
from .core import DocumentRecord, ExtractiveGroundedGenerator, FuzzyEvidence, FuzzyIndex
class RAGState(TypedDict, total=False):
question: str
k: int
expanded: bool
evidence: list[dict[str, Any]]
confidence: float
margin: float
reason: str
route: str
answer: str
citations: list[str]
def _serialize_hits(hits: list[FuzzyEvidence]) -> list[dict[str, Any]]:
return [
{
"id": h.document.id,
"text": h.document.text,
"title": h.document.title,
"source": h.document.source,
"metadata": dict(h.document.metadata),
"raw_relevance": h.raw_relevance,
"fuzzy_score": h.fuzzy_score,
"authority": h.authority,
"freshness": h.freshness,
"reliability": h.reliability,
}
for h in hits
]
def _deserialize_hits(items: list[dict[str, Any]]) -> list[FuzzyEvidence]:
hits: list[FuzzyEvidence] = []
for x in items:
d = DocumentRecord(x["id"], x["text"], x.get("title", ""), x.get("source", ""), x.get("metadata", {}))
hits.append(FuzzyEvidence(d, float(x["raw_relevance"]), float(x["fuzzy_score"]), float(x["authority"]), float(x["freshness"]), float(x["reliability"]), {}))
return hits
class FuzzyRAGGraph:
"""LangGraph orchestration with fuzzy generate/expand/abstain routing."""
def __init__(self, index: FuzzyIndex, initial_k: int = 4, generator=None):
self.index = index
self.initial_k = initial_k
self.generator = ExtractiveGroundedGenerator()
self.external_generator = generator
self.graph = self._build()
def _build(self):
try:
from langgraph.graph import END, START, StateGraph
except Exception as exc: # pragma: no cover
raise RuntimeError("langgraph is required to build FuzzyRAGGraph") from exc
builder = StateGraph(RAGState)
builder.add_node("retrieve", self._retrieve)
builder.add_node("assess", self._assess)
builder.add_node("expand", self._expand)
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",
lambda s: s["route"],
{"generate": "generate", "expand": "expand", "abstain": "abstain"},
)
builder.add_edge("expand", "retrieve")
builder.add_edge("generate", END)
builder.add_edge("abstain", END)
return builder.compile()
def _retrieve(self, state: RAGState) -> dict[str, Any]:
k = int(state.get("k", self.initial_k))
pool = self.index.config.extended_pool if state.get("expanded", False) else self.index.config.candidate_pool
hits = self.index.search(state["question"], k=k, candidate_pool=pool)
return {"evidence": _serialize_hits(hits)}
def _assess(self, state: RAGState) -> dict[str, Any]:
hits = _deserialize_hits(state.get("evidence", []))
d = self.index.decision(hits)
route = d["route"]
# Only one retrieval expansion is permitted. If ambiguity survives the
# wider search, the graph abstains rather than cycling or bluffing.
if route == "expand" and state.get("expanded", False):
route = "abstain"
reason = "ambiguous_after_expansion"
else:
reason = d["reason"]
return {"route": route, "confidence": float(d["confidence"]), "margin": float(d["margin"]), "reason": reason}
def _expand(self, state: RAGState) -> dict[str, Any]:
return {"k": min(self.index.config.extended_pool, max(int(state.get("k", self.initial_k)) * 2, self.initial_k + 2)), "expanded": True}
def _generate(self, state: RAGState) -> dict[str, Any]:
hits = _deserialize_hits(state.get("evidence", []))
if self.external_generator is None:
answer, citations = self.generator.generate(state["question"], hits)
else:
answer = str(self.external_generator(state["question"], hits))
citations = [hits[0].document.id] if hits else []
return {"answer": answer, "citations": citations}
def _abstain(self, state: RAGState) -> dict[str, Any]:
return {"answer": self.generator.abstain_text, "citations": []}
def invoke(self, question: str) -> dict[str, Any]:
out = self.graph.invoke({"question": question, "k": self.initial_k, "expanded": False})
return out
|