# How Needle's Tool Retrieval Works: Why Only the Top 5 Tools Are Passed to the Engine

> Discover how Needle retrieves the top 5 relevant tools for your queries. Learn about embedding, nearest-neighbor search, and efficient inference to keep prompts within context limits.

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

---

**Needle retrieves the most relevant tools for each query by embedding tool descriptions and performing nearest-neighbor search to select only the top S tools, keeping prompts within context limits and improving inference efficiency.**

Needle's **tool retrieval** is a multi-stage pipeline that converts Python functions and Pydantic models into JSON-schema tools, indexes them by embedding, and dynamically selects only the most relevant subset for each query. This article explains exactly how this works in the `cactus-compute/needle` repository and why the system enforces a strict limit on the number of tools passed to the language model.

## How Tool Retrieval Works in Needle

### Tool Registration and Schema Resolution

When you create a `Needle` agent, the `__init__` method accepts a `tools` argument containing Python callables, Pydantic models, or pre-built dictionaries. The private `_resolve` method in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) processes each entry:

```python
def _resolve(self, tools):
    schemas = []
    for entry in tools or []:
        if _is_pydantic_model(entry):
            schema = pydantic_schema(entry)
            self._functions[schema["name"]] = entry
            schemas.append(schema)
        elif callable(entry):
            schema = getattr(entry, "_needle_tool", None) or build_schema(entry)
            self._functions[schema["name"]] = entry
            schemas.append(schema)
        elif isinstance(entry, dict):
            schemas.append(entry)
    return schemas

```

Each resolved schema is stored in `self._functions` by name for later invocation, while the schemas themselves are JSON-encoded and passed to the native C-library via `needle_init`.

### From Schema to Embeddings: The Tool Index

The **tool index** lives in [`needle/agent/fetch.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/fetch.py). This module builds and loads a vector index where each tool description is embedded into a high-dimensional vector space. The index enables fast nearest-neighbor search based on semantic similarity to user queries.

When initializing an agent with `tool_index_path`, the encoded tools and index path are passed to the engine:

```python
self._tools_json = tools_json.encode("utf-8")
_lib().needle_init(self._system, self._tools_json, self._tool_index_path)

```

### Runtime Retrieval: Selecting the Top S Tools

On every inference call, the engine:

1. **Embeds the user query** using the same embedding model used for tool descriptions
2. **Performs nearest-neighbor search** against the tool index
3. **Retrieves exactly S most similar tool schemas**
4. **Injects only these S schemas** into the prompt sent to the LLM

The Python layer does not control S directly; this value is determined by engine configuration and can be overridden via CLI flags like `--top-tools`.

## Why Only the Top 5 Tools? Four Core Reasons

| Reason | Explanation |
|--------|-------------|
| **Context window limits** | LLMs have finite token budgets. Tool schemas are verbose JSON objects; including dozens could exhaust the available context. |
| **Inference latency** | Fewer tools reduce the "function call" search space, speeding up both token generation and tool selection decisions. |
| **Relevance filtering** | Cosine similarity on embeddings ensures only semantically relevant tools reach the model, reducing noise and hallucinated tool calls. |
| **Reproducibility** | A fixed S value (default 5) guarantees deterministic behavior across identical queries, critical for testing and production systems. |

The default value of **S = 5** balances these concerns: enough tools to cover multi-step queries, few enough to stay within typical context limits.

## Code Examples: Working with Tool Retrieval

### Example 1: Registering a Function as a Tool

The `@tool` decorator in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) attaches a JSON schema to any callable:

```python
from needle import Needle, tool

@tool
def add(a: int, b: int) -> int:
    """Add two integers."""
    return a + b

agent = Needle(tools=[add])
print(agent.run("What is 7 plus 5?"))

```

The decorator adds a `_needle_tool` attribute containing the pre-built schema, skipping redundant introspection.

### Example 2: Using Pydantic Models as Tools

Pydantic models automatically convert to structured tool schemas via `pydantic_schema` (lines 52-61 in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py)):

```python
from pydantic import BaseModel
from needle import Needle

class WeatherQuery(BaseModel):
    """Ask for the weather in a city."""
    city: str

agent = Needle(tools=[WeatherQuery])
print(agent.run("Tell me the weather in Berlin."))

```

### Example 3: Pre-Computing a Tool Index

For production deployments, build the tool index ahead of time to skip embedding computation at runtime:

```python

# First, build the index via CLI: needle fetch --build-index

agent = Needle(
    tools=[add, WeatherQuery],
    tool_index_path="my_tools.index"  # Enables top-S retrieval

)
print(agent.run("Add 12 and 30, then give me the weather in Tokyo."))

```

Only the top S tools from this index will be considered for each query, regardless of how many tools are registered.

## Key Implementation Files

| File | Purpose |
|------|---------|
| [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) | Core `Needle` class with `_resolve` method for schema generation |
| [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) | `build_schema`, `pydantic_schema`, and `@tool` decorator implementations |
| [`needle/agent/fetch.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/fetch.py) | Tool index creation/loading and nearest-neighbor search for top-S retrieval |
| [`tests/test_generate.py`](https://github.com/cactus-compute/needle/blob/main/tests/test_generate.py) | Verifies tool deduplication by name in `_collect_tools` |
| [`tests/test_tools.py`](https://github.com/cactus-compute/needle/blob/main/tests/test_tools.py) | Validates `@tool` decorator preserves original callable behavior |

## Summary

- **Tool retrieval in Needle** converts Python callables and Pydantic models into JSON schemas via `_resolve` in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py)
- The [`needle/agent/fetch.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/fetch.py) module maintains an embedding-based index for semantic similarity search
- Only the **top S tools** (default 5) are retrieved per query to respect context limits, reduce latency, and maximize relevance
- S is controlled by engine configuration, not Python code, ensuring consistent behavior across deployments

## Frequently Asked Questions

### How does Needle decide which 5 tools to use for a query?

Needle embeds the user's query into the same vector space as tool descriptions, then performs nearest-neighbor search against the tool index built by [`needle/agent/fetch.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/fetch.py). The 5 tools with highest cosine similarity scores are selected. This happens inside the native engine, not in Python.

### Can I increase the number of tools beyond 5?

Yes. The `S` parameter is configurable via the `--top-tools` CLI flag or engine configuration options. However, increasing this value may impact latency and risk context window exhaustion, especially with verbose tool schemas.

### What happens if I don't provide a `tool_index_path`?

Without a pre-built index, Needle falls back to including all registered tools in every prompt. This bypasses semantic retrieval entirely and should be avoided for production systems with more than a handful of tools.

### Why does Needle use both Python schema generation and a native C library?

The Python layer ([`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py)) handles the flexible, introspection-heavy work of converting arbitrary callables into JSON schemas. The native library manages the performance-critical embedding and retrieval operations, achieving speeds that pure Python cannot match for large tool collections.