224 lines
7.4 KiB
Markdown
224 lines
7.4 KiB
Markdown
|
---
|
|||
|
|
license: apache-2.0
|
|||
|
|
language:
|
|||
|
|
- en
|
|||
|
|
datasets:
|
|||
|
|
- intuit/tool-optimizer-dataset
|
|||
|
|
base_model:
|
|||
|
|
- Qwen/Qwen3-4B-Instruct-2507
|
|||
|
|
pipeline_tag: text-generation
|
|||
|
|
library_name: transformers
|
|||
|
|
tags:
|
|||
|
|
- agents
|
|||
|
|
- tool-use
|
|||
|
|
- sft
|
|||
|
|
- documentation
|
|||
|
|
- text-generation
|
|||
|
|
---
|
|||
|
|
|
|||
|
|
# Agent Tool Optimizer (`intuit/agent-tool-optimizer`)
|
|||
|
|
|
|||
|
|
`intuit/agent-tool-optimizer` is a **supervised fine-tuned (SFT)** model that rewrites **tool / API descriptions** to be more usable by **LLM agents**. Given a tool name, a parameter schema, and a baseline (often human-written) description, the model produces an improved description that helps an agent:
|
|||
|
|
|
|||
|
|
- decide **when to use vs. not use** the tool
|
|||
|
|
- generate **valid parameters** (required vs optional, constraints, defaults)
|
|||
|
|
- avoid common mistakes and likely validation failures
|
|||
|
|
|
|||
|
|
This model is trained to work in a **trace-free** setting at inference time (i.e., **no tool execution traces required**).
|
|||
|
|
|
|||
|
|
For the accompanying codebase (inference + training), see: [Agent Tool Interface Optimizer](https://github.com/intuit-ai-research/tool-optimizer).
|
|||
|
|
|
|||
|
|
---
|
|||
|
|
|
|||
|
|
## What problem does this solve?
|
|||
|
|
|
|||
|
|
Tool interfaces (descriptions + parameter schemas) are the “contract” between agents and tools, but are typically written for humans. When descriptions under-specify **required parameters**, omit **constraints**, or fail to explain **tool boundaries**, agent performance can plateau and can degrade as the number of available tools increases.
|
|||
|
|
|
|||
|
|
We study tool interface improvement as a scalable complement to agent fine-tuning, and propose **Trace-Free+**: a curriculum-learning approach that transfers knowledge learned from trace-rich training to trace-free inference for unseen tools.
|
|||
|
|
|
|||
|
|
---
|
|||
|
|
|
|||
|
|
## Paper (arXiv)
|
|||
|
|
|
|||
|
|
This model is released alongside the preprint:
|
|||
|
|
|
|||
|
|
- **Learning to Rewrite Tool Descriptions for Reliable LLM-Agent Tool Use**
|
|||
|
|
Ruocheng Guo, Kaiwen Dong, Xiang Gao, Kamalika Das
|
|||
|
|
arXiv:2602.20426 (2026) — [paper](https://arxiv.org/abs/2602.20426)
|
|||
|
|
|
|||
|
|
### Citation
|
|||
|
|
|
|||
|
|
```bibtex
|
|||
|
|
@misc{guo2026learningrewritetooldescriptions,
|
|||
|
|
title={Learning to Rewrite Tool Descriptions for Reliable LLM-Agent Tool Use},
|
|||
|
|
author={Ruocheng Guo and Kaiwen Dong and Xiang Gao and Kamalika Das},
|
|||
|
|
year={2026},
|
|||
|
|
eprint={2602.20426},
|
|||
|
|
archivePrefix={arXiv},
|
|||
|
|
primaryClass={cs.AI},
|
|||
|
|
url={https://arxiv.org/abs/2602.20426},
|
|||
|
|
}
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
---
|
|||
|
|
|
|||
|
|
## Recommended prompt (trace-free)
|
|||
|
|
|
|||
|
|
This is the **canonical inference prompt** used for trace-free tool description generation (also available as `tool_prompt.txt` in the `tool-optimizer` repo).
|
|||
|
|
|
|||
|
|
```
|
|||
|
|
You are an API documentation specialist.
|
|||
|
|
|
|||
|
|
Rewrite the API description so an AI agent can:
|
|||
|
|
1) Decide when to use this API
|
|||
|
|
2) Generate valid parameters
|
|||
|
|
|
|||
|
|
Inputs:
|
|||
|
|
- API name: {tool_name}
|
|||
|
|
- Parameter schema: {parameter_json}
|
|||
|
|
- Baseline description: {original_description}
|
|||
|
|
|
|||
|
|
Infer (do not output):
|
|||
|
|
- When to use vs not use this API
|
|||
|
|
- Required vs optional parameters
|
|||
|
|
- Parameter meanings and constraints
|
|||
|
|
- Cross-parameter dependencies or exclusions
|
|||
|
|
- Common parameter mistakes
|
|||
|
|
- no examples are provided, infer from the schema and baseline description only
|
|||
|
|
|
|||
|
|
Write a clear API description that:
|
|||
|
|
- States when to use and NOT use the API
|
|||
|
|
- Does not invent or reference non-provided APIs
|
|||
|
|
- Explains each parameter's meaning, type, required/optional status, constraints, and defaults
|
|||
|
|
- Describes likely validation failures and how to avoid them
|
|||
|
|
- Abstracts patterns into general rules
|
|||
|
|
- Does not restate the full schema verbatim
|
|||
|
|
- Does not mention whether examples were provided
|
|||
|
|
|
|||
|
|
You may replace the baseline description entirely.
|
|||
|
|
|
|||
|
|
Output ONLY valid JSON (no markdown, no code blocks):
|
|||
|
|
{{"description": "<your improved API description here>"}}
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
### Inputs
|
|||
|
|
|
|||
|
|
- **`tool_name`**: the tool/API name
|
|||
|
|
- **`parameter_json`**: a JSON string describing the parameter schema (treat this as authoritative)
|
|||
|
|
- **`original_description`**: the baseline description you want to improve
|
|||
|
|
|
|||
|
|
### Output
|
|||
|
|
|
|||
|
|
The model is trained to output **only valid JSON** with a single field:
|
|||
|
|
|
|||
|
|
- **`description`**: the improved tool description (string)
|
|||
|
|
|
|||
|
|
---
|
|||
|
|
|
|||
|
|
## Prompt variation guidance (SFT-sensitive)
|
|||
|
|
|
|||
|
|
Because this model is SFT to follow a specific prompt and output contract, it can be sensitive to prompt changes. The safest strategy is to treat the prompt as a template and apply only **minimal, well-scoped edits**.
|
|||
|
|
|
|||
|
|
### Prompt invariants (do not change)
|
|||
|
|
|
|||
|
|
- Keep the three input slots exactly: `{tool_name}`, `{parameter_json}`, `{original_description}`
|
|||
|
|
- Keep: **“Output ONLY valid JSON (no markdown, no code blocks)”**
|
|||
|
|
- Keep the output schema exactly: `{"description": "..."}` (same key name; no extra keys)
|
|||
|
|
|
|||
|
|
### Safe, minimal edits (usually OK)
|
|||
|
|
|
|||
|
|
- Add 1–3 bullets under **“Infer (do not output)”** to clarify what to pay attention to
|
|||
|
|
- Add constraints under **“Write a clear API description that:”** as additional bullets
|
|||
|
|
- Add brief reminders about schema authority, parameter-name exactness, or concision
|
|||
|
|
|
|||
|
|
### Risky edits (often break JSON / behavior)
|
|||
|
|
|
|||
|
|
- Reordering or removing the output-format lines
|
|||
|
|
- Asking for examples, multi-part outputs, markdown, or extra keys
|
|||
|
|
- Changing placeholder names or introducing new “inputs” not present during training
|
|||
|
|
|
|||
|
|
### Concrete example: minimal diff that still tends to work
|
|||
|
|
|
|||
|
|
The prompt below is a conservative variation. It adds clarifications without changing the core structure or output contract:
|
|||
|
|
|
|||
|
|
```diff
|
|||
|
|
Infer (do not output):
|
|||
|
|
- Preserve key lexical tokens from the baseline description that may match user queries
|
|||
|
|
- Clarify boundaries if this API could be confused with similar tools
|
|||
|
|
|
|||
|
|
Write a clear API description that:
|
|||
|
|
- Treats the parameter schema as authoritative and does not introduce fields, types, or requirements not defined in it
|
|||
|
|
- Explains each parameter's meaning ... while keeping parameter names exactly as defined in the schema
|
|||
|
|
- Lists REQUIRED parameters before optional ones
|
|||
|
|
- Uses enumerated or candidate values exactly as defined in the schema when applicable
|
|||
|
|
- Describes likely validation failures strictly based on schema-defined constraints ...
|
|||
|
|
- Keeps the description concise and avoids unnecessary verbosity
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
---
|
|||
|
|
|
|||
|
|
## Inference
|
|||
|
|
|
|||
|
|
### Option A: Use the `tool-optimizer` library (recommended)
|
|||
|
|
|
|||
|
|
The open-source repo includes a working CLI that runs this model with either **vLLM** or **Hugging Face Transformers**:
|
|||
|
|
|
|||
|
|
```bash
|
|||
|
|
git clone https://github.com/intuit-ai-research/tool-optimizer
|
|||
|
|
cd tool-optimizer
|
|||
|
|
|
|||
|
|
# Install (one option)
|
|||
|
|
python -m pip install -e .
|
|||
|
|
|
|||
|
|
# Run inference (vLLM default)
|
|||
|
|
python src/agent_tool_optimizer/inference_main.py \
|
|||
|
|
--model_name intuit/agent-tool-optimizer \
|
|||
|
|
--dataset_id intuit/tool-optimizer-dataset
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
Notes:
|
|||
|
|
|
|||
|
|
- `--inference_engine vllm` (default) or `--inference_engine hf`
|
|||
|
|
- The dataset is expected to have a `test` split with a `prompt` field.
|
|||
|
|
|
|||
|
|
### Option B: Transformers (direct)
|
|||
|
|
|
|||
|
|
```python
|
|||
|
|
import json
|
|||
|
|
from transformers import pipeline
|
|||
|
|
import torch
|
|||
|
|
|
|||
|
|
model_id = "intuit/agent-tool-optimizer"
|
|||
|
|
gen = pipeline(
|
|||
|
|
"text-generation",
|
|||
|
|
model=model_id,
|
|||
|
|
torch_dtype=torch.bfloat16,
|
|||
|
|
device_map="auto",
|
|||
|
|
trust_remote_code=True,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
prompt = """<prompt above>"""
|
|||
|
|
|
|||
|
|
out = gen(
|
|||
|
|
[{\"role\": \"user\", \"content\": prompt}],
|
|||
|
|
max_new_tokens=512,
|
|||
|
|
do_sample=True,
|
|||
|
|
temperature=0.6,
|
|||
|
|
top_p=0.95,
|
|||
|
|
top_k=40,
|
|||
|
|
return_full_text=False,
|
|||
|
|
)
|
|||
|
|
result = out[0][\"generated_text\"]
|
|||
|
|
print(result)
|
|||
|
|
|
|||
|
|
# Optional: validate JSON
|
|||
|
|
json.loads(result)
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
---
|
|||
|
|
|
|||
|
|
## Example (Before vs After)
|
|||
|
|
|
|||
|
|
|
|||
|
|
|