Skip to content

Latest commit

 

History

History
378 lines (286 loc) · 11.1 KB

File metadata and controls

378 lines (286 loc) · 11.1 KB

ProtForge Developer Guide

This guide walks you through integrating a new protein design tool into ProtForge.

The starter repository contains no real protein design tool modules. Until you copy tools/_template.py to a new module and implement it, ToolRegistry().list_tools() is expected to return an empty list.

Table of Contents

  1. Prerequisites
  2. Quick Start: Adding Your First Tool
  3. Understanding the ToolBase Contract
  4. Configuration Schema
  5. Input/Output Handling
  6. Error Handling & Timeouts
  7. Registering Your Tool
  8. Testing Your Integration
  9. Common Patterns

Prerequisites

  • Python 3.9+
  • The target protein design tool installed by the user (ProtForge does not install tools)
  • Basic understanding of how the tool is invoked (command-line arguments, input/output formats)

Quick Start: Adding Your First Tool

Let's walk through adding a hypothetical tool called FooFold that predicts protein structures.

Step 1: Copy the Template

cp tools/_template.py tools/foofold.py

Step 2: Fill in Metadata

from protforge.base import ToolBase, ToolInput, ToolOutput, ToolStatus
from protforge.registry import register_tool

@register_tool
class FooFold(ToolBase):
    _tool_name = "foofold"
    _tool_version = "2.1.0"
    _tool_category = "prediction"
    _input_types = ["sequence"]
    _output_types = ["structure"]

Step 3: Define Configuration Schema

    _config_schema = {
        "executable_path": {
            "type": "path",
            "description": "Path to FooFold's main executable",
            "required": True,
            "default": None,
        },
        "model_dir": {
            "type": "path",
            "description": "Directory containing FooFold model weights",
            "required": True,
            "default": None,
        },
        "use_gpu": {
            "type": "boolean",
            "description": "Whether to use GPU acceleration",
            "required": False,
            "default": True,
        },
    }

Step 4: Implement validate_config

    def validate_config(self) -> bool:
        exe = self.config.get("executable_path")
        model_dir = self.config.get("model_dir")

        if not exe or not os.path.isfile(exe):
            return False
        if not model_dir or not os.path.isdir(model_dir):
            return False
        return True

Step 5: Implement run

    def run(self, tool_input: ToolInput) -> ToolOutput:
        # Prepare FASTA input
        temp_dir = self.get_temp_dir()
        fasta_path = temp_dir / "input.fasta"
        self.write_sequence_fasta(tool_input.sequences, fasta_path)

        # Build command
        exe = self.config.get("executable_path")
        model_dir = self.config.get("model_dir")
        cmd = [
            exe,
            "--input", str(fasta_path),
            "--model-dir", model_dir,
            "--output", str(temp_dir / "output.pdb"),
        ]

        if self.config.get("use_gpu", True):
            cmd.append("--gpu")

        # Execute
        try:
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=tool_input.timeout or self.default_timeout,
            )
        except subprocess.TimeoutExpired:
            return ToolOutput(status=ToolStatus.TIMEOUT, error_message="Timed out")
        except Exception as e:
            return ToolOutput(status=ToolStatus.ERROR, error_message=str(e))

        if result.returncode != 0:
            return ToolOutput(
                status=ToolStatus.ERROR,
                error_message=result.stderr,
                stdout=result.stdout,
                stderr=result.stderr,
            )

        # Parse output
        output_pdb = temp_dir / "output.pdb"
        return ToolOutput(
            status=ToolStatus.SUCCESS,
            files={"pdb": str(output_pdb)},
            data={"plddt": self._parse_plddt(result.stdout)},
            stdout=result.stdout,
            stderr=result.stderr,
        )

    def _parse_plddt(self, stdout: str) -> float:
        # Tool-specific parsing logic
        for line in stdout.splitlines():
            if "pLDDT:" in line:
                return float(line.split(":")[1].strip())
        return 0.0

That's it! ProtForge will auto-discover FooFold on the next skill load when the project root (which contains the top-level tools/ package) is available to Python.


Understanding the ToolBase Contract

Every tool must subclass ToolBase and implement:

Method Purpose
validate_config() Check if the tool is installed and configured correctly
run(tool_input) Execute the tool and return standardized output

Class attributes you must define:

Attribute Description
_tool_name Unique identifier (lowercase, no spaces)
_tool_version Version string for tracking
_tool_category Broad category: prediction, design, generation, dynamics, analysis
_input_types List of accepted input types: sequence, structure, msa, embedding
_output_types List of produced output types
_config_schema Dictionary defining configurable parameters

Configuration Schema

The _config_schema drives the interactive configuration UI. Each key maps to a specification dict:

{
    "param_name": {
        "type": "path",          # One of: string, integer, boolean, float, path
        "description": "...",    # Human-readable explanation
        "required": True,        # Whether the tool fails without this
        "default": None,         # Default value (can be None)
    }
}

