From d9d7d9bd7471057ad910d8844c462957f1031ae2 Mon Sep 17 00:00:00 2001 From: Andrea Gariboldi Date: Thu, 3 Sep 2026 17:32:45 +0200 Subject: [PATCH 1/7] Create paper_fetch tool --- envs/environment.yaml | 4 +- src/agentomics/agents/paper_fetch.py | 90 +++++++++++++++++++++ src/agentomics/runtime/run_lifecycle.py | 2 +- src/agentomics/tools/paper_fetch_tool.py | 38 +++++++++ src/agentomics/tools/save_paper_pdf_tool.py | 40 +++++++++ src/agentomics/tools/tool_registry.py | 20 +++-- src/agentomics/utils/config.py | 12 +++ 7 files changed, 199 insertions(+), 7 deletions(-) create mode 100644 src/agentomics/agents/paper_fetch.py create mode 100644 src/agentomics/tools/paper_fetch_tool.py create mode 100644 src/agentomics/tools/save_paper_pdf_tool.py diff --git a/envs/environment.yaml b/envs/environment.yaml index 24d40f18..93cf8fe4 100644 --- a/envs/environment.yaml +++ b/envs/environment.yaml @@ -6,7 +6,9 @@ dependencies: - conda-pack==0.9.1 - pandas==2.3.1 - matplotlib==3.8.4 - - pydantic-ai==1.53.0 + - pydantic-ai==1.77.0 + - markdownify==1.2.3 + - ddgs==9.14.4 - pip - pip: - weave==0.52.14 # Often breaks with new versions, update as necessary diff --git a/src/agentomics/agents/paper_fetch.py b/src/agentomics/agents/paper_fetch.py new file mode 100644 index 00000000..594a40e4 --- /dev/null +++ b/src/agentomics/agents/paper_fetch.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +from pathlib import Path + +from pydantic import BaseModel, Field +from pydantic_ai import Agent, ModelRetry +from pydantic_ai.models import Model +from pydantic_ai.settings import ModelSettings +from pydantic_ai.common_tools.web_fetch import web_fetch_tool + +from agentomics.tools.save_paper_pdf_tool import create_save_paper_pdf_tool, is_pdf +from agentomics.utils.config import Config + + +PAPER_FETCH_SYSTEM_PROMPT = """ +Your goal is to find scientific papers informing a machine learning modeling task. + +Work in three phases: +1. Check the returned search results for relevance to the query. If the same paper appears several times, choose a single host to retrieve it from. +2. Read. Use web_fetch on the web pages the search returned to confirm what each paper actually + is and to find its PDF link. Publisher landing pages state the real PDF URL; read it rather than + assembling one from the article id. Never report a paper you have not fetched. +3. Retrieve. Call save_paper_pdf for each paper relevant to the query. It writes the PDF to disk; you do not + need to read the PDF contents yourself. If a URL returns a download page instead of a PDF, go back + to the landing page for the direct link rather than guessing. + +Report only papers whose PDF you successfully saved. +""" + +class Paper(BaseModel): + title: str = Field( + description="The title of the paper." + ) + pdf_url: str | None = Field( + default=None, + description="The URL the PDF was downloaded from.") + pdf_filename: str | None = Field( + default=None, + description="Filename of the saved PDF, as returned by save_paper_pdf. " + "Empty if the PDF could not be retrieved.") + relevance: str = Field( + description="Why this paper is relevant to the query, and what it contributes to the modeling task.") + +class PaperFetchOutput(BaseModel): + papers: list[Paper] = Field( + description="List of relevant fetched papers. Empty if no relevant papers were retrieved or saved.") + +def create_paper_fetch_agent(config: Config, model: Model) -> Agent[None, PaperFetchOutput]: + agent = Agent( + model=model, + system_prompt=PAPER_FETCH_SYSTEM_PROMPT, + tools=[ + web_fetch_tool(timeout=config.web_fetch_timeout), + create_save_paper_pdf_tool(config), + ], + model_settings=ModelSettings(temperature=config.temperature), + output_type=PaperFetchOutput, + retries=config.max_validation_retries + ) + attach_output_validator(agent, config) + return agent + +def attach_output_validator(agent: Agent[None, PaperFetchOutput], config: Config) -> None: + @agent.output_validator + async def _papers_stored_in_shared_dir(result: PaperFetchOutput) -> PaperFetchOutput: + papers_dir = config.fetched_papers_dir.resolve() + + for paper in result.papers: + if paper.pdf_filename is None: + continue + # save_paper_pdf stores the basename, so accept a full path reported back. + pdf_path = papers_dir / Path(paper.pdf_filename).name + if not pdf_path.is_file(): + raise ModelRetry(f"'{paper.pdf_filename}' is not in {papers_dir}. Save it with save_paper_pdf or clear pdf_filename.") + with pdf_path.open("rb") as pdf_file: + header = pdf_file.read(5) + if not is_pdf(header): + raise ModelRetry(f"'{paper.pdf_filename}' is not a PDF. Remove it and clear pdf_filename.") + + stray = [str(p) for p in config.current_step_dir.resolve().rglob("*.pdf") if p.parent != papers_dir] + if stray: + raise ModelRetry(f"PDFs must live only in {papers_dir}, but these are elsewhere: {stray[:10]}") + return result + +def build_prompt(config: Config, query: str, websearch_result: list[dict]) -> str: + return f"""Based on the query received as input: + {query} + and the results of web search for that query: + {websearch_result} + - Save the pdf of the papers at: {config.fetched_papers_dir}""" diff --git a/src/agentomics/runtime/run_lifecycle.py b/src/agentomics/runtime/run_lifecycle.py index 4c5fa329..6b5dbcc5 100644 --- a/src/agentomics/runtime/run_lifecycle.py +++ b/src/agentomics/runtime/run_lifecycle.py @@ -33,7 +33,7 @@ async def run_agentomics(config: Config, default_model: Model, iteration_plan_model: Model, provider: Provider) -> None: total_iterations = config.iterations initialize_shared_environment(config) - tools = create_tools(config, config.tool_ids) + tools = create_tools(config, config.tool_ids, model=default_model) print(f"Starting training loop with {total_iterations} iterations") #TODO for run continuing this is wrong? while True: diff --git a/src/agentomics/tools/paper_fetch_tool.py b/src/agentomics/tools/paper_fetch_tool.py new file mode 100644 index 00000000..08ef97df --- /dev/null +++ b/src/agentomics/tools/paper_fetch_tool.py @@ -0,0 +1,38 @@ +from typing import Any + +from pydantic_ai import ModelRetry, RunContext, Tool +from pydantic_ai.exceptions import UnexpectedModelBehavior, UsageLimitExceeded +from pydantic_ai.models import Model +from pydantic_ai.usage import UsageLimits +from pydantic_ai.common_tools.duckduckgo import duckduckgo_search_tool +from ddgs.exceptions import DDGSException + +from agentomics.agents.paper_fetch import create_paper_fetch_agent, PaperFetchOutput, build_prompt +from agentomics.utils.config import Config + + +def create_paper_fetch_tool(config: Config, model: Model) -> Tool[Any]: + async def paper_fetch(ctx: RunContext[dict], query: str) -> PaperFetchOutput: + """ + A tool used to fetch scientific papers relevant to the provided query. + The fetched papers are saved in PDF format in the fetched_papers/ , inside the shared folder . + This tool should be called with one query at a time; call this tool multiple times if multiple topics need to be searched. + + Args: + query: A string describing one scientific topic or question for which relevant papers are to be fetched. + """ + try: + websearch_result = await duckduckgo_search_tool(max_results=config.paper_fetch_max_results).function(query=query) + except DDGSException as e: + raise ModelRetry(f"Web search is currently unavailable: {e}. Retry paper_fetch or continue without literature.") from e + agent = create_paper_fetch_agent(config, model) + try: + result = await agent.run( + user_prompt=build_prompt(config, query, websearch_result), + usage_limits=UsageLimits(request_limit=config.paper_fetch_request_limit), + ) + except (UnexpectedModelBehavior, UsageLimitExceeded) as e: + raise ModelRetry(f"paper_fetch tool could not complete the search: {e}. Continue without it or retry with a different query.") from e + return result.output + + return Tool(paper_fetch, name="paper_fetch") diff --git a/src/agentomics/tools/save_paper_pdf_tool.py b/src/agentomics/tools/save_paper_pdf_tool.py new file mode 100644 index 00000000..4b68b746 --- /dev/null +++ b/src/agentomics/tools/save_paper_pdf_tool.py @@ -0,0 +1,40 @@ +from typing import Any +from pathlib import Path + +from pydantic_ai import Tool +from pydantic_ai.messages import BinaryContent +from pydantic_ai.exceptions import ModelRetry +from pydantic_ai.common_tools.web_fetch import web_fetch_tool + +from agentomics.utils.config import Config + + +def create_save_paper_pdf_tool(config: Config) -> Tool[Any]: + fetch = web_fetch_tool(timeout=config.web_fetch_timeout).function + + async def save_paper_pdf(url: str, filename: str) -> str: + """Downloads the PDF at the given URL and saves it in the papers folder. + + Args: + url: Direct URL to the PDF. + filename: Name to save it under, ending in .pdf. Directories are ignored; + the tool decides where the file goes. + """ + response = await fetch(url) + if not isinstance(response, BinaryContent) or not is_pdf(response.data): + raise ModelRetry( + f"{url} did not return a PDF. Some sites serve an HTML download page instead; " + "find the direct PDF link and retry." + ) + + papers_dir = config.fetched_papers_dir + papers_dir.mkdir(parents=True, exist_ok=True) + pdf_path = papers_dir / Path(filename).name + pdf_path.write_bytes(response.data) + return f"Saved {len(response.data)} bytes to {pdf_path}" + + return Tool(save_paper_pdf, name="save_paper_pdf") + +def is_pdf(data: bytes) -> bool: + """A PDF always starts with the bytes `%PDF-`, whatever content type the server claims.""" + return data.startswith(b"%PDF-") diff --git a/src/agentomics/tools/tool_registry.py b/src/agentomics/tools/tool_registry.py index 207db43f..0748cffa 100644 --- a/src/agentomics/tools/tool_registry.py +++ b/src/agentomics/tools/tool_registry.py @@ -5,28 +5,38 @@ import weave from pydantic_ai import Tool +from pydantic_ai.models import Model from .bash_tool import create_bash_tool +from .paper_fetch_tool import create_paper_fetch_tool from .replace_tool import create_replace_tool from .run_python_tool import create_run_python_tool from .write_python_tool import create_write_python_tool from agentomics.utils.config import Config -TOOL_FACTORY_REGISTRY: dict[str, Callable] = { +TOOL_FACTORY_REGISTRY: dict[str, Callable[[Config], Tool[Any]]] = { "bash": create_bash_tool, "write_python": create_write_python_tool, "run_python": create_run_python_tool, "replace": create_replace_tool, } -def create_tool(config: Config, tool_id: str) -> Tool[Any]: - tool = TOOL_FACTORY_REGISTRY[tool_id](config) +# Tools delegating to sub-agents, model definition needed +MODEL_TOOL_FACTORY_REGISTRY: dict[str, Callable[[Config, Model], Tool[Any]]] = { + "paper_fetch": create_paper_fetch_tool, +} + +def create_tool(config: Config, tool_id: str, model: Model | None = None) -> Tool[Any]: + if tool_id in MODEL_TOOL_FACTORY_REGISTRY: + tool = MODEL_TOOL_FACTORY_REGISTRY[tool_id](config, model) + else: + tool = TOOL_FACTORY_REGISTRY[tool_id](config) tool.function = weave.op(tool.function, call_display_name=tool.name) return tool -def create_tools(config: Config, tool_ids: Iterable[str]) -> list[Tool[Any]]: - return [create_tool(config, tool_id) for tool_id in tool_ids] +def create_tools(config: Config, tool_ids: Iterable[str], model: Model | None = None) -> list[Tool[Any]]: + return [create_tool(config, tool_id, model) for tool_id in tool_ids] def get_tool_names(config: Config, tool_ids: Iterable[str]) -> list[str]: return [tool.name for tool in create_tools(config, tool_ids)] diff --git a/src/agentomics/utils/config.py b/src/agentomics/utils/config.py index 2a6a5602..5fe4a54f 100644 --- a/src/agentomics/utils/config.py +++ b/src/agentomics/utils/config.py @@ -24,6 +24,7 @@ class Config: ITERATION_STATE_FILENAME: ClassVar[str] = "iteration_state.json" ENVIRONMENT_DESCRIPTOR_FILENAME: ClassVar[str] = "environment.yml" STEP_OUTPUT_FILENAME: ClassVar[str] = "output.json" + FETCHED_PAPERS_DIRNAME: ClassVar[str] = "fetched_papers" DEFAULT_ITERATIONS: ClassVar[int] = 5 DEFAULT_SPLIT_ALLOWED_ITERATIONS: ClassVar[int] = 1 @@ -40,6 +41,7 @@ class Config: "validation_evaluation", ] DEFAULT_TOOL_IDS: ClassVar[list[str]] = [ + "paper_fetch", "bash", "write_python", "run_python", @@ -48,9 +50,12 @@ class Config: DEFAULT_MAX_STEPS: ClassVar[int] = 100 DEFAULT_TEMPERATURE: ClassVar[float] = 0.7 DEFAULT_MAX_VALIDATION_RETRIES: ClassVar[int] = 5 + DEFAULT_PAPER_FETCH_REQUEST_LIMIT: ClassVar[int] = 10 + DEFAULT_PAPER_FETCH_MAX_RESULTS: ClassVar[int] = 10 DEFAULT_USE_PROXY: ClassVar[bool] = True DEFAULT_LLM_RESPONSE_TIMEOUT: ClassVar[int] = 60 * 10 DEFAULT_BASH_TOOL_TIMEOUT: ClassVar[int] = 60 * 3 + DEFAULT_WEB_FETCH_TIMEOUT: ClassVar[int] = 60 DEFAULT_MAX_TOOL_RETRIES: ClassVar[int] = 5 DEFAULT_RUN_PYTHON_TOOL_TIMEOUT: ClassVar[int] = 60 * 60 * 6 DEFAULT_USER_PROMPT: ClassVar[str] = "Develop a machine learning model that generalizes well to new unseen data." @@ -92,6 +97,9 @@ class Config: use_proxy: bool = DEFAULT_USE_PROXY llm_response_timeout: int = DEFAULT_LLM_RESPONSE_TIMEOUT bash_tool_timeout: int = DEFAULT_BASH_TOOL_TIMEOUT + web_fetch_timeout: int = DEFAULT_WEB_FETCH_TIMEOUT + paper_fetch_request_limit: int = DEFAULT_PAPER_FETCH_REQUEST_LIMIT + paper_fetch_max_results: int = DEFAULT_PAPER_FETCH_MAX_RESULTS max_tool_retries: int = DEFAULT_MAX_TOOL_RETRIES label_to_scalar: dict[str, int] | None = None wandb_run_id: str | None = None @@ -141,6 +149,10 @@ def current_iteration_dir(self) -> Path: def current_step_dir(self) -> Path: return self.current_iteration_dir / "current_step" + @property + def fetched_papers_dir(self) -> Path: + return self.shared_dir / self.FETCHED_PAPERS_DIRNAME + @property def current_iteration_runtime_info_dir(self) -> Path: return self.current_iteration_dir / self.RUNTIME_INFO_DIRNAME From 2cd41836d114f3cff83fed4c01c04b60acecf986 Mon Sep 17 00:00:00 2001 From: Andrea Gariboldi Date: Fri, 4 Sep 2026 11:30:21 +0200 Subject: [PATCH 2/7] Streamline paper_fetch validator --- src/agentomics/agents/paper_fetch.py | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/src/agentomics/agents/paper_fetch.py b/src/agentomics/agents/paper_fetch.py index 594a40e4..f099fdc9 100644 --- a/src/agentomics/agents/paper_fetch.py +++ b/src/agentomics/agents/paper_fetch.py @@ -8,7 +8,7 @@ from pydantic_ai.settings import ModelSettings from pydantic_ai.common_tools.web_fetch import web_fetch_tool -from agentomics.tools.save_paper_pdf_tool import create_save_paper_pdf_tool, is_pdf +from agentomics.tools.save_paper_pdf_tool import create_save_paper_pdf_tool from agentomics.utils.config import Config @@ -68,18 +68,9 @@ async def _papers_stored_in_shared_dir(result: PaperFetchOutput) -> PaperFetchOu for paper in result.papers: if paper.pdf_filename is None: continue - # save_paper_pdf stores the basename, so accept a full path reported back. pdf_path = papers_dir / Path(paper.pdf_filename).name if not pdf_path.is_file(): raise ModelRetry(f"'{paper.pdf_filename}' is not in {papers_dir}. Save it with save_paper_pdf or clear pdf_filename.") - with pdf_path.open("rb") as pdf_file: - header = pdf_file.read(5) - if not is_pdf(header): - raise ModelRetry(f"'{paper.pdf_filename}' is not a PDF. Remove it and clear pdf_filename.") - - stray = [str(p) for p in config.current_step_dir.resolve().rglob("*.pdf") if p.parent != papers_dir] - if stray: - raise ModelRetry(f"PDFs must live only in {papers_dir}, but these are elsewhere: {stray[:10]}") return result def build_prompt(config: Config, query: str, websearch_result: list[dict]) -> str: From 9eb5c6ac92c0be26bcf7e9e3d694b14606ef55d1 Mon Sep 17 00:00:00 2001 From: Andrea Gariboldi Date: Fri, 4 Sep 2026 11:40:17 +0200 Subject: [PATCH 3/7] Revamp paper_fetch prompt --- src/agentomics/agents/paper_fetch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/agentomics/agents/paper_fetch.py b/src/agentomics/agents/paper_fetch.py index f099fdc9..2b533f45 100644 --- a/src/agentomics/agents/paper_fetch.py +++ b/src/agentomics/agents/paper_fetch.py @@ -78,4 +78,4 @@ def build_prompt(config: Config, query: str, websearch_result: list[dict]) -> st {query} and the results of web search for that query: {websearch_result} - - Save the pdf of the papers at: {config.fetched_papers_dir}""" + - Save the pdf of relevant papers at: {config.fetched_papers_dir}""" From 608b76895556406a61eaf96eefb2c25be5e43659 Mon Sep 17 00:00:00 2001 From: Andrea Gariboldi Date: Mon, 12 Oct 2026 08:53:21 +0000 Subject: [PATCH 4/7] Update paper fetch tool --- envs/environment.yaml | 1 + src/agentomics/cli/run.py | 1 + src/agentomics/runtime/run_lifecycle.py | 2 +- src/agentomics/tools/paper_fetch_tool.py | 379 +++++++++++++++++++++-- src/agentomics/tools/tool_registry.py | 18 +- src/agentomics/utils/config.py | 5 +- 6 files changed, 365 insertions(+), 41 deletions(-) diff --git a/envs/environment.yaml b/envs/environment.yaml index 93cf8fe4..76689f0b 100644 --- a/envs/environment.yaml +++ b/envs/environment.yaml @@ -23,3 +23,4 @@ dependencies: - openai==2.16.0 - anthropic==0.79.0 - pytest==9.1.1 + - mcp==1.28.0 diff --git a/src/agentomics/cli/run.py b/src/agentomics/cli/run.py index 4b6d8a5f..c80f36d9 100644 --- a/src/agentomics/cli/run.py +++ b/src/agentomics/cli/run.py @@ -108,6 +108,7 @@ def _build_container_arguments(arguments: argparse.Namespace) -> list[str]: def _docker_environment_arguments(agent_id: str) -> list[str]: docker_arguments = docker_environment_arguments( "OPENAI_API_KEY", "ANTHROPIC_API_KEY", "OPENROUTER_API_KEY", + "ASTA_API_KEY", "S2_API_KEY", "UNPAYWALL_EMAIL", "PROVISIONING_OPENROUTER_API_KEY", "OLLAMA_BASE_URL", "WANDB_API_KEY", "WANDB_PROJECT_NAME", "WANDB_ENTITY", "HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "NO_PROXY", diff --git a/src/agentomics/runtime/run_lifecycle.py b/src/agentomics/runtime/run_lifecycle.py index 6b5dbcc5..4c5fa329 100644 --- a/src/agentomics/runtime/run_lifecycle.py +++ b/src/agentomics/runtime/run_lifecycle.py @@ -33,7 +33,7 @@ async def run_agentomics(config: Config, default_model: Model, iteration_plan_model: Model, provider: Provider) -> None: total_iterations = config.iterations initialize_shared_environment(config) - tools = create_tools(config, config.tool_ids, model=default_model) + tools = create_tools(config, config.tool_ids) print(f"Starting training loop with {total_iterations} iterations") #TODO for run continuing this is wrong? while True: diff --git a/src/agentomics/tools/paper_fetch_tool.py b/src/agentomics/tools/paper_fetch_tool.py index 08ef97df..6992d2c0 100644 --- a/src/agentomics/tools/paper_fetch_tool.py +++ b/src/agentomics/tools/paper_fetch_tool.py @@ -1,38 +1,369 @@ +from __future__ import annotations + +import asyncio +import json +import os +from dataclasses import dataclass +from pathlib import Path from typing import Any +from urllib.parse import quote -from pydantic_ai import ModelRetry, RunContext, Tool -from pydantic_ai.exceptions import UnexpectedModelBehavior, UsageLimitExceeded -from pydantic_ai.models import Model -from pydantic_ai.usage import UsageLimits -from pydantic_ai.common_tools.duckduckgo import duckduckgo_search_tool -from ddgs.exceptions import DDGSException +import httpx +from pydantic import BaseModel +from pydantic_ai import Tool +from pydantic_ai.mcp import MCPServerStreamableHTTP -from agentomics.agents.paper_fetch import create_paper_fetch_agent, PaperFetchOutput, build_prompt from agentomics.utils.config import Config -def create_paper_fetch_tool(config: Config, model: Model) -> Tool[Any]: - async def paper_fetch(ctx: RunContext[dict], query: str) -> PaperFetchOutput: +class RetrievedPaper(BaseModel): + title: str + year: int | None + file_path: str + +@dataclass(frozen=True) +class PaperSearchResult: + paper_id: str + title: str + year: int | None + + @classmethod + def from_asta(cls, paper: dict[str, Any]) -> PaperSearchResult: + return cls( + paper_id=paper["paperId"], + title=paper["title"], + year=paper.get("year"), + ) + +def create_paper_fetch_tool(config: Config) -> Tool[Any]: + async def paper_fetch(query: str) -> list[RetrievedPaper]: """ - A tool used to fetch scientific papers relevant to the provided query. - The fetched papers are saved in PDF format in the fetched_papers/ , inside the shared folder . - This tool should be called with one query at a time; call this tool multiple times if multiple topics need to be searched. + Find and download open-access scientific papers relevant to a research question. + Save available PDFs. Each result includes the paper's title, + publication year, and local file_path. Args: - query: A string describing one scientific topic or question for which relevant papers are to be fetched. + query: One focused scientific topic or question to search for. """ + search_results = await search_asta_papers( + query=query, + limit=config.paper_fetch_max_results, + timeout=config.paper_fetch_timeout, + ) + return await download_papers( + papers=search_results, + destination=config.fetched_papers_dir, + timeout=config.paper_fetch_timeout, + ) + + return Tool(paper_fetch, name="paper_fetch") + +async def search_asta_papers( + query: str, limit: int, timeout: float +) -> list[PaperSearchResult]: + headers = {} + # An optional Asta key enables higher rate limits. + if api_key := os.getenv("ASTA_API_KEY"): + headers["x-api-key"] = api_key + + asta_mcp = MCPServerStreamableHTTP( + url="https://asta-tools.allen.ai/mcp/v1", + headers=headers, + timeout=timeout, + ) + result = await asta_mcp.direct_call_tool( + "search_papers_by_relevance", + { + "keyword": query, + "fields": "title,year", + "limit": limit, + }, + ) + if isinstance(result, str): + result = json.loads(result) + return [PaperSearchResult.from_asta(paper) for paper in result] + +async def download_papers( + papers: list[PaperSearchResult], destination: Path, timeout: float +) -> list[RetrievedPaper]: + destination.mkdir(parents=True, exist_ok=True) + + async with httpx.AsyncClient( + timeout=timeout, + follow_redirects=True, + headers={ + "User-Agent": "Mozilla/5.0 (compatible; Agentomics/1.0; scientific paper retrieval)", + }, + ) as client: + downloaded_papers = [] + papers_to_download = [] + for paper in papers: + pdf_path = destination / f"{paper.paper_id}.pdf" + if _is_pdf_file(pdf_path): + downloaded_papers.append( + RetrievedPaper( + title=paper.title, + year=paper.year, + file_path=str(pdf_path), + ) + ) + continue + + papers_to_download.append(paper) + + metadata_by_paper_id = await _get_semantic_scholar_metadata( + client, + [paper.paper_id for paper in papers_to_download], + ) + + for paper in papers_to_download: + pdf_path = destination / f"{paper.paper_id}.pdf" + candidates = await resolve_pdf_candidates( + client, + paper, + metadata_by_paper_id.get(paper.paper_id, {}), + ) + for candidate in candidates: + try: + response = await client.get(candidate) + response.raise_for_status() + if not response.content.startswith(b"%PDF-"): + continue + pdf_path.write_bytes(response.content) + except (httpx.HTTPError, OSError): + continue + + downloaded_papers.append( + RetrievedPaper( + title=paper.title, + year=paper.year, + file_path=str(pdf_path), + ) + ) + break + + return downloaded_papers + +async def resolve_pdf_candidates( + client: httpx.AsyncClient, + paper: PaperSearchResult, + metadata: dict[str, Any], +) -> list[str]: + candidates: list[str] = [] + + external_ids = metadata.get("externalIds") or {} + metadata_doi = external_ids.get("DOI") + resolved_doi = metadata_doi + arxiv_id = external_ids.get("ArXiv") or _arxiv_id_from_doi(metadata_doi) + pmcid = external_ids.get("PubMedCentral") + pdf_url = (metadata.get("openAccessPdf") or {}).get("url") or None + + candidates.extend(await _unpaywall_pdf_urls(client, metadata_doi)) + + if arxiv_id: + arxiv_id = quote(arxiv_id, safe="/.") + candidates.append(f"https://arxiv.org/pdf/{arxiv_id}") + + pmcid = pmcid or await _find_pmcid(client, metadata_doi) + if pmcid is None: + europe_pmc_doi, pmcid = await _find_europe_pmc_ids(client, paper.title) + resolved_doi = metadata_doi or europe_pmc_doi + # Europe PMC may identify a preprint with a different DOI. + if metadata_doi and europe_pmc_doi and europe_pmc_doi != metadata_doi: + candidates.extend(await _biorxiv_pdf_urls(client, europe_pmc_doi)) + + if pmcid: + pmcid = str(pmcid).upper() + if not pmcid.startswith("PMC"): + pmcid = f"PMC{pmcid}" + candidates.extend( + [ + f"https://europepmc.org/articles/{pmcid}?pdf=render", + "https://europepmc.org/backend/ptpmcrender.fcgi" + f"?accid={pmcid}&blobtype=pdf", + ] + ) + + candidates.extend(await _biorxiv_pdf_urls(client, resolved_doi)) + + if pdf_url: + candidates.append(pdf_url) + + candidates.extend(await _crossref_pdf_urls(client, resolved_doi)) + + return list(dict.fromkeys(candidates)) + +async def _get_semantic_scholar_metadata( + client: httpx.AsyncClient, + paper_ids: list[str], +) -> dict[str, dict[str, Any]]: + if not paper_ids: + return {} + + headers = {"Accept": "application/json"} + if api_key := os.getenv("S2_API_KEY"): + headers["x-api-key"] = api_key + + for attempt in range(3): try: - websearch_result = await duckduckgo_search_tool(max_results=config.paper_fetch_max_results).function(query=query) - except DDGSException as e: - raise ModelRetry(f"Web search is currently unavailable: {e}. Retry paper_fetch or continue without literature.") from e - agent = create_paper_fetch_agent(config, model) + response = await client.post( + "https://api.semanticscholar.org/graph/v1/paper/batch", + params={"fields": "externalIds,openAccessPdf"}, + headers=headers, + json={"ids": paper_ids}, + ) + if response.status_code != 429: + response.raise_for_status() + metadata = response.json() + break + except (httpx.HTTPError, ValueError): + return {} + + if attempt == 2: + return {} + await asyncio.sleep(3) + + return { + paper["paperId"]: paper + for paper in metadata + if paper and paper.get("paperId") + } + +async def _unpaywall_pdf_urls( + client: httpx.AsyncClient, + doi: str | None, +) -> list[str]: + email = os.getenv("UNPAYWALL_EMAIL") + if not doi or not email: + return [] + + try: + response = await client.get( + f"https://api.unpaywall.org/v2/{quote(doi, safe='/')}", + params={"email": email}, + headers={"Accept": "application/json"}, + ) + response.raise_for_status() + data = response.json() + except (httpx.HTTPError, ValueError): + return [] + + locations = [data.get("best_oa_location"), *(data.get("oa_locations") or [])] + return [ + location["url_for_pdf"] + for location in locations + if location and location.get("url_for_pdf") + ] + +def _arxiv_id_from_doi(doi: str | None) -> str | None: + prefix = "10.48550/arxiv." + if doi and doi.lower().startswith(prefix.lower()): + return doi[len(prefix):] + return None + +async def _find_pmcid(client: httpx.AsyncClient, doi: str | None) -> str | None: + if doi is None: + return None + + params = { + "ids": doi, + "idtype": "doi", + "format": "json", + "tool": "agentomics", + } + + try: + response = await client.get( + "https://pmc.ncbi.nlm.nih.gov/tools/idconv/api/v1/articles/", + params=params, + headers={"Accept": "application/json"}, + ) + response.raise_for_status() + records = response.json().get("records") or [] + except (httpx.HTTPError, ValueError): + return None + + return records[0].get("pmcid") if records else None + +async def _find_europe_pmc_ids( + client: httpx.AsyncClient, + title: str, +) -> tuple[str | None, str | None]: + try: + response = await client.get( + "https://www.ebi.ac.uk/europepmc/webservices/rest/search", + params={ + "query": f'TITLE:"{title}"', + "format": "json", + "resultType": "lite", + "pageSize": 1, + }, + headers={"Accept": "application/json"}, + ) + response.raise_for_status() + results = response.json().get("resultList", {}).get("result") or [] + except (httpx.HTTPError, ValueError): + return None, None + + if not results: + return None, None + + result = results[0] + returned_title = result.get("title", "").casefold().rstrip(".") + if returned_title != title.casefold().rstrip("."): + return None, None + + return result.get("doi"), result.get("pmcid") + +async def _biorxiv_pdf_urls( + client: httpx.AsyncClient, + doi: str | None, +) -> list[str]: + if doi is None or not doi.lower().startswith("10.1101/"): + return [] + + encoded_doi = quote(doi, safe="/") + for server in ("biorxiv", "medrxiv"): try: - result = await agent.run( - user_prompt=build_prompt(config, query, websearch_result), - usage_limits=UsageLimits(request_limit=config.paper_fetch_request_limit), + response = await client.get( + f"https://api.biorxiv.org/details/{server}/{encoded_doi}/na/json", + headers={"Accept": "application/json"}, ) - except (UnexpectedModelBehavior, UsageLimitExceeded) as e: - raise ModelRetry(f"paper_fetch tool could not complete the search: {e}. Continue without it or retry with a different query.") from e - return result.output + response.raise_for_status() + if response.json().get("collection"): + return [ + f"https://www.{server}.org/content/{encoded_doi}.full.pdf", + ] + except (httpx.HTTPError, ValueError): + continue - return Tool(paper_fetch, name="paper_fetch") + return [] + +async def _crossref_pdf_urls( + client: httpx.AsyncClient, + doi: str | None, +) -> list[str]: + if not doi: + return [] + + try: + response = await client.get( + f"https://api.crossref.org/works/{quote(doi, safe='')}", + headers={"Accept": "application/json"}, + ) + response.raise_for_status() + metadata = response.json()["message"] + except (httpx.HTTPError, ValueError, KeyError): + return [] + + return [ + link["URL"] + for link in metadata.get("link", []) + if link.get("URL") and link.get("content-type") == "application/pdf" + ] + +def _is_pdf_file(path: Path) -> bool: + if not path.is_file(): + return False + with path.open("rb") as pdf_file: + return pdf_file.read(5) == b"%PDF-" diff --git a/src/agentomics/tools/tool_registry.py b/src/agentomics/tools/tool_registry.py index 0748cffa..fca42a64 100644 --- a/src/agentomics/tools/tool_registry.py +++ b/src/agentomics/tools/tool_registry.py @@ -5,7 +5,6 @@ import weave from pydantic_ai import Tool -from pydantic_ai.models import Model from .bash_tool import create_bash_tool from .paper_fetch_tool import create_paper_fetch_tool @@ -16,27 +15,20 @@ from agentomics.utils.config import Config TOOL_FACTORY_REGISTRY: dict[str, Callable[[Config], Tool[Any]]] = { + "paper_fetch": create_paper_fetch_tool, "bash": create_bash_tool, "write_python": create_write_python_tool, "run_python": create_run_python_tool, "replace": create_replace_tool, } -# Tools delegating to sub-agents, model definition needed -MODEL_TOOL_FACTORY_REGISTRY: dict[str, Callable[[Config, Model], Tool[Any]]] = { - "paper_fetch": create_paper_fetch_tool, -} - -def create_tool(config: Config, tool_id: str, model: Model | None = None) -> Tool[Any]: - if tool_id in MODEL_TOOL_FACTORY_REGISTRY: - tool = MODEL_TOOL_FACTORY_REGISTRY[tool_id](config, model) - else: - tool = TOOL_FACTORY_REGISTRY[tool_id](config) +def create_tool(config: Config, tool_id: str) -> Tool[Any]: + tool = TOOL_FACTORY_REGISTRY[tool_id](config) tool.function = weave.op(tool.function, call_display_name=tool.name) return tool -def create_tools(config: Config, tool_ids: Iterable[str], model: Model | None = None) -> list[Tool[Any]]: - return [create_tool(config, tool_id, model) for tool_id in tool_ids] +def create_tools(config: Config, tool_ids: Iterable[str]) -> list[Tool[Any]]: + return [create_tool(config, tool_id) for tool_id in tool_ids] def get_tool_names(config: Config, tool_ids: Iterable[str]) -> list[str]: return [tool.name for tool in create_tools(config, tool_ids)] diff --git a/src/agentomics/utils/config.py b/src/agentomics/utils/config.py index 5fe4a54f..a7088260 100644 --- a/src/agentomics/utils/config.py +++ b/src/agentomics/utils/config.py @@ -52,10 +52,10 @@ class Config: DEFAULT_MAX_VALIDATION_RETRIES: ClassVar[int] = 5 DEFAULT_PAPER_FETCH_REQUEST_LIMIT: ClassVar[int] = 10 DEFAULT_PAPER_FETCH_MAX_RESULTS: ClassVar[int] = 10 + DEFAULT_PAPER_FETCH_TIMEOUT: ClassVar[int] = 120 DEFAULT_USE_PROXY: ClassVar[bool] = True DEFAULT_LLM_RESPONSE_TIMEOUT: ClassVar[int] = 60 * 10 DEFAULT_BASH_TOOL_TIMEOUT: ClassVar[int] = 60 * 3 - DEFAULT_WEB_FETCH_TIMEOUT: ClassVar[int] = 60 DEFAULT_MAX_TOOL_RETRIES: ClassVar[int] = 5 DEFAULT_RUN_PYTHON_TOOL_TIMEOUT: ClassVar[int] = 60 * 60 * 6 DEFAULT_USER_PROMPT: ClassVar[str] = "Develop a machine learning model that generalizes well to new unseen data." @@ -97,8 +97,7 @@ class Config: use_proxy: bool = DEFAULT_USE_PROXY llm_response_timeout: int = DEFAULT_LLM_RESPONSE_TIMEOUT bash_tool_timeout: int = DEFAULT_BASH_TOOL_TIMEOUT - web_fetch_timeout: int = DEFAULT_WEB_FETCH_TIMEOUT - paper_fetch_request_limit: int = DEFAULT_PAPER_FETCH_REQUEST_LIMIT + paper_fetch_timeout: int = DEFAULT_PAPER_FETCH_TIMEOUT paper_fetch_max_results: int = DEFAULT_PAPER_FETCH_MAX_RESULTS max_tool_retries: int = DEFAULT_MAX_TOOL_RETRIES label_to_scalar: dict[str, int] | None = None From f592f93afecb6bdee29187806b3fcc35eb9bd699 Mon Sep 17 00:00:00 2001 From: Andrea Gariboldi Date: Mon, 12 Oct 2026 09:58:25 +0000 Subject: [PATCH 5/7] Update docs --- .env.example | 7 +++++++ docs/configuration/environment.md | 30 ++++++++++++++++++++++++++++++ docs/user-guide/outputs.md | 14 +++++++++++++- 3 files changed, 50 insertions(+), 1 deletion(-) diff --git a/.env.example b/.env.example index cd546754..d861f23d 100644 --- a/.env.example +++ b/.env.example @@ -6,6 +6,13 @@ OPENROUTER_API_KEY=your_openrouter_api_key_here OPENAI_API_KEY=your_openai_api_key_here # OR See the Providers section in README.md for other supported providers +# Optional: Scientific paper fetching +# Asta and Semantic Scholar use separate API keys. +# ASTA_API_KEY=your_asta_api_key_here +# S2_API_KEY=your_semantic_scholar_api_key_here +# Set a contact email to enable Unpaywall DOI lookups. +# UNPAYWALL_EMAIL=you@your-institution.org + # Optional: Weights & Biases / Weave experiment tracking # Get your API key from: https://wandb.ai/authorize # Uncomment all three WANDB_* variables to enable logging. diff --git a/docs/configuration/environment.md b/docs/configuration/environment.md index 9f3725f8..4b35041d 100644 --- a/docs/configuration/environment.md +++ b/docs/configuration/environment.md @@ -50,6 +50,31 @@ If you use `--use-provisioning-key`, set: OPENROUTER_API_KEY=sk-or-v1-xxxxxxxxxxxx ``` +## Paper Fetching + +The paper-fetch tool reads these optional variables. Export them in your shell +or add them to `.env`; `agentomics-run` forwards them into Docker. + +| Variable | Purpose | +|----------|---------| +| `ASTA_API_KEY` | Authenticate Asta MCP search requests. This is separate from the Semantic Scholar key. | +| `S2_API_KEY` | Authenticate Semantic Scholar batch requests for paper identifiers and PDF metadata. | +| `UNPAYWALL_EMAIL` | Contact email sent with DOI lookups. Unpaywall is skipped if this is unset. | + +Asta and Semantic Scholar requests are attempted without keys when they are +unset, but may encounter rate limits or service failures. Crossref and the +repository resolvers do not require credentials in this implementation. + +```bash +# Optional paper-fetch credentials +# ASTA_API_KEY=your_asta_key +# S2_API_KEY=your_semantic_scholar_key +# UNPAYWALL_EMAIL=you@your-institution.org +``` + +See [Fetching Scientific Papers](../user-guide/running-agent.md#fetching-scientific-papers) +for the retrieval flow and its limitations. + ## Weights & Biases (W&B) Enable experiment tracking and logging: @@ -124,6 +149,11 @@ OPENROUTER_API_KEY=sk-or-v1-xxxxxxxxxxxx # OPENAI_API_KEY=sk-xxxxxxxxxxxx # ANTHROPIC_API_KEY=sk-ant-xxxxxxxxxxxx +# Paper fetching (optional) +# ASTA_API_KEY=your_asta_key +# S2_API_KEY=your_semantic_scholar_key +# UNPAYWALL_EMAIL=you@your-institution.org + # Weights & Biases (optional) WANDB_API_KEY=xxxxxxxxxxxxxxxxxxxxxxxx WANDB_PROJECT_NAME=agentomics diff --git a/docs/user-guide/outputs.md b/docs/user-guide/outputs.md index b0e71008..f8407f38 100644 --- a/docs/user-guide/outputs.md +++ b/docs/user-guide/outputs.md @@ -24,7 +24,8 @@ outputs// │ ├── shared/ │ │ ├── config.json │ │ ├── dataset_metadata.json -│ │ └── splits/ +│ │ ├── splits/ +│ │ └── fetched_papers/ # PDFs downloaded by paper_fetch, when used │ ├── iteration_0/ │ ├── iteration_1/ │ └── ... @@ -91,6 +92,17 @@ run/iteration_N/ └── ... # Other iteration artifacts ``` +## Fetched Papers + +The `paper_fetch` tool saves PDFs to +`outputs//run/shared/fetched_papers/` on the host. Each filename is the +Semantic Scholar paper ID followed by `.pdf`. The returned `file_path` points to +the file inside the run's container workspace. + +The folder is shared across iterations of the run. Existing files with a PDF +signature are reused on later calls for the same paper ID. Only successfully +saved or cached PDFs are returned to the agent; no HTML fallback is used. + ## Reports `reports/best_iteration.md` and `reports/best_iteration.pdf` are copies of the From 5e3cb85ed50ff9331c7d94f096f51fd701692277 Mon Sep 17 00:00:00 2001 From: Andrea Gariboldi Date: Mon, 12 Oct 2026 10:13:16 +0000 Subject: [PATCH 6/7] Add tests --- tests/unit/test_paper_fetch_tool.py | 147 ++++++++++++++++++++++++++++ 1 file changed, 147 insertions(+) create mode 100644 tests/unit/test_paper_fetch_tool.py diff --git a/tests/unit/test_paper_fetch_tool.py b/tests/unit/test_paper_fetch_tool.py new file mode 100644 index 00000000..b38a7154 --- /dev/null +++ b/tests/unit/test_paper_fetch_tool.py @@ -0,0 +1,147 @@ +import asyncio +from io import BytesIO + +import httpx +import pytest +from pypdf import PdfWriter + +from agentomics.tools import paper_fetch_tool +from agentomics.tools.paper_fetch_tool import PaperSearchResult, RetrievedPaper, download_papers + + +@pytest.mark.parametrize( + "external_ids,pdf_url,resolver_url,resolver_response", + [ + pytest.param( + {"ArXiv": "2501.12948"}, "https://arxiv.org/pdf/2501.12948", + None, None, id="arxiv", + ), + pytest.param( + {"PubMedCentral": "123456"}, + "https://europepmc.org/articles/PMC123456?pdf=render", + None, None, id="europe-pmc", + ), + pytest.param( + {"DOI": "10.1101/2024.01.01.123456"}, + "https://www.biorxiv.org/content/10.1101/2024.01.01.123456.full.pdf", + "https://api.biorxiv.org/details/biorxiv/10.1101/2024.01.01.123456/na/json", + {"collection": [{"doi": "10.1101/2024.01.01.123456"}]}, id="biorxiv", + ), + pytest.param( + {"DOI": "10.1101/2024.01.01.123456"}, + "https://www.medrxiv.org/content/10.1101/2024.01.01.123456.full.pdf", + "https://api.biorxiv.org/details/medrxiv/10.1101/2024.01.01.123456/na/json", + {"collection": [{"doi": "10.1101/2024.01.01.123456"}]}, id="medrxiv", + ), + pytest.param( + {}, "https://publisher.example/semantic-scholar.pdf", + None, None, id="semantic-scholar", + ), + pytest.param( + {"DOI": "10.1234/paper"}, "https://repository.example/unpaywall.pdf", + "https://api.unpaywall.org/v2/10.1234/paper", + {"best_oa_location": {"url_for_pdf": "https://repository.example/unpaywall.pdf"}}, + id="unpaywall", + ), + pytest.param( + {"DOI": "10.1234/paper"}, "https://publisher.example/crossref.pdf", + "https://api.crossref.org/works/10.1234%2Fpaper", + {"message": {"link": [{ + "URL": "https://publisher.example/crossref.pdf", + "content-type": "application/pdf", + }]}}, id="crossref", + ), + ], +) +def test_download_papers_saves_pdf_from_resource( + external_ids, pdf_url, resolver_url, resolver_response, tmp_path, monkeypatch, +): + monkeypatch.delenv("S2_API_KEY", raising=False) + monkeypatch.setenv("UNPAYWALL_EMAIL", "researcher@example.org") + paper = PaperSearchResult(paper_id="paper-id", title="Sequence models", year=2025) + metadata = {"paperId": paper.paper_id, "externalIds": external_ids} + if resolver_url is None and not external_ids: + metadata["openAccessPdf"] = {"url": pdf_url} + + pdf = BytesIO() + writer = PdfWriter() + writer.add_blank_page(width=72, height=72) + writer.write(pdf) + pdf_bytes = pdf.getvalue() + + def respond(request): + if request.method == "POST" and request.url.path == "/graph/v1/paper/batch": + return httpx.Response(200, json=[metadata]) + if str(request.url) == pdf_url: + return httpx.Response(200, content=pdf_bytes) + if str(request.url.copy_with(query=None)) == resolver_url: + return httpx.Response(200, json=resolver_response) + return httpx.Response(404) + + original_client = httpx.AsyncClient + monkeypatch.setattr( + paper_fetch_tool.httpx, "AsyncClient", + lambda **kwargs: original_client(transport=httpx.MockTransport(respond), **kwargs), + ) + + destination = tmp_path / "papers" + result = asyncio.run(download_papers([paper], destination, timeout=1)) + + saved_pdf = destination / "paper-id.pdf" + assert result == [RetrievedPaper(title=paper.title, year=paper.year, file_path=str(saved_pdf))] + assert saved_pdf.read_bytes() == pdf_bytes + + +@pytest.mark.parametrize( + "paper,pdf_url,status,body", + [ + pytest.param( + PaperSearchResult( + paper_id="f36f5e1578017e53df01a95ecc002a44906de433", + title="miRBench: novel benchmark datasets for microRNA binding site prediction", + year=2025, + ), + "https://academic.oup.com/bioinformatics/article-pdf/41/Supplement_1/i542/63745693/btaf233.pdf", + 403, b"Forbidden", id="oup-forbidden", + ), + pytest.param( + PaperSearchResult( + paper_id="2c3c17026138fee75aa0bf662dd25bbab601899b", + title="GraphTar: applying word2vec and graph neural networks to miRNA target prediction", + year=2023, + ), + "https://bmcbioinformatics.biomedcentral.com/counter/pdf/10.1186/s12859-023-05564-x", + 200, b"Verify you are human", id="bmc-html-response", + ), + ], +) +def test_download_papers_omits_paper_when_pdf_is_inaccessible( + paper, pdf_url, status, body, tmp_path, monkeypatch, +): + # Reproduce previously observed responses without depending on live availability. + monkeypatch.delenv("S2_API_KEY", raising=False) + monkeypatch.delenv("UNPAYWALL_EMAIL", raising=False) + requested_urls = [] + + def respond(request): + requested_urls.append(str(request.url)) + if request.method == "POST" and request.url.path == "/graph/v1/paper/batch": + return httpx.Response(200, json=[{ + "paperId": paper.paper_id, "openAccessPdf": {"url": pdf_url}, + }]) + if str(request.url) == pdf_url: + return httpx.Response(status, content=body) + return httpx.Response(404) + + original_client = httpx.AsyncClient + monkeypatch.setattr( + paper_fetch_tool.httpx, "AsyncClient", + lambda **kwargs: original_client(transport=httpx.MockTransport(respond), **kwargs), + ) + + destination = tmp_path / "papers" + result = asyncio.run(download_papers([paper], destination, timeout=1)) + + assert result == [] + assert pdf_url in requested_urls + assert list(destination.iterdir()) == [] From e742be4d1693fbec8a1ae11f04bc951e3c50842e Mon Sep 17 00:00:00 2001 From: Andrea Gariboldi Date: Mon, 12 Oct 2026 10:14:03 +0000 Subject: [PATCH 7/7] Use already installed library in tests --- tests/unit/test_paper_fetch_tool.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/unit/test_paper_fetch_tool.py b/tests/unit/test_paper_fetch_tool.py index b38a7154..0e5351cf 100644 --- a/tests/unit/test_paper_fetch_tool.py +++ b/tests/unit/test_paper_fetch_tool.py @@ -3,7 +3,7 @@ import httpx import pytest -from pypdf import PdfWriter +from reportlab.pdfgen.canvas import Canvas from agentomics.tools import paper_fetch_tool from agentomics.tools.paper_fetch_tool import PaperSearchResult, RetrievedPaper, download_papers @@ -64,9 +64,9 @@ def test_download_papers_saves_pdf_from_resource( metadata["openAccessPdf"] = {"url": pdf_url} pdf = BytesIO() - writer = PdfWriter() - writer.add_blank_page(width=72, height=72) - writer.write(pdf) + canvas = Canvas(pdf, pagesize=(72, 72), invariant=1) + canvas.showPage() + canvas.save() pdf_bytes = pdf.getvalue() def respond(request):