| 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: |
| 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"] |
| |
| |
| 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 |
|
|
