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.
- Prerequisites
- Quick Start: Adding Your First Tool
- Understanding the ToolBase Contract
- Configuration Schema
- Input/Output Handling
- Error Handling & Timeouts
- Registering Your Tool
- Testing Your Integration
- Common Patterns
- 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)
Let's walk through adding a hypothetical tool called FooFold that predicts protein structures.
cp tools/_template.py tools/foofold.pyfrom 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"] _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,
},
} 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 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.0That'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.
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 |
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.
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 |
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)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,
},
)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.
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:
- Loading modules from the project-owned top-level
tools/directory under a private namespace - Skipping
tools/_template.py, because it is only a template - Collecting classes decorated with
@register_tool - Checking that
_tool_nameis 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().
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 foofoldRun 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/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"))
...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}"])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)- See
tools/_template.pyfor a fully commented template - See
references/ARCHITECTURE.mdfor framework design decisions - See
references/API-REFERENCE.mdfor complete API documentation