| from __future__ import annotations |
|
|
| from dataclasses import dataclass, field |
| from typing import Any, Iterable, Sequence |
| import math |
| import re |
|
|
| import numpy as np |
| from sklearn.feature_extraction.text import TfidfVectorizer |
| from sklearn.metrics.pairwise import cosine_similarity |
|
|
|
|
| @dataclass(frozen=True) |
| class DocumentRecord: |
| id: str |
| text: str |
| title: str = "" |
| source: str = "" |
| metadata: dict[str, Any] = field(default_factory=dict) |
|
|
|
|
| @dataclass(frozen=True) |
| class FuzzyRetrievalConfig: |
| candidate_pool: int = 10 |
| answer_threshold: float = 0.60 |
| ambiguity_margin: float = 0.015 |
| min_raw_relevance: float = 0.14 |
| expand_threshold: float = 0.48 |
| extended_pool: int = 16 |
|
|
|
|
| @dataclass(frozen=True) |
| class FuzzyEvidence: |
| document: DocumentRecord |
| raw_relevance: float |
| fuzzy_score: float |
| authority: float |
| freshness: float |
| reliability: float |
| memberships: dict[str, float] |
|
|
|
|
| class FuzzyEvidenceController: |
| """Small dependency-free Sugeno fuzzy controller for RAG evidence quality. |
| |
| Inputs are normalized to [0, 1]: lexical/semantic relevance, source authority, |
| freshness, and source reliability. The rules intentionally make relevance a |
| gate: highly trusted but irrelevant evidence must not become high confidence. |
| """ |
|
|
| @staticmethod |
| def _clip(x: float) -> float: |
| return max(0.0, min(1.0, float(x))) |
|
|
| @classmethod |
| def _left(cls, x: float, full_until: float, zero_at: float) -> float: |
| x = cls._clip(x) |
| if x <= full_until: |
| return 1.0 |
| if x >= zero_at: |
| return 0.0 |
| return (zero_at - x) / (zero_at - full_until) |
|
|
| @classmethod |
| def _right(cls, x: float, zero_until: float, full_at: float) -> float: |
| x = cls._clip(x) |
| if x <= zero_until: |
| return 0.0 |
| if x >= full_at: |
| return 1.0 |
| return (x - zero_until) / (full_at - zero_until) |
|
|
| @classmethod |
| def _triangle(cls, x: float, left: float, peak: float, right: float) -> float: |
| x = cls._clip(x) |
| if x <= left or x >= right: |
| return 0.0 |
| if x == peak: |
| return 1.0 |
| if x < peak: |
| return (x - left) / (peak - left) |
| return (right - x) / (right - peak) |
|
|
| def memberships(self, relevance: float, authority: float, freshness: float, reliability: float) -> dict[str, float]: |
| r = self._clip(relevance) |
| a = self._clip(authority) |
| f = self._clip(freshness) |
| q = self._clip(reliability) |
| return { |
| "rel_low": self._left(r, 0.05, 0.30), |
| "rel_med": self._triangle(r, 0.12, 0.42, 0.70), |
| "rel_high": self._right(r, 0.42, 0.72), |
| "auth_low": self._left(a, 0.0, 0.55), |
| "auth_med": self._triangle(a, 0.35, 0.62, 0.84), |
| "auth_high": self._right(a, 0.68, 0.90), |
| "fresh_low": self._left(f, 0.0, 0.42), |
| "fresh_med": self._triangle(f, 0.28, 0.58, 0.82), |
| "fresh_high": self._right(f, 0.68, 0.90), |
| "reliability_low": self._left(q, 0.0, 0.50), |
| "reliability_high": self._right(q, 0.58, 0.88), |
| } |
|
|
| def score(self, relevance: float, authority: float, freshness: float, reliability: float) -> tuple[float, dict[str, float]]: |
| m = self.memberships(relevance, authority, freshness, reliability) |
| |
| rules: list[tuple[float, float]] = [ |
| (min(m["rel_high"], m["auth_high"], m["fresh_high"], m["reliability_high"]), 0.98), |
| (min(m["rel_med"], m["auth_high"], m["fresh_high"], m["reliability_high"]), 0.90), |
| (min(m["rel_high"], m["auth_high"]), 0.92), |
| (min(m["rel_med"], m["auth_high"]), 0.84), |
| (min(m["rel_high"], m["auth_med"], m["fresh_high"]), 0.74), |
| (min(m["rel_med"], m["auth_med"]), 0.54), |
| (min(m["rel_high"], m["fresh_low"]), 0.28), |
| (min(m["rel_high"], m["auth_low"]), 0.24), |
| (min(m["rel_med"], m["auth_low"]), 0.20), |
| (min(m["rel_high"], m["reliability_low"]), 0.26), |
| (m["rel_low"], 0.06), |
| ] |
| denom = sum(w for w, _ in rules) |
| if denom <= 1e-12: |
| |
| base = 0.10 + 0.55 * self._clip(relevance) |
| else: |
| base = sum(w * z for w, z in rules) / denom |
| |
| |
| value = self._clip(base) |
| return value, m |
|
|
|
|
| class FuzzyIndex: |
| def __init__(self, documents: Sequence[DocumentRecord] | None = None, config: FuzzyRetrievalConfig | None = None): |
| self.config = config or FuzzyRetrievalConfig() |
| self.controller = FuzzyEvidenceController() |
| self.documents: list[DocumentRecord] = [] |
| self._vectorizer: TfidfVectorizer | None = None |
| self._matrix = None |
| if documents: |
| self.fit(documents) |
|
|
| @staticmethod |
| def _trust(doc: DocumentRecord) -> tuple[float, float, float]: |
| md = doc.metadata or {} |
| authority = float(md.get("authority", 0.50)) |
| freshness = float(md.get("freshness", 0.50)) |
| if "reliability" in md: |
| reliability = float(md["reliability"]) |
| elif "reliability_alpha" in md or "reliability_beta" in md: |
| a = float(md.get("reliability_alpha", 1.0)) |
| b = float(md.get("reliability_beta", 1.0)) |
| reliability = a / max(a + b, 1e-12) |
| else: |
| reliability = 0.50 |
| return tuple(max(0.0, min(1.0, x)) for x in (authority, freshness, reliability)) |
|
|
| def fit(self, documents: Sequence[DocumentRecord]) -> "FuzzyIndex": |
| self.documents = list(documents) |
| corpus = [f"{d.title}. {d.text}".strip() for d in self.documents] |
| self._vectorizer = TfidfVectorizer(ngram_range=(1, 2), stop_words="english", sublinear_tf=True) |
| self._matrix = self._vectorizer.fit_transform(corpus) if corpus else None |
| return self |
|
|
| def add(self, documents: Iterable[DocumentRecord]) -> "FuzzyIndex": |
| return self.fit([*self.documents, *list(documents)]) |
|
|
| def search(self, query: str, k: int = 4, candidate_pool: int | None = None) -> list[FuzzyEvidence]: |
| if not self.documents or self._vectorizer is None or self._matrix is None: |
| return [] |
| q = self._vectorizer.transform([query]) |
| sims = cosine_similarity(q, self._matrix)[0] |
| pool = min(candidate_pool or self.config.candidate_pool, len(self.documents)) |
| candidate_ids = np.argsort(-sims, kind="stable")[:pool] |
| hits: list[FuzzyEvidence] = [] |
| for i in candidate_ids: |
| doc = self.documents[int(i)] |
| relevance = float(sims[int(i)]) |
| authority, freshness, reliability = self._trust(doc) |
| fuzzy, memberships = self.controller.score(relevance, authority, freshness, reliability) |
| hits.append(FuzzyEvidence(doc, relevance, fuzzy, authority, freshness, reliability, memberships)) |
| hits.sort(key=lambda h: (h.fuzzy_score, h.raw_relevance, h.authority, h.freshness), reverse=True) |
| return hits[: max(1, k)] |
|
|
| def decision(self, hits: Sequence[FuzzyEvidence]) -> dict[str, Any]: |
| if not hits: |
| return {"route": "abstain", "confidence": 0.0, "margin": 0.0, "reason": "no_evidence"} |
| top = hits[0] |
| second = hits[1] if len(hits) > 1 else None |
| if second is None: |
| margin = top.fuzzy_score |
| else: |
| |
| |
| |
| margin = 0.80 * (top.fuzzy_score - second.fuzzy_score) + 0.20 * (top.raw_relevance - second.raw_relevance) |
| if top.raw_relevance < self.config.min_raw_relevance: |
| return {"route": "abstain", "confidence": top.fuzzy_score, "margin": margin, "reason": "low_relevance"} |
| trusted_enough = top.authority >= 0.60 or top.reliability >= 0.75 |
| if top.fuzzy_score >= self.config.answer_threshold and margin >= self.config.ambiguity_margin and trusted_enough: |
| return {"route": "generate", "confidence": top.fuzzy_score, "margin": margin, "reason": "strong_evidence"} |
| if top.fuzzy_score >= self.config.expand_threshold: |
| return {"route": "expand", "confidence": top.fuzzy_score, "margin": margin, "reason": "borderline_or_ambiguous"} |
| return {"route": "abstain", "confidence": top.fuzzy_score, "margin": margin, "reason": "weak_evidence"} |
|
|
| def should_abstain(self, hits: Sequence[FuzzyEvidence]) -> bool: |
| return self.decision(hits)["route"] != "generate" |
|
|
|
|
| class PlainTfidfIndex: |
| """Baseline retriever intentionally unaware of source quality or uncertainty.""" |
|
|
| def __init__(self, documents: Sequence[DocumentRecord]): |
| self.documents = list(documents) |
| corpus = [f"{d.title}. {d.text}".strip() for d in self.documents] |
| self.vectorizer = TfidfVectorizer(ngram_range=(1, 2), stop_words="english", sublinear_tf=True) |
| self.matrix = self.vectorizer.fit_transform(corpus) |
|
|
| def search(self, query: str, k: int = 4) -> list[FuzzyEvidence]: |
| q = self.vectorizer.transform([query]) |
| sims = cosine_similarity(q, self.matrix)[0] |
| ids = np.argsort(-sims, kind="stable")[: min(k, len(self.documents))] |
| out = [] |
| for i in ids: |
| d = self.documents[int(i)] |
| s = float(sims[int(i)]) |
| out.append(FuzzyEvidence(d, s, s, 0.5, 0.5, 0.5, {})) |
| return out |
|
|
|
|
| class ExtractiveGroundedGenerator: |
| """Offline answer synthesizer used for reproducible RAG evaluation. |
| |
| It selects the query-most-relevant sentence from the accepted evidence and |
| returns it with an explicit citation. Production users can replace this with |
| any LangChain Runnable/LLM without changing fuzzy retrieval or routing. |
| """ |
|
|
| abstain_text = "Insufficient reliable evidence to answer." |
|
|
| @staticmethod |
| def _sentences(text: str) -> list[str]: |
| parts = re.split(r"(?<=[.!?])\s+", text.strip()) |
| return [p.strip() for p in parts if p.strip()] |
|
|
| def generate(self, question: str, hits: Sequence[FuzzyEvidence]) -> tuple[str, list[str]]: |
| if not hits: |
| return self.abstain_text, [] |
| top = hits[0] |
| sentences = self._sentences(top.document.text) or [top.document.text] |
| if len(sentences) == 1: |
| best = sentences[0] |
| else: |
| vec = TfidfVectorizer(ngram_range=(1, 2), stop_words="english") |
| X = vec.fit_transform([question, *sentences]) |
| scores = cosine_similarity(X[0:1], X[1:])[0] |
| best = sentences[int(np.argmax(scores))] |
| return f"{best} [{top.document.id}]", [top.document.id] |
|
|
|
|
| class FuzzyRAGEngine: |
| def __init__(self, documents: Sequence[DocumentRecord] | None = None, config: FuzzyRetrievalConfig | None = None): |
| self.index = FuzzyIndex(documents or [], config) |
| self.generator = ExtractiveGroundedGenerator() |
|
|
| def ingest(self, documents: Sequence[DocumentRecord], replace: bool = True) -> int: |
| if replace: |
| self.index.fit(documents) |
| else: |
| self.index.add(documents) |
| return len(self.index.documents) |
|
|
| def query(self, question: str, k: int = 4) -> dict[str, Any]: |
| hits = self.index.search(question, k=k) |
| decision = self.index.decision(hits) |
| expanded = False |
| if decision["route"] == "expand": |
| expanded = True |
| hits = self.index.search(question, k=max(k * 2, 6), candidate_pool=self.index.config.extended_pool) |
| decision = self.index.decision(hits) |
| if decision["route"] == "expand": |
| decision = {**decision, "route": "abstain", "reason": "ambiguous_after_expansion"} |
| if decision["route"] == "generate": |
| answer, citations = self.generator.generate(question, hits) |
| else: |
| answer, citations = self.generator.abstain_text, [] |
| return { |
| "question": question, |
| "answer": answer, |
| "citations": citations, |
| "confidence": round(float(decision["confidence"]), 6), |
| "margin": round(float(decision["margin"]), 6), |
| "decision": decision["route"], |
| "reason": decision["reason"], |
| "expanded": expanded, |
| "evidence": [ |
| { |
| "id": h.document.id, |
| "title": h.document.title, |
| "source": h.document.source, |
| "raw_relevance": round(h.raw_relevance, 6), |
| "fuzzy_score": round(h.fuzzy_score, 6), |
| "authority": round(h.authority, 6), |
| "freshness": round(h.freshness, 6), |
| "reliability": round(h.reliability, 6), |
| } |
| for h in hits |
| ], |
| } |
|
|
