Files
GEWU/scripts/run_retrieval_benchmark.py
vasily 78033bb830
CI / quality (push) Canceled after 0s
feat: online verification, node removals, and knowledge browser
Adds claim-bound online evidence records and citation verification scripts,
node removal transport/adapter, retrieval benchmark and owner-review tooling,
GraphCanvas/SourceDetail knowledge browser UI, and the ingest handoff helper.
Bumps the agents/gewu pack to the online-verification branch.
2026-10-16 20:02:29 +08:00

127 lines
5.1 KiB
Python

#!/usr/bin/env python3
"""Capture scoped MCP retrieval results for a frozen, blind question packet.
This records evidence for independent grading; it does not infer relevance from
keyword overlap or assign semantic scores. Use an authenticated MCP endpoint.
"""
from __future__ import annotations
import argparse
import asyncio
import hashlib
import json
import time
from datetime import UTC, datetime
from pathlib import Path
from mcp import Client
from mcp.client.streamable_http import streamable_http_client
async def run(args: argparse.Namespace) -> None:
raw = args.questions.read_bytes()
questions = json.loads(raw)
if not isinstance(questions, list) or not questions:
raise ValueError("Questions must be a nonempty JSON list")
ids: set[str] = set()
for question in questions:
if not isinstance(question, dict) or not all(
isinstance(question.get(field), str) and question[field].strip()
for field in ("id", "question", "partition")
):
raise ValueError("Each question requires nonempty id, question and partition strings")
if question["id"] in ids:
raise ValueError("Question IDs must be unique")
ids.add(question["id"])
# Exclusive creation prevents accidentally overwriting an earlier scored run.
with args.output.open("x", encoding="utf-8") as output:
output.write(
json.dumps(
{
"kind": "manifest",
"created_utc": datetime.now(UTC).isoformat(),
"questions_sha256": hashlib.sha256(raw).hexdigest(),
"project_id": args.project,
"limit": 5,
"budget_tokens": args.budget_tokens,
"question_count": len(questions),
"collect_top_five": args.collect_top_five,
"scoring": "independent relevance judgments required",
}
)
+ "\n"
)
output.flush()
async with Client(streamable_http_client(args.mcp_url)) as client:
for question in questions:
started = time.monotonic()
responses = []
node_ids: list[str] = []
cursor: str | None = None
seen_cursors: set[str] = set()
is_error = False
pagination_blocked = False
for _page in range(5 if args.collect_top_five else 1):
arguments = {
"query": question["question"],
"projects": [args.project],
"include_shared": True,
"all_projects": False,
"limit": 5,
"budget_tokens": args.budget_tokens,
"include_evidence": False,
}
if cursor is not None:
arguments["cursor"] = cursor
response = await client.call_tool("search", arguments)
is_error = is_error or response.is_error
pages = [
json.loads(item.text) for item in response.content if item.type == "text"
]
responses.extend(pages)
if response.is_error or not pages:
break
envelope = pages[0]
rows = envelope.get("data", [])
node_ids.extend(row["node"]["node_id"] for row in rows)
next_cursor = envelope.get("next_cursor")
if len(node_ids) >= 5 or next_cursor is None:
break
if not rows or next_cursor in seen_cursors:
pagination_blocked = True
break
seen_cursors.add(next_cursor)
cursor = next_cursor
result = {
"kind": "retrieval",
"question": question,
"seconds": round(time.monotonic() - started, 6),
"is_error": is_error,
"response": responses,
"returned_node_ids": node_ids[:5],
"pagination_blocked": pagination_blocked,
}
output.write(json.dumps(result, ensure_ascii=False) + "\n")
output.flush()
print(f"{question['id']}: {'error' if is_error else 'captured'}", flush=True)
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--questions", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--project", required=True)
parser.add_argument("--mcp-url", required=True)
parser.add_argument("--budget-tokens", type=int, default=8000)
parser.add_argument("--collect-top-five", action="store_true")
args = parser.parse_args()
if not 256 <= args.budget_tokens <= 50_000:
parser.error("--budget-tokens must be between 256 and 50000")
asyncio.run(run(args))
if __name__ == "__main__":
main()