File size: 16,543 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 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 | from __future__ import annotations
import json
from pathlib import Path
from typing import Any
from bayesian_langgraph.core import (
BayesianIndex,
BayesianRetrievalConfig,
DocumentRecord,
PlainTfidfIndex,
ExtractiveGroundedGenerator,
)
def _md(kind: str) -> dict[str, Any]:
if kind == "official":
return {"reliability_alpha": 18, "reliability_beta": 2, "authority": 0.98, "freshness": 1.0, "tier": kind}
if kind == "stale":
return {"reliability_alpha": 6, "reliability_beta": 6, "authority": 0.45, "freshness": 0.18, "tier": kind}
if kind == "community":
return {"reliability_alpha": 7, "reliability_beta": 5, "authority": 0.55, "freshness": 0.65, "tier": kind}
if kind == "official_equal":
return {"reliability_alpha": 18, "reliability_beta": 2, "authority": 0.98, "freshness": 1.0, "tier": "official"}
return {"reliability_alpha": 1, "reliability_beta": 1, "authority": 0.5, "freshness": 0.5, "tier": kind}
CONFLICT_DOCS = [
DocumentRecord("LEG-HR-PORTAL", "The current HR portal for leave requests and payslips is WorkSphere.", "Current HR portal login guide", "legacy-cache", _md("stale")),
DocumentRecord("OFF-HR-PORTAL", "All employee self-service tasks, including leave requests and payslips, are now handled in PeopleHub.", "Employee self-service migration completed", "official-handbook", _md("official")),
DocumentRecord("COM-HR-PORTAL", "Some teams used WorkSphere during migration; ask HR if your bookmark no longer works.", "HR portal discussion", "community-wiki", _md("community")),
DocumentRecord("LEG-LEAVE", "Employees receive 15 paid leave days per year under the leave policy.", "Paid leave entitlement FAQ", "legacy-cache", _md("stale")),
DocumentRecord("OFF-LEAVE", "The 2026 employee benefits policy grants 18 paid leave days per calendar year.", "2026 employee benefits policy", "official-handbook", _md("official")),
DocumentRecord("COM-LEAVE", "Leave balances vary by category and should be checked in the HR system.", "Leave balance help", "community-wiki", _md("community")),
DocumentRecord("LEG-REMOTE", "The current remote work policy permits three remote days per week with manager approval.", "Current remote work days", "legacy-cache", _md("stale")),
DocumentRecord("OFF-REMOTE", "Under the active flexible-work policy, employees may work remotely up to two days per week with manager approval.", "Flexible-work policy revision", "official-handbook", _md("official")),
DocumentRecord("COM-REMOTE", "Remote work requires manager approval and may vary during special projects.", "Remote work tips", "community-wiki", _md("community")),
DocumentRecord("LEG-PAYSLIP", "Salary slips are published on the tenth working day of each month.", "Payslip publication date", "legacy-cache", _md("stale")),
DocumentRecord("OFF-PAYSLIP", "Payroll publishes salary slips by the fifth working day of each month.", "Payroll service standard", "official-handbook", _md("official")),
DocumentRecord("COM-PAYSLIP", "Payslip publication can be delayed around holidays.", "Payroll discussion", "community-wiki", _md("community")),
DocumentRecord("LEG-VPN", "The supported corporate VPN protocol is OpenVPN.", "Corporate VPN protocol", "legacy-cache", _md("stale")),
DocumentRecord("OFF-VPN", "The security baseline now requires WireGuard for the standard corporate VPN profile.", "Network security baseline", "official-it", _md("official")),
DocumentRecord("COM-VPN", "OpenVPN profiles may still exist on older laptops.", "VPN migration notes", "community-wiki", _md("community")),
DocumentRecord("LEG-BACKUP", "Production database backups are retained for 14 days.", "Database backup retention policy", "legacy-cache", _md("stale")),
DocumentRecord("OFF-BACKUP", "Production database backups have a retention period of 30 days.", "Disaster recovery control DR-7", "official-it", _md("official")),
DocumentRecord("COM-BACKUP", "Teams should not rely on local snapshots as their only backup.", "Backup practices", "community-wiki", _md("community")),
DocumentRecord("LEG-RATE", "The public API rate limit is 60 requests per minute per token.", "Public API rate limit", "legacy-cache", _md("stale")),
DocumentRecord("OFF-RATE", "The current gateway quota permits 120 requests per minute per token.", "API gateway operating limits", "official-it", _md("official")),
DocumentRecord("COM-RATE", "Burst traffic may still be throttled even below the nominal quota.", "API troubleshooting", "community-wiki", _md("community")),
DocumentRecord("LEG-REGION", "The primary deployment region is us-east-1.", "Primary cloud deployment region", "legacy-cache", _md("stale")),
DocumentRecord("OFF-REGION", "The production primary region is ap-south-1; disaster recovery remains in a separate region.", "Cloud production topology", "official-it", _md("official")),
DocumentRecord("COM-REGION", "Developers may use several regions for non-production experiments.", "Cloud developer notes", "community-wiki", _md("community")),
DocumentRecord("LEG-EXPENSE", "The meal expense cap is INR 3,000 per day on domestic travel.", "Domestic meal expense cap", "legacy-cache", _md("stale")),
DocumentRecord("OFF-EXPENSE", "For domestic business travel, reimbursable meals are capped at INR 4,500 per day.", "Travel and expense policy 2026", "official-finance", _md("official")),
DocumentRecord("COM-EXPENSE", "Itemized receipts are recommended for meal reimbursement.", "Expense claim tips", "community-wiki", _md("community")),
DocumentRecord("LEG-REIMB", "Employees must submit reimbursement claims within 30 days of the expense.", "Reimbursement submission deadline", "legacy-cache", _md("stale")),
DocumentRecord("OFF-REIMB", "Expense claims must be submitted within 15 days of the transaction date.", "Expense reimbursement standard", "official-finance", _md("official")),
DocumentRecord("COM-REIMB", "Finance may request additional receipts during reimbursement review.", "Reimbursement FAQ", "community-wiki", _md("community")),
DocumentRecord("LEG-TRAVEL", "All corporate travel should be booked through TravelPro.", "Corporate travel booking portal", "legacy-cache", _md("stale")),
DocumentRecord("OFF-TRAVEL", "Approved business travel must be booked through TripDesk.", "Business travel booking procedure", "official-finance", _md("official")),
DocumentRecord("COM-TRAVEL", "Direct airline booking can cause reimbursement delays.", "Travel booking discussion", "community-wiki", _md("community")),
DocumentRecord("LEG-PROC", "Purchases above INR 25,000 require procurement review.", "Procurement review threshold", "legacy-cache", _md("stale")),
DocumentRecord("OFF-PROC", "Procurement review is mandatory for purchases above INR 50,000.", "Procurement control PC-12", "official-finance", _md("official")),
DocumentRecord("COM-PROC", "Some categories require procurement review regardless of value.", "Procurement FAQ", "community-wiki", _md("community")),
DocumentRecord("OFF-LAB-A", "The graduate AI lab has moved to room B-204.", "AI lab room notice A", "official-campus", _md("official_equal")),
DocumentRecord("OFF-LAB-B", "The graduate AI lab has moved to room C-101.", "AI lab room notice B", "official-campus", _md("official_equal")),
DocumentRecord("OFF-OFFICE-A", "Professor Sen holds office hours on Tuesday at 3 PM.", "Office hours notice A", "official-campus", _md("official_equal")),
DocumentRecord("OFF-OFFICE-B", "Professor Sen holds office hours on Wednesday at 4 PM.", "Office hours notice B", "official-campus", _md("official_equal")),
]
CONFLICT_CASES = [
{"id":"q1","q":"What is the current HR portal for leave requests and payslips?","gold":"PeopleHub","doc":"OFF-HR-PORTAL","abstain":False},
{"id":"q2","q":"How many paid leave days do employees get per year under the current policy?","gold":"18 paid leave days","doc":"OFF-LEAVE","abstain":False},
{"id":"q3","q":"How many remote work days per week are currently allowed with manager approval?","gold":"two days per week","doc":"OFF-REMOTE","abstain":False},
{"id":"q4","q":"By which working day are salary slips published each month?","gold":"fifth working day","doc":"OFF-PAYSLIP","abstain":False},
{"id":"q5","q":"Which VPN protocol is required by the current corporate security baseline?","gold":"WireGuard","doc":"OFF-VPN","abstain":False},
{"id":"q6","q":"What is the production database backup retention period?","gold":"30 days","doc":"OFF-BACKUP","abstain":False},
{"id":"q7","q":"What is the current public API rate limit per minute per token?","gold":"120 requests per minute","doc":"OFF-RATE","abstain":False},
{"id":"q8","q":"Which cloud region is primary for production?","gold":"ap-south-1","doc":"OFF-REGION","abstain":False},
{"id":"q9","q":"What is the domestic business travel meal expense cap per day?","gold":"INR 4,500 per day","doc":"OFF-EXPENSE","abstain":False},
{"id":"q10","q":"Within how many days must expense reimbursement claims be submitted?","gold":"15 days","doc":"OFF-REIMB","abstain":False},
{"id":"q11","q":"Which portal must be used to book approved business travel?","gold":"TripDesk","doc":"OFF-TRAVEL","abstain":False},
{"id":"q12","q":"Above what amount is procurement review mandatory?","gold":"INR 50,000","doc":"OFF-PROC","abstain":False},
{"id":"q13","q":"What is the company dental insurance annual limit?","gold":None,"doc":None,"abstain":True},
{"id":"q14","q":"What is the cafeteria breakfast price?","gold":None,"doc":None,"abstain":True},
{"id":"q15","q":"What is the emergency satellite phone number?","gold":None,"doc":None,"abstain":True},
{"id":"q16","q":"What is the annual gym membership reimbursement?","gold":None,"doc":None,"abstain":True},
{"id":"q17","q":"Which room is the graduate AI lab currently in?","gold":None,"doc":None,"abstain":True},
{"id":"q18","q":"When are Professor Sen's office hours?","gold":None,"doc":None,"abstain":True},
]
CLEAN_DOCS = [
DocumentRecord("C1", "The observatory opens to visitors at 9 AM every Saturday.", "Observatory hours", "handbook", _md("official")),
DocumentRecord("C2", "The chemistry safety briefing is held in Hall 7.", "Chemistry briefing", "handbook", _md("official")),
DocumentRecord("C3", "Library books may be borrowed for 21 days.", "Library borrowing", "handbook", _md("official")),
DocumentRecord("C4", "The robotics club meets on Thursday evenings.", "Robotics club", "handbook", _md("official")),
DocumentRecord("C5", "Visitor parking is located beside Gate 3.", "Visitor parking", "handbook", _md("official")),
DocumentRecord("C6", "The campus shuttle begins service at 6:30 AM.", "Campus shuttle", "handbook", _md("official")),
DocumentRecord("C7", "The data science seminar carries two academic credits.", "Data science seminar", "handbook", _md("official")),
DocumentRecord("C8", "The emergency assembly point is the North Lawn.", "Emergency procedure", "handbook", _md("official")),
]
CLEAN_CASES = [
{"id":"c1","q":"When does the observatory open on Saturday?","gold":"9 AM","doc":"C1","abstain":False},
{"id":"c2","q":"Where is the chemistry safety briefing held?","gold":"Hall 7","doc":"C2","abstain":False},
{"id":"c3","q":"How long may library books be borrowed?","gold":"21 days","doc":"C3","abstain":False},
{"id":"c4","q":"When does the robotics club meet?","gold":"Thursday evenings","doc":"C4","abstain":False},
{"id":"c5","q":"Where is visitor parking?","gold":"Gate 3","doc":"C5","abstain":False},
{"id":"c6","q":"What time does the campus shuttle begin service?","gold":"6:30 AM","doc":"C6","abstain":False},
{"id":"c7","q":"How many credits is the data science seminar worth?","gold":"two academic credits","doc":"C7","abstain":False},
{"id":"c8","q":"Where is the emergency assembly point?","gold":"North Lawn","doc":"C8","abstain":False},
]
def _eval(name: str, docs: list[DocumentRecord], cases: list[dict[str, Any]], bayesian: bool) -> tuple[dict[str, Any], list[dict[str, Any]]]:
gen = ExtractiveGroundedGenerator()
if bayesian:
searcher = BayesianIndex(docs, BayesianRetrievalConfig(candidate_pool=10, likelihood_temperature=2.0))
else:
searcher = PlainTfidfIndex(docs)
rows: list[dict[str, Any]] = []
rr: list[float] = []
r1: list[float] = []
r3: list[float] = []
briers: list[float] = []
doc_map = {d.id: d for d in docs}
for c in cases:
hits = searcher.search(c["q"], k=4)
top = hits[0] if hits else None
if bayesian:
abstain = searcher.should_abstain(hits)
else:
abstain = (not hits) or top.retrieval_score < 0.06
if abstain:
answer, citations = gen.abstain_text, []
else:
answer, citations = gen.generate(c["q"], hits)
if c["abstain"]:
correct = answer == gen.abstain_text
else:
correct = str(c["gold"]).lower() in answer.lower()
ids = [h.document.id for h in hits]
rank = ids.index(c["doc"]) + 1 if c["doc"] in ids else 0
r1.append(1.0 if rank == 1 else 0.0)
r3.append(1.0 if 0 < rank <= 3 else 0.0)
rr.append(1.0 / rank if rank else 0.0)
p = top.posterior if top else 0.0
y = 1.0 if top and top.document.id == c["doc"] else 0.0
briers.append((p - y) ** 2)
grounded = True
if citations:
sentence = answer.rsplit(" [", 1)[0]
grounded = sentence in doc_map[citations[0]].text
rows.append({
"id": c["id"],
"query": c["q"],
"expected_abstain": c["abstain"],
"top_doc": None if top is None else top.document.id,
"retrieval_score": None if top is None else round(top.retrieval_score, 6),
"top_posterior": None if top is None else round(top.posterior, 6),
"answer": answer,
"citations": citations,
"correct": correct,
"grounded": grounded,
})
n_answerable = sum(not c["abstain"] for c in cases)
n_abstain = sum(c["abstain"] for c in cases)
summary = {
"system": name,
"n_cases": len(cases),
"answerable_cases": n_answerable,
"abstention_cases": n_abstain,
"recall_at_1": sum(r1) / len(r1) if r1 else None,
"recall_at_3": sum(r3) / len(r3) if r3 else None,
"mrr": sum(rr) / len(rr) if rr else None,
"end_to_end_accuracy": sum(r["correct"] for r in rows) / len(rows),
"grounding_accuracy": sum(r["grounded"] for r in rows) / len(rows),
"abstention_accuracy": (sum(r["correct"] for r, c in zip(rows, cases) if c["abstain"]) / n_abstain) if n_abstain else None,
"answer_accuracy_on_answerable": (sum(r["correct"] for r, c in zip(rows, cases) if not c["abstain"]) / n_answerable) if n_answerable else None,
"top_posterior_brier": sum(briers) / len(briers) if briers else None,
}
return summary, rows
def run_benchmark(output_path: str | Path | None = None) -> dict[str, Any]:
clean_b, clean_b_rows = _eval("bayesian", CLEAN_DOCS, CLEAN_CASES, True)
clean_t, clean_t_rows = _eval("plain_tfidf", CLEAN_DOCS, CLEAN_CASES, False)
conflict_b, conflict_b_rows = _eval("bayesian", CONFLICT_DOCS, CONFLICT_CASES, True)
conflict_t, conflict_t_rows = _eval("plain_tfidf", CONFLICT_DOCS, CONFLICT_CASES, False)
result = {
"benchmark": "controlled_rag_v2",
"methodology": (
"Two end-to-end suites. CLEAN checks ordinary retrieval+grounded answer synthesis. "
"CONFLICT adds stale authoritative conflicts, missing evidence, and unresolved equal-authority contradictions. "
"The generator is extractive and citation-grounded so the evaluation isolates retrieval and uncertainty behavior without an external LLM confound."
),
"clean": {"bayesian": clean_b, "baseline": clean_t, "bayesian_cases": clean_b_rows, "baseline_cases": clean_t_rows},
"conflict": {"bayesian": conflict_b, "baseline": conflict_t, "bayesian_cases": conflict_b_rows, "baseline_cases": conflict_t_rows},
}
if output_path is not None:
Path(output_path).write_text(json.dumps(result, indent=2), encoding="utf-8")
return result
if __name__ == "__main__":
print(json.dumps(run_benchmark(), indent=2))
|