Implement Tool

Implement a New Tool — Orchestrator Pipeline

You are implementing a new bioinformatics tool in the proto-tools codebase. This skill orchestrates the full lifecycle using parallel subagents for speed.

Pipeline overview:

Phase 1: Research → Phase 2: Contract → Phase 3: Fan-out (5 parallel agents) → Phase 4: Verify → Phase 4.5: Config Field Audit → Phase 4.6: Temp Integration & Stress → Phase 4.7: Docs Verification → Phase 4.8: Self-Audit & Full PR Review → Phase 5: Ship

The Phase 4.x gates are decisive and non-negotiable — they catch the failure modes that the fast fan-out reliably misses (invented config parameters, untested real workloads, hallucinated citations/licenses, missing license.yaml/links.yaml). Do not skip them to ship faster. Ground every decision in upstream source or a credible reference — never assume.

Key principle: The core tool file (Input/Config/Output + @tool() + run_*()) is the contract everything else depends on. Write it first (sequential), then fan out to subagents (parallel).

CRITICAL: Standalone scripts run in isolated environments and MUST NOT import from proto_tools:

• from proto_tools.utils import ... — will fail at runtime
• from proto_tools.entities import ... — will fail at runtime
• from standalone_helpers import get_subprocess_device_env — published on PYTHONPATH (OK)
• Standard library imports: import json, import subprocess, etc. (OK)
• Dependencies from requirements.txt: import torch, import numpy, etc. (OK)
• NEVER install proto_tools in standalone environments (creates circular dependency, breaks isolation)

Phase 0: Parse Input

The user provides EITHER:

