From 04b06e0e8b7e40dfaeaafa67bff79e0f1d14b89d Mon Sep 17 00:00:00 2001 From: Pankaj Rajdeo Date: Tue, 6 Oct 2026 03:09:51 -0400 Subject: [PATCH 1/2] feat: add paper retrieval tools --- docs/how-it-works/architecture.md | 2 + envs/environment.yaml | 2 + pyproject.toml | 1 + src/agentomics/runtime/paper_retrieval.py | 148 ++++++++++++++++++ src/agentomics/tools/inspect_image_tool.py | 51 ++++++ .../tools/retrieve_paper_chunks_tool.py | 29 ++++ src/agentomics/tools/tool_registry.py | 4 + src/agentomics/utils/config.py | 12 ++ tests/unit/test_paper_retrieval_tools.py | 67 ++++++++ 9 files changed, 316 insertions(+) create mode 100644 src/agentomics/runtime/paper_retrieval.py create mode 100644 src/agentomics/tools/inspect_image_tool.py create mode 100644 src/agentomics/tools/retrieve_paper_chunks_tool.py create mode 100644 tests/unit/test_paper_retrieval_tools.py diff --git a/docs/how-it-works/architecture.md b/docs/how-it-works/architecture.md index feaee42c..0309bd9f 100644 --- a/docs/how-it-works/architecture.md +++ b/docs/how-it-works/architecture.md @@ -133,6 +133,8 @@ The agent has access to these tools: | Tool | Purpose | |------|---------| +| **Retrieve Paper Chunks** | Find relevant passages in fetched scientific papers | +| **Inspect Image** | Inspect paper pages and other images from the current run | | **Bash** | Execute shell commands | | **Write Python** | Create Python files | | **Run Python** | Execute Python scripts | diff --git a/envs/environment.yaml b/envs/environment.yaml index 30d08979..24b0f946 100644 --- a/envs/environment.yaml +++ b/envs/environment.yaml @@ -21,3 +21,5 @@ dependencies: - openai==3.29.0 - anthropic==1.13.0 - pytest==9.1.1 + - pypdf==6.19.0 + - pypdfium2==5.14.0 diff --git a/pyproject.toml b/pyproject.toml index d889a861..f23d523f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,7 @@ dependencies = [ "rich>=14.0", "pytest==9.1.1", "pypdf>=6.0", + "pypdfium2>=5.14", ] license = { file = "LICENSE" } authors = [ diff --git a/src/agentomics/runtime/paper_retrieval.py b/src/agentomics/runtime/paper_retrieval.py new file mode 100644 index 00000000..07189948 --- /dev/null +++ b/src/agentomics/runtime/paper_retrieval.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +import hashlib +import json +from pathlib import Path + +import pypdfium2 as pdfium +from pypdf import PdfReader +from sklearn.feature_extraction.text import TfidfVectorizer +from sklearn.metrics.pairwise import cosine_similarity + +from agentomics.utils.config import Config + + +CHUNK_WORDS = 350 +CHUNK_OVERLAP_WORDS = 50 + + +def retrieve_paper_chunks( + config: Config, + query: str, + max_results: int, +) -> list[dict[str, str | int | float]]: + query = query.strip() + if not query: + raise ValueError("query cannot be empty") + + chunks: list[dict[str, str | int]] = [] + for pdf_path in sorted(config.fetched_papers_dir.glob("*.pdf")): + paper_dir, paper_chunks = _prepare_pdf(config, pdf_path) + chunks.extend( + { + "paper": pdf_path.name, + **chunk, + "markdown_path": str(paper_dir / "document.md"), + } + for chunk in paper_chunks + ) + + if not chunks: + return [] + + vectorizer = TfidfVectorizer( + stop_words="english", + ngram_range=(1, 2), + sublinear_tf=True, + ) + chunk_matrix = vectorizer.fit_transform(str(chunk["text"]) for chunk in chunks) + scores = cosine_similarity(vectorizer.transform([query]), chunk_matrix).ravel() + selected = [index for index in scores.argsort()[::-1] if scores[index] > 0][:max_results] + + results: list[dict[str, str | int | float]] = [] + for index in selected: + chunk = chunks[int(index)] + pdf_path = config.fetched_papers_dir / str(chunk["paper"]) + page_number = int(chunk["page"]) + page_image_path = _page_image_path( + Path(str(chunk["markdown_path"])).parent, + page_number, + ) + if not page_image_path.exists(): + render_pdf_page(pdf_path, page_number, page_image_path) + results.append( + { + **chunk, + "score": float(scores[index]), + "image_path": str(page_image_path), + } + ) + return results + + +def render_pdf_page( + pdf_path: Path, + page_number: int, + image_path: Path, +) -> None: + document = pdfium.PdfDocument(pdf_path) + if not 1 <= page_number <= len(document): + document.close() + raise ValueError(f"page {page_number} is outside {pdf_path.name}") + + page = document[page_number - 1] + try: + bitmap = page.render(scale=2) + try: + image_path.parent.mkdir(parents=True, exist_ok=True) + bitmap.to_pil().save(image_path, format="PNG") + finally: + bitmap.close() + finally: + page.close() + document.close() + + +def _prepare_pdf( + config: Config, + pdf_path: Path, +) -> tuple[Path, list[dict[str, str | int]]]: + with pdf_path.open("rb") as pdf_file: + digest = hashlib.file_digest(pdf_file, "sha256").hexdigest()[:12] + paper_dir = config.processed_papers_dir / f"{pdf_path.stem}-{digest}" + paper_dir.mkdir(parents=True, exist_ok=True) + document_path = paper_dir / "document.md" + chunks_path = paper_dir / "chunks.json" + if document_path.exists() and chunks_path.exists(): + return paper_dir, json.loads(chunks_path.read_text(encoding="utf-8")) + + document_pages: list[str] = [] + chunks: list[dict[str, str | int]] = [] + for page_number, page in enumerate(PdfReader(pdf_path).pages, start=1): + text = ( + (page.extract_text(extraction_mode="layout") or "").strip() + if page.get_contents() is not None + else "" + ) + document_pages.append(f"## Page {page_number}\n\n{text}") + chunks.extend( + { + "page": page_number, + "chunk": chunk_number, + "text": chunk, + } + for chunk_number, chunk in enumerate(_split_text(text), start=1) + ) + + document_path.write_text( + f"# {pdf_path.name}\n\n" + "\n\n".join(document_pages) + "\n", + encoding="utf-8", + ) + chunks_path.write_text( + json.dumps(chunks, ensure_ascii=False), + encoding="utf-8", + ) + return paper_dir, chunks + + +def _split_text(text: str) -> list[str]: + words = text.split() + step = CHUNK_WORDS - CHUNK_OVERLAP_WORDS + return [ + " ".join(words[start : start + CHUNK_WORDS]) + for start in range(0, len(words), step) + ] + + +def _page_image_path(paper_dir: Path, page_number: int) -> Path: + return paper_dir / "pages" / f"page_{page_number:04d}.png" diff --git a/src/agentomics/tools/inspect_image_tool.py b/src/agentomics/tools/inspect_image_tool.py new file mode 100644 index 00000000..36f3832e --- /dev/null +++ b/src/agentomics/tools/inspect_image_tool.py @@ -0,0 +1,51 @@ +from pathlib import Path +from typing import Any + +from pydantic_ai import ModelRetry, Tool +from pydantic_ai.messages import BinaryContent, ToolReturn + +from agentomics.utils.config import Config + + +IMAGE_MEDIA_TYPES = { + ".gif": "image/gif", + ".jpeg": "image/jpeg", + ".jpg": "image/jpeg", + ".png": "image/png", + ".webp": "image/webp", +} +MAX_IMAGE_BYTES = 20 * 1024 * 1024 + + +def create_inspect_image_tool(config: Config) -> Tool[Any]: + def _inspect_image(image_path: str, query: str) -> ToolReturn: + """Send an image to the configured model, which must support image input. + + Args: + image_path: Absolute path to an image inside the run directory. + query: Question to answer using the image. + """ + path = Path(image_path).resolve() + if not path.is_relative_to(config.run_dir.resolve()): + raise ModelRetry(f"Image must be inside {config.run_dir}") + + media_type = IMAGE_MEDIA_TYPES.get(path.suffix.lower()) + if not path.is_file() or media_type is None: + raise ModelRetry(f"Unsupported or missing image: {image_path}") + if path.stat().st_size > MAX_IMAGE_BYTES: + raise ModelRetry(f"Image exceeds {MAX_IMAGE_BYTES // (1024 * 1024)} MB: {image_path}") + + return ToolReturn( + return_value=f"Loaded {path}", + content=[ + query, + BinaryContent(data=path.read_bytes(), media_type=media_type), + ], + ) + + return Tool( + _inspect_image, + name="inspect_image", + max_retries=config.max_tool_retries, + require_parameter_descriptions=True, + ) diff --git a/src/agentomics/tools/retrieve_paper_chunks_tool.py b/src/agentomics/tools/retrieve_paper_chunks_tool.py new file mode 100644 index 00000000..92decf33 --- /dev/null +++ b/src/agentomics/tools/retrieve_paper_chunks_tool.py @@ -0,0 +1,29 @@ +from typing import Annotated, Any + +from pydantic import Field +from pydantic_ai import Tool + +from agentomics.runtime.paper_retrieval import retrieve_paper_chunks +from agentomics.utils.config import Config + + +def create_retrieve_paper_chunks_tool(config: Config) -> Tool[Any]: + def _retrieve_paper_chunks( + query: Annotated[str, Field(min_length=1)], + max_results: Annotated[int, Field(ge=1, le=10)] = 5, + ) -> list[dict[str, str | int | float]]: + """Find passages relevant to a query in the fetched scientific papers. + + Args: + query: Information to find in the fetched papers. + max_results: Maximum number of passages to return. + """ + return retrieve_paper_chunks(config, query, max_results) + + return Tool( + _retrieve_paper_chunks, + name="retrieve_paper_chunks", + max_retries=config.max_tool_retries, + require_parameter_descriptions=True, + sequential=True, + ) diff --git a/src/agentomics/tools/tool_registry.py b/src/agentomics/tools/tool_registry.py index 207db43f..c6353e11 100644 --- a/src/agentomics/tools/tool_registry.py +++ b/src/agentomics/tools/tool_registry.py @@ -7,7 +7,9 @@ from pydantic_ai import Tool from .bash_tool import create_bash_tool +from .inspect_image_tool import create_inspect_image_tool from .replace_tool import create_replace_tool +from .retrieve_paper_chunks_tool import create_retrieve_paper_chunks_tool from .run_python_tool import create_run_python_tool from .write_python_tool import create_write_python_tool @@ -15,6 +17,8 @@ TOOL_FACTORY_REGISTRY: dict[str, Callable] = { "bash": create_bash_tool, + "inspect_image": create_inspect_image_tool, + "retrieve_paper_chunks": create_retrieve_paper_chunks_tool, "write_python": create_write_python_tool, "run_python": create_run_python_tool, "replace": create_replace_tool, diff --git a/src/agentomics/utils/config.py b/src/agentomics/utils/config.py index c33fe762..b2d9eab9 100644 --- a/src/agentomics/utils/config.py +++ b/src/agentomics/utils/config.py @@ -24,6 +24,8 @@ 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" + PROCESSED_PAPERS_DIRNAME: ClassVar[str] = "processed_papers" DEFAULT_ITERATIONS: ClassVar[int] = 5 DEFAULT_SPLIT_ALLOWED_ITERATIONS: ClassVar[int] = 1 @@ -40,6 +42,8 @@ class Config: "prediction_exploration", ] DEFAULT_TOOL_IDS: ClassVar[list[str]] = [ + "retrieve_paper_chunks", + "inspect_image", "bash", "write_python", "run_python", @@ -141,6 +145,14 @@ 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 processed_papers_dir(self) -> Path: + return self.shared_dir / self.PROCESSED_PAPERS_DIRNAME + @property def current_iteration_runtime_info_dir(self) -> Path: return self.current_iteration_dir / self.RUNTIME_INFO_DIRNAME diff --git a/tests/unit/test_paper_retrieval_tools.py b/tests/unit/test_paper_retrieval_tools.py new file mode 100644 index 00000000..d0e60cc8 --- /dev/null +++ b/tests/unit/test_paper_retrieval_tools.py @@ -0,0 +1,67 @@ +from pathlib import Path + +import pytest +from pydantic_ai import ModelRetry +from pydantic_ai.messages import BinaryContent, ToolReturn +from pypdf import PdfWriter +from reportlab.pdfgen import canvas + +from agentomics.tools.tool_registry import create_tool + + +def _write_paper(path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + document = canvas.Canvas(str(path)) + document.drawString(72, 720, "Background information about protein binding.") + document.showPage() + document.drawString(72, 720, "The zebrafish classifier uses a convolutional network.") + document.save() + + +def _write_blank_paper(path: Path) -> None: + document = PdfWriter() + document.add_blank_page(width=612, height=792) + with path.open("wb") as output: + document.write(output) + + +def test_retrieves_and_inspects_relevant_paper_page(default_agent_config): + paper_path = default_agent_config.fetched_papers_dir / "example.pdf" + _write_paper(paper_path) + _write_blank_paper(default_agent_config.fetched_papers_dir / "blank.pdf") + + retrieval_tool = create_tool(default_agent_config, "retrieve_paper_chunks") + results = retrieval_tool.function( + query="Which model is used for the zebrafish classifier?", + max_results=1, + ) + + assert results[0]["paper"] == paper_path.name + assert results[0]["page"] == 2 + assert "convolutional network" in results[0]["text"] + assert Path(results[0]["markdown_path"]).is_file() + assert Path(results[0]["markdown_path"]).with_name("chunks.json").is_file() + assert Path(results[0]["image_path"]).is_file() + + inspect_tool = create_tool(default_agent_config, "inspect_image") + inspected = inspect_tool.function( + image_path=results[0]["image_path"], + query="What model does this page describe?", + ) + + assert isinstance(inspected, ToolReturn) + assert isinstance(inspected.content[1], BinaryContent) + assert inspected.content[1].media_type == "image/png" + assert retrieval_tool.function( + query="zygomorphic quasar upholstery", + max_results=1, + ) == [] + + +def test_inspect_image_rejects_files_outside_the_run(default_agent_config, tmp_path): + image_path = tmp_path / "outside.png" + image_path.write_bytes(b"not an accessible image") + inspect_tool = create_tool(default_agent_config, "inspect_image") + + with pytest.raises(ModelRetry, match="must be inside"): + inspect_tool.function(image_path=str(image_path), query="What is shown?") From 40d21880d7b59de9fdda733a64f7994696c0946d Mon Sep 17 00:00:00 2001 From: Pankaj Rajdeo Date: Tue, 20 Oct 2026 05:34:39 -0400 Subject: [PATCH 2/2] refactor: use agentic paper search --- docs/how-it-works/architecture.md | 4 +- src/agentomics/runtime/paper_retrieval.py | 145 ++++----------- src/agentomics/tools/inspect_image_tool.py | 10 +- src/agentomics/tools/paper_files_tool.py | 167 ++++++++++++++++++ .../tools/retrieve_paper_chunks_tool.py | 29 --- src/agentomics/tools/tool_registry.py | 14 +- src/agentomics/utils/config.py | 6 +- tests/unit/test_paper_retrieval_tools.py | 48 +++-- 8 files changed, 255 insertions(+), 168 deletions(-) create mode 100644 src/agentomics/tools/paper_files_tool.py delete mode 100644 src/agentomics/tools/retrieve_paper_chunks_tool.py diff --git a/docs/how-it-works/architecture.md b/docs/how-it-works/architecture.md index 0309bd9f..e07b8eb1 100644 --- a/docs/how-it-works/architecture.md +++ b/docs/how-it-works/architecture.md @@ -133,8 +133,8 @@ The agent has access to these tools: | Tool | Purpose | |------|---------| -| **Retrieve Paper Chunks** | Find relevant passages in fetched scientific papers | -| **Inspect Image** | Inspect paper pages and other images from the current run | +| **ls / glob / grep / read_file** | Navigate and search fetched scientific papers | +| **Render Paper Page / Inspect Image** | Inspect paper figures and page layouts | | **Bash** | Execute shell commands | | **Write Python** | Create Python files | | **Run Python** | Execute Python scripts | diff --git a/src/agentomics/runtime/paper_retrieval.py b/src/agentomics/runtime/paper_retrieval.py index 07189948..b9509fb3 100644 --- a/src/agentomics/runtime/paper_retrieval.py +++ b/src/agentomics/runtime/paper_retrieval.py @@ -1,84 +1,53 @@ from __future__ import annotations import hashlib -import json from pathlib import Path import pypdfium2 as pdfium from pypdf import PdfReader -from sklearn.feature_extraction.text import TfidfVectorizer -from sklearn.metrics.pairwise import cosine_similarity from agentomics.utils.config import Config -CHUNK_WORDS = 350 -CHUNK_OVERLAP_WORDS = 50 - - -def retrieve_paper_chunks( - config: Config, - query: str, - max_results: int, -) -> list[dict[str, str | int | float]]: - query = query.strip() - if not query: - raise ValueError("query cannot be empty") - - chunks: list[dict[str, str | int]] = [] +def prepare_papers(config: Config) -> Path: for pdf_path in sorted(config.fetched_papers_dir.glob("*.pdf")): - paper_dir, paper_chunks = _prepare_pdf(config, pdf_path) - chunks.extend( - { - "paper": pdf_path.name, - **chunk, - "markdown_path": str(paper_dir / "document.md"), - } - for chunk in paper_chunks + document_path = _paper_dir(config, pdf_path) / "document.md" + if document_path.exists(): + continue + + pages = [] + for page_number, page in enumerate(PdfReader(pdf_path).pages, start=1): + text = ( + (page.extract_text(extraction_mode="layout") or "").strip() + if page.get_contents() is not None + else "" + ) + pages.append(f"## Page {page_number}\n\n{text}") + + document_path.parent.mkdir(parents=True, exist_ok=True) + document_path.write_text( + f"# {pdf_path.name}\n\nSource PDF: {pdf_path.resolve()}\n\n" + + "\n\n".join(pages) + + "\n", + encoding="utf-8", ) - if not chunks: - return [] + return config.processed_papers_dir - vectorizer = TfidfVectorizer( - stop_words="english", - ngram_range=(1, 2), - sublinear_tf=True, - ) - chunk_matrix = vectorizer.fit_transform(str(chunk["text"]) for chunk in chunks) - scores = cosine_similarity(vectorizer.transform([query]), chunk_matrix).ravel() - selected = [index for index in scores.argsort()[::-1] if scores[index] > 0][:max_results] - results: list[dict[str, str | int | float]] = [] - for index in selected: - chunk = chunks[int(index)] - pdf_path = config.fetched_papers_dir / str(chunk["paper"]) - page_number = int(chunk["page"]) - page_image_path = _page_image_path( - Path(str(chunk["markdown_path"])).parent, - page_number, - ) - if not page_image_path.exists(): - render_pdf_page(pdf_path, page_number, page_image_path) - results.append( - { - **chunk, - "score": float(scores[index]), - "image_path": str(page_image_path), - } - ) - return results +def render_pdf_page(config: Config, pdf_path: str, page_number: int) -> Path: + source = Path(pdf_path).resolve() + if not source.is_relative_to(config.fetched_papers_dir.resolve()) or not source.is_file(): + raise ValueError(f"Paper must be inside {config.fetched_papers_dir}: {pdf_path}") + image_path = _paper_dir(config, source) / "pages" / f"page_{page_number:04d}.png" + if image_path.exists(): + return image_path -def render_pdf_page( - pdf_path: Path, - page_number: int, - image_path: Path, -) -> None: - document = pdfium.PdfDocument(pdf_path) + document = pdfium.PdfDocument(source) if not 1 <= page_number <= len(document): document.close() - raise ValueError(f"page {page_number} is outside {pdf_path.name}") + raise ValueError(f"page {page_number} is outside {source.name}") page = document[page_number - 1] try: @@ -91,58 +60,10 @@ def render_pdf_page( finally: page.close() document.close() + return image_path -def _prepare_pdf( - config: Config, - pdf_path: Path, -) -> tuple[Path, list[dict[str, str | int]]]: +def _paper_dir(config: Config, pdf_path: Path) -> Path: with pdf_path.open("rb") as pdf_file: digest = hashlib.file_digest(pdf_file, "sha256").hexdigest()[:12] - paper_dir = config.processed_papers_dir / f"{pdf_path.stem}-{digest}" - paper_dir.mkdir(parents=True, exist_ok=True) - document_path = paper_dir / "document.md" - chunks_path = paper_dir / "chunks.json" - if document_path.exists() and chunks_path.exists(): - return paper_dir, json.loads(chunks_path.read_text(encoding="utf-8")) - - document_pages: list[str] = [] - chunks: list[dict[str, str | int]] = [] - for page_number, page in enumerate(PdfReader(pdf_path).pages, start=1): - text = ( - (page.extract_text(extraction_mode="layout") or "").strip() - if page.get_contents() is not None - else "" - ) - document_pages.append(f"## Page {page_number}\n\n{text}") - chunks.extend( - { - "page": page_number, - "chunk": chunk_number, - "text": chunk, - } - for chunk_number, chunk in enumerate(_split_text(text), start=1) - ) - - document_path.write_text( - f"# {pdf_path.name}\n\n" + "\n\n".join(document_pages) + "\n", - encoding="utf-8", - ) - chunks_path.write_text( - json.dumps(chunks, ensure_ascii=False), - encoding="utf-8", - ) - return paper_dir, chunks - - -def _split_text(text: str) -> list[str]: - words = text.split() - step = CHUNK_WORDS - CHUNK_OVERLAP_WORDS - return [ - " ".join(words[start : start + CHUNK_WORDS]) - for start in range(0, len(words), step) - ] - - -def _page_image_path(paper_dir: Path, page_number: int) -> Path: - return paper_dir / "pages" / f"page_{page_number:04d}.png" + return config.processed_papers_dir / f"{pdf_path.stem}-{digest}" diff --git a/src/agentomics/tools/inspect_image_tool.py b/src/agentomics/tools/inspect_image_tool.py index 36f3832e..a2f0db5b 100644 --- a/src/agentomics/tools/inspect_image_tool.py +++ b/src/agentomics/tools/inspect_image_tool.py @@ -22,12 +22,16 @@ def _inspect_image(image_path: str, query: str) -> ToolReturn: """Send an image to the configured model, which must support image input. Args: - image_path: Absolute path to an image inside the run directory. + image_path: Image inside the current step or processed papers directory. query: Question to answer using the image. """ path = Path(image_path).resolve() - if not path.is_relative_to(config.run_dir.resolve()): - raise ModelRetry(f"Image must be inside {config.run_dir}") + allowed_roots = ( + config.current_step_dir.resolve(), + config.processed_papers_dir.resolve(), + ) + if not any(path.is_relative_to(root) for root in allowed_roots): + raise ModelRetry("Image must be inside the current step or processed papers directory") media_type = IMAGE_MEDIA_TYPES.get(path.suffix.lower()) if not path.is_file() or media_type is None: diff --git a/src/agentomics/tools/paper_files_tool.py b/src/agentomics/tools/paper_files_tool.py new file mode 100644 index 00000000..a7ec1b6b --- /dev/null +++ b/src/agentomics/tools/paper_files_tool.py @@ -0,0 +1,167 @@ +from __future__ import annotations + +import re +from pathlib import Path +from typing import Annotated, Any + +from pydantic import Field +from pydantic_ai import ModelRetry, Tool + +from agentomics.runtime.paper_retrieval import prepare_papers, render_pdf_page +from agentomics.utils.config import Config + + +def _paper_path(config: Config, path: str) -> Path: + root = prepare_papers(config).resolve() + resolved = (root / path).resolve() + if not resolved.is_relative_to(root): + raise ModelRetry(f"Path must be inside the processed papers directory: {root}") + return resolved + + +def create_ls_tool(config: Config) -> Tool[Any]: + def ls( + path: str = ".", + max_results: Annotated[int, Field(ge=1, le=100)] = 100, + ) -> list[dict[str, str | int]]: + """List files and directories containing processed scientific papers. + + Args: + path: Directory path relative to the processed papers directory. + max_results: Maximum number of entries to return. + """ + directory = _paper_path(config, path) + if not directory.is_dir(): + raise ModelRetry(f"Directory not found: {path}") + root = config.processed_papers_dir.resolve() + return [ + { + "path": str(child.relative_to(root)), + "type": "directory" if child.is_dir() else "file", + "size": child.stat().st_size, + } + for child in sorted(directory.iterdir())[:max_results] + ] + + return _tool(config, ls, "ls") + + +def create_glob_tool(config: Config) -> Tool[Any]: + def glob( + pattern: str, + max_results: Annotated[int, Field(ge=1, le=100)] = 100, + ) -> list[str]: + """Find processed paper files matching a glob pattern. + + Args: + pattern: Pattern relative to the processed papers directory, such as **/*.md. + max_results: Maximum number of paths to return. + """ + if Path(pattern).is_absolute() or ".." in Path(pattern).parts: + raise ModelRetry("Glob pattern must stay inside the processed papers directory") + root = prepare_papers(config).resolve() + return [ + str(path.relative_to(root)) + for path in sorted(root.glob(pattern))[:max_results] + if path.resolve().is_relative_to(root) + ] + + return _tool(config, glob, "glob") + + +def create_grep_tool(config: Config) -> Tool[Any]: + def grep( + pattern: str, + path: str = ".", + max_results: Annotated[int, Field(ge=1, le=100)] = 50, + ) -> list[dict[str, str | int]]: + """Search processed paper text with a regular expression. + + Args: + pattern: Case-insensitive regular expression to search for. + path: File or directory path relative to the processed papers directory. + max_results: Maximum number of matching lines to return. + """ + try: + expression = re.compile(pattern, re.IGNORECASE) + except re.error as error: + raise ModelRetry(f"Invalid regular expression: {error}") from error + + target = _paper_path(config, path) + if not target.exists(): + raise ModelRetry(f"Path not found: {path}") + root = config.processed_papers_dir.resolve() + files = [target] if target.is_file() else sorted(target.rglob("*.md")) + matches = [] + for file_path in files: + for line_number, line in enumerate( + file_path.read_text(encoding="utf-8", errors="replace").splitlines(), + start=1, + ): + if expression.search(line): + matches.append( + { + "path": str(file_path.relative_to(root)), + "line": line_number, + "text": line[:1000], + } + ) + if len(matches) == max_results: + return matches + return matches + + return _tool(config, grep, "grep") + + +def create_read_file_tool(config: Config) -> Tool[Any]: + def read_file( + path: str, + offset: Annotated[int, Field(ge=1)] = 1, + limit: Annotated[int, Field(ge=1, le=500)] = 200, + ) -> str: + """Read numbered lines from a processed scientific paper. + + Args: + path: File path relative to the processed papers directory. + offset: First line to read, starting at 1. + limit: Maximum number of lines to return. + """ + file_path = _paper_path(config, path) + if not file_path.is_file(): + raise ModelRetry(f"File not found: {path}") + lines = file_path.read_text(encoding="utf-8", errors="replace").splitlines() + return "\n".join( + f"{line_number}: {lines[line_number - 1]}" + for line_number in range(offset, min(offset + limit, len(lines) + 1)) + ) + + return _tool(config, read_file, "read_file") + + +def create_render_paper_page_tool(config: Config) -> Tool[Any]: + def render_paper_page( + pdf_path: str, + page_number: Annotated[int, Field(ge=1)], + ) -> str: + """Render one paper page as an image for inspection. + + Args: + pdf_path: Source PDF path shown at the top of its processed Markdown file. + page_number: One-based page number to render. + """ + try: + return str(render_pdf_page(config, pdf_path, page_number)) + except ValueError as error: + raise ModelRetry(str(error)) from error + + return _tool(config, render_paper_page, "render_paper_page") + + +def _tool(config: Config, function: Any, name: str) -> Tool[Any]: + return Tool( + function, + name=name, + max_retries=config.max_tool_retries, + require_parameter_descriptions=True, + sequential=True, + ) diff --git a/src/agentomics/tools/retrieve_paper_chunks_tool.py b/src/agentomics/tools/retrieve_paper_chunks_tool.py deleted file mode 100644 index 92decf33..00000000 --- a/src/agentomics/tools/retrieve_paper_chunks_tool.py +++ /dev/null @@ -1,29 +0,0 @@ -from typing import Annotated, Any - -from pydantic import Field -from pydantic_ai import Tool - -from agentomics.runtime.paper_retrieval import retrieve_paper_chunks -from agentomics.utils.config import Config - - -def create_retrieve_paper_chunks_tool(config: Config) -> Tool[Any]: - def _retrieve_paper_chunks( - query: Annotated[str, Field(min_length=1)], - max_results: Annotated[int, Field(ge=1, le=10)] = 5, - ) -> list[dict[str, str | int | float]]: - """Find passages relevant to a query in the fetched scientific papers. - - Args: - query: Information to find in the fetched papers. - max_results: Maximum number of passages to return. - """ - return retrieve_paper_chunks(config, query, max_results) - - return Tool( - _retrieve_paper_chunks, - name="retrieve_paper_chunks", - max_retries=config.max_tool_retries, - require_parameter_descriptions=True, - sequential=True, - ) diff --git a/src/agentomics/tools/tool_registry.py b/src/agentomics/tools/tool_registry.py index c6353e11..7f88c106 100644 --- a/src/agentomics/tools/tool_registry.py +++ b/src/agentomics/tools/tool_registry.py @@ -8,8 +8,14 @@ from .bash_tool import create_bash_tool from .inspect_image_tool import create_inspect_image_tool +from .paper_files_tool import ( + create_glob_tool, + create_grep_tool, + create_ls_tool, + create_read_file_tool, + create_render_paper_page_tool, +) from .replace_tool import create_replace_tool -from .retrieve_paper_chunks_tool import create_retrieve_paper_chunks_tool from .run_python_tool import create_run_python_tool from .write_python_tool import create_write_python_tool @@ -17,8 +23,12 @@ TOOL_FACTORY_REGISTRY: dict[str, Callable] = { "bash": create_bash_tool, + "glob": create_glob_tool, + "grep": create_grep_tool, "inspect_image": create_inspect_image_tool, - "retrieve_paper_chunks": create_retrieve_paper_chunks_tool, + "ls": create_ls_tool, + "read_file": create_read_file_tool, + "render_paper_page": create_render_paper_page_tool, "write_python": create_write_python_tool, "run_python": create_run_python_tool, "replace": create_replace_tool, diff --git a/src/agentomics/utils/config.py b/src/agentomics/utils/config.py index b2d9eab9..bceef639 100644 --- a/src/agentomics/utils/config.py +++ b/src/agentomics/utils/config.py @@ -42,7 +42,11 @@ class Config: "prediction_exploration", ] DEFAULT_TOOL_IDS: ClassVar[list[str]] = [ - "retrieve_paper_chunks", + "ls", + "glob", + "grep", + "read_file", + "render_paper_page", "inspect_image", "bash", "write_python", diff --git a/tests/unit/test_paper_retrieval_tools.py b/tests/unit/test_paper_retrieval_tools.py index d0e60cc8..be52ece4 100644 --- a/tests/unit/test_paper_retrieval_tools.py +++ b/tests/unit/test_paper_retrieval_tools.py @@ -25,43 +25,53 @@ def _write_blank_paper(path: Path) -> None: document.write(output) -def test_retrieves_and_inspects_relevant_paper_page(default_agent_config): +def test_searches_reads_and_inspects_papers(default_agent_config): paper_path = default_agent_config.fetched_papers_dir / "example.pdf" _write_paper(paper_path) _write_blank_paper(default_agent_config.fetched_papers_dir / "blank.pdf") - retrieval_tool = create_tool(default_agent_config, "retrieve_paper_chunks") - results = retrieval_tool.function( - query="Which model is used for the zebrafish classifier?", - max_results=1, + entries = create_tool(default_agent_config, "ls").function() + assert len(entries) == 2 + + documents = create_tool(default_agent_config, "glob").function(pattern="**/*.md") + assert len(documents) == 2 + + matches = create_tool(default_agent_config, "grep").function( + pattern="convolutional network", + ) + assert len(matches) == 1 + text = create_tool(default_agent_config, "read_file").function( + path=matches[0]["path"], + offset=matches[0]["line"], + limit=1, ) + assert "convolutional network" in text - assert results[0]["paper"] == paper_path.name - assert results[0]["page"] == 2 - assert "convolutional network" in results[0]["text"] - assert Path(results[0]["markdown_path"]).is_file() - assert Path(results[0]["markdown_path"]).with_name("chunks.json").is_file() - assert Path(results[0]["image_path"]).is_file() + image_path = create_tool(default_agent_config, "render_paper_page").function( + pdf_path=str(paper_path), + page_number=2, + ) + assert Path(image_path).is_file() inspect_tool = create_tool(default_agent_config, "inspect_image") inspected = inspect_tool.function( - image_path=results[0]["image_path"], + image_path=image_path, query="What model does this page describe?", ) assert isinstance(inspected, ToolReturn) assert isinstance(inspected.content[1], BinaryContent) assert inspected.content[1].media_type == "image/png" - assert retrieval_tool.function( - query="zygomorphic quasar upholstery", - max_results=1, - ) == [] -def test_inspect_image_rejects_files_outside_the_run(default_agent_config, tmp_path): +def test_paper_tools_reject_files_outside_their_directories(default_agent_config, tmp_path): image_path = tmp_path / "outside.png" image_path.write_bytes(b"not an accessible image") - inspect_tool = create_tool(default_agent_config, "inspect_image") + with pytest.raises(ModelRetry, match="processed papers"): + create_tool(default_agent_config, "read_file").function(path=str(image_path)) with pytest.raises(ModelRetry, match="must be inside"): - inspect_tool.function(image_path=str(image_path), query="What is shown?") + create_tool(default_agent_config, "inspect_image").function( + image_path=str(image_path), + query="What is shown?", + )