# How Does Needle 2 Perform Tool Retrieval and Select the Top 5 Tools Per Turn?

> Discover how Needle 2 retrieves tools using CLI arguments and selects top 5 tool calls per turn by ranking log-probabilities during inference.

- Repository: [Cactus Compute, Inc./needle](https://github.com/cactus-compute/needle)
- Tags: deep-dive
- Published: 2026-08-25

---

**Needle 2 retrieves available tools by loading a JSON schema catalog via CLI arguments, then selects the top 5 most probable tool calls per turn by ranking them based on cumulative token log-probabilities generated during inference.**

Needle 2 implements a lightweight function-calling architecture that allows language models to invoke external Python functions during text generation. According to the cactus-compute/needle source code, the framework handles tool exposure, runtime retrieval, and probabilistic selection through a three-stage pipeline involving schema registration, system prompt injection, and token-level scoring.

## Understanding the Tool Registration Pipeline

### Schema Generation with `@tool`

Every Python callable exposed to the model must be decorated with `@tool`. This decorator, defined in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py), automatically constructs a JSON schema from the function’s type hints, default values, and docstring.

```python
def tool(fn: Callable) -> Callable:
    fn._needle_tool = build_schema(fn)   # stores JSON description

    return fn

```

The `build_schema` function inspects the function signature and extracts typing information including `Enum`, `typing.Literal`, and Pydantic models. It generates a schema compliant with the OpenAI/OpenRouter function-calling specification, attaching it to the function as `fn._needle_tool`. This schema contains the tool name, description, and parameter definitions required by the model.

## How Needle 2 Retrieves Tools at Runtime

### CLI Tool Loading

Tool retrieval begins at the command line. When invoking Needle 2 with the `--tools` flag (implemented in [`needle/cli.py`](https://github.com/cactus-compute/needle/blob/main/needle/cli.py)), the framework expects a path to a JSON file containing an array of tool schemas.

```bash
needle run --checkpoint model.ckpt --tools tools.json --query "Search for quantum computing news"

```

The CLI parses this JSON file and forwards the catalog to the generation orchestration layer. This design separates tool definition from model execution, allowing dynamic reconfiguration without code changes.

### System Prompt Injection

Before inference begins, Needle 2 injects the loaded tool schemas into the model’s system prompt. The model receives the complete catalog of available functions as structured JSON, enabling it to emit function-call tokens when it determines a tool invocation is appropriate for the user’s query.

## Ranking and Selecting the Top S Tools Per Turn

### Token-Level Probability Scoring

During each generation step, Needle 2 computes per-token probabilities in [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py) via the `_step_stats` function. This module uses JAX operations to calculate log-softmax distributions over the vocabulary.

```python
logp = jax.nn.log_softmax(step_logits.astype(jnp.float32), axis=-1)
p = jnp.exp(logp)
chosen = jnp.max(p, axis=-1)          # highest token probability

```

After the model emits a complete turn, Needle 2 identifies all syntactically valid tool calls within the generated output. Each candidate call receives a score equal to the sum of log-probabilities for the tokens that constituted the tool name and arguments. Higher scores indicate greater model confidence in that specific tool selection.

### The Top-S Selection Algorithm

Once scoring completes, Needle 2 sorts all valid tool calls in descending order by their cumulative log-probability. The framework then applies a configurable threshold **S** (defaulting to 5) to retain only the most likely candidates:

```python
candidates.sort(reverse=True)       # highest score first

top_s = [call for _, call in candidates[:S]]

```

The value of **S** is controlled via the `--max-tools` CLI argument. If the user specifies `--max-tools 3`, only the three highest-probability tool calls proceed to execution. This mechanism prevents the model from invoking low-confidence tools while maintaining flexibility for multi-tool workflows.

## Practical Implementation Examples

**Defining a searchable tool:**

```python
from needle.agent.tools import tool, Field

@tool
def search_web(query: str, *, top_k: int = Field(default=5,
                                                description="Maximum number of results")) -> list[str]:
    """Search the internet and return the top `top_k` URLs for `query`."""
    ...

```

**Exporting tool schemas for runtime use:**

```python
import json
from needle.agent.tools import tool
from my_toolkit import search_web, calculate_metric

# Extract schemas from decorated functions

schemas = [fn._needle_tool for fn in (search_web, calculate_metric)]

# Write catalog for CLI consumption

with open('tools.json', 'w') as f:
    json.dump(schemas, f, indent=2)

```

**Running with top-S selection:**

```bash

# Retrieve and execute only the top 3 most probable tools per turn

needle run --checkpoint model.ckpt \
           --tools tools.json \
           --query "Analyze recent AI benchmarks" \
           --max-tools 3

```

## Summary

- **Tool Registration:** The `@tool` decorator in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) attaches JSON schemas to functions via `build_schema`, extracting type hints and documentation automatically.
- **Runtime Retrieval:** The `--tools` flag in [`needle/cli.py`](https://github.com/cactus-compute/needle/blob/main/needle/cli.py) loads external schema catalogs and injects them into the system prompt.
- **Probabilistic Ranking:** The `_step_stats` function in [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py) computes token log-probabilities using `jax.nn.log_softmax` to score candidate tool calls.
- **Top-S Selection:** Valid tool calls are sorted by cumulative log-probability and filtered to the top **S** entries (default 5), configurable via `--max-tools`.

## Frequently Asked Questions

### How does Needle 2 calculate confidence scores for tool calls?

Needle 2 calculates confidence by summing the log-probabilities of tokens generated for each tool call. During inference, `_step_stats` in [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py) computes `jax.nn.log_softmax` over model logits, then aggregates these values across the tool name and argument tokens to produce a final ranking score.

### Can I change the number of tools selected per turn from the default 5?

Yes. The selection limit **S** is configurable via the `--max-tools` CLI flag when running `needle run`. For example, `--max-tools 10` allows the model to execute up to ten tool calls per turn, while `--max-tools 1` restricts it to the single most probable call.

### What format does the tool schema file need to follow?

The tools file must be a JSON array containing schemas generated by `build_schema` in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py). Each object includes standard fields like `name`, `description`, and `parameters`, matching the OpenAI function-calling specification. You generate this file by collecting `fn._needle_tool` attributes from your `@tool`-decorated functions.

### Does Needle 2 execute all valid tool calls or only the top S?

Needle 2 executes only the top **S** calls as determined by log-probability ranking. While the model may generate multiple syntactically valid function calls during a turn, the framework filters these to the highest-confidence subset before execution, preventing noisy or uncertain invocations from affecting the workflow.