• A GitHub issue URL/number (e.g., #<issue-number> or https://github.com/evo-design/proto-tools/issues/<issue-number>)
• A tool name + source repo (e.g., "ESM-IF from https://github.com/facebookresearch/esm")

If a GitHub issue is provided:

1. Fetch the issue with gh issue view <number> --json title,body,labels
2. Extract from the issue body: tool name, source repo URLs, paper links, category, operations
3. Present the extracted info to the user for confirmation

If direct tool info is provided:

1. Confirm tool name, category, and source repo with the user

Output of Phase 0: A clear specification:

• Tool name (e.g., esmif)
• Category (e.g., inverse_folding)
• Operations (e.g., sample, score)
• Source repo URL(s)
• Paper DOI (if available)
• Whether it uses shared data models from the category

Phase 1: Research

Goal: Understand the tool's API, inputs, outputs, and dependencies.

Steps:

1. Fetch source repo README — Use WebFetch or gh to read the README of the research repo. Extract:

   • Installation instructions (pip packages, dependencies)
   • Python API or CLI interface
   • Input/output formats
   • License — read the repo's actual LICENSE/COPYING file (not just the README badge) and the weights license if it differs. Record the SPDX id (or note it's non-standard) and the file URL. Feeds license.yaml and the README License callout.
   • Canonical links — the repo URL, project homepage/docs, paper DOI or preprint URL, HuggingFace model page, affiliated orgs. Feeds links.yaml. Verify each resolves.
   • The full upstream parameter surface — which parameters the upstream CLI/Python API actually exposes to end users, with their defaults and valid ranges. This is the ground truth for Phase 4.5's config audit; capture it now while reading the source.
   • Model weights location (HuggingFace, GitHub releases, etc.) — determines which weight management pattern to use:
     • HuggingFace from_pretrained → no weight code needed (HF_HOME set automatically). Example: tools/masked_models/esm2/
     • Direct download in setup.sh → must implement PROTO_MODEL_CACHE pattern. Example: tools/inverse_folding/fampnn/standalone/setup.sh
     • Foundry install → must implement PROTO_MODEL_CACHE pattern. Example: tools/inverse_folding/ligandmpnn/standalone/setup.sh
     • Cannot be auto-downloaded (the asset is too large to auto-download, or there's no programmatic fetch path at all — e.g. license requires a manual approval flow before the user receives a one-time download link). Distinct from gated-but-fetchable assets like HF gated repos, where setup.sh can still download once a token is present (use proto_check_gated_hf_repo for those). For not-fetchable assets, the user must place files on disk themselves; setup.sh doesn't fetch. → precheck at the top of setup.sh: when the asset is missing, print [proto-tools] ASSET_NOT_AVAILABLE: <toolkit>:<asset_kind> and exit 64. Tests skip cleanly on hosts without weights; misconfigured PROTO_<TOOLKIT>_WEIGHTS_DIR still fails (exit 1). Example: tools/structure_prediction/x3dna/standalone/setup.sh. See PATTERNS.md → "standalone/setup.sh for tools whose assets are NOT automatically downloaded".
       • A restrictive license is not by itself a reason to use this pattern. If the weights sit at a URL anyone can fetch, download them in setup.sh and record the restrictions in license.yaml (commercial_use, redistribution, weights.text) instead of weights.access. tools/structure_prediction/alphafold3/standalone/setup.sh does this: public download, non-commercial-only terms.

2. Read key source files — Find and read the research repo's inference/prediction scripts to understand:

   • How the model is loaded
   • What inputs it expects (sequences, structures, files)
   • What outputs it produces (scores, sequences, structures)
   • Device handling (CUDA/CPU)

3. Read existing tools in the same category — Find the reference implementation:

   ls proto_tools/tools/{category}/

   Read the reference tool's:

   • Main tool file (the @tool() decorated function)
   • standalone/inference.py (or run.py)
   • standalone/setup.sh
   • standalone/requirements.txt
   • standalone/python_version.txt (required — the Python version pin for the env)
   • __init__.py (at tool and category level)
   • README.md
   • cite.bib
   • Test file in tests/{category}_tests/

4. Check for shared data models — If the category has a shared_data_models.py, read it. The new tool should extend these base classes rather than creating new ones. Input and Output classes may be reused directly across sibling tools. The Config must be per-tool: each tool needs its own config class so it carries its own tool_key (stamped at registration and checked by test_config_carries_its_tool_key). If the shape is identical to a sibling's, subclass it with a one-line docstring rather than registering the shared class twice.

Output of Phase 1: Mental model of:

• The tool's Python API (function signatures, class names)
• Its dependencies and installation requirements
• How it maps to the existing tool patterns
• Which reference tool to use as the structural template

Phase 2: Contract (Write Core Tool File)

Goal: Write the core tool file that defines the API surface. Everything else depends on this.

This phase is sequential — no subagents. The orchestrator writes this directly.

Steps:

Placeholder glossary (used throughout this skill — all distinct concepts, not interchangeable):

• {toolkit} — snake_case directory name shared by all tools in the family (e.g., evo2, pyrosetta). Strict — drives the directory path and the dispatch/worker identifier.
• {tool_key} — kebab-case registration key, {toolkit}-{suffix} (e.g., evo2-sample, pyrosetta-energy). Strict — the value passed to @tool(key=...).
• {tool_key_snake} — snake_case form of {tool_key} (e.g., evo2_sample). Strict — the core tool file name, the run_* function, and the test file.
• {ToolName} — PascalCase class-name prefix for this tool's Input / Config / Output (e.g., Evo2Sample, ESMFoldPrediction). Developer's choice — typically the PascalCase of {tool_key}, but pick whatever reads cleanly (e.g., ESMFold over Esmfold) as long as it's specific to this tool.
• {tool_display_name} — human-readable label (e.g., "Evo 2"). Used in docstrings, commit messages, and PR titles.

1. Create the tool directory structure:

   mkdir -p proto_tools/tools/{category}/{toolkit}/standalone
   mkdir -p proto_tools/tools/{category}/{toolkit}/examples

   Target file tree:

   tools/{category}/
   +-- shared_data_models.py   # Shared Input/Config/Output base classes (if category has 2+ tools)
   +-- {toolkit}/
   |   +-- __init__.py
   |   +-- {tool_key_snake}.py # Core tool file (Input/Config/Output + @tool + run_*); one per registered tool
   |   +-- cite.bib            # BibTeX citation (required)
   |   +-- license.yaml        # License metadata (required — enforced by test_license_consistency.py)
   |   +-- links.yaml          # Canonical upstream links (required — github/website/paper/huggingface/...)
   |   +-- README.md
   |   +-- examples/
   |   |   +-- example.ipynb   # Working example notebook (required)
   |   +-- helpers.py            # [optional] Self-contained plain-type file-format conversion helpers
   |   +-- standalone/
   |   |   +-- setup.sh
   |   |   +-- run.py OR inference.py  # run.py for CPU tools, inference.py for AI models
   |   |   +-- requirements.txt
   |   |   +-- binary_config.py   # [optional]
   +-- __init__.py

2. Write the core tool file proto_tools/tools/{category}/{toolkit}/{tool_key_snake}.py with:

   • Proper imports
   • Input class extending BaseToolInput (or shared base) with Field() — extra="forbid"
   • Config class extending BaseConfig (or shared base) with ConfigField() — extra="forbid". Use reload_on_change=True on fields that require worker restart (model checkpoint, etc.). Use include_in_key=False on fields that don't affect computation results (device, verbose, timeout are already excluded on BaseConfig; tool-level overrides of device must also set include_in_key=False). include_in_key defaults to True. Use xor_group="<slug>" to mark mutually exclusive sibling fields (see "Mutual-exclusion fields (XOR groups)" below).
   • Output class extending BaseToolOutput (or shared base) with Field() — extra="ignore" (set by BaseToolOutput; the base also emits a warning for unexpected fields so computed-field JSON round-trips can validate cleanly)
   • @tool() decorator with the 7 required kwargs (key, label, category, input_class, config_class, output_class, description) plus conventionally set uses_gpu. Optional: example_input, device_count, cacheable, stochastic, iterable_input_fields (a list — ["sequences"] for an ordinary iterable, or an index-parallel group like ["complexes", "msas"]), iterable_output_field, metrics_class, gpu_only, post_process_iterable
   • run_*() function that calls ToolInstance.dispatch()
   • If category has shared data models, use type aliases (e.g., ToolInput = InverseFoldingInput)
   • Metrics: if the tool emits scalar metric-like values (plDDT, perplexity, scores, etc.), route them through a Metrics subclass (from proto_tools/utils/tool_io.py). Prefer a shared per-category class (MaskedModelScoringMetrics, InverseFoldingScoringMetrics, CausalModelScoringMetrics, etc.) when the metric set matches siblings. Otherwise declare a per-tool <Tool>Metrics(Metrics) subclass in the tool file with a metric_spec: ClassVar[dict[str, MetricSpec]] documenting type/range and a better_values_are direction ("higher", "lower", "in-range", or "context-dependent") for each metric — test_metric_spec_consistency.py requires a valid better_values_are on every metric. Verification lives in each tool's e2e test via assert_metrics_in_spec(result) (helper at tests/tool_infra_tests/_metric_helpers.py).

Critical conventions:

• Tool registry key ({tool_key}): {toolkit}-{suffix} kebab-case (e.g., "esmif-sample")
• Run function: run_{tool_key_snake} (e.g., run_esmif_sample)
• Classes: {ToolName}Input/Config/Output, PascalCase (e.g., ESMIFSampleInput)
• batch_size defaults to 1 for GPU tools
• No try/except — @tool decorator handles errors
• Use logging.getLogger(__name__), never print()
• Output must implement output_format_options, output_format_default, _export_output()
• Iterable input/output cardinality must be 1:1: when iterable_input_fields and iterable_output_field are both set, len(output.{iterable_output_field}) must equal the (primary) input list length, by position. The framework's per-item cache stitching, dedup, and diversification testing all rely on this. If the tool produces N samples per input (e.g. a num_sequences_per_structure or num_designs config parameter), bundle the N samples inside a single per-input wrapper class — do NOT flatten N×inputs into the output list. Canonical patterns to follow: ProteinMPNNSequences / LigandMPNNSequences (per-structure sequence bundles, num_sequences_per_structure config) and RFdiffusion3Designs (per-spec design bundles, n_batches * diffusion_batch_size per spec). See PATTERNS.md → "Iterable cardinality" for the full rule and code shape.
• Input = per-item, Config = shared: model each field by "is there one of these per item, or one for the whole batch?" Shared whole-batch params/collections → Config (which also keeps the per-item cache key correct by construction). Per-item values → Input, either bundled inside the iterable element or as a parallel sibling in iterable_input_fields. A per-item value that is usually uniform can be a broadcastable field — InputField(..., broadcastable=True), declared as list[...] | None, added to iterable_input_fields, normalized via broadcast_parallel_fields() (one value broadcasts to all, N is strict 1:1). Broadcast is opt-in per field and only when a single value applied to every item is genuinely correct (safe for binder_chain; wrong for msas). See PATTERNS.md → "Input = per-item, Config = shared".
• cpus_per_instance opt-in: BaseConfig.cpus_per_instance defaults to None — every CPU tool stays off ToolPool's CPU scheduler and runs as a single direct call. Most CPU tools should leave this alone. Only opt in (override to a positive int) when per-call work is heavy enough to amortize spinning up N persistent worker subprocesses — each holds its own venv in RAM and pays a startup tax, so cheap tools (short per-item compute, internal threading, network IO) lose more than they gain. The canonical opt-in is PyRosetta (heavy init, multi-second per pose, embarrassingly parallel poses → cpus_per_instance = 1). GPU tools (gpus_per_instance > 0) ignore cpus_per_instance entirely.

Inherited field audit: When reusing a shared base config (e.g., InverseFoldingConfig) or base input, enumerate every inherited field and verify the target model can implement it. For each unsupported field, either implement support (e.g., logit masking for excluded_amino_acids) or override the field with a validator that raises ValueError("'{field_name}' is not supported by {tool_display_name}") when a non-default value is provided. Do not silently inherit fields that the model ignores.

No config fields for internal/ephemeral plumbing: Don't expose a config field that doesn't change the computation result — e.g., an internal job name, a scratch/output directory, or intermediate-file labels. The wrapper returns results via the Output object, and the export API (_export_output() / output.export(...)) owns user-facing output files and naming — so such fields add API surface without behavior, and (unless marked include_in_key=False) wrongly split the cache so identical calls miss each other. Hardcode the internal value instead. Contributors reflexively add these; drop them. Example: opendde removed its name (job label → ephemeral temp-file names) and root_dir (duplicated the managed weights-cache resolution) fields for exactly this reason.

Remote fail-fast for local-resource configs: If a config setting needs a local resource that can't be staged to the hosted service — a local database/index, an on-disk weights/artifact directory, or a local file path — override BaseConfig.remote_unsupported_reason(self, device: str) -> str | None so device='proto' fails fast at dispatch with a clear message instead of a late runtime error (a missing-file crash, or silently ignoring the override). Return a user-facing reason when the offending setting is active, None otherwise (the default means remote-compatible). Guard on the setting actually being non-default so ordinary remote runs are untouched:

def remote_unsupported_reason(self, device: str) -> str | None:
    """A local weights directory (``local_path``) isn't present on a hosted worker."""
    if self.local_path:  # optional override; unset by default → remote-OK
        return "local_path points to a local weights directory not available on device='proto'. Unset it, or run locally with device='cpu'."
    return None

A tool whose only mode needs a local file (e.g. a required local DB/HMM input, like blast-search local mode or pyhmmer-hmmscan) returns the reason unconditionally. Scope this to intrinsically local configs — do not encode deployment-specific hosting knowledge here (e.g. which model checkpoints are pre-warmed on the hosted service); that stays in the private hosting layer. The router invokes this on any remote device path, after the local_cpu no-op and the license gate. Branch on device when a restriction reflects what Proto chooses to host rather than an intrinsically local resource: a caller running on their own compute is not bound by our hosting limits.

Mutual-exclusion fields (XOR groups)

For "pick one" siblings: make each Optional with default=None, tag with the same xor_group="<slug>", and add a @model_validator(mode="after") to enforce at runtime.

mmseqs_db: str | None = InputField(default=None, xor_group="target",
    description="Target DB (path/slug/AssetRef).")
target_sequences: list[str] | None = InputField(default=None, xor_group="target",
    description="Inline target protein sequences.")

@model_validator(mode="after")
def exactly_one_target(self) -> "Mmseqs2SearchProteinsInput":
    if (self.mmseqs_db is None) == (self.target_sequences is None):
        raise ValueError("provide exactly one of `mmseqs_db` or `target_sequences`")
    return self

Model / checkpoint selection: one model_checkpoint field

When a tool exposes both a fixed set of known/bundled checkpoints and the ability to load an arbitrary checkpoint file, collapse them into a single model_checkpoint: str field — do not ship separate model_name (an enum) and checkpoint_path fields. Contributors reflexively add both; standardize on one. The single field takes either a bundled model name or a path to a checkpoint file:

• Bundled names are the tool's known presets. They must be auto-downloaded by setup.sh into the managed weights cache (proto_resolve_weights_dir → PROTO_MODEL_CACHE), so a user never fetches anything to use them. A small map is the single source of truth for "which names exist," each resolving to a file under the weights dir (or None = the framework's own default weights).
• Anything else is treated as an explicit path to the user's own checkpoint, used as-is with a clear not-found error (which also guards typos like "opendde_v2").

