diff --git a/apps/studio/lib/ai/tool-filter.test.ts b/apps/studio/lib/ai/tool-filter.test.ts index c419c119e47..5ca60c9b7f3 100644 --- a/apps/studio/lib/ai/tool-filter.test.ts +++ b/apps/studio/lib/ai/tool-filter.test.ts @@ -8,7 +8,6 @@ import { filterToolsByOptInLevel, createPrivacyMessageTool, toolSetValidationSchema, - transformToolResult, } from './tool-filter' describe('TOOL_CATEGORY_MAP', () => { @@ -237,56 +236,6 @@ describe('createPrivacyMessageTool', () => { }) }) -describe('transformToolResult', () => { - it('should wrap a tool with a result transformation function', async () => { - const originalResult = { data: 'original' } - - const mockTool = { - description: 'Test tool', - execute: vitest.fn().mockResolvedValue(originalResult), - } as unknown as Tool - - const transformFn = vitest.fn((result: typeof originalResult) => ({ - data: `${result.data} - transformed`, - })) - - const transformedTool = transformToolResult(mockTool, transformFn) - - // Tool properties should be preserved - expect(transformedTool.description).toBe(mockTool.description) - - // Execute the transformed tool - const args = { key: 'value' } - const options = {} as any - - if (!transformedTool.execute) { - throw new Error('Transformed tool does not have an execute function') - } - - const result = await transformedTool.execute(args, options) - - // Original tool should have been called with the same arguments - expect(mockTool.execute).toHaveBeenCalledWith(args, options) - - // Transform function should have been called with the original result - expect(transformFn).toHaveBeenCalledWith(originalResult) - - // Final result should be the transformed value - expect(result).toEqual({ data: 'original - transformed' }) - }) - - it('should throw an error if tool is null', () => { - expect(() => transformToolResult(null as any, () => ({}))).toThrow('Tool is required') - }) - - it('should throw an error if tool does not have an execute function', () => { - const invalidTool = { name: 'invalid' } as any - expect(() => transformToolResult(invalidTool, () => ({}))).toThrow( - 'Tool does not have an execute function' - ) - }) -}) - describe('toolSetValidationSchema', () => { it('should accept subset of known tools', () => { const validSubset = { diff --git a/apps/studio/lib/ai/tool-filter.ts b/apps/studio/lib/ai/tool-filter.ts index 33d5f07c7fa..a344b3f03e0 100644 --- a/apps/studio/lib/ai/tool-filter.ts +++ b/apps/studio/lib/ai/tool-filter.ts @@ -42,31 +42,6 @@ export const toolSetValidationSchema = z.record( basicToolSchema ) -/** - * Transforms the result of a tool execution to a new output. - */ -export function transformToolResult( - tool: Tool, - execute: (result: OriginalResult) => NewResult -): Tool { - if (!tool) { - throw new Error('Tool is required') - } - - if (!tool.execute) { - throw new Error('Tool does not have an execute function') - } - - // Intercept the tool to add a custom execute function - return { - ...tool, - execute: async (args: any, options: any) => { - const result = await tool.execute!(args, options) - return execute(result) - }, - } as unknown as Tool -} - /** * Tool categories based on the data they access */