tools

tools give the model the ability to take actions during training: search a corpus, call an API, execute code, read a file. each tool has a schema that the model sees (name, description, input parameters) and an implementation that runs when the model calls it.

during training, the model learns when to use each tool and how to call it with the right arguments to maximize its reward.

defining a tool

a tool is an OpenAI-format Tool dict with a name, description, and JSON schema for its inputs:

from benchmax.envs import Tool

search_tool: Tool = {
    "type": "function",
    "function": {
        "name": "search",
        "description": "Search the knowledge base.",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "the search query"},
                "limit": {"type": "integer", "description": "max results to return"},
            },
            "required": ["query"],
        },
    },
}

the description and parameters are what the model sees when deciding whether and how to call the tool. clear descriptions lead to better tool usage.

registering and implementing tools

implement list_tools to declare available tools and run_tool to handle calls:

class MyEnv(BaseEnv):
    async def list_tools(self):
        return [search_tool]

    async def run_tool(self, rollout_id: str, tool_name: str, **tool_args):
        if tool_name == "search":
            return await self._search(tool_args["query"], tool_args.get("limit", 10))

run_tool receives the tool name and arguments exactly as the model provided them. return a string; this becomes the tool result the model sees in its context. tools can also return a list of OpenAI content parts: the geo3k example returns an image crop as an image_url part, which vision-language models receive as real image tokens. see multimodal training.

multiple tools

when your environment has several tools, use a dispatch pattern:

async def list_tools(self):
    return [search_tool, execute_tool, summarize_tool]

async def run_tool(self, rollout_id: str, tool_name: str, **tool_args):
    handler = {
        "search": self._search,
        "execute": self._execute,
        "summarize": self._summarize,
    }[tool_name]
    return await handler(**tool_args)

this is how the postgres-search example’s SearchEnv works: it registers a search tool over your corpus and dispatches calls to the search backend.

rollout lifecycle hooks

for tools that need per-example setup (e.g. copying a file into a workspace before the model can edit it), override the rollout_context hook. it’s an async context manager that wraps each rollout: acquire resources before yield, release them after.

from contextlib import asynccontextmanager

@asynccontextmanager
async def rollout_context(self, rollout_id, example):
    # example.payload holds the dataset row for this rollout
    self._workspace[rollout_id] = setup_workspace(example.payload["file_path"])
    try:
        yield
    finally:
        cleanup(self._workspace.pop(rollout_id, None))

per-example setup data rides in the example payload built by your create_dataset row mapper; every payload field other than prompt_messages is available here and in rollout.example_args. see dataset.

rollout_context is for per-rollout resources. shared resources the environment owns across rollouts (an http client, a database pool) belong in the constructor and are closed in aclose, which runs once when the environment is done.

full example

here’s a complete environment for a code-editing task. the model can read files, write patches, and run tests; after it finishes, compute_reward checks whether the tests pass.

import subprocess
from contextlib import asynccontextmanager
from pathlib import Path

from benchmax.envs import (
    BaseEnv,
    BaseRollout,
    Example,
    JsonlDataset,
    Tool,
    canonical_example_id,
)

read_tool: Tool = {
    "type": "function",
    "function": {
        "name": "read_file",
        "description": "Read the contents of a file in the workspace.",
        "parameters": {
            "type": "object",
            "properties": {
                "path": {"type": "string", "description": "relative path to the file"},
            },
            "required": ["path"],
        },
    },
}

write_tool: Tool = {
    "type": "function",
    "function": {
        "name": "write_file",
        "description": "Overwrite a file in the workspace with new content.",
        "parameters": {
            "type": "object",
            "properties": {
                "path": {"type": "string", "description": "relative path to the file"},
                "content": {"type": "string", "description": "new file content"},
            },
            "required": ["path", "content"],
        },
    },
}

run_tests_tool: Tool = {
    "type": "function",
    "function": {
        "name": "run_tests",
        "description": "Run the test suite and return stdout + exit code.",
        "parameters": {"type": "object", "properties": {}},
    },
}


class CodeEditEnv(BaseEnv):
    max_turns = 20

    def __init__(self):
        super().__init__()
        self._workspaces: dict[str, str] = {}  # rollout_id -> tmp dir path

    async def list_tools(self):
        return [read_tool, write_tool, run_tests_tool]

    @asynccontextmanager
    async def rollout_context(self, rollout_id, example):
        # copy the repo snapshot for this example into an isolated tmp dir
        self._workspaces[rollout_id] = setup_workspace(example.payload["repo_path"])
        try:
            yield
        finally:
            cleanup(self._workspaces.pop(rollout_id, None))

    async def run_tool(self, rollout_id: str, tool_name: str, **tool_args):
        workspace = self._workspaces[rollout_id]
        if tool_name == "read_file":
            full_path = f"{workspace}/{tool_args['path']}"
            return open(full_path).read()
        elif tool_name == "write_file":
            full_path = f"{workspace}/{tool_args['path']}"
            open(full_path, "w").write(tool_args["content"])
            return "ok"
        elif tool_name == "run_tests":
            result = subprocess.run(
                ["pytest", "--tb=short", "-q"],
                cwd=workspace,
                capture_output=True,
                text=True,
            )
            return f"exit {result.returncode}\n{result.stdout}{result.stderr}"

    async def compute_reward(self, rollout: BaseRollout) -> dict[str, float]:
        workspace = self._workspaces[rollout.rollout_id]
        result = subprocess.run(
            ["pytest", "-q"], cwd=workspace, capture_output=True, text=True
        )
        return {"tests_pass": 1.0 if result.returncode == 0 else 0.0}

    async def create_dataset(self, split, base_dir: Path, *, max_examples=None):
        return JsonlDataset(
            base_dir / f"{split}.jsonl",
            row_to_example=self._example,
            max_examples=max_examples,
        )

    def _example(self, row: dict) -> Example:
        payload = {
            "prompt_messages": [
                {
                    "role": "user",
                    "content": f"Fix the failing tests in this repo.\n\n{row['task']}",
                }
            ],
            "repo_path": row["repo_path"],
        }
        return Example(id=canonical_example_id(payload), payload=payload)