diff --git a/docs/how-it-works/architecture.md b/docs/how-it-works/architecture.md index feaee42c..e07b8eb1 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 | |------|---------| +| **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/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..b9509fb3 --- /dev/null +++ b/src/agentomics/runtime/paper_retrieval.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +import hashlib +from pathlib import Path + +import pypdfium2 as pdfium +from pypdf import PdfReader + +from agentomics.utils.config import Config + + +def prepare_papers(config: Config) -> Path: + for pdf_path in sorted(config.fetched_papers_dir.glob("*.pdf")): + 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", + ) + + return config.processed_papers_dir + + +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 + + document = pdfium.PdfDocument(source) + if not 1 <= page_number <= len(document): + document.close() + raise ValueError(f"page {page_number} is outside {source.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() + return image_path + + +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] + 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 new file mode 100644 index 00000000..a2f0db5b --- /dev/null +++ b/src/agentomics/tools/inspect_image_tool.py @@ -0,0 +1,55 @@ +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: Image inside the current step or processed papers directory. + query: Question to answer using the image. + """ + path = Path(image_path).resolve() + 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: + 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/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/tool_registry.py b/src/agentomics/tools/tool_registry.py index 207db43f..7f88c106 100644 --- a/src/agentomics/tools/tool_registry.py +++ b/src/agentomics/tools/tool_registry.py @@ -7,6 +7,14 @@ from pydantic_ai import Tool 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 .run_python_tool import create_run_python_tool from .write_python_tool import create_write_python_tool @@ -15,6 +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, + "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 c33fe762..bceef639 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,12 @@ class Config: "prediction_exploration", ] DEFAULT_TOOL_IDS: ClassVar[list[str]] = [ + "ls", + "glob", + "grep", + "read_file", + "render_paper_page", + "inspect_image", "bash", "write_python", "run_python", @@ -141,6 +149,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..be52ece4 --- /dev/null +++ b/tests/unit/test_paper_retrieval_tools.py @@ -0,0 +1,77 @@ +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_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") + + 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 + + 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=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" + + +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") + + 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"): + create_tool(default_agent_config, "inspect_image").function( + image_path=str(image_path), + query="What is shown?", + )