universal-ai-brain
Version:
🧠UNIVERSAL AI BRAIN 3.3 - The world's most advanced cognitive architecture with 24 specialized systems, MongoDB 8.1 $rankFusion hybrid search, latest Voyage 3.5 embeddings, and framework-agnostic design. Works with Mastra, Vercel AI, LangChain, OpenAI A
291 lines (251 loc) • 9.89 kB
text/typescript
import { ToolExecutor, ToolDefinition, ToolExecutionContext } from '../agent/ToolExecutor';
import { MongoDataStore } from '../persistance/MongoDataStore';
import { setupTestDb, teardownTestDb, getTestDb } from './setup';
describe('ToolExecutor', () => {
let toolExecutor: ToolExecutor;
let toolStore: MongoDataStore<ToolDefinition>;
let executionStore: MongoDataStore<any>;
beforeAll(async () => {
await setupTestDb();
const db = getTestDb();
toolStore = new MongoDataStore(db, 'agent_tools');
executionStore = new MongoDataStore(db, 'tool_executions');
toolExecutor = new ToolExecutor(toolStore, executionStore);
});
afterAll(async () => {
await teardownTestDb();
});
beforeEach(async () => {
// Clear collections before each test
await getTestDb().collection('agent_tools').deleteMany({});
await getTestDb().collection('tool_executions').deleteMany({});
// Create built-in tool definitions
await toolExecutor.createTool({
tool_id: 'tavily_web_search',
name: 'Tavily Web Search',
description: 'A tool for searching the web.',
version: '1.0.0',
status: 'active',
input_schema: {},
output_schema: {},
execute: async () => ({})
});
await toolExecutor.createTool({
tool_id: 'internal_db_lookup',
name: 'Internal DB Lookup',
description: 'A tool for looking up data in the internal database.',
version: '1.0.0',
status: 'active',
input_schema: {},
output_schema: {},
execute: async () => ({})
});
await toolExecutor.createTool({
tool_id: 'text_analysis',
name: 'Text Analysis',
description: 'A tool for analyzing text.',
version: '1.0.0',
status: 'active',
input_schema: {},
output_schema: {},
execute: async () => ({})
});
});
describe('built-in tools', () => {
it('should execute tavily_web_search tool', async () => {
const input = {
query: 'test search query',
max_results: 3,
search_depth: 'advanced'
};
const context: ToolExecutionContext = {
agent_id: 'test-agent',
workflow_id: 'test-workflow',
timeout_ms: 5000
};
const result = await toolExecutor.execute('tavily_web_search', input, context);
expect(result).toBeDefined();
expect(result.results).toBeDefined();
expect(Array.isArray(result.results)).toBe(true);
expect(result.answer).toContain('test search query');
expect(result.confidence).toBeGreaterThan(0);
});
it('should execute internal_db_lookup tool', async () => {
const input = {
collection: 'test_collection',
filter: { status: 'active' },
limit: 5
};
const result = await toolExecutor.execute('internal_db_lookup', input, { agent_id: 'test-agent', workflow_id: 'test-workflow', timeout_ms: 5000 });
expect(result).toBeDefined();
expect(result.results).toBeDefined();
expect(result.count).toBe(0); // Mock returns empty results
expect(result.message).toContain('test_collection');
});
it('should execute text_analysis tool', async () => {
const input = {
text: 'This is a positive test message',
analysis_type: 'sentiment'
};
const result = await toolExecutor.execute('text_analysis', input, { agent_id: 'test-agent', workflow_id: 'test-workflow', timeout_ms: 5000 });
expect(result).toBeDefined();
expect(result.analysis_type).toBe('sentiment');
expect(result.sentiment).toBe('positive');
expect(result.confidence).toBeGreaterThan(0);
expect(result.key_phrases).toBeDefined();
});
});
describe('tool registration and execution', () => {
it('should register and execute custom tool', async () => {
// Register a custom tool
const customTool = async (input: Record<string, any>) => {
return {
message: `Processed: ${input.data}`,
timestamp: new Date().toISOString()
};
};
await toolExecutor.createTool({
tool_id: 'custom_tool',
name: 'Custom Tool',
description: 'A custom tool.',
version: '1.0.0',
status: 'active',
input_schema: {},
output_schema: {},
execute: customTool
});
const input = { data: 'test data' };
const result = await toolExecutor.execute('custom_tool', input, { agent_id: 'test-agent', workflow_id: 'test-workflow', timeout_ms: 5000 });
expect(result.message).toBe('Processed: test data');
expect(result.timestamp).toBeDefined();
});
it('should throw error for non-existent tool', async () => {
await expect(
toolExecutor.execute('non_existent_tool', {}, { agent_id: 'test-agent', workflow_id: 'test-workflow', timeout_ms: 5000 })
).rejects.toThrow("Tool 'non_existent_tool' not found or inactive");
});
});
describe('tool definition management', () => {
it('should create tool definition', async () => {
const toolData = {
tool_id: 'test_tool',
name: 'Test Tool',
description: 'A tool for testing',
version: '1.0.0',
status: 'active' as const,
input_schema: {
type: 'object',
properties: {
input: { type: 'string' }
},
required: ['input']
},
output_schema: {
type: 'object',
properties: {
output: { type: 'string' }
}
}
};
const created = await toolExecutor.createTool({ ...toolData, execute: async () => ({}) });
expect(created._id).toBeDefined();
expect(created.tool_id).toBe('test_tool');
expect(created.performance_stats).toBeDefined();
expect(created.performance_stats!.total_calls).toBe(0);
});
it('should get tool definition', async () => {
// First create a tool
const toolData = {
tool_id: 'lookup_tool',
name: 'Lookup Tool',
description: 'A lookup tool',
version: '1.0.0',
status: 'active' as const,
input_schema: { type: 'object' },
output_schema: { type: 'object' }
};
await toolExecutor.createTool({ ...toolData, execute: async () => ({}) });
const retrieved = await toolStore.findOne({ tool_id: 'lookup_tool' });
expect(retrieved).toBeTruthy();
expect(retrieved!.tool_id).toBe('lookup_tool');
});
});
describe('execution logging', () => {
it('should log tool execution', async () => {
const input = { query: 'test' };
const context: ToolExecutionContext = {
agent_id: 'test-agent',
workflow_id: 'test-workflow',
timeout_ms: 5000
};
await toolExecutor.execute('tavily_web_search', input, context);
// Check that execution was logged
const executions = await executionStore.find({ tool_id: 'tavily_web_search' });
expect(executions.length).toBe(1);
const execution = executions[0];
expect(execution.agent_id).toBe('test-agent');
expect(execution.workflow_id).toBe('test-workflow');
expect(execution.input).toEqual(input);
expect(execution.performance.success).toBe(true);
expect(execution.performance.execution_time_ms).toBeGreaterThan(0);
});
it('should get execution history', async () => {
// Execute tool multiple times
await toolExecutor.execute('text_analysis', { text: 'test 1' }, { agent_id: 'test-agent', workflow_id: 'test-workflow', timeout_ms: 5000 });
await toolExecutor.execute('text_analysis', { text: 'test 2' }, { agent_id: 'test-agent', workflow_id: 'test-workflow', timeout_ms: 5000 });
const history = await executionStore.find({ tool_id: 'text_analysis' });
expect(history.length).toBe(2);
});
});
describe('error handling', () => {
it('should handle tool execution errors', async () => {
// Register a tool that throws an error
const errorTool = async () => {
throw new Error('Tool execution failed');
};
await toolExecutor.createTool({
tool_id: 'error_tool',
name: 'Error Tool',
description: 'A tool that throws an error.',
version: '1.0.0',
status: 'active',
input_schema: {},
output_schema: {},
execute: errorTool
});
await expect(
toolExecutor.execute('error_tool', {}, { agent_id: 'test-agent', workflow_id: 'test-workflow', timeout_ms: 5000 })
).rejects.toThrow('Tool execution failed');
// Check that error was logged
const executions = await executionStore.find({ tool_id: 'error_tool' });
expect(executions.length).toBe(1);
expect(executions[0].performance.success).toBe(false);
expect(executions[0].error).toBeDefined();
});
});
describe('rate limiting', () => {
it('should handle rate limiting', async () => {
// Create a tool with strict rate limits
const toolData = {
tool_id: 'rate_limited_tool',
name: 'Rate Limited Tool',
description: 'A tool with rate limits',
version: '1.0.0',
status: 'active' as const,
input_schema: { type: 'object' },
output_schema: { type: 'object' },
rate_limits: {
calls_per_minute: 1,
calls_per_hour: 5,
calls_per_day: 10
}
};
await toolExecutor.createTool({ ...toolData, execute: async () => ({}) });
// Register the tool function
// Rate limiting is not implemented in the new ToolExecutor
// First call should succeed
const result1 = await toolExecutor.execute('rate_limited_tool', {}, { agent_id: 'test-agent', workflow_id: 'test-workflow', timeout_ms: 5000 });
expect(result1.success).toBe(true);
});
});
});