AI Tool Retriever

A lightweight, extensible library for dynamically selecting the most relevant tools for AI SDK-powered agent based on user queries.

Uses semantic search to find the best tools for the job, ensuring that models receives only the necessary tools saving context space and improving accuracy.

Benefits

• Reduce Context Window Usage: Only send relevant tool descriptions to the AI model
• Improve Response Accuracy: Models perform better with focused, relevant tool sets
• Scale to Hundreds of Tools: Semantic search handles large tool collections efficiently
• Bring Your Own Stack: Pluggable architecture lets you use any embedding model or vector store
• Zero Config Start: Default local providers work out of the box with @xenova/transformers

Features

• Dynamic Tool Selection: Find relevant tools based on semantic similarity to user queries
• Pluggable Architecture: Explicitly choose your embedding model and vector store
• Default Local Providers: Start with in-memory embeddings and vector store—no external services required
• Custom Provider Support: Implement EmbeddingProvider and ToolStore interfaces for any service (OpenAI, Supabase, Pinecone, etc.)
• Search Term Enhancement: Add custom searchTerms to tool definitions for better matching
• Explicit Tool Forcing: Force tool inclusion with [toolName] syntax in queries
• Configurable Retrieval: Fine-tune with similarity thresholds, match counts, and strict mode
• Idempotent Syncing: Content hashing utilities for efficient vector store synchronization
• TypeScript Native: Full type definitions with comprehensive JSDoc documentation

Quick Start

npm install ai-tool-retriever

import { ToolRetriever, ToolDefinition } from 'ai-tool-retriever'
import { TransformersEmbeddingProvider } from 'ai-tool-retriever/providers/embedding/transformers'
import { InMemoryStore } from 'ai-tool-retriever/providers/store/in-memory'
import { z } from 'zod'
import { tool } from 'ai'

const allMyTools: ToolDefinition[] = [
	{
		name: 'getWeather',
		tool: tool({
			description: 'Fetches the weather for a given location.',
			parameters: z.object({ city: z.string() }),
		}),
		searchTerms: ['forecast', 'temperature', 'climate', 'rain', 'sun'],
	},
	{
		name: 'searchFinancialNews',
		tool: tool({
			description: 'Searches for financial news articles about a company.',
			parameters: z.object({ ticker: z.string() }),
		}),
		searchTerms: ['stocks', 'market', 'earnings', 'sec filings', 'investing'],
	},
]

const embeddingProvider = await TransformersEmbeddingProvider.create()
const store = new InMemoryStore()

const retriever = await ToolRetriever.create({
	tools: allMyTools,
	embeddingProvider,
	store,
})

const relevantTools = await retriever.retrieve('What are the latest earnings for TSLA?')
// Returns only the searchFinancialNews tool

Advanced Usage

Custom Embedding Provider (OpenAI Example)

import type { EmbeddingProvider } from 'ai-tool-retriever'
import OpenAI from 'openai'

const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY })

export class OpenAIEmbeddingProvider implements EmbeddingProvider {
	public readonly dimensions = 1536

	async getFloatEmbeddingsBatch(texts: string[]): Promise<number[][]> {
		const response = await openai.embeddings.create({
			model: 'text-embedding-3-small',
			input: texts,
		})
		return response.data.map((d) => d.embedding)
	}
}

Retrieve Options

const relevantTools = await retriever.retrieve('Query with [forcedTool]', {
	matchCount: 5, // Max tools to return
	matchThreshold: 0.75, // Minimum similarity score
	strict: true, // Throw if forced tool not found
})

Persistent Vector Stores

The library supports any vector database through the ToolStore interface. See the core package README for a complete Supabase implementation example with idempotent syncing.

Packages

Package	Description
ai-tool-retriever	Core library with semantic tool retrieval, pluggable providers, and utilities
benchmark	Performance benchmarks for embedding throughput and search scalability
e2e	End-to-end tests validating real-world behavior

Documentation

• Core Package README - Complete guide with installation, advanced usage, custom providers, and CI/CD setup
• AI SDK Documentation - Build AI-powered applications with TypeScript

License

Licensed under the MIT License.