| 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)) |
|
|
