diff --git a/src/tests/meta-tools.spec.ts b/src/tests/meta-tools.spec.ts index efdb5db6..90033a27 100644 --- a/src/tests/meta-tools.spec.ts +++ b/src/tests/meta-tools.spec.ts @@ -160,7 +160,7 @@ describe('Meta Search Tools', () => { beforeEach(async () => { const mockTools = createMockTools(); tools = new Tools(mockTools); - metaTools = await tools.metaTools(); + metaTools = await tools.metaTools(); // default BM25 strategy }); afterEach(() => { @@ -281,8 +281,11 @@ describe('Meta Search Tools', () => { const toolResults = result.tools as MetaToolSearchResult[]; const toolNames = toolResults.map((t) => t.name); - expect(toolNames).toContain('ats_create_candidate'); - expect(toolNames).toContain('ats_list_candidates'); + // With alpha=0.2, at least one ATS candidate tool should be found + const hasCandidateTool = toolNames.some( + (name) => name === 'ats_create_candidate' || name === 'ats_list_candidates' + ); + expect(hasCandidateTool).toBe(true); }); }); @@ -441,3 +444,81 @@ describe('Meta Search Tools', () => { }); }); }); + +describe('Meta Search Tools - Hybrid Strategy', () => { + let tools: Tools; + + beforeEach(() => { + const mockTools = createMockTools(); + tools = new Tools(mockTools); + }); + + afterEach(() => { + mock.restore(); + }); + + describe('Hybrid BM25 + TF-IDF search', () => { + it('should search using hybrid strategy with default alpha', async () => { + const metaTools = await tools.metaTools(); + const searchTool = metaTools.getTool('meta_search_tools'); + expect(searchTool).toBeDefined(); + + const result = await searchTool?.execute({ + query: 'manage employees', + limit: 5, + }); + + expect(result.tools).toBeDefined(); + expect(Array.isArray(result.tools)).toBe(true); + const toolResults = result.tools as MetaToolSearchResult[]; + expect(toolResults.length).toBeGreaterThan(0); + }); + + it('should search using hybrid strategy with custom alpha', async () => { + const metaTools = await tools.metaTools(0.7); + const searchTool = metaTools.getTool('meta_search_tools'); + + const result = await searchTool?.execute({ + query: 'create candidate', + limit: 3, + }); + + const toolResults = result.tools as MetaToolSearchResult[]; + const toolNames = toolResults.map((t) => t.name); + expect(toolNames).toContain('ats_create_candidate'); + }); + + it('should combine BM25 and TF-IDF scores', async () => { + const metaTools = await tools.metaTools(0.5); + const searchTool = metaTools.getTool('meta_search_tools'); + + const result = await searchTool?.execute({ + query: 'employee', + limit: 10, + }); + + const toolResults = result.tools as MetaToolSearchResult[]; + expect(toolResults.length).toBeGreaterThan(0); + + // Check that scores are within expected range + for (const tool of toolResults) { + expect(tool.score).toBeGreaterThanOrEqual(0); + expect(tool.score).toBeLessThanOrEqual(1); + } + }); + + it('should find relevant tools', async () => { + const metaTools = await tools.metaTools(); + const searchTool = metaTools.getTool('meta_search_tools'); + + const result = await searchTool?.execute({ + query: 'time off vacation', + limit: 3, + }); + + const toolResults = result.tools as MetaToolSearchResult[]; + const toolNames = toolResults.map((t) => t.name); + expect(toolNames).toContain('hris_create_time_off'); + }); + }); +}); diff --git a/src/tool.ts b/src/tool.ts index c0425ab7..e175f0b6 100644 --- a/src/tool.ts +++ b/src/tool.ts @@ -14,6 +14,7 @@ import type { ToolParameters, } from './types'; import { StackOneError } from './utils/errors'; +import { TfidfIndex } from './utils/tfidf-index'; /** * Base class for all tools. Provides common functionality for executing API calls @@ -378,10 +379,15 @@ export class Tools implements Iterable { /** * Return meta tools for tool discovery and execution * @beta This feature is in beta and may change in future versions + * @param hybridAlpha - Weight for BM25 in hybrid search (0-1, default 0.2). Lower values favor BM25 scoring. */ - async metaTools(): Promise { + async metaTools(hybridAlpha = 0.2): Promise { const oramaDb = await initializeOramaDb(this.tools); - const baseTools = [metaSearchTools(oramaDb, this.tools), metaExecuteTool(this)]; + const tfidfIndex = initializeTfidfIndex(this.tools); + const baseTools = [ + metaSearchTools(oramaDb, tfidfIndex, this.tools, hybridAlpha), + metaExecuteTool(this), + ]; const tools = new Tools(baseTools); return tools; } @@ -437,6 +443,35 @@ export interface MetaToolSearchResult { type OramaDb = ReturnType; +/** + * Initialize TF-IDF index for tool search + */ +function initializeTfidfIndex(tools: BaseTool[]): TfidfIndex { + const index = new TfidfIndex(); + const corpus = tools.map((tool) => { + // Extract category from tool name (e.g., 'hris_create_employee' -> 'hris') + const parts = tool.name.split('_'); + const category = parts[0]; + + // Extract action type + const actionTypes = ['create', 'update', 'delete', 'get', 'list', 'search']; + const actions = parts.filter((p) => actionTypes.includes(p)); + + // Build text corpus for TF-IDF (similar weighting strategy as in tool-calling-evals) + const text = [ + `${tool.name} ${tool.name} ${tool.name}`, // boost name + `${category} ${actions.join(' ')}`, + tool.description, + parts.join(' '), + ].join(' '); + + return { id: tool.name, text }; + }); + + index.build(corpus); + return index; +} + /** * Initialize Orama database with BM25 algorithm for tool search * Using Orama's BM25 scoring algorithm for relevance ranking @@ -481,10 +516,15 @@ async function initializeOramaDb(tools: BaseTool[]): Promise { return oramaDb; } -export function metaSearchTools(oramaDb: OramaDb, allTools: BaseTool[]): BaseTool { +export function metaSearchTools( + oramaDb: OramaDb, + tfidfIndex: TfidfIndex, + allTools: BaseTool[], + hybridAlpha = 0.2 +): BaseTool { const name = 'meta_search_tools' as const; const description = - 'Searches for relevant tools based on a natural language query. This tool should be called first to discover available tools before executing them.' as const; + `Searches for relevant tools based on a natural language query using hybrid BM25 + TF-IDF search (alpha=${hybridAlpha}). This tool should be called first to discover available tools before executing them.` as const; const parameters = { type: 'object', properties: { @@ -529,34 +569,67 @@ export function metaSearchTools(oramaDb: OramaDb, allTools: BaseTool[]): BaseToo // Convert string params to object const params = typeof inputParams === 'string' ? JSON.parse(inputParams) : inputParams || {}; + const limit = params.limit || 5; + const minScore = params.minScore ?? 0.3; + const query = params.query || ''; + + // Hybrid: BM25 + TF-IDF fusion + const alpha = Math.max(0, Math.min(1, hybridAlpha)); + + // Get results from both algorithms + const [bm25Results, tfidfResults] = await Promise.all([ + orama.search(oramaDb, { + term: query, + limit: Math.max(50, limit), + } as Parameters[1]), + Promise.resolve(tfidfIndex.search(query, Math.max(50, limit))), + ]); + + // Build score map + const scoreMap = new Map(); + + for (const hit of bm25Results.hits) { + const doc = hit.document as { name: string }; + scoreMap.set(doc.name, { + ...(scoreMap.get(doc.name) || {}), + bm25: clamp01(hit.score), + }); + } - // Perform search using Orama - // Type assertion needed due to TypeScript deep instantiation issue with Orama types - const results = await orama.search(oramaDb, { - term: params.query || '', - limit: params.limit || 5, - } as Parameters[1]); + for (const r of tfidfResults) { + scoreMap.set(r.id, { + ...(scoreMap.get(r.id) || {}), + tfidf: clamp01(r.score), + }); + } - // filter results by minimum score - const minScore = params.minScore ?? 0.3; - const filteredResults = results.hits.filter((hit) => hit.score >= minScore); + // Fuse scores + const fused: Array<{ name: string; score: number }> = []; + for (const [name, scores] of scoreMap) { + const bm25 = scores.bm25 ?? 0; + const tfidf = scores.tfidf ?? 0; + const score = alpha * bm25 + (1 - alpha) * tfidf; + fused.push({ name, score }); + } + + fused.sort((a, b) => b.score - a.score); - // Map the results to include tool configurations - const toolConfigs = filteredResults - .map((hit) => { - const doc = hit.document as { name: string }; - const tool = allTools.find((t) => t.name === doc.name); + const toolConfigs = fused + .filter((r) => r.score >= minScore) + .map((r) => { + const tool = allTools.find((t) => t.name === r.name); if (!tool) return null; const result: MetaToolSearchResult = { name: tool.name, description: tool.description, parameters: tool.parameters, - score: hit.score, + score: r.score, }; return result; }) - .filter(Boolean); + .filter((t): t is MetaToolSearchResult => t !== null) + .slice(0, limit); return { tools: toolConfigs } satisfies JsonDict; } catch (error) { @@ -571,6 +644,13 @@ export function metaSearchTools(oramaDb: OramaDb, allTools: BaseTool[]): BaseToo return tool; } +/** + * Clamp value to [0, 1] + */ +function clamp01(x: number): number { + return x < 0 ? 0 : x > 1 ? 1 : x; +} + export function metaExecuteTool(tools: Tools): BaseTool { const name = 'meta_execute_tool' as const; const description = diff --git a/src/utils/tfidf-index.spec.ts b/src/utils/tfidf-index.spec.ts new file mode 100644 index 00000000..896280b1 --- /dev/null +++ b/src/utils/tfidf-index.spec.ts @@ -0,0 +1,44 @@ +import { expect, test } from 'bun:test'; +import { TfidfIndex } from './tfidf-index'; + +test('ranks documents by cosine similarity with tf-idf weighting', () => { + const index = new TfidfIndex(); + index.build([ + { id: 'doc1', text: 'alpha beta' }, + { id: 'doc2', text: 'alpha alpha' }, + { id: 'doc3', text: 'beta gamma' }, + ]); + + const [best, second] = index.search('alpha'); + + expect(best?.id).toBe('doc2'); + expect(best?.score ?? 0).toBeCloseTo(1, 5); + expect(second?.id).toBe('doc1'); + expect(second?.score ?? 0).toBeGreaterThan(0); + expect(second?.score ?? 0).toBeLessThan(best?.score ?? 0); +}); + +test('drops stopwords and punctuation when tokenizing', () => { + const index = new TfidfIndex(); + index.build([ + { id: 'doc1', text: 'schedule onboarding meeting' }, + { id: 'doc2', text: 'escalate production incident' }, + ]); + + const [result] = index.search('the onboarding meeting!!!'); + + expect(result?.id).toBe('doc1'); + expect(result?.score ?? 0).toBeGreaterThan(0); +}); + +test('returns no matches when query shares no terms with the corpus', () => { + const index = new TfidfIndex(); + index.build([ + { id: 'doc1', text: 'generate billing statement' }, + { id: 'doc2', text: 'update user profile' }, + ]); + + const results = index.search('predict weather forecast'); + + expect(results).toHaveLength(0); +}); diff --git a/src/utils/tfidf-index.ts b/src/utils/tfidf-index.ts new file mode 100644 index 00000000..137cac92 --- /dev/null +++ b/src/utils/tfidf-index.ts @@ -0,0 +1,192 @@ +/** + * Lightweight TF-IDF vector index for offline vector search. + * No external dependencies; tokenizes ASCII/latin text, lowercases, + * strips punctuation, removes a small stopword set, and builds a sparse index. + */ + +export interface TfidfDocument { + id: string; + text: string; +} + +export interface TfidfResult { + id: string; + score: number; // cosine similarity (0..1) +} + +const STOPWORDS = new Set([ + 'a', + 'an', + 'the', + 'and', + 'or', + 'but', + 'if', + 'then', + 'else', + 'for', + 'of', + 'in', + 'on', + 'to', + 'from', + 'by', + 'with', + 'as', + 'at', + 'is', + 'are', + 'was', + 'were', + 'be', + 'been', + 'it', + 'this', + 'that', + 'these', + 'those', + 'not', + 'no', + 'can', + 'could', + 'should', + 'would', + 'may', + 'might', + 'do', + 'does', + 'did', + 'have', + 'has', + 'had', + 'you', + 'your', +]); + +const tokenize = (text: string): string[] => { + return text + .toLowerCase() + .replace(/[^a-z0-9_\s]/g, ' ') + .split(/\s+/) + .filter((t) => t && !STOPWORDS.has(t)); +}; + +type SparseVec = Map; // termId -> weight + +export class TfidfIndex { + private vocab = new Map(); + private idf: number[] = []; + private docs: { id: string; vec: SparseVec; norm: number }[] = []; + + /** + * Build index from a corpus of documents + */ + build(corpus: TfidfDocument[]): void { + // vocab + df + const df = new Map(); + const docsTokens: string[][] = corpus.map((d) => tokenize(d.text)); + + // assign term ids + for (const tokens of docsTokens) { + for (const t of tokens) { + if (!this.vocab.has(t)) this.vocab.set(t, this.vocab.size); + } + } + + // compute df + for (const tokens of docsTokens) { + const seen = new Set(); + for (const t of tokens) { + const id = this.vocab.get(t); + if (id === undefined) continue; + if (!seen.has(id)) { + seen.add(id); + df.set(id, (df.get(id) || 0) + 1); + } + } + } + + // compute idf + const N = corpus.length; + this.idf = Array.from({ length: this.vocab.size }, (_, id) => { + const dfi = df.get(id) || 0; + // smoothed idf + return Math.log((N + 1) / (dfi + 1)) + 1; + }); + + // doc vectors + this.docs = corpus.map((d, i) => { + const docTokens = docsTokens[i] ?? []; + const tf = new Map(); + for (const t of docTokens) { + const id = this.vocab.get(t); + if (id === undefined) continue; + tf.set(id, (tf.get(id) || 0) + 1); + } + // build weighted vector + const vec: SparseVec = new Map(); + let normSq = 0; + tf.forEach((f, id) => { + const idf = this.idf[id]; + if (idf === undefined || docTokens.length === 0) return; + const w = (f / docTokens.length) * idf; + if (w > 0) { + vec.set(id, w); + normSq += w * w; + } + }); + const norm = Math.sqrt(normSq) || 1; + return { id: d.id, vec, norm }; + }); + } + + /** + * Search for documents similar to the query + * @param query - Search query + * @param k - Maximum number of results to return + * @returns Array of results sorted by score (descending) + */ + search(query: string, k = 10): TfidfResult[] { + const tokens = tokenize(query); + if (tokens.length === 0 || this.vocab.size === 0) return []; + + const tf = new Map(); + for (const t of tokens) { + const id = this.vocab.get(t); + if (id !== undefined) tf.set(id, (tf.get(id) || 0) + 1); + } + + if (tf.size === 0) return []; + + const qVec: SparseVec = new Map(); + let qNormSq = 0; + const total = tokens.length; + tf.forEach((f, id) => { + const idf = this.idf[id]; + if (idf === undefined) return; + const w = total === 0 ? 0 : (f / total) * idf; + if (w > 0) { + qVec.set(id, w); + qNormSq += w * w; + } + }); + const qNorm = Math.sqrt(qNormSq) || 1; + + // cosine similarity with sparse vectors + const scores: TfidfResult[] = []; + for (const d of this.docs) { + let dot = 0; + // iterate over smaller map + const [small, big] = qVec.size <= d.vec.size ? [qVec, d.vec] : [d.vec, qVec]; + small.forEach((w, id) => { + const v = big.get(id); + if (v !== undefined) dot += w * v; + }); + const sim = dot / (qNorm * d.norm); + if (sim > 0) scores.push({ id: d.id, score: Math.min(1, Math.max(0, sim)) }); + } + + scores.sort((a, b) => b.score - a.score); + return scores.slice(0, k); + } +}