Teaching Enterprise AI to Choose the Right Tool

An AI assistant can only act well if it chooses the right tool: the API that looks up a particular record, updates a system, or handles a request. In a catalog of thousands of enterprise APIs, descriptions may overlap. The deciding detail might be an internal rule about which product version or date range belongs to which endpoint—information a short tool description may not capture.

The study introduces TRACE, a method for training a language model to retrieve tools from a catalog while preserving its learned knowledge about them. The authors, researchers at SAP Labs, ask whether a model can recognize tools and apply business rules when selecting among them, without relying on the slower constrained decoding used by earlier parametric retrieval methods.

A two-stage training recipe

In parametric retrieval, each tool is assigned a special token, and the model learns to produce that token for the desired tool. Earlier approaches can use constrained beam search, which checks candidate outputs against a list of valid tool tokens. The paper reports that retrieval training of this kind can weaken the model’s ability to answer questions about tools, while beam search adds latency.

TRACE separates training into two stages. First, the model learns associations between tool information and special tokens through multiple formats. Next, it trains on user queries paired with a reasoning trace and a JSON list of selected tool tokens. Some examples target general retrieval; others are designed around rules curated by domain experts, such as routing constraints for overlapping APIs. Tool names in the training traces are replaced with their special tokens, connecting the explanation to the output.

At inference, the model generates a trace and then its tool-token list using greedy decoding, rather than searching through constrained beams. The authors test it on 8,283 tools across HR and Finance. Their evaluation includes retrieval, questions probing tool knowledge, and a separate multiple-choice knowledge test authored by experts. The production retrieval benchmark and tool catalogs are proprietary and are not released.

Results depend on the domain

The paper reports that training only to map queries to tool tokens sharply reduced scores on tool-knowledge probes, in some cases bringing them close to chance. Reasoning-based training recovered much of that performance. In the headline single-beam setup, TRACE achieved 85.5% recall on Domain A and 60.2% on Domain B. The authors also report embedding-baseline top-10 recall of 27.5% and 52.7%. These are not like-for-like figures: the baseline returns up to ten candidates, while TRACE’s reported result measures recall over its generated tool set. The numbers therefore do not establish a direct comparison under identical output conditions.

The paper also evaluates constrained-decoding variants and an embedding-plus-LLM-reranking pipeline. Under single-beam decoding, TRACE’s recall exceeded that pipeline’s reported results on both domains, with a much larger margin in Domain A. The authors associate weaker Domain B performance with gaps in the rules covering confusing tool groups. Adding rule-targeted training data improved results in Domain A, but the paper does not establish a general formula for how many rules are sufficient.

The speed result is notable but specific to the tested setup. On one H200 GPU, the authors measured about 1.9 seconds per query for single-beam reasoning and about 19 seconds for constrained beam-10 decoding. At concurrency 32, they report 11.2 queries per second versus 0.05. These are measurements from the authors’ configuration, not guarantees for other hardware, models, or workloads.

What remains uncertain

The findings come from one base model, one proprietary catalog spanning two domains, and benchmarks that outside researchers cannot inspect. Domain experts wrote the business rules, and the study does not test automatic rule extraction. The authors also acknowledge that single-beam retrieval trails constrained decoding in the larger Finance catalog. A small share of generated tokens do not correspond to catalog tools; the system discards them, but this shows that unconstrained output is not error-free.

TRACE is described in an arXiv preprint, which may not have undergone peer review. The results are best treated as evidence about a promising approach in a specific enterprise setting, not proof that it will transfer broadly. Independent replication, comparable accessible benchmarks, and tests across models and organizations would help establish how robust the gains are and how much expert rule-writing they require.