When the agent calls configure_tool("foofold"), ProtForge returns this schema so the agent can ask the user for each value.


Input/Output Handling

Standardized Data Types

ProtForge uses these canonical types for data exchange:

Type Python Representation File Formats
sequence Dict[str, str] (chain_id -> sequence) FASTA
structure Dict[str, str] (name -> file path) PDB, mmCIF
msa Dict[str, str] (name -> file path) A3M, FASTA
embedding Dict[str, Any] NPZ, PT
trajectory Dict[str, str] (name -> file path) XTC, DCD
score Dict[str, float] JSON, CSV

Writing Input Files

Use built-in utilities from ToolBase:

# Write sequences to FASTA
fasta_path = self.get_temp_dir() / "input.fasta"
self.write_sequence_fasta(tool_input.sequences, fasta_path)

# Write PDB content
pdb_path = self.get_temp_dir() / "input.pdb"
self.write_pdb(tool_input.structures["model"], pdb_path)

Returning Output Files

Always return file paths in ToolOutput.files with canonical type keys:

return ToolOutput(
    status=ToolStatus.SUCCESS,
    files={
        "pdb": "/path/to/output.pdb",      # primary structure output
        "json": "/path/to/scores.json",    # auxiliary data
    },
    data={
        "plddt": 92.5,
        "ptm": 0.89,
    },
)

Error Handling & Timeouts

Always catch exceptions and return a ToolOutput with an appropriate ToolStatus:

except subprocess.TimeoutExpired:
    return ToolOutput(status=ToolStatus.TIMEOUT, error_message="Exceeded time limit")
except FileNotFoundError:
    return ToolOutput(status=ToolStatus.ERROR, error_message="Executable not found")
except Exception as e:
    return ToolOutput(status=ToolStatus.ERROR, error_message=str(e))

Respect the tool_input.timeout parameter; fall back to self.default_timeout if not provided. New integrations should prefer protforge.runner.CommandRunner for consistent subprocess handling, timeout conversion, stdout/stderr capture, and command redaction.


Registering Your Tool

Use the @register_tool decorator on your customized class. The starter tools/_template.py keeps this decorator commented out so accidentally importing the template does not register a fake tool:

from protforge.registry import register_tool

@register_tool
class FooFold(ToolBase):
    ...

The registry discovers tools by:

  1. Loading modules from the project-owned top-level tools/ directory under a private namespace
  2. Skipping tools/_template.py, because it is only a template
  3. Collecting classes decorated with @register_tool
  4. Checking that _tool_name is defined

No additional registration step is needed once your integration module lives in the project tools/ directory. For embedded or packaged use cases, you can also register a pre-instantiated tool with ToolRegistry.register_tool_instance().


Testing Your Integration

Create a test in tests/:

import unittest
from protforge import ToolRegistry, ConfigManager

class TestFooFold(unittest.TestCase):
    def setUp(self):
        self.cfg = ConfigManager()
        self.registry = ToolRegistry(self.cfg)
        self.tool = self.registry.get_tool("foofold")

    def test_tool_discovered(self):
        self.assertIsNotNone(self.tool)

    def test_not_configured_initially(self):
        self.assertFalse(self.tool.validate_config())

    def test_configurable(self):
        self.cfg.set_tool_config("foofold", {
            "executable_path": "/usr/bin/true",  # mock
            "model_dir": "/tmp/models",
        }, schema=self.tool._config_schema)
        self.tool.config = self.cfg.get_tool_config("foofold")
        # Note: validate_config may still fail if /usr/bin/true doesn't act like foofold

Run tests:

cd /path/to/protforge && python -m unittest -v tests.test_framework
# or, after installing test extras:
python -m pip install -e '.[test]'
python -m pytest tests/

Common Patterns

Conda Environment Activation

Many bioinformatics tools require a specific conda environment. Handle this in _build_command:

def _build_command(self, input_files, parameters):
    cmd = []
    env = self.config.get("conda_env")
    if env:
        cmd.extend(["conda", "run", "-n", env, "python"])
    cmd.append(self.config.get("executable_path"))
    ...

Multi-Sequence Input

Some tools accept multiple chains or homomers. Iterate over tool_input.sequences:

for chain_id, seq in tool_input.sequences.items():
    cmd.extend(["--chain", f"{chain_id}={seq}"])

Parsing Complex Output

For tools that write multiple files, glob the output directory:

from pathlib import Path

output_dir = self.get_temp_dir() / "results"
output_dir.mkdir(exist_ok=True)

cmd.extend(["--out-dir", str(output_dir)])
# ... run ...

files = {}
for pdb_file in output_dir.glob("*.pdb"):
    files[pdb_file.stem] = str(pdb_file)

return ToolOutput(status=ToolStatus.SUCCESS, files=files)

Next Steps

  • See tools/_template.py for a fully commented template
  • See references/ARCHITECTURE.md for framework design decisions
  • See references/API-REFERENCE.md for complete API documentation