Metadata-Version: 2.4
Name: tool-trainer
Version: 0.4.1
Summary: Train LLMs to use web search as a tool — trace generation, fine-tuning, and evaluation with Tavily.
Project-URL: Homepage, https://github.com/Galhadarr/tool-trainer
Project-URL: Documentation, https://github.com/Galhadarr/tool-trainer#readme
Project-URL: Repository, https://github.com/Galhadarr/tool-trainer
Project-URL: Issues, https://github.com/Galhadarr/tool-trainer/issues
Project-URL: Changelog, https://github.com/Galhadarr/tool-trainer/blob/master/CHANGELOG.md
Author: Gal Hadar
License-Expression: Apache-2.0
License-File: LICENSE
License-File: NOTICE
Keywords: agent,browsecomp,fine-tuning,llm,sft,tavily,tool-use,web-search
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Python: >=3.10
Requires-Dist: datasets>=2.19
Requires-Dist: openai>=1.40
Requires-Dist: pydantic-settings>=2.2
Requires-Dist: pydantic>=2.6
Requires-Dist: pyyaml>=6
Requires-Dist: rich>=13
Requires-Dist: tavily-python>=0.5
Requires-Dist: typer>=0.12
Provides-Extra: dev
Requires-Dist: mypy>=1.10; extra == 'dev'
Requires-Dist: pytest-mock>=3.12; extra == 'dev'
Requires-Dist: pytest>=8; extra == 'dev'
Requires-Dist: ruff>=0.5; extra == 'dev'
Requires-Dist: twine>=5; extra == 'dev'
Provides-Extra: train
Requires-Dist: accelerate>=0.33; extra == 'train'
Requires-Dist: peft>=0.12; extra == 'train'
Requires-Dist: torch>=2.3; extra == 'train'
Requires-Dist: transformers>=4.45; extra == 'train'
Requires-Dist: trl>=0.11; extra == 'train'
Description-Content-Type: text/markdown

# 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](https://tavily.com).

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](https://openai.com/index/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](./CHANGELOG.md).

## Install

```bash
pip install tool-trainer
```

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

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

## Quickstart

```bash
# 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:

```json
{
  "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.

```bash
# 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.

```bash
# 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](./ROADMAP.md). See [CLAUDE.md](./CLAUDE.md) for the original MVP spec.

## License

[Apache-2.0](./LICENSE)
