126 lines
5.1 KiB
Python
126 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()
|