tool-trainer 0.4.1

Train LLMs to use web search as a tool — trace generation, fine-tuning, and evaluation with Tavily.

tool-trainer

Train LLMs to use web search as a tool.

Generate tool-use trajectories, fine-tune models on them, and evaluate search behavior — all in one pipeline, powered by Tavily.

The goal isn't to fine-tune on web data. It's to teach models when to search, how to phrase queries, and how to ground answers in retrieved evidence. Future evals target tool-use benchmarks like BrowseComp, GAIA, and SimpleQA-style factuality — not generic QA accuracy.

Status: v0.4 — trace generation, multi-step critique loop, eval on BrowseComp/SimpleQA, and LoRA SFT training all shipped. Pre-release (first public release is v1.0); see CHANGELOG.md.

Install

pip install tool-trainer

You'll need API keys for Tavily (search) and OpenAI (teacher model):

export TAVILY_API_KEY=tvly-...
export OPENAI_API_KEY=sk-...

Quickstart

# 1. Scaffold a project
tt init --path ./my-run
cd my-run

# 2. Put your questions in data/train.jsonl (JSONL with {"id", "question"} rows)

# 3. Validate config + data
tt validate --config config.yaml --data data/train.jsonl

# 4. Generate tool-use trajectories
tt generate \
  --config config.yaml \
  --input data/train.jsonl \
  --output outputs/traces.jsonl

# 5. Inspect what the teacher produced
tt inspect outputs/traces.jsonl

# 6. Re-run deterministically from the cache (no live Tavily calls)
tt replay \
  --config config.yaml \
  --input data/train.jsonl \
  --cache outputs/caches/search.sqlite \
  --output outputs/replayed.jsonl

Model availability: the scaffolded config.yaml defaults to gpt-4o-mini (broadly available). If you copy examples/config.minimal.yaml instead it uses gpt-4.1-mini; if your OpenAI account doesn't have access, edit the file's teacher.model to gpt-4o-mini (or whichever chat model your account does have).

Trace format

Each trajectory is a tool-use sequence suitable for SFT:

{
  "id": "q1",
  "messages": [
    {"role": "user", "content": "Who is the current president of Argentina?"},
    {"role": "assistant", "tool_call": {"name": "web_search", "arguments": {"query": "current president of Argentina 2026", "max_results": 5}}},
    {"role": "tool", "name": "web_search", "content": {"results": [...]}},
    {"role": "assistant", "content": "..."}
  ],
  "labels": {"search_used": true, "search_reason": "freshness", "num_search_steps": 1}
}

What's in v0.4

• tt init — project scaffolding
• tt validate — config + schema + API-key checks
• tt generate --mode sft --max-search-steps N — agentic search loop (router → query_writer → Tavily → critic → re-query → … → answer)
• tt inspect — rich-rendered trajectory viewer
• tt replay — deterministic re-runs from the SQLite cache
• tt eval — measure the pipeline on tool-use benchmarks (HTML + Markdown + JSONL, --concurrency for parallel runs)
• tt train --method sft — fine-tune a base model with LoRA on generated trajectories (multi-step SFT supervision built in)
• --teacher-model hf:<path> — run any stage against a local HF checkpoint

Fine-tuning a model

Hardware: the default base model (Qwen/Qwen2.5-7B-Instruct) needs a CUDA GPU with ≥24 GB VRAM (single A100 / 4090 is fine). For a quick CPU/Mac smoke without real training value, swap in --base-model HuggingFaceTB/SmolLM2-135M-Instruct — that loads in seconds and runs on a laptop.

# One-time: install the training extra (torch, transformers, trl, peft, accelerate).
pip install 'tool-trainer[train]'

# Generate a bunch of trajectories from a larger training set.
tt generate --config config.yaml --input data/train.jsonl --output outputs/traces.jsonl

# Fine-tune Qwen2.5-7B-Instruct with LoRA SFT.
tt train \
  --config config.yaml \
  --method sft \
  --data outputs/traces.jsonl \
  --base-model Qwen/Qwen2.5-7B-Instruct \
  --output outputs/checkpoints/sft-001 \
  --epochs 2

# Evaluate the fine-tuned checkpoint (uses OpenAI as judge by default).
tt eval \
  --config config.yaml \
  --benchmark simpleqa \
  --teacher-model hf:Qwen/Qwen2.5-7B-Instruct+lora:outputs/checkpoints/sft-001 \
  --output outputs/eval-sft-001 \
  --limit 100

Evaluating on a benchmark

Cost: every example consumes Tavily search credits + OpenAI tokens for router → query_writer → critic → answerer → judge. With gpt-4o-mini as both teacher and judge, a 50-row run is roughly $0.05–0.15, a full SimpleQA (4,326 rows) is ~$5–8, and BrowseComp (1,266 rows × multi-step) is ~$10–15. Always start with --limit 20 to calibrate before unleashing a full sweep.

# Your own eval set:
tt eval --config config.yaml \
  --benchmark local:data/eval.jsonl \
  --output outputs/eval-001

# OpenAI SimpleQA (4,326 factoid Qs, public) — start small:
tt eval --config config.yaml --benchmark simpleqa --output outputs/sqa --limit 50

# OpenAI BrowseComp (1,266 agentic browsing Qs) — start small:
tt eval --config config.yaml --benchmark browsecomp --output outputs/bcomp --limit 20

The output directory gets summary.json (aggregate metrics), failures.jsonl (per-example traces + scores), report.md, and report.html (with collapsible per-failure drill-downs).

Current metrics: answer accuracy (LLM judge, CORRECT/INCORRECT/NOT_ATTEMPTED), search-decision accuracy, searches per example, latency, token cost. Groundedness + query-quality judges land in v0.2.x.

Roadmap

• v0.2 (shipped) — tt eval with SimpleQA/BrowseComp adapters and tool-calling metrics
• v0.3 (shipped) — tt train --method sft with LoRA + HFInferenceClient for evaluating local checkpoints
• v0.4 (shipped) — multi-step critique loop (router → query_writer → search → critic → re-query → answer)
• v0.4.1 — tt serve (terminal REPL + FastAPI)
• v0.5 — DPO + preferences mode
• v0.6 — Anthropic teacher, vLLM serving

Full details in ROADMAP.md. See CLAUDE.md for the original MVP spec.

License

Apache-2.0