from __future__ import annotations import json import pytest from pwc_cli.cli import UsageError from pwc_cli.transport import Response, ResponseError from pwc_mcp.catalog import ( CatalogClient, PublicPaperMiss, cache_ttl, near_title, normalize_title, ) class StubTransport: def __init__(self, routes): self.routes = routes self.calls = [] def get(self, path, params=None): self.calls.append((path, params)) value = self.routes[path] if isinstance(value, str): return Response(value.encode(), {"content-type": "text/markdown"}) return Response( json.dumps(value).encode(), {"content-type": "application/json"} ) def test_catalog_normalizes_paper_urls_and_caches_identical_reads(): transport = StubTransport( { "papers/1706.03762": {"id": "755", "title": "Attention Is All You Need"}, } ) catalog = CatalogClient(transport=transport) first = catalog.query( ("paper", "info"), {"paper": "https://arxiv.org/pdf/1706.03762v7.pdf", "include_resources": True}, ) second = catalog.query( ("paper", "info"), { "paper": "https://huggingface.co/papers/1706.03762", "include_resources": True, }, ) assert first == second assert transport.calls == [("papers/1706.03762", {"include_resources": True})] def test_catalog_reads_one_version_bound_markdown_chunk_per_call(): version = "a" * 64 class ChunkTransport: def __init__(self): self.calls = [] def get(self, path, params=None): self.calls.append((path, params)) if params["offset"] == 0: return Response( "α first".encode(), { "x-pwc-truncated": "1", "x-pwc-content-version": version, "x-pwc-next-offset": "8", }, ) return Response( b" terminal", { "x-pwc-truncated": "0", "x-pwc-content-version": version, }, ) transport = ChunkTransport() catalog = CatalogClient(transport=transport) first = catalog.read_paper_chunk("1706.03762", limit=65_536) second = catalog.read_paper_chunk( "1706.03762", offset=first.next_offset, content_version=first.content_version, limit=65_536, resolved=True, ) assert first.markdown == "α first" assert first.next_offset == 8 assert second.markdown == " terminal" assert second.next_offset is None assert transport.calls == [ ( "research/papers/1706.03762/read", {"offset": 0, "limit": 65_536, "content_version": None}, ), ( "research/papers/1706.03762/read", {"offset": 8, "limit": 65_536, "content_version": version}, ), ] def test_catalog_caches_versioned_chunks_and_rejects_nonadvancing_offsets(): version = "b" * 64 class ChunkTransport: def __init__(self): self.calls = 0 def get(self, _path, _params=None): self.calls += 1 return Response( b"data", { "x-pwc-truncated": "1", "x-pwc-content-version": version, "x-pwc-next-offset": "4", }, ) transport = ChunkTransport() catalog = CatalogClient(transport=transport) with pytest.raises(ResponseError, match="did not advance"): catalog.read_paper_chunk( "1706.03762", offset=4, content_version=version, resolved=True, ) assert transport.calls == 1 def test_catalog_versioned_chunk_cache_avoids_repeated_upstream_reads(): version = "c" * 64 class ChunkTransport: def __init__(self): self.calls = 0 def get(self, _path, _params=None): self.calls += 1 return Response( b"done", { "x-pwc-truncated": "0", "x-pwc-content-version": version, }, ) transport = ChunkTransport() catalog = CatalogClient(transport=transport) first = catalog.read_paper_chunk( "1706.03762", offset=8, content_version=version, resolved=True, ) second = catalog.read_paper_chunk( "1706.03762", offset=8, content_version=version, resolved=True, ) assert first == second assert transport.calls == 1 def test_catalog_resolves_exact_titles_and_rejects_ambiguous_titles(): transport = StubTransport( { "papers/search": { "results": [ {"id": "1", "arxiv_id": "1111.11111", "title": "Same Title"}, {"id": "2", "arxiv_id": "2222.22222", "title": "Same Title"}, ] } } ) catalog = CatalogClient(transport=transport) with pytest.raises(ResponseError, match="ambiguous"): catalog.query(("paper", "lineage", "list"), {"paper": "Same Title"}) def test_catalog_resolves_pwc_urls_and_dotted_legacy_arxiv_ids(): transport = StubTransport( { "papers/search": { "results": [ { "id": "755", "arxiv_id": "1706.03762", "title": "Attention Is All You Need", } ] }, "papers/1706.03762": {"id": "755"}, "papers/math.GT%2F0309136": {"id": "900"}, } ) catalog = CatalogClient(transport=transport) assert catalog.query( ("paper", "info"), {"paper": "https://paperswithcode.co/paper/attention-is-all-you-need"}, ) == {"id": "755"} assert catalog.query(("paper", "info"), {"paper": "math.GT/0309136"}) == { "id": "900" } def test_catalog_checks_later_search_pages_for_ambiguous_exact_titles(): class PaginatedTransport: def get(self, path, params=None): assert path == "papers/search" page = params["page"] payload = { "results": [ { "id": str(page), "arxiv_id": f"1111.1111{page}", "title": "Same Title", } ], "next_page": 2 if page == 1 else None, } return Response(json.dumps(payload).encode(), {}) catalog = CatalogClient(transport=PaginatedTransport()) with pytest.raises(ResponseError, match="ambiguous"): catalog.query(("paper", "lineage", "list"), {"paper": "Same Title"}) def test_catalog_fails_closed_when_paper_filters_are_not_confirmed(): transport = StubTransport({"papers/": {"results": [], "applied_filters": {}}}) catalog = CatalogClient(transport=transport) with pytest.raises(ResponseError, match="did not confirm"): 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 == [] def test_catalog_canonicalises_numeric_ids_of_arxiv_papers(): transport = StubTransport( { "papers/86122": {"id": "86122", "arxiv_id": "2505.18132v1"}, "papers/2505.18132": {"id": "86122", "title": "BiggerGait"}, "papers/4242": {"id": "4242", "arxiv_id": None, "source": "external"}, } ) catalog = CatalogClient(transport=transport) # List tools hand out numeric IDs for arXiv papers too; the Markdown route # only knows arXiv papers by arXiv ID, so numeric references canonicalise. assert catalog.resolve_paper("86122") == "2505.18132" assert catalog.query(("paper", "info"), {"paper": "86122"}) == { "id": "86122", "title": "BiggerGait", } # External papers have no arXiv ID and keep their numeric reference. assert catalog.resolve_paper("4242") == "4242" assert transport.calls[0] == ("papers/86122", {}) def test_catalog_rejects_unsupported_urls_without_searching(): transport = StubTransport({}) catalog = CatalogClient(transport=transport) for url in ( "https://doi.org/10.1007/s10479-024-06277-x", "https://aclanthology.org/2026.acl-long.200/", ): with pytest.raises(ResponseError, match="Paper URL not supported"): catalog.resolve_paper(url) assert transport.calls == [] def test_catalog_distinguishes_missing_papers_from_missing_markdown(): from pwc_cli.transport import HTTPStatusError class MissingTransport: def __init__(self, paper_exists): self.paper_exists = paper_exists def get(self, path, params=None): if path.endswith("/read"): raise HTTPStatusError(404, "Paper Markdown was not found") if self.paper_exists: return Response(b'{"id": "1"}', {}) raise HTTPStatusError(404, "No paper found") with pytest.raises(ResponseError, match="Paper not found: 2308.10195"): CatalogClient(transport=MissingTransport(False)).read_paper_chunk( "2308.10195", resolved=True ) with pytest.raises(HTTPStatusError) as missing_markdown: CatalogClient(transport=MissingTransport(True)).read_paper_chunk( "2505.18132", resolved=True ) assert missing_markdown.value.status == 404 def test_catalog_reports_titles_missing_from_an_empty_search_as_not_found(): # The public API answers an unknown title with an empty result list; that # must surface as "title not found", not as a malformed response. transport = StubTransport( { "papers/search": { "next_page": None, "previous_page": None, "results": [], "applied_filters": {}, } } ) catalog = CatalogClient(transport=transport) with pytest.raises(ResponseError, match="Paper title not found: Dropout"): catalog.resolve_paper("Dropout: A Simple Way to Prevent Overfitting") # One phrase query, then one plain keyword query as the fallback. assert [params["q"] for _path, params in transport.calls] == [ '"Dropout: A Simple Way to Prevent Overfitting"', "Dropout: A Simple Way to Prevent Overfitting", ] EXACT_TITLE_ROW = {"id": "755", "arxiv_id": "1706.03762", "title": "Attention Is All You Need"} class TitleSearchTransport: """Keyword search whose plain results never end; the phrase query is small.""" def __init__(self, *, phrase_results, plain_has_exact): self.phrase_results = phrase_results self.plain_has_exact = plain_has_exact self.calls = [] def get(self, path, params=None): params = dict(params or {}) self.calls.append((path, params)) assert path == "papers/search" page = params["page"] if params["q"].startswith('"'): body = {"results": self.phrase_results, "next_page": None} else: rows = [ {"id": str(page * 100 + i), "arxiv_id": f"2{page:03d}.{i:05d}", "title": f"Attention {i}"} for i in range(100) ] if self.plain_has_exact and page == 1: rows[0] = EXACT_TITLE_ROW body = {"results": rows, "next_page": page + 1} return Response(json.dumps(body).encode(), {"content-type": "application/json"}) def test_catalog_resolves_common_word_titles_through_one_phrase_query(): # "Attention Is All You Need" has over a thousand keyword hits; paging # through them used to end in "Too many results" although the exact title # was the first result. transport = TitleSearchTransport( phrase_results=[ EXACT_TITLE_ROW, {"id": "9", "arxiv_id": "2010.13154", "title": "Attention is All You Need in Speech Separation"}, ], plain_has_exact=True, ) catalog = CatalogClient(transport=transport) assert catalog.resolve_paper("Attention Is All You Need") == "1706.03762" assert [params["q"] for _path, params in transport.calls] == ['"Attention Is All You Need"'] def test_catalog_keeps_exact_matches_found_before_the_page_budget_runs_out(): transport = TitleSearchTransport(phrase_results=[], plain_has_exact=True) catalog = CatalogClient(transport=transport) assert catalog.resolve_paper("Attention Is All You Need") == "1706.03762" assert len(transport.calls) == 11 inconclusive = CatalogClient( transport=TitleSearchTransport(phrase_results=[], plain_has_exact=False) ) with pytest.raises(ResponseError, match="Too many results"): inconclusive.resolve_paper("Attention Is All You Need") def test_catalog_resolves_a_unique_short_name_before_the_title_colon(): # Production chat: get_paper_info("MuseTalk") failed although the catalog # holds "MuseTalk: Real-Time High-Fidelity Video Dubbing ..." (2410.10122). musetalk = { "id": "90", "arxiv_id": "2410.10122", "title": "MuseTalk: Real-Time High-Fidelity Video Dubbing via Spatio-Temporal Sampling", } citing = {"id": "91", "arxiv_id": "2501.00001", "title": "Beyond MuseTalk: Faster Dubbing"} transport = TitleSearchTransport(phrase_results=[musetalk, citing], plain_has_exact=False) catalog = CatalogClient(transport=transport) assert catalog.resolve_paper("MuseTalk") == "2410.10122" assert catalog.resolve_paper("musetalk") == "2410.10122" twins = [musetalk, {**musetalk, "id": "92", "arxiv_id": "2502.00002"}] ambiguous = CatalogClient( transport=TitleSearchTransport(phrase_results=twins, plain_has_exact=False) ) with pytest.raises(ResponseError, match="Paper title is ambiguous: MuseTalk"): ambiguous.resolve_paper("MuseTalk") def test_catalog_resolves_a_unique_near_match_of_a_full_title(): # Production chat: get_paper_info("BEIR: A Heterogeneous Benchmark ...") # failed although the catalog holds the arXiv spelling "Heterogenous" # (2104.08663), which keyword search returns first. beir = { "id": "5505", "arxiv_id": "2104.08663", "title": "BEIR: A Heterogenous Benchmark for Zero-shot Evaluation of Information Retrieval Models", } beir_nl = { "id": "93", "arxiv_id": "2412.08329", "title": "BEIR-NL: Zero-shot Information Retrieval Benchmark for the Dutch Language", } catalog = CatalogClient( transport=TitleSearchTransport(phrase_results=[beir, beir_nl], plain_has_exact=False) ) assert catalog.resolve_paper( "BEIR: A Heterogeneous Benchmark for Zero-shot Evaluation of Information Retrieval Models" ) == "2104.08663" twins = [beir, {**beir, "id": "94", "arxiv_id": "2502.00003"}] ambiguous = CatalogClient( transport=TitleSearchTransport(phrase_results=twins, plain_has_exact=False) ) with pytest.raises(ResponseError, match="Paper title is ambiguous"): ambiguous.resolve_paper( "BEIR: A Heterogeneous Benchmark for Zero-shot Evaluation of Information Retrieval Models" ) def test_near_title_rejects_short_names_and_other_editions(): assert near_title( normalize_title("BEIR: A Heterogenous Benchmark for Zero-shot Evaluation"), normalize_title("BEIR: A Heterogeneous Benchmark for Zero-shot Evaluation"), ) assert not near_title(normalize_title("MuseTalks"), normalize_title("MuseTalk")) assert not near_title( normalize_title("The Llama 3 Herd of Language Foundation Models"), normalize_title("The Llama 4 Herd of Language Foundation Models"), ) assert not near_title( normalize_title("Deep Residual Learning for Image Recognition"), normalize_title("Deep Residual Learning for Video Recognition"), ) def test_catalog_reads_public_metadata_and_caches_lookups_including_misses(): from pwc_cli.transport import HTTPStatusError class PublicTransport: def __init__(self): self.calls = [] def get(self, path, params=None): self.calls.append((path, params)) if path == "papers/2610.18869/public-metadata": return Response( json.dumps({"in_catalog": False, "title": "Found"}).encode(), {} ) if path == "papers/public-metadata/search": title = params["title"] if title == "Down Title": raise HTTPStatusError(503, "unavailable") if title == "BERT": raise HTTPStatusError(400, "title too short") if title == "Some Title": return Response( json.dumps({"in_catalog": False, "title": "Some title"}).encode(), {}, ) raise HTTPStatusError(404, "no exact title") if path.startswith("papers/Down/"): raise HTTPStatusError(503, "unavailable") if "%20" in path or path == "papers/BERT/public-metadata": raise HTTPStatusError(400, "titles are not identifiers") raise HTTPStatusError(404, "Paper not found") transport = PublicTransport() catalog = CatalogClient(transport=transport) assert catalog.public_paper(" 2610.18869 ")["title"] == "Found" assert catalog.public_paper("2610.18869")["title"] == "Found" assert catalog.public_paper("10.1/x") == PublicPaperMiss("identifier") assert catalog.public_paper("10.1/x") == PublicPaperMiss("identifier") # Titles the identifier route rejects are searched as exact titles. assert catalog.public_paper("Some Title") == { "in_catalog": False, "title": "Some title", "requested_title": "Some Title", } assert catalog.public_paper("Some Title")["title"] == "Some title" assert catalog.public_paper("Unknown Title") == PublicPaperMiss("title") assert catalog.public_paper("Unknown Title") == PublicPaperMiss("title") assert catalog.public_paper("BERT") == PublicPaperMiss("unsupported") # Upstream failures are best effort and not cached as misses. assert catalog.public_paper("Down") == PublicPaperMiss("unavailable") assert catalog.public_paper("Down") == PublicPaperMiss("unavailable") assert catalog.public_paper("Down Title") == PublicPaperMiss("unavailable") assert catalog.public_paper("Down Title") == PublicPaperMiss("unavailable") search = "papers/public-metadata/search" assert transport.calls == [ ("papers/2610.18869/public-metadata", None), ("papers/10.1/x/public-metadata", None), ("papers/Some%20Title/public-metadata", None), (search, {"title": "Some Title"}), ("papers/Unknown%20Title/public-metadata", None), (search, {"title": "Unknown Title"}), ("papers/BERT/public-metadata", None), (search, {"title": "BERT"}), ("papers/Down/public-metadata", None), ("papers/Down/public-metadata", None), ("papers/Down%20Title/public-metadata", None), (search, {"title": "Down Title"}), ("papers/Down%20Title/public-metadata", None), (search, {"title": "Down Title"}), ]