| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
| from typing import Any |
|
|
| from bayesian_langgraph import BayesianIndex, BayesianRAGGraph |
| from evaluation.rag_benchmark import CLEAN_CASES, CLEAN_DOCS, CONFLICT_CASES, CONFLICT_DOCS |
|
|
|
|
| def _case_ok(out: dict[str, Any], case: dict[str, Any]) -> bool: |
| if case["abstain"]: |
| return bool(out.get("abstained")) and not out.get("citations") |
| return ( |
| not out.get("abstained") |
| and case["doc"] in out.get("citations", []) |
| and str(case["gold"]).lower() in str(out.get("answer", "")).lower() |
| ) |
|
|
|
|
| def _run_suite(docs, cases) -> dict[str, Any]: |
| rag = BayesianRAGGraph(BayesianIndex(docs), k=4) |
| rows = [] |
| for case in cases: |
| out = rag.invoke(case["q"]) |
| rows.append({ |
| "id": case["id"], |
| "ok": _case_ok(out, case), |
| "answer": out["answer"], |
| "citations": out["citations"], |
| "abstained": out["abstained"], |
| "top_doc": out["hits"][0]["id"] if out["hits"] else None, |
| }) |
| return { |
| "n_cases": len(rows), |
| "passed": sum(r["ok"] for r in rows), |
| "accuracy": sum(r["ok"] for r in rows) / len(rows), |
| "rows": rows, |
| } |
|
|
|
|
| def run_framework_benchmark(output_path: str | Path | None = None) -> dict[str, Any]: |
| rag = BayesianRAGGraph(BayesianIndex(CLEAN_DOCS), k=4) |
| graph = rag.graph.get_graph() |
| result = { |
| "framework": "langgraph", |
| "stategraph_nodes": sorted(graph.nodes.keys()), |
| "stategraph_edges": sorted([e.source, e.target] for e in graph.edges), |
| "clean": _run_suite(CLEAN_DOCS, CLEAN_CASES), |
| "conflict": _run_suite(CONFLICT_DOCS, CONFLICT_CASES), |
| } |
| 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_framework_benchmark(), indent=2)) |
|
|
