Implementation:Neuml Txtai EmbeddingsTool Init
| Knowledge Sources | |
|---|---|
| Domains | NLP, Agent |
| Last Updated | 2026-02-09 00:00 GMT |
Overview
Concrete tool for wrapping an Embeddings instance as an agent-invocable semantic search tool, provided by the txtai library.
Description
EmbeddingsTool is a subclass of smolagents.Tool that adapts a txtai Embeddings vector index into the standardised tool interface expected by smolagents-based agents. The constructor (__init__) accepts a configuration dictionary that specifies the tool's name, description, and the source of the embeddings index. The forward method executes a semantic search query against the index and returns the top 5 results.
The constructor performs the following steps:
- Extracts name and description from the configuration dictionary. The description is augmented with a note explaining that results are returned as a list of dictionaries with id, text, and score keys.
- Sets inputs to a fixed schema: a single query parameter of type string.
- Sets output_type to "any".
- Calls the load method, which either reuses an existing Embeddings instance from config["target"] or creates a new one by calling Embeddings().load(**config) with the provided path and/or container parameters.
- Calls super().__init__() to validate parameters and complete smolagents Tool initialisation.
Usage
Use EmbeddingsTool when you need to register a semantic search index as a tool in an agent's toolset. It is typically created indirectly through the ToolFactory (when a dict with an Embeddings target or path/container keys is encountered in the tools configuration), but can also be instantiated directly.
Code Reference
Source Location
- Repository: txtai
- File:
src/python/txtai/agent/tool/embeddings.py - Lines: 15-69
Signature
class EmbeddingsTool(Tool):
def __init__(self, config):
...
def forward(self, query):
return self.embeddings.search(query, 5)
def load(self, config):
if "target" in config:
return config["target"]
embeddings = Embeddings()
embeddings.load(**config)
return embeddings
Import
from txtai.agent import Agent # typically used via Agent config, not directly imported
from txtai.agent.tool import EmbeddingsTool # direct import
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| config | dict | Yes | Configuration dictionary for the embeddings tool. |
| config["name"] | str | Yes | Short identifier for the tool, used by the LLM to reference it. |
| config["description"] | str | Yes | Natural-language description of the tool's purpose, injected into the agent prompt. |
| config["target"] | Embeddings | No | An existing Embeddings instance to reuse. If provided, path and container are ignored. |
| config["path"] | str | No | Path to load an Embeddings index from disk. Required if target is not provided. |
| config["container"] | str | No | Container name for loading embeddings from cloud storage or a specific sub-path. |
Outputs
| Name | Type | Description |
|---|---|---|
| instance | EmbeddingsTool | A smolagents.Tool subclass with a forward(query) method. |
| forward(query) return | list[dict] | List of up to 5 result dictionaries, each with keys id (str), text (str), and score (float). |
Usage Examples
Basic Example
from txtai.agent import Agent
from txtai.embeddings import Embeddings
# Create and index an embeddings instance
embeddings = Embeddings(path="sentence-transformers/all-MiniLM-L6-v2")
embeddings.index(
[
{"id": "0", "text": "The sun is a star at the center of our solar system."},
{"id": "1", "text": "Water freezes at 0 degrees Celsius under standard pressure."},
{"id": "2", "text": "Python is a popular programming language for data science."},
]
)
# Create an agent with the embeddings as a tool
agent = Agent(
tools=[
{
"name": "knowledge_base",
"description": "Search a knowledge base of science facts",
"target": embeddings,
}
],
model="huggingface-hub/Meta-Llama-3-8B-Instruct",
)
# The agent can now autonomously search the knowledge base
result = agent("What is the sun?")
Loading from Disk
from txtai.agent import Agent
# Create an agent with embeddings loaded from a saved index
agent = Agent(
tools=[
{
"name": "documents",
"description": "Search indexed documents",
"path": "/data/my-index",
}
],
model="huggingface-hub/Meta-Llama-3-8B-Instruct",
)