Skip to content

Write a custom tool#

Goal: add a new Python function that agents can call. The schema is auto-generated from type hints + docstring, so you just write the function. Time: 15 minutes.

When to write a tool vs use an agent#

A tool is a function. An agent is a stateful ReAct loop that calls tools. Use a tool when: - It's a single callable operation with clear input / output - It doesn't need its own LLM reasoning - It can be reused across many agents

If it needs reasoning ("pick the best algorithm for this data"), that's an agent, not a tool.

1. Write the function#

Add to ml_team/tools/<your_module>.py. Example — a tool that geocodes an Indian PIN code:

# ml_team/tools/geocoding.py
from ml_team.core.tool_executor import tool
import httpx

@tool(
    description="Geocode an Indian PIN code to {city, state, district} via India Post API.",
    parallel_safe=True,      # can run concurrently with other parallel_safe tools
)
def geocode_pin(pin: str) -> dict:
    """Given a 6-digit Indian PIN code, return city, state, district.

    Args:
        pin: 6-digit numeric PIN code (string).

    Returns:
        dict with keys: city (str), state (str), district (str), pin (str).

    Raises:
        ValueError: if pin is not 6 digits.
        RuntimeError: if India Post API unavailable.
    """
    if not (pin.isdigit() and len(pin) == 6):
        raise ValueError(f"PIN must be 6 digits, got {pin!r}")

    resp = httpx.get(f"https://api.postalpincode.in/pincode/{pin}", timeout=5)
    resp.raise_for_status()
    data = resp.json()[0]
    if data.get("Status") != "Success":
        raise RuntimeError(f"India Post returned: {data.get('Status')}")
    post_office = data["PostOffice"][0]
    return {
        "city": post_office["Taluk"] or post_office["Name"],
        "state": post_office["State"],
        "district": post_office["District"],
        "pin": pin,
    }

2. Register the tool#

Tools auto-register via the @tool decorator + the module being imported. Update ml_team/tools/__init__.py to import your module:

# ml_team/tools/__init__.py
from .geocoding import geocode_pin    # triggers registration

Verify:

swarm agents tools list | grep geocode_pin
# geocode_pin | parallel_safe | ml_team.tools.geocoding

3. Add to agent allowlists#

An agent can only call tools on its tools=[] list. Give it to relevant agents:

# ml_team/config/agent_defs.py
AgentConfig(
    name="data_enricher",
    tools=["load_dataset", "geocode_pin", "execute_python"],   # added here
    ...
)

4. Write a test#

# ml_team/tests/test_tool_geocoding.py
import pytest
from unittest.mock import patch
from ml_team.tools.geocoding import geocode_pin

def test_valid_pin_returns_expected_keys():
    with patch("httpx.get") as mock_get:
        mock_get.return_value.json.return_value = [{
            "Status": "Success",
            "PostOffice": [{
                "Name": "Koramangala", "Taluk": "Bangalore South",
                "State": "Karnataka", "District": "Bengaluru Urban",
            }],
        }]
        mock_get.return_value.raise_for_status = lambda: None
        result = geocode_pin("560034")
    assert result == {
        "city": "Bangalore South",
        "state": "Karnataka",
        "district": "Bengaluru Urban",
        "pin": "560034",
    }

def test_invalid_pin_raises():
    with pytest.raises(ValueError, match="6 digits"):
        geocode_pin("12345")

def test_non_numeric_pin_raises():
    with pytest.raises(ValueError):
        geocode_pin("56OO34")

Run:

.venv/bin/pytest ml_team/tests/test_tool_geocoding.py -v

5. Use it in a pipeline#

Agents now pick it up automatically — the LLM sees it in the tool catalogue. Verify:

swarm pipelines run \
  --problem "Enrich customer records with their city from PIN" \
  --dataset customers.csv \
  --template fast_prototype

Watch the agent journal:

jq -c '.[] | select(.kind == "tool_call" and .tool == "geocode_pin")' \
  pipeline_runs/<run_id>/conversations/data_enricher.jsonl

What the LLM sees#

The @tool decorator generates this OpenAI-format schema from your function:

{
  "type": "function",
  "function": {
    "name": "geocode_pin",
    "description": "Geocode an Indian PIN code to {city, state, district} via India Post API.",
    "parameters": {
      "type": "object",
      "properties": {
        "pin": {
          "type": "string",
          "description": "6-digit numeric PIN code (string)."
        }
      },
      "required": ["pin"]
    }
  }
}

So your docstring arg descriptions are what the LLM sees. Write them clearly.

Type annotation → JSON schema#

Python type JSON schema
str {"type": "string"}
int {"type": "integer"}
float {"type": "number"}
bool {"type": "boolean"}
list[T] {"type": "array", "items": T}
dict {"type": "object"}
Optional[T] type + nullable
Literal["a", "b"] enum
pydantic.BaseModel full object schema (recursive)

If you need a richer schema (e.g. regex pattern on a string), pass it explicitly:

@tool(
    description="...",
    schema={
        "pin": {"type": "string", "pattern": "^\\d{6}$"},
    },
)
def geocode_pin(pin: str) -> dict:
    ...

Security model#

Tools are gated by the permission engine. Specifically:

  • Per-agent allowlists deny your tool for agents that don't list it
  • Feature flags can disable the tool platform-wide (@tool(feature_flag="geocoding_enabled"))
  • HITL gates — if sensitive (@require_approval("geocoding")) it prompts an operator before running
  • YAML policies — organisation-wide deny rules

Write tools as if anyone might call them; let the engine decide who actually can.

Parallel-safe#

If a tool is side-effect-free and thread-safe, mark parallel_safe=True. The tool executor will run it concurrently with other parallel-safe tools within the same agent turn — big speedup on multi-tool turns.

Default: parallel_safe=False (safe default).

Error handling contract#

  • Raise ValueError for bad user input → returned as {"error": "validation", "message": "..."}
  • Raise ApprovalRequired → escalates to HITL flow
  • Raise RuntimeError or Exception → returned as {"error": "tool_failure", "message": "..."}
  • Never swallow ApprovalRequired — always let it propagate

Next#