model_checkpoint: str = ConfigField(
    title="Model / Checkpoint", default="<default_name>",
    description="Bundled model name ('<name_a>' or '<name_b>') or a path to a custom .pt checkpoint.",
)

• remote_unsupported_reason rejects only a non-bundled value (a local path isn't on a hosted worker); bundled names are remote-OK because they're provisioned.
• Keep the bundled-name map in sync with setup.sh's download list, and add a test asserting setup.sh fetches every bundled checkpoint.
• Reference implementation: tools/structure_prediction/opendde (_BUNDLED_MODELS in the tool, _BUNDLED_CHECKPOINTS in standalone/inference.py).

Code Style Conventions

These ensure consistent formatting across all generated tools:

1. Module docstrings — Use one-liner format: """Tool description.""". No multi-line module docstrings.
2. No unused imports — Only import what's actually used. Don't speculatively include Path, Optional, numpy, etc.
3. Type alias labels — When using type aliases for shared data models, label them with comments:

   # Input:
   ToolSampleInput = InverseFoldingInput
   # Output:
   ToolSampleOutput = InverseFoldingOutput

4. ToolInstance import at top level — Import ToolInstance at module level, not lazily inside the run function:

   from proto_tools.utils.tool_instance import ToolInstance

5. logger.debug() before main loop — Add a status message before the processing loop:

   logger.debug("Using local venv for {toolkit} {operation}")

6. Don't redefine inherited fields — Fields like verbose, timeout, and device are inherited from BaseConfig. Don't redeclare them in tool-specific Config classes unless overriding the default value.
7. __init__.py files — Sort __all__ alphabetically.

For the complete tool file template, refer to .claude/skills/implement-tool/TEMPLATES.md. For implementation patterns (CPU/GPU/compile-from-source), refer to .claude/skills/implement-tool/PATTERNS.md.

Phase 3: Fan-Out (5 Parallel Subagents)

Goal: Generate all supporting files in parallel. Each subagent gets the contract (tool file) + one reference file + a focused prompt.

IMPORTANT: Launch all 5 Task() calls in a single message. If sent in separate messages, the first blocks and the rest wait.

Before launching subagents:

1. Read the contract tool file you just wrote (Phase 2)
2. Read the matching reference file for each subagent from the reference tool

Then launch all 5 subagents simultaneously:

Subagent 1: Standalone Environment

What it produces: standalone/inference.py (or run.py), standalone/setup.sh, standalone/requirements.txt, standalone/python_version.txt, optionally standalone/env_vars.txt

Prompt template:

You are implementing the standalone execution environment for a bioinformatics tool.

## Your Task
Create the standalone files for the {tool_display_name} tool in:
  proto_tools/tools/{category}/{toolkit}/standalone/

## Contract (the tool file this standalone serves)
<paste full tool file content here>

## Reference Implementation
Here is the standalone from {reference_tool} to use as a structural template:
<paste reference inference.py content>
<paste reference setup.sh content>
<paste reference requirements.txt content>

## Research: Source Repo API
<paste key findings about how the research code works — model loading, inference API, etc.>

## Files to Create

### 1. inference.py (for GPU/AI tools) or run.py (for CPU tools)
Structure:
- Module docstring
- Imports (json, sys, os at top; heavy imports inside functions)
- **Module-level logger** (REQUIRED — see "Logger Convention" below):
  ```python
  from standalone_helpers import get_logger

  logger = get_logger(__name__)

• Model class with lazy loading pattern:
  • __init__: set self._loaded = False
  • load(device): import heavy deps, load model weights
  • Core method(s) matching the operations from the contract
  • serialize_output() from standalone_helpers for tensor → list conversion
• dispatch(input_dict) function as entry point:
  • Global model instance (lazy-initialized)
  • Routes by input_dict["operation"] to model methods
  • Returns JSON-serializable dict
• if __name__ == "__main__": block reading/writing JSON files

Logger Convention (enforced by tests/style_consistency_tests/test_standalone_logger_consistency.py)

Every .py file under proto_tools/tools/*/*/standalone/ — inference.py, run.py, binary_config.py, model.py, helper modules, etc. — must declare a module-level logger:

from standalone_helpers import get_logger

logger = get_logger(__name__)

Why: standalones run inside isolated micromamba subprocesses where proto_tools is NOT importable. get_logger is shipped via standalone_helpers/proto_logging.py (published on PYTHONPATH). Records emitted through it are bridged back to the parent process for re-emission under proto_tools.worker.{toolkit}.*. Plain logging.getLogger(__name__) produces a logger outside the bridge namespace and its records are silently dropped.

__init__.py is the only exemption (no log calls there). The consistency test enforces this on every other .py file in the standalone tree.

For status updates that should drive the parent's spinner subtitle (and never clutter console output), use the dedicated method:

logger.update_status("Loading checkpoint")  # spinner subtitle update; not shown in console
logger.info("...")                           # normal log line

If the tool has file-format conversion helpers (e.g., writing MSA to Parquet, converting complexes to YAML/FASTA), implement them as plain-type functions in helpers.py at the tool directory level. These functions must be self-contained (no proto_tools imports) and take only plain types (str, list, dict). The tool layer imports and wraps them with typed signatures. Keeping them dependency-free lets the hosted execution backend reuse the same files, preserving a single source of truth. See esmfold, chai1, boltz2 for examples.

CRITICAL RULES:

• Every .py file in standalone/ (except __init__.py) MUST declare logger = get_logger(__name__) from standalone_helpers — see "Logger Convention" above. Plain logging.getLogger(__name__) is forbidden and enforced by a consistency test.
• Heavy imports (torch, model libraries) ONLY inside methods, never at module level
• The dispatch() function is the entry point for both persistent-worker and one-shot execution
• All tensor/array outputs must be converted to Python lists via serialize_output() from standalone_helpers
• Match the operation names used in the contract's ToolInstance.dispatch() calls
• Device handling: accept device from input_dict, pass to model.load()
• Seed handling: pass config.seed (raw Optional[int]) through the dispatch dict. In inference.py, call set_torch_seed(seed) unconditionally — the helper's None-check gates the expensive cuDNN flags. For any downstream sampler that needs a concrete int, do sampling_seed = seed if seed is not None else get_random_int() using the helper from standalone_helpers.
• stochastic=True: set when outputs depend on config.seed. Cacheable unseeded calls skip cache; iterable dispatches skip dedup so duplicate items diverge via per-item RNG advancement. The framework does not unroll multi-item dispatches — per-item seed handling is the tool's responsibility. Do not set for tools that accept but ignore the seed.
• Audit hardcoded values in the reference/research code (chain IDs, model paths, default parameters, assumed dimensions). If a value could vary across valid inputs, parameterize it — accept it from the dispatch input dict rather than hardcoding. For example, chain_id="A" should come from the input, not be a constant.

Required Device Management Protocol Functions

Sync note: These patterns are duplicated in PATTERNS.md (templates + reference section) because subagents can't see PATTERNS.md. If the protocol changes, update both files.

All standalone scripts MUST implement two module-level protocol functions for DeviceManager integration. Add these AFTER the dispatch() function (or after operation functions for CPU tools), BEFORE if __name__.

1. to_device(device: str) -> dict

Enables DeviceManager to move models between GPUs and CPU.

Persistent PyTorch tools (global _model kept loaded):

def to_device(device: str) -> dict:
    global _model
    if _model is not None and _model._loaded:
        _model.to_device(device)
        return {"success": True, "device": device}
    else:
        return {"success": True, "device": device, "note": "model not loaded yet"}

Non-persistent / CLI tools (no persistent model state):

def to_device(device: str) -> dict:
    return {"success": True, "device": device, "note": "CLI tool, auto-unloads"}

2. get_memory_stats() -> dict

Enables DeviceManager to query GPU memory usage. Import the helper from standalone_helpers (published on PYTHONPATH).

PyTorch tools:

def get_memory_stats() -> dict:
    from standalone_helpers import get_pytorch_memory_stats
    global _model
    device = _model.device if _model and hasattr(_model, "device") else 0
    return get_pytorch_memory_stats(device)

JAX tools:

def get_memory_stats() -> dict:
    from standalone_helpers import get_jax_memory_stats
    return get_jax_memory_stats(device_index=0)

CPU-only tools:

def get_memory_stats() -> dict:
    return {"available": False, "framework": "cpu", "note": "CPU tool"}

GPU CLI tools (e.g., Boltz2, RFDiffusion3 — spawn subprocesses that use GPU):

def get_memory_stats() -> dict:
    from standalone_helpers import get_pytorch_memory_stats
    return get_pytorch_memory_stats(device=0)

Key rules:

• Both functions MUST be at module level (not inside classes)
• to_device() returns {"success": bool, "device": str}
• get_memory_stats() returns {"available": bool, "framework": str, ...}
• Import standalone_helpers (published on PYTHONPATH) — do NOT import from proto_tools

2. setup.sh

All setup.sh scripts source standalone_helpers.sh, resolved off PATH inside the build subprocess, for shared infrastructure functions. Reference: tools/inverse_folding/fampnn/standalone/setup.sh.

uv is already installed. ToolInstance creates every tool env with pip and uv pinned to UV_VERSION (proto_tools/utils/tool_instance.py) before setup.sh runs. Call uv pip install directly; never pip install uv or guard on command -v uv. Only if the tool's build breaks on that uv, pin another with an optional standalone/uv_version.txt (default: <version>; see "uv Version Override" in notes/tool-environments.md).

#!/bin/bash
set -euo pipefail
source standalone_helpers.sh

# PyTorch tools (add extras like torchvision if needed):
proto_install_pytorch
# proto_install_pytorch "" torchvision  # torch + version-matched torchvision

# JAX tools (with optional tool-specific override prefix):
# proto_install_jax MYTOOL

# If CUDA toolkit needed (for JIT compilation):
# proto_install_cuda_toolkit

uv pip install -r requirements.txt

# Non-HF tools that download weights:
proto_resolve_weights_dir {toolkit}
if [ ! -f "$WEIGHTS_DIR/model.pt" ]; then
    wget -q -O "$WEIGHTS_DIR/model.pt" "https://example.com/model.pt"
fi

# Gated HF models (validates access before pip install):
# proto_check_gated_hf_repo "org/model-name" "https://huggingface.co/org/model-name"

Available functions (see utils/standalone_helpers_source/standalone_helpers.sh):

• proto_install_pytorch [spec] [extras...] — install PyTorch via RECOMMENDED_TORCH_SPEC, with optional version-matched extras (e.g., torchvision)
• proto_install_jax [TOOL_PREFIX] — install JAX with tool-specific overrides
• proto_install_cuda_toolkit [constraint] [extras...] — micromamba CUDA toolkit
• proto_resolve_weights_dir <toolkit> — set $WEIGHTS_DIR via PROTO_MODEL_CACHE
• proto_check_gated_hf_repo <repo_id> <license_url> [probe_file] — validate HF access

For the standalone inference.py, non-HF tools must call resolve_weights_dir() to find weights at runtime:

from standalone_helpers import resolve_weights_dir
weights_dir = resolve_weights_dir("{toolkit}")
model_path = os.path.join(weights_dir, "model.pt")

3. requirements.txt

List all Python dependencies with version pins. Do NOT include torch/jax — those are installed by setup.sh using hardware-aware detection.

4. python_version.txt

Required. Every tool must pin its Python version — ToolInstance raises FileNotFoundError at setup time for a tool without this file, and the style consistency tests fail if any tool is missing it. Always use the keyed format with a required default key. Comments (# to end of line) and blank lines are allowed.

default: 3.12

Per-platform overrides (rarely needed): use only when an upstream dependency is unavailable for the default Python on a specific platform. Three-tier lookup, most specific wins:

1. {system}-{machine} (e.g. linux-aarch64)
2. {system} (e.g. linux)
3. default (required catch-all)

default: 3.11
linux-aarch64: 3.10    # only when forced by dep availability — see PyRosetta

Pick default to match the reference tool unless the upstream library has explicit Python-version constraints. The full spec lives in notes/tool-environments.md. Validation: every python_version.txt is checked by tests/style_consistency_tests/test_python_version_consistency.py, which enforces the format and the required default key on every shipped file.