diff --git a/README.md b/README.md index 263ef35..f4d1eb9 100644 --- a/README.md +++ b/README.md @@ -8,8 +8,9 @@ coding agents and a hosted Model Context Protocol server. The [`mcp_server`](mcp_server) package exposes an anonymous, read-only, agent-native interface over Streamable HTTP. It supports MCP `2026-07-28` and -legacy 2025-era clients from one endpoint, with tools for papers, related work, -lineage, tasks, methods, benchmarks, and leaderboard evaluations. +legacy 2025-era clients from one endpoint, with one tool per `pwc` research +command: papers, related work, lineage, tasks, methods, conferences, +organizations, frameworks, benchmarks, and filtered leaderboard evaluations. Run it locally with: diff --git a/mcp_server/README.md b/mcp_server/README.md index 41176c6..38b38b3 100644 --- a/mcp_server/README.md +++ b/mcp_server/README.md @@ -22,22 +22,51 @@ curl http://127.0.0.1:7860/health ## Tools -- `search_papers` -- `list_papers` -- `get_paper_info` -- `read_paper` -- `get_related_papers` -- `get_paper_lineage` -- `get_task` -- `get_method` -- `list_benchmarks` -- `get_benchmark` +Every read-only `pwc` CLI research command is one tool, and every research flag +of that command is a tool parameter, so results match `pwc ... --json`. Each +result carries `schema_version` and `data`, the complete CLI JSON payload, next +to typed projections such as `items` or `evaluations`. + +| Tool | CLI command | +| --- | --- | +| `search_papers` | `pwc search` | +| `get_paper_info` | `pwc paper info` | +| `read_paper` | `pwc paper read` (64 KiB chunks with a continuation cursor) | +| `list_papers` | `pwc paper list` | +| `list_recent_papers` | `pwc paper recent` | +| `list_trending_papers` | `pwc paper trending` | +| `get_related_papers` | `pwc paper related` | +| `get_paper_lineage` | `pwc paper lineage list` | +| `get_task` | `pwc task --name` | +| `list_tasks` | `pwc task list` | +| `get_method` | `pwc method --name` | +| `list_methods` | `pwc method list` | +| `get_conference` | `pwc conference --name` | +| `list_conferences` | `pwc conference list` | +| `get_organization` | `pwc organization --name` | +| `list_organizations` | `pwc organization list` | +| `get_framework` | `pwc framework --name` | +| `list_frameworks` | `pwc framework list` | +| `get_benchmark` | `pwc benchmark --name` | +| `list_benchmarks` | `pwc benchmark list` | + +Parameter names follow the CLI flags except for the established MCP names +`published_after`/`published_before` (`--start-date`/`--end-date`), `limit` +(`--page-size`), `authors` (`--author`), `order_direction` (`--order-dir`), +`minimum_evaluations` (`--min-eval-count`), and `include_evaluations` +(`--include-evals`). Terminal-only flags (`--json`, +`--implementation-coverage`, `--flat`) have no parameter because MCP output is +always structured. Hosted differences from the CLI: `limit` is capped at 25, +`search_papers` defaults to `keyword` mode, `get_paper_info` includes +resources by default, and `read_paper` is chunked. +`tests/test_parity.py` fails when the CLI parser and the tool schemas drift. All tools are annotated read-only and idempotent. Search is deterministic; -the caller controls keyword or semantic mode. `read_paper` fetches at most one -64 KiB catalog chunk per call and returns a signed, one-hour continuation cursor -when more Markdown remains. Continuations stay pinned to the resolved paper and -content version, so a changed paper fails with an explicit restart response. +the caller controls keyword, hybrid, or semantic mode. `read_paper` fetches at +most one 64 KiB catalog chunk per call and returns a signed, one-hour +continuation cursor when more Markdown remains. Continuations stay pinned to +the resolved paper and content version, so a changed paper fails with an +explicit restart response. ## Resources @@ -60,10 +89,13 @@ content version, so a changed paper fails with an explicit restart response. | `LOG_LEVEL` | Content-free operational log level | Native clients may omit `Origin`. Browser requests must match the configured -allowlist. The server does not log queries, paper references, request bodies, +allowlist. A first-party client on the same host may send +`X-PwC-MCP-Client: ` to name its rate-limit identity (for example one +hashed chat session); the header counts only on a direct loopback connection +without `X-Forwarded-For`, so proxied public traffic cannot use it. The server does not log queries, paper references, request bodies, raw IP addresses, or authorization headers. -The hosted defaults allow 60 total requests and 10 semantic searches per minute +The hosted defaults allow 60 total requests and 10 semantic or hybrid searches per minute per client IP, with at most four concurrent requests per IP and 32 globally. Tool inputs cap list results at 25, catalog calls time out after 25 seconds, and request, upstream, and serialized MCP response bodies are bounded to 2 MiB. diff --git a/mcp_server/SKILL.md b/mcp_server/SKILL.md index a3059e0..d4cc5a7 100644 --- a/mcp_server/SKILL.md +++ b/mcp_server/SKILL.md @@ -1,23 +1,29 @@ --- name: pwc-mcp -description: "Papers With Code MCP tools for searching and reading AI/ML papers, discovering recent research, finding related work and paper lineage, browsing tasks and methods, and inspecting benchmark leaderboards through the public Papers With Code catalog. Use whenever the user asks to find papers, survey literature, compare research, inspect an arXiv paper, explore AI/ML taxonomy, discover benchmarks or state-of-the-art models, or mentions Papers With Code, the PwC MCP server, or paperswithcode.co/mcp." +description: "Papers With Code MCP tools for searching and reading AI/ML papers, discovering recent and trending research, finding related work and paper lineage, browsing tasks, methods, conferences, organizations, and frameworks, and inspecting benchmark leaderboards with model-size and metric filters through the public Papers With Code catalog. Use whenever the user asks to find papers, survey literature, compare research, inspect an arXiv paper, explore AI/ML taxonomy or conferences, discover benchmarks or state-of-the-art models, or mentions Papers With Code, the PwC MCP server, or paperswithcode.co/mcp." compatibility: "Requires an MCP client connected to https://paperswithcode.co/mcp with the Papers With Code tools available." --- -Generated for `pwc-mcp v0.1.0` and MCP protocol `2026-07-28`. +Generated for `pwc-mcp v0.2.0` and MCP protocol `2026-07-28`. The tools query the public [Papers With Code](https://paperswithcode.co) catalog -anonymously and are read-only. If live tool discovery and this skill disagree, +anonymously and are read-only. Every tool runs the matching `pwc` CLI research +command, and every research flag of that command is a tool parameter, so +results match `pwc ... --json`. If live tool discovery and this skill disagree, the discovered tool schemas are authoritative. -Use each tool's structured output directly for filtering, joining, or other -schema-dependent processing. +Every result carries `schema_version` and `data`, the complete CLI JSON payload +for the call. Typed fields such as `items`, `paper`, `task`, `evaluations` are +projections of `data`; read `data` when a projection lacks a detail. Use +structured output directly for filtering, joining, or other schema-dependent +processing. When the user identifies an author, prefer repeatable structured `list_papers` calls with the `authors` argument. Add `search` for stated topic terms and explicit `order_by: "date_published"` and `order_direction: "desc"` -for newest or recent work. Author references accept an exact normalized name, -numeric ID, or `@HF_USERNAME`; multiple authors use AND semantics. +for newest or recent work; the default order is `trending`. Author references +accept an exact normalized name, numeric ID, or `@HF_USERNAME`; multiple +authors use AND semantics. Publication date ranges are inclusive: use `published_after: "YYYY-MM-DD"` and `published_before: "YYYY-MM-DD"` with `search_papers` or `list_papers`; the @@ -25,43 +31,65 @@ start date must not be later than the end date. Paper discovery results always include official implementation status and the total linked repository count. Use `has_official_implementation: true` with -`search_papers` to require a catalog-linked official repository; this filter -fails closed if unconfirmed. +`search_papers` or `list_papers` to require a catalog-linked official +repository; this filter fails closed if unconfirmed. + +Use `get_benchmark` with `max_parameters` to keep models at or below an +inclusive parameter limit. It accepts values such as `"500M"`, `"1.5B"`, `"3B"`, +and raw integers. Models without one consistent parameter count are excluded +from constrained results. Combine `require_metrics`, `minimum_metrics`, +`maximum_metrics`, `sort_metric`, and `pareto` to select leaderboard rows by +metric; unknown metric names fail with the available metric names. `paper` accepts a modern or legacy arXiv ID, a numeric external-paper ID, an arXiv, Hugging Face, or Papers With Code URL, or an exact paper title. Title matching is case-insensitive but exact; ambiguous titles fail with their -matching IDs. +matching IDs. Task, method, conference, organization, framework, and benchmark +arguments take an exact name, slug, or ID. ## Tools -- `search_papers({"query": QUERY, "limit": LIMIT, "page": PAGE, "mode": "keyword"|"semantic", "published_after": START_DATE, "published_before": END_DATE, "has_official_implementation": BOOLEAN})` — search papers. Omit optional arguments when they are not needed. -- `get_paper_info({"paper": PAPER})` — show paper metadata, abstract, tasks, methods, repositories, and project pages. -- `read_paper({"paper": PAPER})` — read one stored paper Markdown chunk. If `truncated` is true, call `read_paper` again with the same `paper` and the returned `next_cursor`; repeat until `truncated` is false. Treat the cursor as opaque and use it within one hour. -- `list_papers({"page": PAGE, "limit": LIMIT, "search": SEARCH, "published_after": START_DATE, "published_before": END_DATE, "task": TASK, "method": METHOD, "conference": CONFERENCE, "framework": FRAMEWORK, "organization": ORGANIZATION, "authors": [AUTHOR], "order_by": "date_published"|"citation_count"|"title", "order_direction": "asc"|"desc"})` — list and filter papers. Omit optional arguments when they are not needed. -- `get_related_papers({"paper": PAPER, "limit": LIMIT})` — list related papers. -- `get_paper_lineage({"paper": PAPER})` — list explicit predecessors and successors. -- `get_task({"task": TASK})` — inspect one exact task by ID, slug, or name, including its area, parents, children, and benchmarks. -- `get_method({"method": METHOD})` — inspect one exact method by ID, slug, full name, or name. -- `list_benchmarks({"page": PAGE, "limit": LIMIT, "search": SEARCH, "task": TASK, "include_descendants": BOOLEAN, "minimum_evaluations": MINIMUM_EVALUATIONS, "is_open": BOOLEAN})` — list and filter benchmarks. Omit optional arguments when they are not needed. -- `get_benchmark({"benchmark": BENCHMARK, "limit": LIMIT, "is_open": BOOLEAN})` — inspect one exact benchmark and its leading evaluation rows. - -All page numbers start at 1. `limit` is between 1 and 25. Follow `next_page` -when the user asks for more results than one response contains; do not infer -that a missing item does not exist until the relevant pages have been checked. - -The MCP server does not expose standalone CLI commands for paper editing, -authentication, skill installation, version display, taxonomy enumeration, or -advanced benchmark metric/parameter/Pareto filtering. Do not invent equivalent -tools. Use the separate `pwc` CLI only when it is available and the user needs -one of those capabilities. +- `search_papers({"query": QUERY, "mode": "hybrid"|"keyword"|"semantic", "page": PAGE, "limit": LIMIT, "published_after": START_DATE, "published_before": END_DATE, "has_official_implementation": BOOLEAN})` — search papers by title, topic, author, or arXiv ID (`pwc search`). Omit optional arguments when they are not needed. +- `get_paper_info({"paper": PAPER, "include_resources": BOOLEAN, "include_evaluations": BOOLEAN})` — show paper metadata, abstract, tasks, methods, lineage, repositories, project pages, and Hugging Face model, dataset, and Space artifacts; `include_evaluations: true` adds every benchmark evaluation of the paper (`pwc paper info`). +- `read_paper({"paper": PAPER, "cursor": CURSOR})` — read one stored paper Markdown chunk (`pwc paper read`). If `truncated` is true, call `read_paper` again with the same `paper` and the returned `next_cursor`; repeat until `truncated` is false. Treat the cursor as opaque and use it within one hour. +- `list_papers({"search": SEARCH, "task": TASK, "method": METHOD, "conference": CONFERENCE, "framework": FRAMEWORK, "organization": ORGANIZATION, "authors": [AUTHOR], "published_after": START_DATE, "published_before": END_DATE, "all_versions": BOOLEAN, "order_by": "trending"|"date_published"|"citation_count", "order_direction": "asc"|"desc", "include_resources": BOOLEAN, "has_official_implementation": BOOLEAN, "page": PAGE, "limit": LIMIT})` — list and filter papers by exact catalog associations (`pwc paper list`). Omit optional arguments when they are not needed. +- `list_recent_papers({"limit": LIMIT})` — list the most recently added papers (`pwc paper recent`). +- `list_trending_papers({"limit": LIMIT, "max_age_days": DAYS, "min_velocity": VELOCITY})` — list trending papers by repository velocity (`pwc paper trending`). +- `get_related_papers({"paper": PAPER, "limit": LIMIT})` — list related papers (`pwc paper related`); `limit` is at most 20. +- `get_paper_lineage({"paper": PAPER})` — list explicit predecessors and successors (`pwc paper lineage list`). +- `get_task({"task": TASK})` — inspect one exact task, including its area, hierarchy, sister tasks, ranked benchmarks, common methods, recommended frameworks, and trending papers (`pwc task --name`). +- `list_tasks({"area": AREA, "level": LEVEL, "visible_only": BOOLEAN, "group_by_area": BOOLEAN, "order_by": "name"|"created_at"|"level"|"paper_count", "order_direction": "asc"|"desc", "page": PAGE, "limit": LIMIT})` — list and filter research tasks, or set `group_by_area: true` for the complete visible top-level taxonomy without pagination (`pwc task list`). +- `get_method({"method": METHOD})` — inspect one exact method with its area (`pwc method --name`). +- `list_methods({"area": AREA, "introduced_year": YEAR, "order_by": "name"|"full_name"|"introduced_year"|"created_at"|"paper_count", "order_direction": "asc"|"desc", "page": PAGE, "limit": LIMIT})` — list and filter research methods (`pwc method list`). +- `get_conference({"conference": CONFERENCE})` — inspect one exact conference (`pwc conference --name`). +- `list_conferences({"year": YEAR})` — list conferences with imported papers (`pwc conference list`). +- `get_organization({"organization": ORGANIZATION})` — inspect one exact research organization (`pwc organization --name`). +- `list_organizations({"featured_only": BOOLEAN})` — list research organizations (`pwc organization list`). +- `get_framework({"framework": FRAMEWORK})` — inspect one exact research framework (`pwc framework --name`). +- `list_frameworks({"domain": DOMAIN, "category": CATEGORY, "platform": PLATFORM})` — list research frameworks (`pwc framework list`). +- `get_benchmark({"benchmark": BENCHMARK, "limit": LIMIT, "is_open": BOOLEAN, "max_parameters": SIZE, "require_metrics": [METRIC], "minimum_metrics": {METRIC: VALUE}, "maximum_metrics": {METRIC: VALUE}, "sort_metric": "METRIC:asc|desc", "pareto": ["METRIC:higher", "METRIC:lower"]})` — inspect one exact benchmark leaderboard with model-size, metric threshold, sort, and Pareto selection (`pwc benchmark --name`). `matched_count` reports rows that satisfied the filters before `limit`. +- `list_benchmarks({"search": SEARCH, "task": TASK, "include_descendants": BOOLEAN, "minimum_evaluations": COUNT, "is_open": BOOLEAN, "group_by_area": BOOLEAN, "area": AREA, "benchmarks_per_task": COUNT, "order_by": "trending"|"name"|"full_name"|"created_at"|"paper_count", "order_direction": "asc"|"desc", "page": PAGE, "limit": LIMIT})` — list and filter benchmarks (`pwc benchmark list`). With `task`, results are ranked by trend unless `order_by` is set; `group_by_area` or `area` returns top benchmarks under each visible task. + +All page numbers start at 1. `limit` is between 1 and 25 and defaults to the +CLI page size capped at 25. Follow `next_page` when the user asks for more +results than one response contains; do not infer that a missing item does not +exist until the relevant pages have been checked. + +The MCP server does not expose the standalone CLI's paper editing, +authentication, skill installation, or version commands, and `pwc paper read` +is served as `read_paper` chunks rather than one document. +Do not invent equivalent tools. Use the separate `pwc` CLI only when it is +available and the user needs one of those capabilities. ## Research workflow 1. Use `list_benchmarks({"task": TASK})` to discover active benchmarks, then - `get_benchmark({"benchmark": NAME})` to inspect a leaderboard. + `get_benchmark({"benchmark": NAME})` to inspect a leaderboard. Add + `max_parameters` when model size is part of the request and `sort_metric` + or `minimum_metrics` when a specific metric matters. 2. Use `get_paper_info({"paper": PAPER})` to inspect promising results. Its - response includes repositories and project pages. + response includes repositories, project pages, and Hugging Face artifacts; + add `include_evaluations: true` to compare one paper across benchmarks. 3. Use exact `list_papers` `authors`, `task`, `method`, `conference`, `framework`, and `organization` arguments for known identities or catalog associations. Combine them to require every association; do not substitute @@ -70,24 +98,28 @@ one of those capabilities. evidence. Follow every returned continuation cursor needed for the user's question. 5. Expand the literature with `get_related_papers` and use - `get_paper_lineage` when model or method ancestry matters. + `get_paper_lineage` when model or method ancestry matters. Use `get_task`, + `get_method`, `list_conferences`, `list_organizations`, and + `list_frameworks` to orient within the taxonomy. 6. Preserve paper titles, identifiers, and URLs so claims remain traceable. ## Output and limits - Tool results use stable, versioned structured output with compact text fallbacks. Prefer structured fields over parsing the text fallback. -- Search mode is deterministic: choose `keyword` by default and use `semantic` - when conceptual similarity is more useful. The MCP server does not support - the CLI's `hybrid` mode. +- Search mode is deterministic: `keyword` is the default; `hybrid` adds dense + retrieval and `semantic` uses it alone. Both `hybrid` and `semantic` count + toward the semantic search rate limit. - Catalog-filtered paper lists fail closed unless the server confirms every requested filter; never treat results from an older server as filtered. + Parameter-filtered leaderboards fail closed unless the server confirms + parameter-filter support and every returned model satisfies the limit. - `read_paper` returns at most one 64 KiB chunk per call. A continuation stays pinned to the resolved paper and content version; if the paper changes, start reading again from the beginning. - The hosted service limits each client to 60 total requests and 10 semantic - searches per minute, with at most four concurrent requests. Respect retry - metadata instead of trying to work around limits. + or hybrid searches per minute, with at most four concurrent requests. Respect + retry metadata instead of trying to work around limits. - The tools contain no authentication, catalog mutation, ingestion, publication, image, embedding, CRON, or infrastructure-maintenance operations. diff --git a/mcp_server/SPEC.md b/mcp_server/SPEC.md index 1af6141..665cef4 100644 --- a/mcp_server/SPEC.md +++ b/mcp_server/SPEC.md @@ -18,18 +18,37 @@ server. ## Public contract -Expose exactly these tools: +Expose exactly one tool per read-only `pwc` CLI research command, with every +research flag of that command as a tool parameter (`tests/test_parity.py` +enforces this against the CLI parser): -- `search_papers` -- `list_papers` -- `get_paper_info` -- `read_paper` -- `get_related_papers` -- `get_paper_lineage` -- `get_task` -- `get_method` -- `list_benchmarks` -- `get_benchmark` +- `search_papers` (`pwc search`) +- `get_paper_info` (`pwc paper info`) +- `read_paper` (`pwc paper read`) +- `list_papers` (`pwc paper list`) +- `list_recent_papers` (`pwc paper recent`) +- `list_trending_papers` (`pwc paper trending`) +- `get_related_papers` (`pwc paper related`) +- `get_paper_lineage` (`pwc paper lineage list`) +- `get_task` (`pwc task --name`) +- `list_tasks` (`pwc task list`) +- `get_method` (`pwc method --name`) +- `list_methods` (`pwc method list`) +- `get_conference` (`pwc conference --name`) +- `list_conferences` (`pwc conference list`) +- `get_organization` (`pwc organization --name`) +- `list_organizations` (`pwc organization list`) +- `get_framework` (`pwc framework --name`) +- `list_frameworks` (`pwc framework list`) +- `get_benchmark` (`pwc benchmark --name`) +- `list_benchmarks` (`pwc benchmark list`) + +Tools run the CLI handlers in-process through the shared cached transport, so +validation, fail-closed filter confirmation, and the JSON payload are the +CLI's. Every result includes that payload as `data` beside typed projections. +Terminal-only flags have no parameter. The hosted service caps `limit` at 25, +defaults `search_papers` to keyword mode, includes paper resources by default, +and serves `read_paper` in chunks. Expose these resource templates and no prompts: @@ -53,7 +72,7 @@ than selecting one result. - Require a strict configurable browser Origin allowlist; native clients may omit Origin. Never configure a wildcard Origin. -- Enforce, per IP, 60 total requests/minute, 10 semantic searches/minute, and +- Enforce, per IP, 60 total requests/minute, 10 semantic or hybrid searches/minute, and four concurrent requests, plus a global ceiling of 32 concurrent requests. List tools return at most 25 rows. - Set catalog timeouts to 25 seconds and bound HTTP request and upstream response @@ -66,7 +85,8 @@ than selecting one result. Cache immutable, versioned Markdown chunks for one hour within both a 256-entry and 16 MiB ceiling. - Bind to loopback on the VPS, trust forwarded identity only from an exact - loopback peer, and expose cached catalog readiness without making `/health` + loopback peer, let a direct loopback client without `X-Forwarded-For` name + its own rate-limit identity with `X-PwC-MCP-Client`, and expose cached catalog readiness without making `/health` wait on an upstream call. ## Release diff --git a/mcp_server/pyproject.toml b/mcp_server/pyproject.toml index 896cd32..a1ccae7 100644 --- a/mcp_server/pyproject.toml +++ b/mcp_server/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "pwc-mcp" -version = "0.1.0" +version = "0.2.0" description = "Read-only Papers With Code MCP server" readme = "README.md" requires-python = ">=3.10" @@ -8,7 +8,7 @@ license = {text = "Apache-2.0"} dependencies = [ "mcp>=2.1,<3", "pydantic>=2.11,<3", - "pwc-cli>=0.3.1,<0.4", + "pwc-cli>=0.4.1,<0.5", "starlette>=0.47,<1", "uvicorn>=0.35,<1", ] diff --git a/mcp_server/src/pwc_mcp/__init__.py b/mcp_server/src/pwc_mcp/__init__.py index ffd884d..cab4aac 100644 --- a/mcp_server/src/pwc_mcp/__init__.py +++ b/mcp_server/src/pwc_mcp/__init__.py @@ -1,3 +1,3 @@ """Read-only Papers With Code MCP server.""" -__version__ = "0.1.0" +__version__ = "0.2.0" diff --git a/mcp_server/src/pwc_mcp/app.py b/mcp_server/src/pwc_mcp/app.py index 1967f42..6c84b7a 100644 --- a/mcp_server/src/pwc_mcp/app.py +++ b/mcp_server/src/pwc_mcp/app.py @@ -6,6 +6,7 @@ import json import logging import os +import re import secrets import threading import time @@ -23,7 +24,7 @@ from pwc_mcp import __version__ from pwc_mcp.catalog import CatalogClient -from pwc_mcp.server import Catalog, build_server +from pwc_mcp.server import TOOL_COMMANDS, Catalog, build_server LOGGER = logging.getLogger("pwc_mcp.requests") # MCP SDK diagnostics can include peer-supplied tool names and resource URIs. @@ -32,18 +33,7 @@ PROTOCOL_VERSION = "2026-07-28" MAX_REQUEST_BODY_SIZE = 2 * 1024 * 1024 MAX_RESPONSE_BODY_SIZE = 2 * 1024 * 1024 -KNOWN_TOOLS = { - "search_papers", - "list_papers", - "get_paper_info", - "read_paper", - "get_related_papers", - "get_paper_lineage", - "get_task", - "get_method", - "list_benchmarks", - "get_benchmark", -} +KNOWN_TOOLS = frozenset(TOOL_COMMANDS) KNOWN_PROTOCOLS = { PROTOCOL_VERSION, "2025-11-25", @@ -105,6 +95,14 @@ async def health(request: Request) -> JSONResponse: ) +# A first-party client on the same host (the chat gateway) may name the +# rate-limit identity of a request, for example one hashed chat session. The +# header counts only on a direct loopback connection that carries no proxy +# header: nginx always adds X-Forwarded-For, so public traffic cannot use it. +CLIENT_IDENTITY_HEADER = "x-pwc-mcp-client" +CLIENT_IDENTITY = re.compile(r"[A-Za-z0-9][A-Za-z0-9._:-]{0,127}") + + def _client_address(scope: Scope, headers: Headers, trust_proxy_headers: bool) -> str: client = scope.get("client") direct_address = str(client[0]) if client else "unknown" @@ -116,6 +114,9 @@ def _client_address(scope: Scope, headers: Headers, trust_proxy_headers: bool) - forwarded = headers.get("x-forwarded-for") if forwarded: return forwarded.split(",", 1)[0].strip() + identity = headers.get(CLIENT_IDENTITY_HEADER, "") + if CLIENT_IDENTITY.fullmatch(identity): + return f"client:{identity}" return direct_address @@ -206,7 +207,8 @@ def _tool_and_semantic(headers: Headers, body: bytes) -> tuple[str | None, bool] mode = arguments.get("mode") tool = body_tool or header_tool tool = tool if tool in KNOWN_TOOLS else None - return tool, tool == "search_papers" and mode == "semantic" + # Hybrid retrieval also runs the dense embedding search upstream. + return tool, tool == "search_papers" and mode in {"semantic", "hybrid"} def _protocol_label(headers: Headers) -> str: diff --git a/mcp_server/src/pwc_mcp/catalog.py b/mcp_server/src/pwc_mcp/catalog.py index c5cacc6..02730b3 100644 --- a/mcp_server/src/pwc_mcp/catalog.py +++ b/mcp_server/src/pwc_mcp/catalog.py @@ -4,10 +4,12 @@ import threading import time from collections import OrderedDict +from collections.abc import Mapping from dataclasses import dataclass from typing import Any, Protocol from urllib.parse import quote, urlparse +from pwc_cli import queries from pwc_cli.transport import Client, HTTPStatusError, Response, ResponseError PAPER_ID = re.compile(r"(?:\d{4}\.\d{4,5}|[a-z][a-z0-9.-]*/\d{7}|\d+)", re.IGNORECASE) @@ -17,6 +19,27 @@ MARKDOWN_CACHE_ENTRIES = 256 MARKDOWN_CACHE_BYTES = 16 * 1024 * 1024 MARKDOWN_CACHE_SECONDS = 3600 +# Search and paper listings change quickly; taxonomy is stable; the rest sits +# between (SPEC.md: one minute, ten minutes, and five minutes respectively). +VOLATILE_PATHS = frozenset( + {"papers/", "papers/search", "papers/recent", "papers/trending"} +) +TAXONOMY_PREFIXES = ( + "tasks/", + "methods/", + "areas/", + "conferences/", + "organizations/", + "frameworks/", +) + + +def cache_ttl(path: str) -> int: + if path in VOLATILE_PATHS: + return 60 + if path.startswith(TAXONOMY_PREFIXES): + return 600 + return 300 @dataclass(frozen=True) @@ -79,9 +102,9 @@ def put(self, key: object, value: object, ttl_seconds: int) -> None: class _MarkdownChunkCache: def __init__(self): - self._values: OrderedDict[ - object, tuple[float, int, PaperMarkdownChunk] - ] = OrderedDict() + self._values: OrderedDict[object, tuple[float, int, PaperMarkdownChunk]] = ( + OrderedDict() + ) self._bytes = 0 self._lock = threading.Lock() @@ -129,6 +152,21 @@ def __init__(self, transport: Transport | None = None): self.cache = _TTLCache() self.markdown_cache = _MarkdownChunkCache() + def _response( + self, + path: str, + params: Mapping[str, object | None] | None = None, + *, + ttl: int, + ) -> Response: + key = ("response", path, _freeze(dict(params or {}))) + cached = self.cache.get(key) + if isinstance(cached, Response): + return cached + response = self.transport.get(path, dict(params or {})) + self.cache.put(key, response, ttl) + return response + def _json( self, path: str, @@ -136,14 +174,9 @@ def _json( *, ttl: int, ) -> dict[str, Any]: - key = ("json", path, _freeze(params or {})) - cached = self.cache.get(key) - if isinstance(cached, dict): - return cached - payload = self.transport.get(path, params).json() + payload = self._response(path, params, ttl=ttl).json() if not isinstance(payload, dict): raise ResponseError("API returned an unexpected response shape") - self.cache.put(key, payload, ttl) return payload def check_readiness(self) -> bool: @@ -263,65 +296,6 @@ def _resolve_paper(self, reference: str) -> str: def resolve_paper(self, paper: str) -> str: return self._resolve_paper(paper) - @staticmethod - def _exact( - reference: str, - items: list[dict[str, Any]], - label: str, - fields: tuple[str, ...] = ("name", "slug", "id"), - ) -> dict[str, Any]: - target = reference.strip().casefold() - for field in fields: - for item in items: - if str(item.get(field) or "").strip().casefold() == target: - return item - raise ResponseError(f"{label} not found: {reference}") - - def search_papers( - self, - *, - query: str, - mode: str = "keyword", - page: int = 1, - limit: int = 10, - published_after: str | None = None, - published_before: str | None = None, - has_official_implementation: bool = False, - ) -> dict[str, Any]: - params: dict[str, object | None] = { - "q": query, - "mode": mode, - "page": page, - "page_size": limit, - "start_date": published_after, - "end_date": published_before, - "has_official_implementation": ( - True if has_official_implementation else None - ), - } - payload = self._json("papers/search", params, ttl=60) - if has_official_implementation and (payload.get("applied_filters") or {}).get( - "has_official_implementation" - ) not in {True, "true"}: - raise ResponseError( - "Papers API did not confirm has_official_implementation" - ) - return payload - - def get_paper_info(self, paper: str, *, include_resources: bool) -> dict[str, Any]: - reference = self._resolve_paper(paper) - return self._json( - f"papers/{quote(reference, safe='.')}", - {"include_resources": include_resources}, - ttl=300, - ) - - def read_paper(self, paper: str) -> str: - reference = self._resolve_paper(paper) - return self._text( - f"research/papers/{quote(reference, safe='.')}/read", ttl=3600 - ) - def read_paper_chunk( self, paper: str, @@ -361,7 +335,9 @@ def read_paper_chunk( markdown = response.body.decode("utf-8") next_offset = int(next_text) if next_text is not None else None except (UnicodeDecodeError, ValueError) as error: - raise ResponseError("Papers API returned an invalid Markdown chunk") from error + raise ResponseError( + "Papers API returned an invalid Markdown chunk" + ) from error if ( CONTENT_VERSION.fullmatch(returned_version) is None or truncated not in {"0", "1"} @@ -369,10 +345,7 @@ def read_paper_chunk( or (truncated == "1" and next_offset is None) or (truncated == "0" and next_offset is not None) or (next_offset is not None and next_offset <= offset) - or ( - next_offset is not None - and next_offset - offset != len(response.body) - ) + or (next_offset is not None and next_offset - offset != len(response.body)) ): detail = ( "Papers API Markdown offset did not advance" @@ -387,166 +360,30 @@ def read_paper_chunk( content_version=returned_version, next_offset=next_offset, ) - self.markdown_cache.put( - (reference, returned_version, offset, limit), result - ) + self.markdown_cache.put((reference, returned_version, offset, limit), result) return result - def list_papers( - self, - *, - search: str | None = None, - task: str | None = None, - method: str | None = None, - conference: str | None = None, - framework: str | None = None, - organization: str | None = None, - authors: list[str] | None = None, - published_after: str | None = None, - published_before: str | None = None, - order_by: str = "date_published", - order_direction: str = "desc", - page: int = 1, - limit: int = 10, - ) -> dict[str, Any]: - requested = { - key: value - for key, value in { - "task": task, - "method": method, - "conference": conference, - "framework": framework, - "organization": organization, - "start_date": published_after, - "end_date": published_before, - }.items() - if value is not None - } - params: dict[str, object | None] = { - "search": search, - **requested, - "author": authors or None, - "latest_only": True, - "order_by": order_by, - "order_dir": order_direction, - "page": page, - "page_size": limit, - } - payload = self._json("papers/", params, ttl=60) - applied = payload.get("applied_filters") - if requested and ( - not isinstance(applied, dict) - or any( - str(applied.get(key, "")).casefold() != str(value).casefold() - for key, value in requested.items() - ) - ): - raise ResponseError("Papers API did not confirm requested catalog filters") - return payload + def query(self, command: tuple[str, ...], options: Mapping[str, Any]) -> Any: + """Run one read-only ``pwc`` command in-process against the cached catalog. - def get_related_papers(self, paper: str, *, limit: int) -> dict[str, Any]: - reference = self._resolve_paper(paper) - return self._json( - f"papers/{quote(reference, safe='.')}/related", - {"limit": limit}, - ttl=300, - ) + Paper references are resolved first so URLs, legacy IDs, and exact + titles behave exactly as they do for ``read_paper``; the CLI then + receives the canonical identifier and applies its own validation, + fail-closed filter checks, and JSON payload shape. + """ + resolved = dict(options) + if resolved.get("paper") is not None: + resolved["paper"] = self._resolve_paper(str(resolved["paper"])) + return queries.query(tuple(command), resolved, _CachedTransport(self)) - def get_paper_lineage(self, paper: str) -> dict[str, Any]: - reference = self._resolve_paper(paper) - return self._json( - f"research/papers/{quote(reference, safe='.')}/lineage", ttl=300 - ) - def get_task(self, task: str) -> dict[str, Any]: - if task.strip().isdigit(): - task_id = task.strip() - else: - candidates = self._rows( - self._json( - "tasks/", - {"q": task, "page": 1, "page_size": 100}, - ttl=600, - ) - ) - task_id = str(self._exact(task, candidates, "Task").get("id")) - return self._json(f"tasks/{quote(task_id, safe='')}/page", ttl=600) +class _CachedTransport: + """``pwc_cli.transport.Client`` stand-in that serves CLI handlers from the cache.""" - def get_method(self, method: str) -> dict[str, Any]: - if method.strip().isdigit(): - method_id = method.strip() - else: - candidates = self._rows( - self._json( - "methods/", - {"q": method, "page": 1, "page_size": 100}, - ttl=600, - ) - ) - matched = self._exact( - method, - candidates, - "Method", - fields=("name", "full_name", "slug", "id"), - ) - method_id = str(matched.get("id") or matched.get("slug")) - return self._json(f"methods/{quote(method_id, safe='')}", ttl=600) - - def list_benchmarks( - self, - *, - search: str | None = None, - task: str | None = None, - include_descendants: bool = False, - minimum_evaluations: int | None = None, - is_open: bool | None = None, - page: int = 1, - limit: int = 10, - ) -> dict[str, Any]: - return self._json( - "datasets/", - { - "q": search, - "task": task, - "include_descendants": include_descendants, - "min_eval_count": minimum_evaluations, - "is_open": is_open, - "ordering": "-paper_count", - "page": page, - "page_size": limit, - }, - ttl=300, - ) + def __init__(self, catalog: CatalogClient): + self.catalog = catalog - def get_benchmark( - self, benchmark: str, *, limit: int, is_open: bool | None - ) -> dict[str, Any]: - candidates = self._rows( - self._json( - "datasets/", - {"q": benchmark, "page": 1, "page_size": 100}, - ttl=300, - ) - ) - matched = self._exact( - benchmark, - candidates, - "Benchmark", - fields=("name", "full_name", "slug", "id"), - ) - benchmark_id = str(matched.get("id")) - evaluations = self._json( - f"datasets/{quote(benchmark_id, safe='')}/evaluations/", - { - "page": 1, - "page_size": limit, - "ordering": "best_rank", - "is_open": is_open, - }, - ttl=300, - ) - return { - "benchmark": matched, - "count": evaluations.get("count") or 0, - "results": evaluations.get("results") or [], - } + def get( + self, path: str, params: Mapping[str, object | None] | None = None + ) -> Response: + return self.catalog._response(path, params, ttl=cache_ttl(path)) diff --git a/mcp_server/src/pwc_mcp/models.py b/mcp_server/src/pwc_mcp/models.py index 71bffbd..5fdf0af 100644 --- a/mcp_server/src/pwc_mcp/models.py +++ b/mcp_server/src/pwc_mcp/models.py @@ -2,13 +2,23 @@ from typing import Any, Literal -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, Field class OutputModel(BaseModel): model_config = ConfigDict(extra="forbid") +class QueryResult(OutputModel): + """The exact ``data`` document that ``pwc --json`` prints.""" + + schema_version: Literal["v1"] = "v1" + data: Any = Field( + default=None, + description="Complete CLI JSON payload for this call; typed fields are projections of it.", + ) + + class PaperSummary(OutputModel): id: str arxiv_id: str | None = None @@ -21,8 +31,7 @@ class PaperSummary(OutputModel): code_repository_count: int -class PaperPage(OutputModel): - schema_version: Literal["v1"] = "v1" +class PaperPage(QueryResult): items: list[PaperSummary] next_page: int | None = None @@ -52,11 +61,29 @@ class PaperDetail(OutputModel): methods: list[CatalogReference] repositories: list[RepositoryReference] project_pages: list[str] + hf_models: list[str] = [] + hf_datasets: list[str] = [] + hf_spaces: list[str] = [] -class PaperInfoResult(OutputModel): - schema_version: Literal["v1"] = "v1" +class PaperEvaluation(OutputModel): + id: str + benchmark: str | None = None + task: str | None = None + model_name: str + harness: str | None = None + metrics: dict[str, float | int | str | None] + best_metric: str | None = None + best_rank: int | None = None + is_open: bool + num_parameters: int | None = None + source_url: str | None = None + + +class PaperInfoResult(QueryResult): paper: PaperDetail + evaluation_count: int | None = None + evaluations: list[PaperEvaluation] | None = None class PaperReadResult(OutputModel): @@ -73,8 +100,7 @@ class PaperReference(OutputModel): title: str -class PaperLineageResult(OutputModel): - schema_version: Literal["v1"] = "v1" +class PaperLineageResult(QueryResult): paper: PaperReference predecessors: list[PaperReference] successors: list[PaperReference] @@ -107,8 +133,7 @@ class TaskDetail(OutputModel): benchmarks: list[BenchmarkSummary] -class TaskResult(OutputModel): - schema_version: Literal["v1"] = "v1" +class TaskResult(QueryResult): task: TaskDetail @@ -125,13 +150,11 @@ class MethodDetail(OutputModel): paper_count: int -class MethodResult(OutputModel): - schema_version: Literal["v1"] = "v1" +class MethodResult(QueryResult): method: MethodDetail -class BenchmarkPage(OutputModel): - schema_version: Literal["v1"] = "v1" +class BenchmarkPage(QueryResult): items: list[BenchmarkSummary] next_page: int | None = None @@ -139,35 +162,57 @@ class BenchmarkPage(OutputModel): class Evaluation(OutputModel): id: str model_name: str + harness: str | None = None metrics: dict[str, float | int | str | None] + best_metric: str | None = None best_rank: int | None = None + task: str | None = None paper_id: str | None = None paper_title: str | None = None paper_arxiv_id: str | None = None + paper_published: str | None = None is_open: bool num_parameters: int | None = None -class BenchmarkResult(OutputModel): - schema_version: Literal["v1"] = "v1" +class BenchmarkResult(QueryResult): benchmark: BenchmarkSummary evaluation_count: int + matched_count: int | None = None evaluations: list[Evaluation] +def _text(value: Any) -> str | None: + return str(value) if value not in (None, "") else None + + +def _int(value: Any) -> int | None: + if isinstance(value, bool) or value is None: + return None + try: + return int(value) + except (TypeError, ValueError): + return None + + +def _urls(values: Any) -> list[str]: + urls = [] + for value in values or []: + url = value.get("url") if isinstance(value, dict) else value + if url: + urls.append(str(url)) + return urls + + def paper_summary(item: dict[str, Any]) -> PaperSummary: return PaperSummary( id=str(item.get("id") or ""), - arxiv_id=str(item["arxiv_id"]) if item.get("arxiv_id") else None, + arxiv_id=_text(item.get("arxiv_id")), title=str(item.get("title") or "Untitled paper"), authors=[str(author) for author in item.get("authors") or []], - published=str(item["published"]) if item.get("published") else None, - citation_count=( - int(item["citation_count"]) - if item.get("citation_count") is not None - else None - ), - url=str(item.get("url_abs") or item.get("source_url") or "") or None, + published=_text(item.get("published") or item.get("date_published")), + citation_count=_int(item.get("citation_count")), + url=_text(item.get("url_abs") or item.get("source_url")), has_official_implementation=item.get("has_official_implementation") is True, code_repository_count=int(item.get("code_repository_count") or 0), ) @@ -177,7 +222,7 @@ def catalog_reference(item: dict[str, Any]) -> CatalogReference: return CatalogReference( id=str(item.get("id") or ""), name=str(item.get("name") or item.get("slug") or "Unknown"), - slug=str(item["slug"]) if item.get("slug") else None, + slug=_text(item.get("slug")), ) @@ -191,25 +236,16 @@ def paper_detail(item: dict[str, Any]) -> PaperDetail: is_official=repository.get("is_official") is True, ) ) - project_pages = [] - for page in item.get("project_pages") or []: - url = page.get("url") if isinstance(page, dict) else page - if url: - project_pages.append(str(url)) return PaperDetail( id=str(item.get("id") or ""), - arxiv_id=str(item["arxiv_id"]) if item.get("arxiv_id") else None, + arxiv_id=_text(item.get("arxiv_id")), title=str(item.get("title") or "Untitled paper"), - abstract=str(item["abstract"]) if item.get("abstract") else None, + abstract=_text(item.get("abstract")), authors=[str(author) for author in item.get("authors") or []], - published=str(item["published"]) if item.get("published") else None, - citation_count=( - int(item["citation_count"]) - if item.get("citation_count") is not None - else None - ), - url=str(item.get("url_abs") or item.get("source_url") or "") or None, - pdf_url=str(item["url_pdf"]) if item.get("url_pdf") else None, + published=_text(item.get("published")), + citation_count=_int(item.get("citation_count")), + url=_text(item.get("url_abs") or item.get("source_url")), + pdf_url=_text(item.get("url_pdf")), tasks=[ catalog_reference(task) for task in item.get("tasks") or [] @@ -221,53 +257,75 @@ def paper_detail(item: dict[str, Any]) -> PaperDetail: if isinstance(method, dict) ], repositories=repositories, - project_pages=project_pages, + project_pages=_urls(item.get("project_pages")), + hf_models=_urls(item.get("hf_models")), + hf_datasets=_urls(item.get("hf_datasets")), + hf_spaces=_urls(item.get("hf_spaces")), + ) + + +def _metrics(item: dict[str, Any]) -> dict[str, float | int | str | None]: + metrics = item.get("metrics") + if not isinstance(metrics, dict): + return {} + return {str(key): value for key, value in metrics.items()} + + +def paper_evaluation(item: dict[str, Any]) -> PaperEvaluation: + return PaperEvaluation( + id=str(item.get("id") or ""), + benchmark=_text(item.get("dataset_name")), + task=_text(item.get("task_name")), + model_name=str(item.get("model_name") or "Unknown model"), + harness=_text(item.get("harness")), + metrics=_metrics(item), + best_metric=_text(item.get("best_metric")), + best_rank=_int(item.get("best_rank")), + is_open=item.get("is_open") is not False, + num_parameters=_int(item.get("num_parameters")), + source_url=_text(item.get("result_url") or item.get("source_url")), ) def paper_reference(item: dict[str, Any]) -> PaperReference: + reference = item.get("reference") or item.get("arxiv_id") return PaperReference( id=str(item.get("id") or ""), - reference=( - str(item.get("reference") or item.get("arxiv_id")) - if item.get("reference") or item.get("arxiv_id") - else None - ), + reference=str(reference) if reference else None, title=str(item.get("title") or item.get("reference") or "Untitled paper"), ) def benchmark_summary(item: dict[str, Any]) -> BenchmarkSummary: + count = ( + item.get("paper_count") + if item.get("paper_count") is not None + else item.get("all_time_paper_count", item.get("evaluation_count")) + ) return BenchmarkSummary( id=str(item.get("id") or ""), name=str(item.get("name") or item.get("slug") or "Unknown benchmark"), - slug=str(item["slug"]) if item.get("slug") else None, - full_name=str(item["full_name"]) if item.get("full_name") else None, - description=str(item["description"]) if item.get("description") else None, - hf_url=str(item["hf_url"]) if item.get("hf_url") else None, - paper_count=int(item.get("paper_count") or 0), + slug=_text(item.get("slug")), + full_name=_text(item.get("full_name")), + description=_text(item.get("description")), + hf_url=_text(item.get("hf_url")), + paper_count=_int(count) or 0, ) def evaluation(item: dict[str, Any]) -> Evaluation: - metrics = item.get("metrics") return Evaluation( id=str(item.get("id") or ""), model_name=str(item.get("model_name") or "Unknown model"), - metrics={str(key): value for key, value in metrics.items()} - if isinstance(metrics, dict) - else {}, - best_rank=int(item["best_rank"]) if item.get("best_rank") is not None else None, - paper_id=str(item["paper_id"]) if item.get("paper_id") else None, - paper_title=str(item["paper_title"]) if item.get("paper_title") else None, - paper_arxiv_id=( - str(item["paper_arxiv_id"]) if item.get("paper_arxiv_id") else None - ), + harness=_text(item.get("harness")), + metrics=_metrics(item), + best_metric=_text(item.get("best_metric")), + best_rank=_int(item.get("best_rank")), + task=_text(item.get("task_name")), + paper_id=_text(item.get("paper_id")), + paper_title=_text(item.get("paper_title")), + paper_arxiv_id=_text(item.get("paper_arxiv_id")), + paper_published=_text(item.get("paper_published_date")), is_open=item.get("is_open") is not False, - num_parameters=( - int(item["num_parameters"]) - if isinstance(item.get("num_parameters"), int) - and not isinstance(item.get("num_parameters"), bool) - else None - ), + num_parameters=_int(item.get("num_parameters")), ) diff --git a/mcp_server/src/pwc_mcp/server.py b/mcp_server/src/pwc_mcp/server.py index ea30566..f2d9d7e 100644 --- a/mcp_server/src/pwc_mcp/server.py +++ b/mcp_server/src/pwc_mcp/server.py @@ -1,14 +1,27 @@ +"""MCP tools that mirror every read-only ``pwc`` research command. + +Each tool maps one CLI command and exposes each of its research flags as a +typed parameter, then runs the CLI handler in-process through the catalog. +``TOOL_COMMANDS``, ``PARAMETER_NAMES``, ``ENTITY_PARAMETERS`` and +``MCP_ONLY_PARAMETERS`` are the parity contract that ``tests/test_parity.py`` +checks against the CLI parser. +""" + from __future__ import annotations import os import time +from collections.abc import Mapping from datetime import date +from functools import cache from typing import Annotated, Any, Literal, Protocol from mcp.server.caching import CacheHint from mcp.server.mcpserver import MCPServer from mcp.server.mcpserver.exceptions import ToolError from mcp.types import ToolAnnotations +from pwc_cli import queries +from pwc_cli.cli import UsageError from pwc_cli.transport import ResponseError, TransportError from pydantic import Field @@ -30,12 +43,14 @@ PaperLineageResult, PaperPage, PaperReadResult, + QueryResult, TaskDetail, TaskResult, benchmark_summary, catalog_reference, evaluation, paper_detail, + paper_evaluation, paper_reference, paper_summary, ) @@ -46,11 +61,166 @@ idempotent_hint=True, open_world_hint=True, ) -Page = Annotated[int, Field(ge=1, le=100)] -Limit = Annotated[int, Field(ge=1, le=25)] -Reference = Annotated[str, Field(min_length=1, max_length=500)] +# Hosted ceiling on rows per response (SPEC.md); the CLI allows up to 100. +MAX_ROWS = 25 +# CLI lookup failures the caller can act on (an unknown or ambiguous name, an +# unconfirmed filter). Transport, HTTP, and response-shape failures stay generic. +CLIENT_FACING_ERRORS = ( + "Task not found", + "Method not found", + "Benchmark not found", + "Conference not found", + "Organization not found", + "Framework not found", + "Area not found", + "Paper title not found", + "Paper title is ambiguous", + "Paper reference cannot be empty", + "Too many results to resolve paper title", + "Papers API did not confirm", +) +GENERIC_ERROR = "the Papers With Code catalog request failed" + + +def catalog_error_message(error: Exception) -> str: + """Return the CLI's own lookup message when it is actionable, else a generic one.""" + if isinstance(error, ResponseError) and not isinstance(error, TransportError): + message = str(error) + if message.startswith(CLIENT_FACING_ERRORS): + return message + return GENERIC_ERROR + + +# One read-only CLI command per tool. +TOOL_COMMANDS: dict[str, tuple[str, ...]] = { + "search_papers": ("search",), + "get_paper_info": ("paper", "info"), + "read_paper": ("paper", "read"), + "list_papers": ("paper", "list"), + "list_recent_papers": ("paper", "recent"), + "list_trending_papers": ("paper", "trending"), + "get_related_papers": ("paper", "related"), + "get_paper_lineage": ("paper", "lineage", "list"), + "get_task": ("task",), + "list_tasks": ("task", "list"), + "get_method": ("method",), + "list_methods": ("method", "list"), + "get_conference": ("conference",), + "list_conferences": ("conference", "list"), + "get_organization": ("organization",), + "list_organizations": ("organization", "list"), + "get_framework": ("framework",), + "list_frameworks": ("framework", "list"), + "get_benchmark": ("benchmark",), + "list_benchmarks": ("benchmark", "list"), +} +# CLI destinations that keep their established MCP parameter name. +PARAMETER_NAMES: dict[str, str] = { + "start_date": "published_after", + "end_date": "published_before", + "page_size": "limit", + "author": "authors", + "order_dir": "order_direction", + "min_eval_count": "minimum_evaluations", + "include_evals": "include_evaluations", +} +# ``--name`` selects the tool's entity and is exposed under the entity's name. +ENTITY_PARAMETERS: dict[str, str] = { + "get_task": "task", + "get_method": "method", + "get_conference": "conference", + "get_organization": "organization", + "get_framework": "framework", + "get_benchmark": "benchmark", +} +# Parameters with no CLI flag; ``read_paper`` continues with a signed cursor. +MCP_ONLY_PARAMETERS: dict[str, frozenset[str]] = {"read_paper": frozenset({"cursor"})} +_CLI_DESTINATIONS = {mcp: cli for cli, mcp in PARAMETER_NAMES.items()} + + +@cache +def _cli_destinations(tool: str) -> frozenset[str]: + return frozenset(queries.query_options(TOOL_COMMANDS[tool])) + + +def cli_options(tool: str, **parameters: Any) -> dict[str, Any]: + """Translate MCP parameters into the CLI destinations of ``tool``. + + A parameter keeps its name when the command has that flag (``--limit``), + and otherwise follows ``PARAMETER_NAMES`` (``limit`` -> ``--page-size``). + """ + entity = ENTITY_PARAMETERS.get(tool) + destinations = _cli_destinations(tool) + options: dict[str, Any] = {} + for name, value in parameters.items(): + if name in MCP_ONLY_PARAMETERS.get(tool, frozenset()): + continue + if name == entity: + destination = "name" + elif name in destinations: + destination = name + else: + destination = _CLI_DESTINATIONS.get(name, name) + options[destination] = value + return options + + +Page = Annotated[int, Field(ge=1, le=100, description="Result page, starting at 1.")] +Limit = Annotated[ + int, Field(ge=1, le=MAX_ROWS, description="Rows per response, at most 25.") +] +Reference = Annotated[ + str, + Field( + min_length=1, + max_length=500, + description="arXiv ID, numeric PwC ID, arXiv/Hugging Face/PwC URL, or exact title.", + ), +] Query = Annotated[str, Field(min_length=1, max_length=500)] -AuthorList = Annotated[list[str], Field(max_length=10)] +Entity = Annotated[str, Field(min_length=1, max_length=500)] +IsoDate = Annotated[ + str, Field(max_length=10, description="Inclusive publication date, YYYY-MM-DD.") +] +AuthorList = Annotated[ + list[str], + Field( + max_length=10, + description="Exact author names, numeric IDs, or @HF_USERNAME; every author must match.", + ), +] +Area = Annotated[str, Field(description="Case-insensitive exact area name or area ID.")] +Direction = Literal["asc", "desc"] +ParameterSize = Annotated[ + str, + Field( + max_length=32, + description="Inclusive model-size limit such as 500M, 1.5B, 3B, or a raw integer.", + ), +] +MetricNames = Annotated[ + list[str], + Field( + min_length=1, description="Metric names that every returned row must report." + ), +] +MetricBounds = Annotated[ + dict[str, float], Field(description="Metric name to numeric threshold.") +] +SortMetric = Annotated[ + str, + Field( + max_length=200, + description="METRIC or METRIC:asc|desc; default direction is desc.", + ), +] +ParetoObjectives = Annotated[ + list[str], + Field( + min_length=2, + description="Two or more METRIC:higher or METRIC:lower objectives; keeps the Pareto frontier.", + ), +] def _validate_date_range(start: str | None, end: str | None) -> None: @@ -63,31 +233,40 @@ def _validate_date_range(start: str | None, end: str | None) -> None: raise ToolError("published_after must be on or before published_before") -def _catalog_call(function: Any, *args: Any, **kwargs: Any) -> Any: - """Turn upstream failures into deliberately generic, non-content-bearing errors.""" - try: - return function(*args, **kwargs) - except (ResponseError, TransportError) as error: - raise ToolError("the Papers With Code catalog request failed") from error +def _dicts(values: Any) -> list[dict[str, Any]]: + return [value for value in values or [] if isinstance(value, dict)] -class Catalog(Protocol): - def search_papers( - self, - *, - query: str, - mode: str = "keyword", - page: int = 1, - limit: int = 10, - published_after: str | None = None, - published_before: str | None = None, - has_official_implementation: bool = False, - ) -> dict[str, Any]: ... +def _next_page(data: Any) -> int | None: + value = data.get("next_page") if isinstance(data, dict) else None + return int(value) if value is not None else None - def get_paper_info( - self, paper: str, *, include_resources: bool - ) -> dict[str, Any]: ... +def _paper_page(data: Any) -> PaperPage: + """Project a paper listing; recent/trending endpoints return a bare list.""" + if isinstance(data, list): + rows: Any = data + elif isinstance(data, dict): + rows = data.get("results") + else: + raise TypeError("paper listing did not contain a result document") + return PaperPage( + items=[paper_summary(item) for item in _dicts(rows)], + next_page=_next_page(data), + data=data, + ) + + +def _grouped_benchmarks(areas: list[dict[str, Any]]) -> list[dict[str, Any]]: + return [ + benchmark + for area in areas + for task in _dicts(area.get("tasks")) + for benchmark in _dicts(task.get("benchmarks")) + ] + + +class Catalog(Protocol): def resolve_paper(self, paper: str) -> str: ... def read_paper_chunk( @@ -100,47 +279,7 @@ def read_paper_chunk( resolved: bool = False, ) -> PaperMarkdownChunk: ... - def list_papers( - self, - *, - search: str | None = None, - task: str | None = None, - method: str | None = None, - conference: str | None = None, - framework: str | None = None, - organization: str | None = None, - authors: list[str] | None = None, - published_after: str | None = None, - published_before: str | None = None, - order_by: str = "date_published", - order_direction: str = "desc", - page: int = 1, - limit: int = 10, - ) -> dict[str, Any]: ... - - def get_related_papers(self, paper: str, *, limit: int) -> dict[str, Any]: ... - - def get_paper_lineage(self, paper: str) -> dict[str, Any]: ... - - def get_task(self, task: str) -> dict[str, Any]: ... - - def get_method(self, method: str) -> dict[str, Any]: ... - - def list_benchmarks( - self, - *, - search: str | None = None, - task: str | None = None, - include_descendants: bool = False, - minimum_evaluations: int | None = None, - is_open: bool | None = None, - page: int = 1, - limit: int = 10, - ) -> dict[str, Any]: ... - - def get_benchmark( - self, benchmark: str, *, limit: int, is_open: bool | None - ) -> dict[str, Any]: ... + def query(self, command: tuple[str, ...], options: Mapping[str, Any]) -> Any: ... def build_server( @@ -156,7 +295,10 @@ def build_server( server = MCPServer( "pwc", title="Papers With Code", - description="Read-only access to papers, tasks, methods, and benchmarks.", + description=( + "Read-only access to papers, tasks, methods, conferences, organizations, " + "frameworks, and benchmarks; every pwc CLI research command and flag." + ), version=__version__, website_url="https://paperswithcode.co", cache_hints={ @@ -167,57 +309,84 @@ def build_server( }, ) + def run(tool: str, **parameters: Any) -> Any: + """Run the tool's CLI command; surface usage errors, hide upstream detail.""" + try: + return catalog.query(TOOL_COMMANDS[tool], cli_options(tool, **parameters)) + except UsageError as error: + raise ToolError(str(error)) from error + except (ResponseError, TransportError) as error: + raise ToolError(catalog_error_message(error)) from error + @server.tool(annotations=READ_ONLY, structured_output=True) def search_papers( query: Query, - mode: Literal["keyword", "semantic"] = "keyword", + mode: Literal["hybrid", "keyword", "semantic"] = "keyword", page: Page = 1, limit: Limit = 10, - published_after: str | None = None, - published_before: str | None = None, + published_after: IsoDate | None = None, + published_before: IsoDate | None = None, has_official_implementation: bool = False, ) -> PaperPage: - """Search papers by title, topic, author, or arXiv ID.""" + """Search papers by title, topic, author, or arXiv ID (`pwc search`). Broad discovery only: for best, top, or state-of-the-art model questions start with get_task, list_benchmarks, and get_benchmark, which return leaderboard evidence that search cannot.""" _validate_date_range(published_after, published_before) - payload = _catalog_call( - catalog.search_papers, - query=query, - mode=mode, - page=page, - limit=limit, - published_after=published_after, - published_before=published_before, - has_official_implementation=has_official_implementation, + return _paper_page( + run( + "search_papers", + query=query, + mode=mode, + page=page, + limit=limit, + published_after=published_after, + published_before=published_before, + has_official_implementation=has_official_implementation, + ) ) - return PaperPage( - items=[ - paper_summary(item) - for item in payload.get("results") or [] - if isinstance(item, dict) - ], - next_page=( - int(payload["next_page"]) - if payload.get("next_page") is not None + + @server.tool(annotations=READ_ONLY, structured_output=True) + def get_paper_info( + paper: Reference, + include_resources: bool = True, + include_evaluations: bool = False, + ) -> PaperInfoResult: + """Get paper metadata, abstract, tasks, methods, lineage, repositories, Hugging Face artifacts, and optionally every evaluation (`pwc paper info`).""" + data = run( + "get_paper_info", + paper=paper, + include_resources=include_resources, + include_evaluations=include_evaluations, + ) + if not isinstance(data, dict): + raise TypeError("paper response did not contain a paper") + evaluations = data.get("evaluations") + return PaperInfoResult( + paper=paper_detail(data), + evaluation_count=( + int(evaluations.get("count") or 0) + if isinstance(evaluations, dict) else None ), + evaluations=( + [paper_evaluation(item) for item in _dicts(evaluations.get("results"))] + if isinstance(evaluations, dict) + else None + ), + data=data, ) - @server.tool(annotations=READ_ONLY, structured_output=True) - def get_paper_info(paper: Reference) -> PaperInfoResult: - """Get metadata for an arXiv ID, PwC ID, URL, or exact paper title.""" - payload = _catalog_call(catalog.get_paper_info, paper, include_resources=True) - return PaperInfoResult(paper=paper_detail(payload)) - @server.tool(annotations=READ_ONLY, structured_output=True) def read_paper(paper: Reference, cursor: str | None = None) -> PaperReadResult: - """Read stored paper Markdown, continuing oversized documents with a cursor.""" + """Read stored paper Markdown, continuing oversized documents with a cursor (`pwc paper read`).""" reference = paper.strip() try: state = codec.decode(cursor, reference=reference) if cursor else None except ValueError as error: raise ToolError(str(error)) from error if state is None: - canonical = _catalog_call(catalog.resolve_paper, reference) + try: + canonical = catalog.resolve_paper(reference) + except (ResponseError, TransportError) as error: + raise ToolError(catalog_error_message(error)) from error offset = 0 content_version = None limit = read_chunk_bytes @@ -239,9 +408,11 @@ def read_paper(paper: Reference, cursor: str | None = None) -> PaperReadResult: resolved=True, ) except PaperVersionMismatchError as error: - raise ToolError("paper changed; restart reading from the beginning") from error + raise ToolError( + "paper changed; restart reading from the beginning" + ) from error except (ResponseError, TransportError) as error: - raise ToolError("the Papers With Code catalog request failed") from error + raise ToolError(catalog_error_message(error)) from error if chunk.paper != canonical or chunk.source != source: raise ToolError("the Papers With Code catalog request failed") next_cursor = None @@ -273,96 +444,94 @@ def list_papers( framework: str | None = None, organization: str | None = None, authors: AuthorList | None = None, - published_after: str | None = None, - published_before: str | None = None, - order_by: Literal[ - "date_published", "citation_count", "title" - ] = "date_published", - order_direction: Literal["asc", "desc"] = "desc", + published_after: IsoDate | None = None, + published_before: IsoDate | None = None, + all_versions: bool = False, + order_by: Literal["trending", "date_published", "citation_count"] = "trending", + order_direction: Direction = "desc", + include_resources: bool = False, + has_official_implementation: bool = False, page: Page = 1, - limit: Limit = 10, + limit: Limit = 20, ) -> PaperPage: - """List and filter papers in a deterministic catalog order.""" + """List and filter papers by exact catalog associations in a deterministic order (`pwc paper list`).""" _validate_date_range(published_after, published_before) - payload = _catalog_call( - catalog.list_papers, - search=search, - task=task, - method=method, - conference=conference, - framework=framework, - organization=organization, - authors=authors or [], - published_after=published_after, - published_before=published_before, - order_by=order_by, - order_direction=order_direction, - page=page, - limit=limit, - ) - return PaperPage( - items=[ - paper_summary(item) - for item in payload.get("results") or [] - if isinstance(item, dict) - ], - next_page=( - int(payload["next_page"]) - if payload.get("next_page") is not None - else None - ), + return _paper_page( + run( + "list_papers", + search=search, + task=task, + method=method, + conference=conference, + framework=framework, + organization=organization, + authors=authors or None, + published_after=published_after, + published_before=published_before, + all_versions=all_versions, + order_by=order_by, + order_direction=order_direction, + include_resources=include_resources, + has_official_implementation=has_official_implementation, + page=page, + limit=limit, + ) ) @server.tool(annotations=READ_ONLY, structured_output=True) - def get_related_papers(paper: Reference, limit: Limit = 10) -> PaperPage: - """Find catalog papers related to one paper.""" - payload = _catalog_call(catalog.get_related_papers, paper, limit=limit) - return PaperPage( - items=[ - paper_summary(item) - for item in payload.get("results") or [] - if isinstance(item, dict) - ], - next_page=None, + def list_recent_papers(limit: Limit = 10) -> PaperPage: + """List the most recently added papers (`pwc paper recent`).""" + return _paper_page(run("list_recent_papers", limit=limit)) + + @server.tool(annotations=READ_ONLY, structured_output=True) + def list_trending_papers( + limit: Limit = 20, + max_age_days: Annotated[int, Field(ge=1, le=365)] = 180, + min_velocity: float | None = None, + ) -> PaperPage: + """List trending papers by repository velocity (`pwc paper trending`).""" + return _paper_page( + run( + "list_trending_papers", + limit=limit, + max_age_days=max_age_days, + min_velocity=min_velocity, + ) ) + @server.tool(annotations=READ_ONLY, structured_output=True) + def get_related_papers( + paper: Reference, limit: Annotated[int, Field(ge=1, le=20)] = 4 + ) -> PaperPage: + """Find catalog papers related to one paper (`pwc paper related`).""" + return _paper_page(run("get_related_papers", paper=paper, limit=limit)) + @server.tool(annotations=READ_ONLY, structured_output=True) def get_paper_lineage(paper: Reference) -> PaperLineageResult: - """Get explicit predecessor and successor relationships for a paper.""" - payload = _catalog_call(catalog.get_paper_lineage, paper) - current = payload.get("paper") + """Get explicit predecessor and successor papers (`pwc paper lineage list`).""" + data = run("get_paper_lineage", paper=paper) + current = data.get("paper") if isinstance(data, dict) else None if not isinstance(current, dict): raise TypeError("lineage response did not contain a paper") return PaperLineageResult( paper=paper_reference(current), predecessors=[ - paper_reference(item) - for item in payload.get("predecessors") or [] - if isinstance(item, dict) + paper_reference(item) for item in _dicts(data.get("predecessors")) ], successors=[ - paper_reference(item) - for item in payload.get("successors") or [] - if isinstance(item, dict) + paper_reference(item) for item in _dicts(data.get("successors")) ], + data=data, ) @server.tool(annotations=READ_ONLY, structured_output=True) - def get_task(task: Reference) -> TaskResult: - """Get an exact task by ID, slug, or name, including its benchmarks.""" - payload = _catalog_call(catalog.get_task, task) - item = payload.get("task") + def get_task(task: Entity) -> TaskResult: + """Get one exact task by name, slug, or ID with its hierarchy, ranked benchmarks, common methods, and trending papers (`pwc task --name`). Start here for any question about the best or state-of-the-art models for a task, then inspect a leaderboard with get_benchmark.""" + data = run("get_task", task=task) + item = data.get("task") if isinstance(data, dict) else None if not isinstance(item, dict): raise TypeError("task response did not contain a task") - area_item = payload.get("area") - area = ( - AreaReference( - id=str(area_item.get("id") or ""), - name=str(area_item.get("name") or "Unknown area"), - ) - if isinstance(area_item, dict) - else None - ) + area_item = data.get("area") return TaskResult( task=TaskDetail( id=str(item.get("id") or ""), @@ -371,30 +540,59 @@ def get_task(task: Reference) -> TaskResult: description=( str(item["description"]) if item.get("description") else None ), - paper_count=int(item.get("paper_count") or 0), - area=area, - parents=[ - catalog_reference(value) - for value in payload.get("parents") or [] - if isinstance(value, dict) - ], - children=[ - catalog_reference(value) - for value in payload.get("children") or [] - if isinstance(value, dict) - ], + paper_count=int( + item.get("paper_count") or data.get("paper_count") or 0 + ), + area=( + AreaReference( + id=str(area_item.get("id") or ""), + name=str(area_item.get("name") or "Unknown area"), + ) + if isinstance(area_item, dict) + else None + ), + parents=[catalog_reference(v) for v in _dicts(data.get("parents"))], + children=[catalog_reference(v) for v in _dicts(data.get("children"))], benchmarks=[ - benchmark_summary(value) - for value in payload.get("benchmarks") or [] - if isinstance(value, dict) + benchmark_summary(v) for v in _dicts(data.get("benchmarks")) ], + ), + data=data, + ) + + @server.tool(annotations=READ_ONLY, structured_output=True) + def list_tasks( + area: Area | None = None, + level: int | None = None, + visible_only: bool = False, + group_by_area: bool = False, + order_by: Literal["name", "created_at", "level", "paper_count"] = "name", + order_direction: Direction = "asc", + page: Page = 1, + limit: Limit | None = None, + ) -> QueryResult: + """List and filter research tasks, or group the visible top-level taxonomy by area (`pwc task list`).""" + return QueryResult( + data=run( + "list_tasks", + area=area, + level=level, + visible_only=visible_only, + group_by_area=group_by_area, + order_by=order_by, + order_direction=order_direction, + page=page, + limit=limit if limit is not None or group_by_area else MAX_ROWS, ) ) @server.tool(annotations=READ_ONLY, structured_output=True) - def get_method(method: Reference) -> MethodResult: - """Get an exact method by ID, slug, full name, or name.""" - item = _catalog_call(catalog.get_method, method) + def get_method(method: Entity) -> MethodResult: + """Get one exact method by name, full name, slug, or ID (`pwc method --name`).""" + data = run("get_method", method=method) + item = data.get("method") if isinstance(data, dict) else None + if not isinstance(item, dict): + raise TypeError("method response did not contain a method") return MethodResult( method=MethodDetail( id=str(item.get("id") or ""), @@ -419,9 +617,109 @@ def get_method(method: Reference) -> MethodResult: str(item["source_title"]) if item.get("source_title") else None ), paper_count=int(item.get("paper_count") or 0), + ), + data=data, + ) + + @server.tool(annotations=READ_ONLY, structured_output=True) + def list_methods( + area: Area | None = None, + introduced_year: int | None = None, + order_by: Literal[ + "name", "full_name", "introduced_year", "created_at", "paper_count" + ] = "name", + order_direction: Direction = "asc", + page: Page = 1, + limit: Limit = MAX_ROWS, + ) -> QueryResult: + """List and filter research methods (`pwc method list`).""" + return QueryResult( + data=run( + "list_methods", + area=area, + introduced_year=introduced_year, + order_by=order_by, + order_direction=order_direction, + page=page, + limit=limit, + ) + ) + + @server.tool(annotations=READ_ONLY, structured_output=True) + def get_conference(conference: Entity) -> QueryResult: + """Get one exact conference by name, slug, or ID (`pwc conference --name`).""" + return QueryResult(data=run("get_conference", conference=conference)) + + @server.tool(annotations=READ_ONLY, structured_output=True) + def list_conferences(year: int | None = None) -> QueryResult: + """List conferences with imported papers (`pwc conference list`).""" + return QueryResult(data=run("list_conferences", year=year)) + + @server.tool(annotations=READ_ONLY, structured_output=True) + def get_organization(organization: Entity) -> QueryResult: + """Get one exact research organization by name, slug, or ID (`pwc organization --name`).""" + return QueryResult(data=run("get_organization", organization=organization)) + + @server.tool(annotations=READ_ONLY, structured_output=True) + def list_organizations(featured_only: bool = False) -> QueryResult: + """List research organizations (`pwc organization list`).""" + return QueryResult(data=run("list_organizations", featured_only=featured_only)) + + @server.tool(annotations=READ_ONLY, structured_output=True) + def get_framework(framework: Entity) -> QueryResult: + """Get one exact research framework by name, slug, or ID (`pwc framework --name`).""" + return QueryResult(data=run("get_framework", framework=framework)) + + @server.tool(annotations=READ_ONLY, structured_output=True) + def list_frameworks( + domain: str | None = None, + category: str | None = None, + platform: str | None = None, + ) -> QueryResult: + """List research frameworks by domain, category, or platform (`pwc framework list`).""" + return QueryResult( + data=run( + "list_frameworks", domain=domain, category=category, platform=platform ) ) + @server.tool(annotations=READ_ONLY, structured_output=True) + def get_benchmark( + benchmark: Entity, + limit: Limit = 20, + is_open: bool | None = None, + max_parameters: ParameterSize | None = None, + require_metrics: MetricNames | None = None, + minimum_metrics: MetricBounds | None = None, + maximum_metrics: MetricBounds | None = None, + sort_metric: SortMetric | None = None, + pareto: ParetoObjectives | None = None, + ) -> BenchmarkResult: + """Get one exact benchmark and its leaderboard (`pwc benchmark --name`). Use max_parameters (for example "4B") to keep models at or below a size, sort_metric to rank by a metric, and minimum_metrics, maximum_metrics, require_metrics, or pareto to select rows; matched_count reports how many rows passed before limit.""" + data = run( + "get_benchmark", + benchmark=benchmark, + limit=limit, + is_open=is_open, + max_parameters=max_parameters, + require_metrics=require_metrics, + minimum_metrics=minimum_metrics, + maximum_metrics=maximum_metrics, + sort_metric=sort_metric, + pareto=pareto, + ) + item = data.get("benchmark") if isinstance(data, dict) else None + if not isinstance(item, dict): + raise TypeError("benchmark response did not contain a benchmark") + matched = data.get("matched_count") + return BenchmarkResult( + benchmark=benchmark_summary(item), + evaluation_count=int(data.get("count") or 0), + matched_count=int(matched) if matched is not None else None, + evaluations=[evaluation(value) for value in _dicts(data.get("results"))], + data=data, + ) + @server.tool(annotations=READ_ONLY, structured_output=True) def list_benchmarks( search: str | None = None, @@ -429,54 +727,41 @@ def list_benchmarks( include_descendants: bool = False, minimum_evaluations: int | None = None, is_open: bool | None = None, + group_by_area: bool = False, + area: Area | None = None, + benchmarks_per_task: Annotated[int, Field(ge=1, le=10)] = 3, + order_by: Literal["trending", "name", "full_name", "created_at", "paper_count"] + | None = None, + order_direction: Direction = "asc", page: Page = 1, - limit: Limit = 10, + limit: Limit | None = None, ) -> BenchmarkPage: - """List benchmark datasets with optional task and availability filters.""" - payload = _catalog_call( - catalog.list_benchmarks, + """List benchmarks for a task ranked by trend, filter them, or group them by area and task (`pwc benchmark list`). Follow with get_benchmark on the most relevant leaderboard.""" + grouped = group_by_area or area is not None + data = run( + "list_benchmarks", search=search, task=task, include_descendants=include_descendants, minimum_evaluations=minimum_evaluations, is_open=is_open, + group_by_area=group_by_area, + area=area, + benchmarks_per_task=benchmarks_per_task, + order_by=order_by, + order_direction=order_direction, page=page, - limit=limit, + limit=limit if limit is not None or grouped else MAX_ROWS, ) + if not isinstance(data, dict): + raise TypeError("benchmark listing did not contain a result document") + rows = _dicts(data.get("results")) + if grouped: + rows = _grouped_benchmarks(rows) return BenchmarkPage( - items=[ - benchmark_summary(item) - for item in payload.get("results") or [] - if isinstance(item, dict) - ], - next_page=( - int(payload["next_page"]) - if payload.get("next_page") is not None - else None - ), - ) - - @server.tool(annotations=READ_ONLY, structured_output=True) - def get_benchmark( - benchmark: Reference, - limit: Limit = 10, - is_open: bool | None = None, - ) -> BenchmarkResult: - """Get an exact benchmark and its top evaluation rows.""" - payload = _catalog_call( - catalog.get_benchmark, benchmark, limit=limit, is_open=is_open - ) - item = payload.get("benchmark") - if not isinstance(item, dict): - raise TypeError("benchmark response did not contain a benchmark") - return BenchmarkResult( - benchmark=benchmark_summary(item), - evaluation_count=int(payload.get("count") or 0), - evaluations=[ - evaluation(value) - for value in payload.get("results") or [] - if isinstance(value, dict) - ], + items=[benchmark_summary(item) for item in rows], + next_page=_next_page(data), + data=data, ) @server.resource( diff --git a/mcp_server/tests/test_app.py b/mcp_server/tests/test_app.py index b446b01..f79ad8b 100644 --- a/mcp_server/tests/test_app.py +++ b/mcp_server/tests/test_app.py @@ -43,7 +43,7 @@ def test_health_and_browser_origin_policy_are_explicit(): assert health.json() == { "status": "ok", "service": "pwc-mcp", - "version": "0.1.0", + "version": "0.2.0", "protocol": "2026-07-28", } assert rejected.status_code == 403 @@ -206,18 +206,38 @@ def test_one_http_endpoint_serves_modern_and_legacy_protocol_eras(): def test_proxy_identity_trusts_only_an_exact_loopback_peer(): headers = Headers({"x-forwarded-for": "203.0.113.9"}) - assert _client_address( - {"client": ("127.0.0.1", 1234)}, headers, True - ) == "203.0.113.9" - assert _client_address( - {"client": ("::1", 1234)}, headers, True - ) == "203.0.113.9" - assert _client_address( - {"client": ("10.0.0.2", 1234)}, headers, True - ) == "10.0.0.2" - assert _client_address( - {"client": ("192.168.1.2", 1234)}, headers, True - ) == "192.168.1.2" + assert ( + _client_address({"client": ("127.0.0.1", 1234)}, headers, True) == "203.0.113.9" + ) + assert _client_address({"client": ("::1", 1234)}, headers, True) == "203.0.113.9" + assert _client_address({"client": ("10.0.0.2", 1234)}, headers, True) == "10.0.0.2" + assert ( + _client_address({"client": ("192.168.1.2", 1234)}, headers, True) + == "192.168.1.2" + ) + + +def test_loopback_first_party_clients_may_name_their_rate_limit_identity(): + tagged = Headers({"x-pwc-mcp-client": "chat-0123abcd"}) + assert ( + _client_address({"client": ("127.0.0.1", 1)}, tagged, True) + == "client:chat-0123abcd" + ) + assert ( + _client_address({"client": ("::1", 1)}, tagged, True) == "client:chat-0123abcd" + ) + # Proxied traffic always carries X-Forwarded-For, which wins over the tag. + proxied = Headers( + {"x-pwc-mcp-client": "chat-0123abcd", "x-forwarded-for": "203.0.113.9"} + ) + assert _client_address({"client": ("127.0.0.1", 1)}, proxied, True) == "203.0.113.9" + # Remote peers, disabled trust, and malformed tags fall back to the address. + assert _client_address({"client": ("10.0.0.2", 1)}, tagged, True) == "10.0.0.2" + assert _client_address({"client": ("127.0.0.1", 1)}, tagged, False) == "127.0.0.1" + bad = Headers({"x-pwc-mcp-client": "spaces are/not ok"}) + assert _client_address({"client": ("127.0.0.1", 1)}, bad, True) == "127.0.0.1" + long = Headers({"x-pwc-mcp-client": "a" * 129}) + assert _client_address({"client": ("127.0.0.1", 1)}, long, True) == "127.0.0.1" def test_serialized_mcp_response_limit_fails_closed(): @@ -232,3 +252,27 @@ async def oversized(_request): assert response.status_code == 503 assert response.json() == {"error": "response_too_large"} + + +def test_hybrid_search_counts_toward_the_semantic_limit(): + app = create_app( + StubCatalog(), + allowed_hosts=["testserver"], + semantic_limit=1, + ) + body = { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "search_papers", + "arguments": {"query": "attention", "mode": "hybrid"}, + }, + } + + with TestClient(app) as client: + first = client.post("/mcp", json=body) + limited = client.post("/mcp", json=body) + + assert first.status_code != 429 + assert limited.status_code == 429 diff --git a/mcp_server/tests/test_catalog.py b/mcp_server/tests/test_catalog.py index 6c137ba..4d90414 100644 --- a/mcp_server/tests/test_catalog.py +++ b/mcp_server/tests/test_catalog.py @@ -3,8 +3,9 @@ import json import pytest +from pwc_cli.cli import UsageError from pwc_cli.transport import Response, ResponseError -from pwc_mcp.catalog import CatalogClient +from pwc_mcp.catalog import CatalogClient, cache_ttl class StubTransport: @@ -30,11 +31,16 @@ def test_catalog_normalizes_paper_urls_and_caches_identical_reads(): ) catalog = CatalogClient(transport=transport) - first = catalog.get_paper_info( - "https://arxiv.org/pdf/1706.03762v7.pdf", include_resources=True + first = catalog.query( + ("paper", "info"), + {"paper": "https://arxiv.org/pdf/1706.03762v7.pdf", "include_resources": True}, ) - second = catalog.get_paper_info( - "https://huggingface.co/papers/1706.03762", include_resources=True + second = catalog.query( + ("paper", "info"), + { + "paper": "https://huggingface.co/papers/1706.03762", + "include_resources": True, + }, ) assert first == second @@ -174,7 +180,7 @@ def test_catalog_resolves_exact_titles_and_rejects_ambiguous_titles(): catalog = CatalogClient(transport=transport) with pytest.raises(ResponseError, match="ambiguous"): - catalog.get_paper_lineage("Same Title") + catalog.query(("paper", "lineage", "list"), {"paper": "Same Title"}) def test_catalog_resolves_pwc_urls_and_dotted_legacy_arxiv_ids(): @@ -195,11 +201,11 @@ def test_catalog_resolves_pwc_urls_and_dotted_legacy_arxiv_ids(): ) catalog = CatalogClient(transport=transport) - assert catalog.get_paper_info( - "https://paperswithcode.co/paper/attention-is-all-you-need", - include_resources=False, + assert catalog.query( + ("paper", "info"), + {"paper": "https://paperswithcode.co/paper/attention-is-all-you-need"}, ) == {"id": "755"} - assert catalog.get_paper_info("math.GT/0309136", include_resources=False) == { + assert catalog.query(("paper", "info"), {"paper": "math.GT/0309136"}) == { "id": "900" } @@ -224,7 +230,7 @@ def get(self, path, params=None): catalog = CatalogClient(transport=PaginatedTransport()) with pytest.raises(ResponseError, match="ambiguous"): - catalog.get_paper_lineage("Same Title") + catalog.query(("paper", "lineage", "list"), {"paper": "Same Title"}) def test_catalog_fails_closed_when_paper_filters_are_not_confirmed(): @@ -232,4 +238,45 @@ def test_catalog_fails_closed_when_paper_filters_are_not_confirmed(): catalog = CatalogClient(transport=transport) with pytest.raises(ResponseError, match="did not confirm"): - catalog.list_papers(task="image-classification", page=1, limit=10) + catalog.query(("paper", "list"), {"task": "image-classification"}) + + +def test_catalog_query_runs_cli_commands_with_per_path_cache_lifetimes(): + transport = StubTransport( + { + "datasets/": {"results": [{"id": "72", "name": "ImageNet-1k"}]}, + "datasets/72/evaluations/": { + "count": 1, + "results": [{"id": "1", "model_name": "A", "metrics": {"top1": 80}}], + }, + } + ) + catalog = CatalogClient(transport=transport) + + first = catalog.query(("benchmark",), {"name": "ImageNet-1k", "limit": 5}) + second = catalog.query(("benchmark",), {"name": "imagenet-1k", "limit": 5}) + + assert first == second + assert first["results"][0]["model_name"] == "A" + assert first["count"] == 1 + # The second lookup differs only in the search text; the identical + # evaluation request is served from the shared cache. + assert [path for path, _params in transport.calls] == [ + "datasets/", + "datasets/72/evaluations/", + "datasets/", + ] + assert cache_ttl("papers/search") == 60 + assert cache_ttl("tasks/1/page") == 600 + assert cache_ttl("datasets/72/evaluations/") == 300 + + +def test_catalog_query_reports_cli_usage_errors_without_upstream_calls(): + transport = StubTransport({}) + catalog = CatalogClient(transport=transport) + + with pytest.raises(UsageError, match="--limit must be 1-100"): + catalog.query(("search",), {"query": "x", "limit": 0}) + with pytest.raises(UsageError, match="not a read-only"): + catalog.query(("skills", "add"), {}) + assert transport.calls == [] diff --git a/mcp_server/tests/test_parity.py b/mcp_server/tests/test_parity.py new file mode 100644 index 0000000..c25541b --- /dev/null +++ b/mcp_server/tests/test_parity.py @@ -0,0 +1,94 @@ +"""Every read-only CLI command is a tool, and every research flag a parameter.""" + +from __future__ import annotations + +import argparse +import asyncio + +from mcp.client import Client +from pwc_cli import queries +from pwc_mcp.server import ( + ENTITY_PARAMETERS, + MCP_ONLY_PARAMETERS, + PARAMETER_NAMES, + TOOL_COMMANDS, + build_server, + cli_options, +) +from test_server import StubCatalog + +# Deliberate default divergences from the CLI, documented in README.md. +DEFAULT_EXCEPTIONS = {("get_paper_info", "include_resources"): True} + + +def _tools(): + async def exercise(): + async with Client(build_server(StubCatalog())) as client: + return {tool.name: tool for tool in (await client.list_tools()).tools} + + return asyncio.run(exercise()) + + +def _enum_values(schema, root): + if "$ref" in schema: + schema = root["$defs"][schema["$ref"].rsplit("/", 1)[1]] + values = list(schema.get("enum") or []) + for branch in schema.get("anyOf") or []: + values.extend(_enum_values(branch, root)) + return values + + +def _mcp_name(tool: str, destination: str) -> str: + if destination == "name": + return ENTITY_PARAMETERS[tool] + return PARAMETER_NAMES.get(destination, destination) + + +def test_every_read_only_cli_command_has_exactly_one_tool(): + assert sorted(TOOL_COMMANDS.values()) == sorted(queries.READ_ONLY_COMMANDS) + assert len(set(TOOL_COMMANDS.values())) == len(TOOL_COMMANDS) + + +def test_every_cli_research_flag_is_a_tool_parameter_and_vice_versa(): + tools = _tools() + + assert set(tools) == set(TOOL_COMMANDS) + for tool, command in TOOL_COMMANDS.items(): + cli = queries.query_options(command) + schema = tools[tool].input_schema + parameters = set(schema["properties"]) - MCP_ONLY_PARAMETERS.get( + tool, frozenset() + ) + mapped = cli_options(tool, **dict.fromkeys(parameters)) + + assert set(mapped) == set(cli), tool + for destination, action in cli.items(): + parameter = _mcp_name(tool, destination) + property_schema = schema["properties"][parameter] + if isinstance(action, argparse._StoreTrueAction): + assert property_schema.get("type") == "boolean", (tool, parameter) + expected = DEFAULT_EXCEPTIONS.get((tool, parameter), False) + assert property_schema.get("default") is expected, (tool, parameter) + elif action.choices and set(action.choices) != {"true", "false"}: + assert set(action.choices) <= set( + _enum_values(property_schema, schema) + ), ( + tool, + parameter, + ) + if not action.option_strings or destination == "name": + assert parameter in schema["required"], (tool, parameter) + + +def test_paper_references_and_entity_names_are_required_everywhere(): + tools = _tools() + + for tool, entity in ENTITY_PARAMETERS.items(): + assert tools[tool].input_schema["required"] == [entity] + for tool in ( + "get_paper_info", + "read_paper", + "get_related_papers", + "get_paper_lineage", + ): + assert tools[tool].input_schema["required"] == ["paper"] diff --git a/mcp_server/tests/test_server.py b/mcp_server/tests/test_server.py index 07b0ee2..5f12a98 100644 --- a/mcp_server/tests/test_server.py +++ b/mcp_server/tests/test_server.py @@ -4,64 +4,169 @@ import logging from mcp.client import Client -from pwc_cli.transport import ResponseError +from pwc_cli.cli import UsageError +from pwc_cli.transport import HTTPStatusError, ResponseError from pwc_mcp.catalog import PaperMarkdownChunk -from pwc_mcp.server import build_server +from pwc_mcp.server import TOOL_COMMANDS, build_server + +PAPER_ROW = { + "id": "755", + "arxiv_id": "1706.03762", + "title": "Attention Is All You Need", + "authors": ["Ashish Vaswani"], + "published": "2017-06-12", + "citation_count": 190_373, + "url_abs": "https://arxiv.org/abs/1706.03762v7", + "has_official_implementation": True, + "code_repository_count": 595, +} +PAPER_PAGE = {"next_page": 2, "results": [PAPER_ROW]} +PAPER_INFO = { + **PAPER_ROW, + "abstract": "A transformer architecture.", + "url_pdf": "https://arxiv.org/pdf/1706.03762v7.pdf", + "tasks": [ + {"id": "6", "name": "Machine Translation", "slug": "machine-translation"} + ], + "methods": [{"id": "1", "name": "Transformer", "slug": "transformer"}], + "repositories": [ + {"url": "https://github.com/tensorflow/tensor2tensor", "is_official": True} + ], + "project_pages": ["https://example.test/transformer"], + "hf_models": ["https://huggingface.co/google-t5/t5-base"], + "hf_datasets": [], + "hf_spaces": [], +} +PAPER_EVALUATION = { + "id": "40", + "dataset_name": "WMT2014 English-German", + "task_name": "Machine Translation", + "model_name": "Transformer Big", + "metrics": {"BLEU score": 28.4}, + "best_metric": "BLEU score", + "best_rank": 3, + "is_open": True, + "num_parameters": 213_000_000, + "result_url": "https://example.test/result", +} +EVALUATION_ROW = { + "id": "10", + "model_name": "ExampleNet", + "harness": "timm", + "metrics": {"Accuracy": 90.1}, + "best_metric": "Accuracy", + "best_rank": 1, + "task_name": "Image Classification", + "paper_id": "755", + "paper_title": "Attention Is All You Need", + "paper_arxiv_id": "1706.03762", + "paper_published_date": "2017-06-12", + "is_open": True, + "num_parameters": 1000, +} +BENCHMARK = { + "id": "72", + "name": "ImageNet-1k", + "slug": "imagenet-1k", + "paper_count": 124, +} +TASK = { + "id": "1", + "name": "Image Classification", + "slug": "image-classification", + "description": "Assign a class to an image.", + "paper_count": 2343, +} +METHOD = { + "id": "2", + "name": "Transformer", + "slug": "transformer", + "full_name": "Transformer", + "description": "Attention-based architecture.", + "introduced_year": 2017, + "source_paper_id": "755", + "source_url": "/paper/1706.03762", + "source_title": "Attention Is All You Need", + "paper_count": 13505, +} +CONFERENCE = {"slug": "cvpr-2025", "name": "CVPR 2025", "year": 2025, "paper_count": 3} +ORGANIZATION = {"id": "4", "slug": "nvidia", "name": "NVIDIA", "paper_count": 900} +FRAMEWORK = {"id": "7", "slug": "vllm", "name": "vLLM", "platforms": ["gpu"]} +GROUPED_BENCHMARKS = { + "results": [ + { + "id": "1", + "name": "Vision", + "tasks": [ + { + "slug": "image-classification", + "benchmarks": [ + { + "id": "72", + "name": "ImageNet-1k", + "slug": "imagenet-1k", + "evaluation_count": 124, + } + ], + } + ], + } + ] +} + +PAYLOADS = { + ("search",): PAPER_PAGE, + ("paper", "list"): PAPER_PAGE, + # The recent and trending endpoints return a bare list of papers. + ("paper", "recent"): [PAPER_ROW], + ("paper", "trending"): [PAPER_ROW], + ("paper", "related"): {"results": [PAPER_ROW]}, + ("paper", "lineage", "list"): { + "paper": { + "id": 755, + "reference": "1706.03762", + "title": "Attention Is All You Need", + }, + "predecessors": [], + "successors": [{"id": 900, "reference": "2001.00001", "title": "A Follow-up"}], + }, + ("task",): { + "task": TASK, + "area": {"id": "1", "name": "Vision"}, + "parents": [], + "children": [], + "benchmarks": [BENCHMARK], + "common_methods": [METHOD], + "papers": [PAPER_ROW], + "paper_count": 2343, + }, + ("task", "list"): {"count": 1, "results": [TASK]}, + ("method",): {"method": METHOD, "area": {"id": "1", "name": "Vision"}}, + ("method", "list"): {"count": 1, "results": [METHOD]}, + ("conference",): CONFERENCE, + ("conference", "list"): {"count": 1, "results": [CONFERENCE]}, + ("organization",): ORGANIZATION, + ("organization", "list"): {"count": 1, "results": [ORGANIZATION]}, + ("framework",): FRAMEWORK, + ("framework", "list"): {"count": 1, "results": [FRAMEWORK]}, + ("benchmark",): { + "benchmark": {"id": "72", "name": "ImageNet-1k", "slug": "imagenet-1k"}, + "count": 1, + "matched_count": 1, + "results": [EVALUATION_ROW], + }, + ("benchmark", "list"): {"next_page": None, "results": [BENCHMARK]}, +} class StubCatalog: + """Records the CLI command and options each tool requests.""" + def __init__(self): + self.queries: list[tuple[tuple[str, ...], dict]] = [] self.resolve_calls = 0 self.read_calls = [] - def search_papers(self, **_kwargs): - return { - "next_page": 2, - "results": [ - { - "id": "755", - "arxiv_id": "1706.03762", - "title": "Attention Is All You Need", - "authors": ["Ashish Vaswani"], - "published": "2017-06-12", - "citation_count": 190_373, - "url_abs": "https://arxiv.org/abs/1706.03762v7", - "has_official_implementation": True, - "code_repository_count": 595, - } - ], - } - - def get_paper_info(self, paper: str, *, include_resources: bool): - assert paper == "1706.03762" - assert include_resources is True - return { - "id": "755", - "arxiv_id": "1706.03762", - "title": "Attention Is All You Need", - "abstract": "A transformer architecture.", - "authors": ["Ashish Vaswani"], - "published": "2017-06-12", - "citation_count": 190_373, - "url_abs": "https://arxiv.org/abs/1706.03762v7", - "url_pdf": "https://arxiv.org/pdf/1706.03762v7.pdf", - "tasks": [ - { - "id": "6", - "name": "Machine Translation", - "slug": "machine-translation", - } - ], - "methods": [{"id": "1", "name": "Transformer", "slug": "transformer"}], - "repositories": [ - { - "url": "https://github.com/tensorflow/tensor2tensor", - "is_official": True, - } - ], - "project_pages": ["https://example.test/transformer"], - } - def resolve_paper(self, paper: str): self.resolve_calls += 1 assert paper == "1706.03762" @@ -83,7 +188,9 @@ def read_paper_chunk( self.read_calls.append((offset, content_version, limit)) raw = b"abcdefgh" markdown = raw[offset : offset + limit].decode() - next_offset = offset + len(markdown) if offset + len(markdown) < len(raw) else None + next_offset = ( + offset + len(markdown) if offset + len(markdown) < len(raw) else None + ) return PaperMarkdownChunk( paper=paper, source="arxiv", @@ -92,105 +199,39 @@ def read_paper_chunk( next_offset=next_offset, ) - def list_papers(self, **_kwargs): - return self.search_papers() - - def get_related_papers(self, paper: str, *, limit: int): - assert paper == "1706.03762" - assert limit == 2 - return self.search_papers() - - def get_paper_lineage(self, paper: str): - assert paper == "1706.03762" - return { - "paper": { - "id": 755, - "reference": "1706.03762", - "title": "Attention Is All You Need", - }, - "predecessors": [], - "successors": [ - {"id": 900, "reference": "2001.00001", "title": "A Follow-up"} - ], - } + def query(self, command, options): + self.queries.append((tuple(command), dict(options))) + if command == ("paper", "info"): + payload = dict(PAPER_INFO) + if options.get("include_evals"): + payload["evaluations"] = {"count": 1, "results": [PAPER_EVALUATION]} + return payload + if command == ("benchmark", "list") and ( + options.get("group_by_area") or options.get("area") + ): + return GROUPED_BENCHMARKS + return PAYLOADS[tuple(command)] - def get_task(self, task: str): - assert task == "image-classification" - return { - "task": { - "id": "1", - "name": "Image Classification", - "slug": "image-classification", - "description": "Assign a class to an image.", - "paper_count": 2343, - }, - "area": {"id": "1", "name": "Vision"}, - "parents": [], - "children": [], - "benchmarks": [ - { - "id": "72", - "name": "ImageNet-1k", - "slug": "imagenet-1k", - "paper_count": 124, - } - ], - } + def options(self, command): + return next(options for called, options in self.queries if called == command) - def get_method(self, method: str): - assert method == "transformer" - return { - "id": "2", - "name": "Transformer", - "slug": "transformer", - "full_name": "Transformer", - "description": "Attention-based architecture.", - "introduced_year": 2017, - "source_paper_id": "755", - "source_url": "/paper/1706.03762", - "source_title": "Attention Is All You Need", - "paper_count": 13505, - } - def list_benchmarks(self, **_kwargs): - return { - "next_page": None, - "results": [ - { - "id": "72", - "name": "ImageNet-1k", - "slug": "imagenet-1k", - "paper_count": 124, - } - ], - } +def _call(catalog, requests, **server_options): + async def exercise(): + async with Client(build_server(catalog, **server_options)) as client: + results = [] + for name, arguments in requests: + results.append(await client.call_tool(name, arguments)) + return results - def get_benchmark(self, benchmark: str, *, limit: int, is_open: bool | None): - assert benchmark == "imagenet-1k" - assert limit in {5, 10} - assert is_open in {True, None} - return { - "benchmark": {"id": "72", "name": "ImageNet-1k", "slug": "imagenet-1k"}, - "count": 1, - "results": [ - { - "id": "10", - "model_name": "ExampleNet", - "metrics": {"Accuracy": 90.1}, - "best_rank": 1, - "paper_id": "755", - "paper_title": "Attention Is All You Need", - "paper_arxiv_id": "1706.03762", - "is_open": True, - "num_parameters": 1000, - } - ], - } + return asyncio.run(exercise()) def test_search_papers_is_a_read_only_structured_tool(): + catalog = StubCatalog() + async def exercise(): - async with Client(build_server(StubCatalog())) as client: + async with Client(build_server(catalog)) as client: tools = {tool.name: tool for tool in (await client.list_tools()).tools} result = await client.call_tool( "search_papers", @@ -202,6 +243,12 @@ async def exercise(): assert tools["search_papers"].annotations.read_only_hint is True assert tools["search_papers"].input_schema["properties"]["limit"]["maximum"] == 25 + assert tools["search_papers"].input_schema["properties"]["mode"]["enum"] == [ + "hybrid", + "keyword", + "semantic", + ] + assert result.is_error is False assert result.structured_content == { "schema_version": "v1", "items": [ @@ -218,14 +265,30 @@ async def exercise(): } ], "next_page": 2, + "data": PAPER_PAGE, } - assert result.is_error is False + assert catalog.queries == [ + ( + ("search",), + { + "query": "transformer", + "mode": "keyword", + "page": 1, + "limit": 1, + "start_date": None, + "end_date": None, + "has_official_implementation": False, + }, + ) + ] def test_search_rejects_invalid_ranges_before_calling_the_catalog(): - async def exercise(): - async with Client(build_server(StubCatalog())) as client: - return await client.call_tool( + catalog = StubCatalog() + (result,) = _call( + catalog, + [ + ( "search_papers", { "query": "transformer", @@ -233,29 +296,28 @@ async def exercise(): "published_before": "2026-08-01", }, ) - - result = asyncio.run(exercise()) + ], + ) assert result.is_error is True assert result.content[0].text == ( "Error executing tool search_papers: " "published_after must be on or before published_before" ) + assert catalog.queries == [] def test_catalog_failures_do_not_expose_or_log_user_queries(caplog): secret_query = "private unreleased project heliotrope" class FailingCatalog(StubCatalog): - def search_papers(self, **_kwargs): - raise ResponseError(f"Paper title not found: {secret_query}") - - async def exercise(): - async with Client(build_server(FailingCatalog())) as client: - return await client.call_tool("search_papers", {"query": secret_query}) + def query(self, command, options): + raise ResponseError(f"API returned invalid JSON for {secret_query}") with caplog.at_level(logging.INFO): - result = asyncio.run(exercise()) + (result,) = _call( + FailingCatalog(), [("search_papers", {"query": secret_query})] + ) assert result.is_error is True assert result.content[0].text == ( @@ -265,12 +327,62 @@ async def exercise(): assert secret_query not in caplog.text +def test_lookup_failures_return_the_cli_hint_without_logging_it(caplog): + class MissingCatalog(StubCatalog): + def query(self, command, options): + raise ResponseError( + "Task not found: language-modelling; closest results: Language Modeling" + ) + + class TransportCatalog(StubCatalog): + def query(self, command, options): + raise HTTPStatusError(404, "not found: language-modelling") + + with caplog.at_level(logging.INFO): + (missing,) = _call( + MissingCatalog(), [("get_task", {"task": "language-modelling"})] + ) + (transport,) = _call( + TransportCatalog(), [("get_task", {"task": "language-modelling"})] + ) + + assert missing.content[0].text == ( + "Error executing tool get_task: " + "Task not found: language-modelling; closest results: Language Modeling" + ) + assert transport.content[0].text == ( + "Error executing tool get_task: the Papers With Code catalog request failed" + ) + assert "language-modelling" not in caplog.text + + +def test_cli_usage_errors_are_returned_verbatim(): + class StrictCatalog(StubCatalog): + def query(self, command, options): + raise UsageError("unknown metric(s): latency; available metrics: Accuracy") + + (result,) = _call( + StrictCatalog(), + [("get_benchmark", {"benchmark": "imagenet-1k", "sort_metric": "latency"})], + ) + + assert result.is_error is True + assert result.content[0].text == ( + "Error executing tool get_benchmark: " + "unknown metric(s): latency; available metrics: Accuracy" + ) + + def test_paper_info_and_reading_use_stable_schemas_and_opaque_continuation(): catalog = StubCatalog() async def exercise(): async with Client(build_server(catalog, read_chunk_bytes=5)) as client: info = await client.call_tool("get_paper_info", {"paper": "1706.03762"}) + evaluated = await client.call_tool( + "get_paper_info", + {"paper": "1706.03762", "include_evaluations": True}, + ) first = await client.call_tool("read_paper", {"paper": "1706.03762"}) second = await client.call_tool( "read_paper", @@ -279,15 +391,42 @@ async def exercise(): "cursor": first.structured_content["next_cursor"], }, ) - return info, first, second + return info, evaluated, first, second - info, first, second = asyncio.run(exercise()) + info, evaluated, first, second = asyncio.run(exercise()) - assert info.structured_content["paper"]["title"] == "Attention Is All You Need" - assert info.structured_content["paper"]["tasks"] == [ + assert catalog.queries[0] == ( + ("paper", "info"), + {"paper": "1706.03762", "include_resources": True, "include_evals": False}, + ) + paper = info.structured_content["paper"] + assert paper["title"] == "Attention Is All You Need" + assert paper["tasks"] == [ {"id": "6", "name": "Machine Translation", "slug": "machine-translation"} ] - assert info.structured_content["paper"]["repositories"][0]["is_official"] is True + assert paper["repositories"][0]["is_official"] is True + assert paper["hf_models"] == ["https://huggingface.co/google-t5/t5-base"] + assert info.structured_content["evaluations"] is None + assert info.structured_content["data"]["abstract"] == "A transformer architecture." + + assert catalog.queries[1][1]["include_evals"] is True + assert evaluated.structured_content["evaluation_count"] == 1 + assert evaluated.structured_content["evaluations"] == [ + { + "id": "40", + "benchmark": "WMT2014 English-German", + "task": "Machine Translation", + "model_name": "Transformer Big", + "harness": None, + "metrics": {"BLEU score": 28.4}, + "best_metric": "BLEU score", + "best_rank": 3, + "is_open": True, + "num_parameters": 213_000_000, + "source_url": "https://example.test/result", + } + ] + assert first.structured_content["markdown"] == "abcde" assert first.structured_content["truncated"] is True assert first.structured_content["next_cursor"] @@ -303,13 +442,10 @@ async def exercise(): def test_read_paper_rejects_invalid_continuation_as_an_expected_error(): - async def exercise(): - async with Client(build_server(StubCatalog())) as client: - return await client.call_tool( - "read_paper", {"paper": "1706.03762", "cursor": "%%%private%%%"} - ) - - result = asyncio.run(exercise()) + (result,) = _call( + StubCatalog(), + [("read_paper", {"paper": "1706.03762", "cursor": "%%%private%%%"})], + ) assert result.is_error is True assert result.content[0].text == ( @@ -318,22 +454,56 @@ async def exercise(): def test_paper_listing_related_work_and_lineage_are_composable(): - async def exercise(): - async with Client(build_server(StubCatalog())) as client: - listed = await client.call_tool( - "list_papers", {"task": "machine-translation", "limit": 1} - ) - related = await client.call_tool( - "get_related_papers", {"paper": "1706.03762", "limit": 2} - ) - lineage = await client.call_tool( - "get_paper_lineage", {"paper": "1706.03762"} - ) - return listed, related, lineage - - listed, related, lineage = asyncio.run(exercise()) + catalog = StubCatalog() + listed, recent, trending, related, lineage = _call( + catalog, + [ + ( + "list_papers", + { + "task": "machine-translation", + "authors": ["Ashish Vaswani"], + "limit": 1, + }, + ), + ("list_recent_papers", {"limit": 3}), + ("list_trending_papers", {"max_age_days": 30, "min_velocity": 1.5}), + ("get_related_papers", {"paper": "1706.03762", "limit": 2}), + ("get_paper_lineage", {"paper": "1706.03762"}), + ], + ) + assert catalog.options(("paper", "list")) == { + "search": None, + "task": "machine-translation", + "method": None, + "conference": None, + "framework": None, + "organization": None, + "author": ["Ashish Vaswani"], + "start_date": None, + "end_date": None, + "all_versions": False, + "order_by": "trending", + "order_dir": "desc", + "include_resources": False, + "has_official_implementation": False, + "page": 1, + "page_size": 1, + } + assert catalog.options(("paper", "recent")) == {"limit": 3} + assert catalog.options(("paper", "trending")) == { + "limit": 20, + "max_age_days": 30, + "min_velocity": 1.5, + } + assert catalog.options(("paper", "related")) == {"paper": "1706.03762", "limit": 2} + assert catalog.options(("paper", "lineage", "list")) == {"paper": "1706.03762"} assert listed.structured_content["items"][0]["arxiv_id"] == "1706.03762" + assert recent.structured_content["next_page"] is None + assert recent.structured_content["items"][0]["id"] == "755" + assert recent.structured_content["data"] == [PAPER_ROW] + assert trending.structured_content["items"][0]["id"] == "755" assert related.structured_content["items"][0]["id"] == "755" assert lineage.structured_content["paper"]["reference"] == "1706.03762" assert lineage.structured_content["successors"] == [ @@ -342,8 +512,10 @@ async def exercise(): def test_taxonomy_and_benchmark_tools_return_stable_catalog_entities(): + catalog = StubCatalog() + async def exercise(): - async with Client(build_server(StubCatalog())) as client: + async with Client(build_server(catalog)) as client: tool_names = {tool.name for tool in (await client.list_tools()).tools} task = await client.call_tool("get_task", {"task": "image-classification"}) method = await client.call_tool("get_method", {"method": "transformer"}) @@ -352,29 +524,156 @@ async def exercise(): ) benchmark = await client.call_tool( "get_benchmark", - {"benchmark": "imagenet-1k", "limit": 5, "is_open": True}, + { + "benchmark": "imagenet-1k", + "limit": 5, + "is_open": True, + "max_parameters": "4B", + "minimum_metrics": {"Accuracy": 80}, + "sort_metric": "Accuracy:desc", + }, ) return tool_names, task, method, benchmarks, benchmark tool_names, task, method, benchmarks, benchmark = asyncio.run(exercise()) - assert tool_names == { - "search_papers", - "list_papers", - "get_paper_info", - "read_paper", - "get_related_papers", - "get_paper_lineage", - "get_task", - "get_method", - "list_benchmarks", - "get_benchmark", + assert tool_names == set(TOOL_COMMANDS) + assert len(tool_names) == 20 + assert catalog.options(("task",)) == {"name": "image-classification"} + assert catalog.options(("method",)) == {"name": "transformer"} + assert catalog.options(("benchmark", "list")) == { + "search": None, + "task": "image-classification", + "include_descendants": False, + "min_eval_count": None, + "is_open": None, + "group_by_area": False, + "area": None, + "benchmarks_per_task": 3, + "order_by": None, + "order_dir": "asc", + "page": 1, + "page_size": 25, + } + assert catalog.options(("benchmark",)) == { + "name": "imagenet-1k", + "limit": 5, + "is_open": True, + "max_parameters": "4B", + "require_metrics": None, + "minimum_metrics": {"Accuracy": 80.0}, + "maximum_metrics": None, + "sort_metric": "Accuracy:desc", + "pareto": None, } assert task.structured_content["task"]["area"] == {"id": "1", "name": "Vision"} + assert task.structured_content["task"]["benchmarks"][0]["slug"] == "imagenet-1k" + assert task.structured_content["data"]["common_methods"][0]["name"] == "Transformer" assert method.structured_content["method"]["introduced_year"] == 2017 assert benchmarks.structured_content["items"][0]["slug"] == "imagenet-1k" - assert benchmark.structured_content["evaluations"][0]["metrics"] == { - "Accuracy": 90.1 + assert benchmark.structured_content["matched_count"] == 1 + assert benchmark.structured_content["evaluations"] == [ + { + "id": "10", + "model_name": "ExampleNet", + "harness": "timm", + "metrics": {"Accuracy": 90.1}, + "best_metric": "Accuracy", + "best_rank": 1, + "task": "Image Classification", + "paper_id": "755", + "paper_title": "Attention Is All You Need", + "paper_arxiv_id": "1706.03762", + "paper_published": "2017-06-12", + "is_open": True, + "num_parameters": 1000, + } + ] + + +def test_grouped_listings_omit_pagination_and_flatten_benchmarks(): + catalog = StubCatalog() + tasks, grouped_tasks, benchmarks = _call( + catalog, + [ + ("list_tasks", {"area": "Vision", "level": 1}), + ("list_tasks", {"group_by_area": True}), + ("list_benchmarks", {"area": "Vision", "benchmarks_per_task": 2}), + ], + ) + + assert catalog.queries[0] == ( + ("task", "list"), + { + "area": "Vision", + "level": 1, + "visible_only": False, + "group_by_area": False, + "order_by": "name", + "order_dir": "asc", + "page": 1, + "page_size": 25, + }, + ) + assert catalog.queries[1][1]["group_by_area"] is True + assert catalog.queries[1][1]["page_size"] is None + assert catalog.queries[2][1]["page_size"] is None + assert catalog.queries[2][1]["benchmarks_per_task"] == 2 + assert tasks.structured_content == { + "schema_version": "v1", + "data": {"count": 1, "results": [TASK]}, + } + assert grouped_tasks.structured_content["data"] == {"count": 1, "results": [TASK]} + assert benchmarks.structured_content["items"] == [ + { + "id": "72", + "name": "ImageNet-1k", + "slug": "imagenet-1k", + "full_name": None, + "description": None, + "hf_url": None, + "paper_count": 124, + } + ] + assert benchmarks.structured_content["data"] == GROUPED_BENCHMARKS + + +def test_new_catalog_tools_return_the_cli_json_document(): + catalog = StubCatalog() + requests = [ + ("list_methods", {"introduced_year": 2017, "order_by": "paper_count"}), + ("get_conference", {"conference": "CVPR 2025"}), + ("list_conferences", {"year": 2025}), + ("get_organization", {"organization": "NVIDIA"}), + ("list_organizations", {"featured_only": True}), + ("get_framework", {"framework": "vLLM"}), + ("list_frameworks", {"platform": "gpu"}), + ] + results = _call(catalog, requests) + + for (name, _arguments), result in zip(requests, results): + assert result.is_error is False, name + assert result.structured_content == { + "schema_version": "v1", + "data": PAYLOADS[TOOL_COMMANDS[name]], + } + assert catalog.options(("method", "list")) == { + "area": None, + "introduced_year": 2017, + "order_by": "paper_count", + "order_dir": "asc", + "page": 1, + "page_size": 25, + } + assert catalog.options(("conference",)) == {"name": "CVPR 2025"} + assert catalog.options(("conference", "list")) == {"year": 2025} + assert catalog.options(("organization",)) == {"name": "NVIDIA"} + assert catalog.options(("organization", "list")) == {"featured_only": True} + assert catalog.options(("framework",)) == {"name": "vLLM"} + assert catalog.options(("framework", "list")) == { + "domain": None, + "category": None, + "platform": "gpu", } diff --git a/mcp_server/tests/test_skill.py b/mcp_server/tests/test_skill.py index 0ccb341..54cfdaa 100644 --- a/mcp_server/tests/test_skill.py +++ b/mcp_server/tests/test_skill.py @@ -3,19 +3,10 @@ import re from pathlib import Path +from pwc_mcp.server import TOOL_COMMANDS + SKILL = Path(__file__).parents[1] / "SKILL.md" -EXPECTED_TOOLS = { - "search_papers", - "list_papers", - "get_paper_info", - "read_paper", - "get_related_papers", - "get_paper_lineage", - "get_task", - "get_method", - "list_benchmarks", - "get_benchmark", -} +EXPECTED_TOOLS = set(TOOL_COMMANDS) def test_mcp_skill_has_valid_agent_skills_frontmatter() -> None: diff --git a/mcp_server/uv.lock b/mcp_server/uv.lock index f9c4aca..a620cde 100644 --- a/mcp_server/uv.lock +++ b/mcp_server/uv.lock @@ -438,12 +438,12 @@ wheels = [ [[package]] name = "pwc-cli" -version = "0.4.0" +version = "0.4.1" source = { editable = "../standalone_cli" } [[package]] name = "pwc-mcp" -version = "0.1.0" +version = "0.2.0" source = { editable = "." } dependencies = [ { name = "mcp" }, diff --git a/standalone_cli/SKILL.md b/standalone_cli/SKILL.md index ec9221e..f1ebe33 100644 --- a/standalone_cli/SKILL.md +++ b/standalone_cli/SKILL.md @@ -3,7 +3,7 @@ name: pwc-cli description: "Papers With Code CLI (`pwc`) for searching and reading AI/ML papers, discovering recent and trending research, finding related work and paper lineage, browsing tasks, methods, conferences, organizations, frameworks, and benchmark leaderboards, and submitting authenticated paper edits through the public Papers With Code catalog. Use whenever the user asks to find papers, survey literature, compare research, inspect an arXiv paper, explore AI/ML taxonomy or conferences, discover benchmarks or state-of-the-art models, or mentions Papers With Code, `pwc`, or `pwc-cli`." --- -Generated with `pwc v0.4.0`. Run `pwc skills add --force` to regenerate. +Generated with `pwc v0.4.1`. Run `pwc skills add --force` to regenerate. Research commands query the public [Papers With Code](https://paperswithcode.co) catalog anonymously. Paper editing requires explicit browser authorization through `pwc auth login --paper PAPER`. diff --git a/standalone_cli/pyproject.toml b/standalone_cli/pyproject.toml index c68da7c..281940b 100644 --- a/standalone_cli/pyproject.toml +++ b/standalone_cli/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "pwc-cli" -version = "0.4.0" +version = "0.4.1" description = "Papers With Code research and paper-editing CLI" readme = "README.md" requires-python = ">=3.10" diff --git a/standalone_cli/src/pwc_cli/__init__.py b/standalone_cli/src/pwc_cli/__init__.py index 9257051..6bc0ac0 100644 --- a/standalone_cli/src/pwc_cli/__init__.py +++ b/standalone_cli/src/pwc_cli/__init__.py @@ -1,4 +1,4 @@ """Standalone Papers With Code client.""" -__version__ = "0.4.0" +__version__ = "0.4.1" API_CONTRACT_VERSION = "v1" diff --git a/standalone_cli/src/pwc_cli/cli.py b/standalone_cli/src/pwc_cli/cli.py index 20d8cd1..9e9961f 100644 --- a/standalone_cli/src/pwc_cli/cli.py +++ b/standalone_cli/src/pwc_cli/cli.py @@ -305,20 +305,34 @@ def _paper_rows( _print_table(tuple(headers), rows, right_align=tuple(right_align)) +def _emit_json(args: argparse.Namespace, data: Any) -> int: + """Emit one versioned JSON document, or hand it to an in-process caller. + + ``pwc_cli.queries`` sets ``args.result_sink`` so embedders such as the MCP + server receive exactly the ``data`` that ``--json`` would print, without + touching stdout. + """ + sink = getattr(args, "result_sink", None) + if sink is not None: + sink.append(data) + return 0 + print( + json.dumps( + {"schema_version": API_CONTRACT_VERSION, "data": data}, + ensure_ascii=False, + separators=(",", ":"), + ) + ) + return 0 + + def _emit_page( payload: Any, args: argparse.Namespace, renderer: Callable[[list[dict[str, Any]]], None], ) -> int: if args.json: - print( - json.dumps( - {"schema_version": API_CONTRACT_VERSION, "data": payload}, - ensure_ascii=False, - separators=(",", ":"), - ) - ) - return 0 + return _emit_json(args, payload) items, total = _rows(payload) renderer(items) page = getattr(args, "page", 1) @@ -386,7 +400,8 @@ def search(args: argparse.Namespace, client: Client) -> int: def paper_info(args: argparse.Namespace, client: Client) -> int: paper = _resolve_paper(args.paper, client) payload = client.get( - f"papers/{paper}", {"include_resources": args.include_resources} + f"papers/{quote(paper, safe='.')}", + {"include_resources": args.include_resources}, ).json() evaluations = None if args.include_evals: @@ -396,14 +411,7 @@ def paper_info(args: argparse.Namespace, client: Client) -> int: evaluations = _paper_evaluations(client, paper_id) payload = {**payload, "evaluations": evaluations} if args.json: - print( - json.dumps( - {"schema_version": API_CONTRACT_VERSION, "data": payload}, - ensure_ascii=False, - separators=(",", ":"), - ) - ) - return 0 + return _emit_json(args, payload) organizations = ", ".join( str(item.get("name") or item.get("slug")) for item in payload.get("organizations") or [] @@ -513,20 +521,10 @@ def _paper_info_lineage_markdown(item: dict[str, Any]) -> str: def paper_read(args: argparse.Namespace, client: Client) -> int: paper = _resolve_paper(args.paper, client) - response = client.get(f"research/papers/{paper}/read") + response = client.get(f"research/papers/{quote(paper, safe='.')}/read") markdown = response.body.decode("utf-8", errors="replace") if args.json: - print( - json.dumps( - { - "schema_version": API_CONTRACT_VERSION, - "data": {"paper": paper, "markdown": markdown}, - }, - ensure_ascii=False, - separators=(",", ":"), - ) - ) - return 0 + return _emit_json(args, {"paper": paper, "markdown": markdown}) sys.stdout.write(markdown) if markdown and not markdown.endswith("\n"): sys.stdout.write("\n") @@ -678,7 +676,9 @@ def paper_trending(args: argparse.Namespace, client: Client) -> int: def paper_related(args: argparse.Namespace, client: Client) -> int: paper = _resolve_paper(args.paper, client) return _emit_paper_page( - client.get(f"papers/{paper}/related", {"limit": args.limit}).json(), + client.get( + f"papers/{quote(paper, safe='.')}/related", {"limit": args.limit} + ).json(), args, ) @@ -699,16 +699,11 @@ def _lineage_markdown(item: dict[str, Any]) -> str: def paper_lineage(args: argparse.Namespace, client: Client) -> int: paper_reference = _resolve_paper(args.paper, client) - payload = client.get(f"research/papers/{paper_reference}/lineage").json() + payload = client.get( + f"research/papers/{quote(paper_reference, safe='.')}/lineage" + ).json() if args.json: - print( - json.dumps( - {"schema_version": API_CONTRACT_VERSION, "data": payload}, - ensure_ascii=False, - separators=(",", ":"), - ) - ) - return 0 + return _emit_json(args, payload) paper = payload.get("paper") or {} print("# Paper lineage") for title, items in ( @@ -761,7 +756,7 @@ def _ordering(field: str, direction: str) -> str: def _task_list_grouped(args: argparse.Namespace, client: Client) -> int: if args.page != 1 or args.page_size != 50 or args.level is not None: - raise ResponseError( + raise UsageError( "--group-by-area cannot be combined with pagination or --level; " "use --flat for the complete task endpoint" ) @@ -810,14 +805,7 @@ def sort_key(item: dict[str, Any]): grouped.append({**area, "tasks": tasks}) data = {"results": grouped} if args.json: - print( - json.dumps( - {"schema_version": API_CONTRACT_VERSION, "data": data}, - ensure_ascii=False, - separators=(",", ":"), - ) - ) - return 0 + return _emit_json(args, data) for area_index, area in enumerate(grouped): if area_index: @@ -1047,14 +1035,7 @@ def task_detail(args: argparse.Namespace, client: Client) -> int: "papers": papers[:10], } if args.json: - print( - json.dumps( - {"schema_version": API_CONTRACT_VERSION, "data": data}, - ensure_ascii=False, - separators=(",", ":"), - ) - ) - return 0 + return _emit_json(args, data) interactive = sys.stdout.isatty() markdown = not interactive @@ -1238,14 +1219,7 @@ def method_detail(args: argparse.Namespace, client: Client) -> int: ) data = {"method": method, "area": area} if args.json: - print( - json.dumps( - {"schema_version": API_CONTRACT_VERSION, "data": data}, - ensure_ascii=False, - separators=(",", ":"), - ) - ) - return 0 + return _emit_json(args, data) _entity_heading(method.get("name") or method.get("slug")) metadata = [] @@ -1314,14 +1288,7 @@ def conference_detail(args: argparse.Namespace, client: Client) -> int: slug = str(summary.get("slug") or summary.get("id")) conference = client.get(f"conferences/{quote(slug, safe='')}").json() if args.json: - print( - json.dumps( - {"schema_version": API_CONTRACT_VERSION, "data": conference}, - ensure_ascii=False, - separators=(",", ":"), - ) - ) - return 0 + return _emit_json(args, conference) _entity_heading(conference.get("name") or conference.get("slug")) metadata = [] @@ -1392,14 +1359,7 @@ def organization_detail(args: argparse.Namespace, client: Client) -> int: slug = str(summary.get("slug") or summary.get("id")) organization = client.get(f"organizations/{quote(slug, safe='')}").json() if args.json: - print( - json.dumps( - {"schema_version": API_CONTRACT_VERSION, "data": organization}, - ensure_ascii=False, - separators=(",", ":"), - ) - ) - return 0 + return _emit_json(args, organization) _entity_heading(organization.get("name") or organization.get("slug")) metadata = [] @@ -1488,14 +1448,7 @@ def framework_detail(args: argparse.Namespace, client: Client) -> int: items = _framework_catalog_items(client.get("frameworks/").json()) framework = _exact_entity_match(args.name, items, label="Framework") if args.json: - print( - json.dumps( - {"schema_version": API_CONTRACT_VERSION, "data": framework}, - ensure_ascii=False, - separators=(",", ":"), - ) - ) - return 0 + return _emit_json(args, framework) _entity_heading(framework.get("name") or framework.get("slug")) metadata = [ @@ -1695,7 +1648,7 @@ def render_trends(rows: list[dict[str, Any]]) -> None: return _emit_page(payload, args, render_trends) if args.order_by == "trending": - raise ResponseError("--order-by trending requires --task") + raise UsageError("--order-by trending requires --task") ordering = f"-{args.order_by}" if args.order_dir == "desc" else args.order_by payload = client.get( "datasets/", @@ -1768,14 +1721,7 @@ def _benchmark_list_grouped(args: argparse.Namespace, client: Client) -> int: grouped = {"results": areas} if args.json: - print( - json.dumps( - {"schema_version": API_CONTRACT_VERSION, "data": grouped}, - ensure_ascii=False, - separators=(",", ":"), - ) - ) - return 0 + return _emit_json(args, grouped) for area_index, area in enumerate(areas): if area_index: @@ -2232,14 +2178,7 @@ def benchmark_detail(args: argparse.Namespace, client: Client) -> int: "results": rows, } if args.json: - print( - json.dumps( - {"schema_version": API_CONTRACT_VERSION, "data": data}, - ensure_ascii=False, - separators=(",", ":"), - ) - ) - return 0 + return _emit_json(args, data) interactive = sys.stdout.isatty() title = ( @@ -2352,8 +2291,11 @@ def _implementation_coverage(parser: argparse.ArgumentParser) -> None: ) -def build_parser() -> argparse.ArgumentParser: - parser = Parser( +def build_parser( + *, parser_class: type[argparse.ArgumentParser] = Parser +) -> argparse.ArgumentParser: + """Build the pwc parser; subparsers inherit ``parser_class``.""" + parser = parser_class( prog="pwc", description="Papers With Code research and paper-editing CLI", epilog=( diff --git a/standalone_cli/src/pwc_cli/queries.py b/standalone_cli/src/pwc_cli/queries.py new file mode 100644 index 0000000..5fc9306 --- /dev/null +++ b/standalone_cli/src/pwc_cli/queries.py @@ -0,0 +1,178 @@ +"""In-process access to the CLI's read-only research queries. + +Embedders such as the Papers With Code MCP server run the exact CLI handlers +through the CLI's own parser. Every flag, default, validation rule, and JSON +payload therefore stays identical to ``pwc ... --json`` without a subprocess, +stdout capture, or a second implementation of the query logic. + +Each call owns its parser and result sink, so concurrent callers never share +state. Only read-only research commands are reachable; authentication, paper +editing, skill installation, and version display are refused. +""" + +from __future__ import annotations + +import argparse +from collections.abc import Mapping +from datetime import date +from typing import Any + +from pwc_cli.cli import Parser, UsageError, build_parser +from pwc_cli.transport import ResponseError + +Command = tuple[str, ...] + +READ_ONLY_COMMANDS: tuple[Command, ...] = ( + ("search",), + ("paper", "info"), + ("paper", "read"), + ("paper", "list"), + ("paper", "recent"), + ("paper", "trending"), + ("paper", "related"), + ("paper", "lineage", "list"), + ("task",), + ("task", "list"), + ("method",), + ("method", "list"), + ("conference",), + ("conference", "list"), + ("organization",), + ("organization", "list"), + ("framework",), + ("framework", "list"), + ("benchmark",), + ("benchmark", "list"), +) + +# Flags that only shape terminal rendering. JSON output ignores them, so they +# have no in-process equivalent. +PRESENTATION_FLAGS = frozenset({"json", "implementation_coverage", "flat"}) + +_HIDDEN_ACTIONS = ( + argparse._SubParsersAction, + argparse._HelpAction, + argparse._VersionAction, +) + + +class _StrictParser(Parser): + """Raise ``UsageError`` for invalid arguments instead of exiting.""" + + def error(self, message: str) -> None: # type: ignore[override] + raise UsageError(message) + + +def _command_label(command: Command) -> str: + return " ".join(("pwc", *command)) + + +def command_parser( + command: Command, parser: argparse.ArgumentParser | None = None +) -> argparse.ArgumentParser: + """Return the subparser that handles ``command``.""" + current = parser or build_parser(parser_class=_StrictParser) + for name in command: + subparsers = next( + ( + action + for action in current._actions + if isinstance(action, argparse._SubParsersAction) + ), + None, + ) + if subparsers is None or name not in subparsers.choices: + raise KeyError(f"unknown command: {_command_label(command)}") + current = subparsers.choices[name] + return current + + +def query_options(command: Command) -> dict[str, argparse.Action]: + """Research options of one command keyed by destination. + + Presentation-only flags, help, version, and nested command selectors are + excluded, leaving exactly the options an embedder must expose. + """ + parser = command_parser(command) + return { + action.dest: action + for action in parser._actions + if not isinstance(action, _HIDDEN_ACTIONS) + and action.dest not in PRESENTATION_FLAGS + } + + +def _text(value: Any) -> str: + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, date): + return value.isoformat() + return str(value) + + +def build_argv(command: Command, options: Mapping[str, Any]) -> list[str]: + """Translate structured options into the argv the CLI parser expects. + + ``None`` leaves a flag at its CLI default. Booleans toggle ``store_true`` + flags. Lists repeat ``append`` flags or join comma-separated values, and + mappings expand to repeated ``NAME=VALUE`` entries. + """ + actions = query_options(command) + unknown = sorted(set(options) - set(actions)) + if unknown: + raise UsageError( + f"unknown option(s) for {_command_label(command)}: {', '.join(unknown)}" + ) + argv = list(command) + positionals: list[str] = [] + for dest, action in actions.items(): + value = options.get(dest) + if value is None: + continue + if not action.option_strings: + positionals.append(_text(value)) + continue + flag = action.option_strings[0] + if isinstance(action, argparse._StoreTrueAction): + if value is True: + argv.append(flag) + elif value is not False: + raise UsageError(f"{flag} expects a boolean") + continue + if isinstance(value, Mapping): + value = [f"{name}={_text(item)}" for name, item in value.items()] + values = list(value) if isinstance(value, (list, tuple, set)) else [value] + if isinstance(action, argparse._AppendAction): + for item in values: + argv.extend((flag, _text(item))) + elif values: + argv.extend((flag, ",".join(_text(item) for item in values))) + if positionals: + argv.append("--") + argv.extend(positionals) + return argv + + +def query(command: Command, options: Mapping[str, Any], client: Any) -> Any: + """Run one read-only CLI command in-process and return its JSON ``data``. + + ``client`` must provide ``get(path, params) -> Response`` like + :class:`pwc_cli.transport.Client`. Invalid options raise ``UsageError``; + catalog failures raise the CLI's ``TransportError`` or ``ResponseError``. + """ + command = tuple(command) + if command not in READ_ONLY_COMMANDS: + raise UsageError( + f"{_command_label(command)} is not a read-only research command" + ) + argv = build_argv(command, options) + args = build_parser(parser_class=_StrictParser).parse_args(argv) + sink: list[Any] = [] + args.json = True + args.result_sink = sink + args.handler(args, client) + if len(sink) != 1: + raise ResponseError( + f"{_command_label(command)} did not produce one JSON result" + ) + return sink[0] diff --git a/standalone_cli/tests/test_cli.py b/standalone_cli/tests/test_cli.py index 6a8747d..617d46e 100644 --- a/standalone_cli/tests/test_cli.py +++ b/standalone_cli/tests/test_cli.py @@ -85,7 +85,7 @@ def test_generated_skill_matches_installed_cli_version_and_commands(): skill = build_skill_md() assert "name: pwc-cli" in skill - assert "Generated with `pwc v0.4.0`" in skill + assert "Generated with `pwc v0.4.1`" in skill assert "`pwc search QUERY" in skill assert "--include-evals" in skill assert "[--organization ORGANIZATION]" in skill @@ -2255,7 +2255,7 @@ def test_top_level_version_is_offline_and_stable(): build_parser().parse_args(["--version"]) except SystemExit as error: assert error.code == 0 - assert output.getvalue() == "pwc 0.4.0\tapi v1\n" + assert output.getvalue() == "pwc 0.4.1\tapi v1\n" def test_search_default_output_is_compact_deterministic_tsv(monkeypatch): diff --git a/standalone_cli/tests/test_queries.py b/standalone_cli/tests/test_queries.py new file mode 100644 index 0000000..6b8cef8 --- /dev/null +++ b/standalone_cli/tests/test_queries.py @@ -0,0 +1,213 @@ +"""The in-process query layer must mirror ``pwc ... --json`` exactly.""" + +from __future__ import annotations + +import argparse +import io +import json +from contextlib import redirect_stdout + +import pytest +from pwc_cli import queries +from pwc_cli.cli import UsageError, build_parser, main +from pwc_cli.transport import Response + +NON_RESEARCH_COMMANDS = { + ("skills", "add"), + ("version",), + ("auth", "login"), + ("auth", "status"), + ("auth", "logout"), + ("paper", "edit", "export"), + ("paper", "edit", "preview"), + ("paper", "edit", "submit"), +} + +EVALUATION_ROWS = [ + { + "id": "1", + "model_name": "Small", + "metrics": {"mAP": 50.0}, + "best_rank": 2, + "num_parameters": 100_000_000, + "paper_id": "1", + "task_id": "t", + "dataset_id": "9", + }, + { + "id": "2", + "model_name": "Smaller", + "metrics": {"mAP": 60.0}, + "best_rank": 1, + "num_parameters": 50_000_000, + "paper_id": "2", + "task_id": "t", + "dataset_id": "9", + }, +] + +BENCHMARK_ROUTES = { + "datasets/": {"results": [{"id": "9", "name": "COCO", "slug": "coco"}]}, + "datasets/9/evaluations/": {"count": 2, "results": EVALUATION_ROWS}, + "evaluations/": { + "count": 2, + "next_page": None, + "parameter_coverage_known": 2, + "parameter_coverage_total": 2, + "results": EVALUATION_ROWS, + }, +} + + +class StubClient: + def __init__(self, routes): + self.routes = routes + self.calls = [] + + def get(self, path, params=None): + self.calls.append((path, dict(params or {}))) + return Response(json.dumps(self.routes[path]).encode(), {}) + + +def _leaf_commands(parser, prefix=()): + for action in parser._actions: + if not isinstance(action, argparse._SubParsersAction): + continue + for name, child in action.choices.items(): + path = (*prefix, name) + if "handler" in child._defaults: + yield path + yield from _leaf_commands(child, path) + + +def test_read_only_commands_cover_every_research_command(): + research = set(_leaf_commands(build_parser())) - NON_RESEARCH_COMMANDS + + assert set(queries.READ_ONLY_COMMANDS) == research + for command in queries.READ_ONLY_COMMANDS: + assert "handler" in queries.command_parser(command)._defaults + + +def test_query_options_exclude_presentation_flags_and_selectors(): + options = queries.query_options(("benchmark", "list")) + + assert set(options) == { + "page", + "page_size", + "search", + "task", + "group_by_area", + "area", + "benchmarks_per_task", + "include_descendants", + "min_eval_count", + "is_open", + "order_by", + "order_dir", + } + assert queries.query_options(("search",))["query"].option_strings == [] + + +def test_build_argv_translates_structured_options(): + argv = queries.build_argv( + ("benchmark",), + { + "name": "COCO", + "limit": 5, + "is_open": True, + "require_metrics": ["mAP", "FPS"], + "minimum_metrics": {"mAP": 40}, + "maximum_metrics": {"FPS": 10.5}, + "pareto": ["mAP:higher", "FPS:lower"], + "max_parameters": None, + }, + ) + + assert argv == [ + "benchmark", + "--name", + "COCO", + "--limit", + "5", + "--is-open", + "true", + "--require-metrics", + "mAP,FPS", + "--min", + "mAP=40", + "--max", + "FPS=10.5", + "--pareto", + "mAP:higher,FPS:lower", + ] + assert queries.build_argv( + ("paper", "list"), + {"author": ["Ada", "@ada"], "all_versions": True, "include_resources": False}, + ) == ["paper", "list", "--author", "Ada", "--author", "@ada", "--all-versions"] + + argv = queries.build_argv(("search",), {"query": "-attention", "limit": 3}) + assert argv == ["search", "--limit", "3", "--", "-attention"] + assert build_parser().parse_args(argv).query == "-attention" + + +def test_query_returns_exactly_the_cli_json_data(monkeypatch, capsys): + monkeypatch.setattr("pwc_cli.cli.Client", lambda: StubClient(BENCHMARK_ROUTES)) + output = io.StringIO() + with redirect_stdout(output): + assert ( + main( + [ + "benchmark", + "--name", + "COCO", + "--max-parameters", + "500M", + "--sort", + "mAP", + "--limit", + "1", + "--json", + ] + ) + == 0 + ) + printed = json.loads(output.getvalue()) + + client = StubClient(BENCHMARK_ROUTES) + data = queries.query( + ("benchmark",), + {"name": "COCO", "max_parameters": "500M", "sort_metric": "mAP", "limit": 1}, + client, + ) + + assert {"schema_version": "v1", "data": data} == printed + assert data["matched_count"] == 2 + assert [row["model_name"] for row in data["results"]] == ["Smaller"] + assert client.calls[1][1]["max_parameters_exclusive"] == 500_000_001 + assert capsys.readouterr().out == "" + + +def test_query_raises_usage_errors_instead_of_exiting(): + client = StubClient({}) + + with pytest.raises(UsageError, match="--limit must be 1-100"): + queries.query(("search",), {"query": "x", "limit": 500}, client) + with pytest.raises(UsageError, match="unknown option"): + queries.query(("search",), {"query": "x", "page_size": 5}, client) + with pytest.raises(UsageError, match="invalid"): + queries.query(("search",), {"query": "x", "start_date": "2026-13-01"}, client) + with pytest.raises(UsageError, match="cannot be combined"): + queries.query(("task", "list"), {"group_by_area": True, "page": 2}, client) + with pytest.raises(UsageError, match="unknown metric"): + queries.query( + ("benchmark",), + {"name": "COCO", "sort_metric": "Latency"}, + StubClient(BENCHMARK_ROUTES), + ) + assert client.calls == [] + + +def test_query_refuses_commands_that_are_not_read_only(): + for command in (("skills", "add"), ("version",), ("paper", "edit", "export")): + with pytest.raises(UsageError, match="not a read-only"): + queries.query(command, {}, StubClient({}))