# How Needle 2 Handles Tool Retrieval for Large Tool Catalogs

> Needle 2 efficiently retrieves tools from large catalogs using contrastive embeddings and top-5 relevance filtering to ensure low latency and respect context limits.

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

---

**Needle 2 indexes every tool schema using a built-in contrastive embedding head at initialization, then retrieves only the top-5 most relevant tools per query to maintain low latency and strict context window limits.**

Needle 2 is designed to work with hundreds or thousands of tools without overwhelming the language model's context window. The cactus-compute/needle repository implements an efficient **tool retrieval for large tool catalogs** that embeds schemas once and performs similarity-based selection on every turn.

## Embedding the Tool Catalog at Initialization

When you instantiate a `Needle` object, the framework immediately processes your entire tool catalog through a native engine containing a built-in contrastive head. According to the source code in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) (lines 54-66), the `Needle.__init__` method passes each tool's JSON schema to `needle_init`, which **embeds every tool schema a single time** using the engine's contrastive head.

This one-time embedding operation converts tool definitions into dense vectors that capture their semantic meaning and functional signatures. The [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) file provides the underlying schema generation via the `@tool` decorator and utility functions `build_schema` and `pydantic_schema`, which transform Python callables into JSON-schema definitions required by the embedding layer.

## Query-Driven Retrieval and the Top-5 Rule

For every user turn, Needle 2 follows a strict retrieval protocol to manage context window constraints. As documented in [`doc/apis.md`](https://github.com/cactus-compute/needle/blob/main/doc/apis.md) under the "Tool retrieval" section:

> "Five or fewer declared tools render directly. Above that, retrieval engages: at init every tool schema is embedded once … each turn embeds the query, and **only the five highest-scoring tools enter the context** … an unselected tool is unreachable, not merely unlikely."

The retrieval pipeline executes outside the language model through the following steps:

1. **Embed the user query** using the same contrastive head used for tools
2. **Compute similarity scores** between the query embedding and the pre-computed tool embeddings
3. **Select exactly the top-5 scoring tools** based on these similarity metrics
4. **Inject only these five tools** into the prompt grammar for that specific turn

Because this retrieval step occurs **outside the LLM**, the model never sees the full catalog, keeping token consumption predictable and latency consistent regardless of total tool count.

## Persisting Embeddings with tool_index_path

Computing embeddings for massive catalogs can be computationally expensive. To optimize repeated runs, `Needle.__init__` accepts a `tool_index_path` argument that enables disk-based caching of embeddings.

When provided, the engine stores embeddings keyed by a fingerprint combining the tool schemas and model version. Subsequent instantiations with matching fingerprints load instantly from disk, while any schema change triggers selective re-embedding of only the modified entries. As noted in the API documentation:

> "`tool_index_path` persists the embeddings on disk, keyed by a fingerprint over the schemas and the model; a matching fingerprint loads instantly, a changed schema re‑embeds only what changed."

## Working with Large Tool Catalogs: Code Examples

The following example demonstrates declaring a catalog with twenty tools and persisting the embedding index to disk:

```python
from needle import Needle, tool

# 1️⃣ Declare a large catalog (more than five tools)

big_catalog = [
    # simple functions decorated with @tool

    *[lambda x, i=i: f"Result {i}" for i in range(20)],   # 20 dummy tools

]

# 2️⃣ Persist embeddings so the first run does the heavy work only once

agent = Needle(
    tools=big_catalog,
    tool_index_path="my_tools.idx",   # ← disk cache for embeddings

)

# 3️⃣ Ask a question that only a few tools are relevant to

response = agent.run(
    "Which tool returns the result for i=7?",
    max_steps=1,
)

print(response["function_calls"])

# → will contain a call only to the tool whose embedding best matches the query

```

To reuse the pre-computed index in a new process without repeating the embedding step:

```python

# Re‑using the same index in a new process (instant load)

agent2 = Needle(
    tools=big_catalog,
    tool_index_path="my_tools.idx",   # loads the pre‑computed embeddings

)

print(agent2.run("Give me the result for i=13")["function_calls"])

```

## Summary

- **One-time indexing**: `Needle.__init__` in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) passes tool schemas to the native engine's `needle_init`, which embeds every tool using a built-in contrastive head
- **Hard context limit**: Only the **top-5 most relevant tools** enter the prompt context per turn, making unselected tools unreachable for that step
- **Persistent caching**: The `tool_index_path` parameter stores embeddings on disk, keyed by fingerprints of the schemas and model version
- **Schema generation**: [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) provides the `@tool` decorator and `build_schema`/`pydantic_schema` functions that create the JSON schemas required for embedding
- **External retrieval**: The similarity computation and selection happen outside the LLM, ensuring predictable token usage and latency

## Frequently Asked Questions

### How does Needle 2 handle more than five tools?

When you declare more than five tools, Needle 2 automatically engages its retrieval system. The framework embeds the user query and compares it against pre-computed tool embeddings, then selects exactly the five highest-scoring tools to include in the prompt. This hard limit ensures the language model maintains low latency and stays within context window constraints.

### What is the purpose of the tool_index_path parameter?

The `tool_index_path` parameter specifies a filesystem location where Needle 2 caches tool embeddings. This persistence layer avoids re-embedding unchanged catalogs by using a fingerprint of the tool schemas and model version. If the fingerprint matches, embeddings load instantly; if schemas changed, only the modified tools get re-embedded.

### Why does Needle 2 limit retrieval to exactly five tools?

The number five represents a design trade-off between context window limitations and functional coverage. According to the [`doc/apis.md`](https://github.com/cactus-compute/needle/blob/main/doc/apis.md) documentation, this count fits comfortably within the model's context while providing enough tool variety for meaningful selection. Since retrieval happens outside the LLM, this cap keeps token consumption predictable regardless of total catalog size.

### Where does the embedding logic reside in the source code?

The embedding initialization occurs in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) within the `Needle.__init__` method (lines 54-66), which passes JSON schemas to the native engine's `needle_init` function containing the contrastive head. The schema generation logic that creates these JSON definitions lives in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py), specifically in the `build_schema` and `pydantic_schema` functions used by the `@tool` decorator.