diff --git a/.changeset/six-stars-eat.md b/.changeset/six-stars-eat.md new file mode 100644 index 0000000..e89bf76 --- /dev/null +++ b/.changeset/six-stars-eat.md @@ -0,0 +1,5 @@ +--- +"@vantige-ai/typescript-sdk": patch +--- + +Adding support for Vercel AI SDK v4. Users can now specify if they are using v4 of the AI SDK which will use the 'parameters' key instead of the newer 'inputSchema' key when creating tools. diff --git a/examples/ai-sdk-tools-usage.ts b/examples/ai-sdk-tools-usage.ts index 611ca48..224e9db 100644 --- a/examples/ai-sdk-tools-usage.ts +++ b/examples/ai-sdk-tools-usage.ts @@ -6,6 +6,13 @@ import { z } from 'zod'; /** * Example showing how to create AI SDK tools from available knowledge bases + * + * This example demonstrates how to create tools compatible with different versions + * of Vercel's AI SDK: + * - AI SDK v4: Uses 'parameters' property for input schema + * - AI SDK v5+: Uses 'inputSchema' property (MCP-aligned) + * + * You can specify the version as the third parameter to any tool creation function. */ async function exampleAISDKToolsUsage() { // Initialize the Vantige client @@ -33,7 +40,8 @@ async function exampleAISDKToolsUsage() { console.log('\n=== Creating Simple AI SDK Tools ==='); const simpleTools = createSimpleKnowledgeBaseTools( availableResponse.knowledgeBases, - client + client, + 'v5' // Default to v5, but you can specify 'v4' for legacy support ); console.log('Simple tools created:', Object.keys(simpleTools)); @@ -60,7 +68,8 @@ async function exampleAISDKToolsUsage() { console.log('\n=== Creating Full-Featured AI SDK Tools ==='); const fullTools = createKnowledgeBaseTools( availableResponse.knowledgeBases, - client + client, + 'v5' // Default to v5, but you can specify 'v4' for legacy support ); console.log('Full tools created:', Object.keys(fullTools)); @@ -72,6 +81,7 @@ async function exampleAISDKToolsUsage() { client, { simplified: true, // Use simple interface + version: 'v5', // Specify AI SDK version ('v4' or 'v5') keyGenerator: (kb) => `search-${kb.name.toLowerCase().replace(/\s+/g, '-')}`, descriptionGenerator: (kb) => `Search the ${kb.name} knowledge base for relevant information. ${kb.description || ''}`, } @@ -79,7 +89,33 @@ async function exampleAISDKToolsUsage() { console.log('Custom tools created:', Object.keys(customTools)); - // ===== EXAMPLE 4: Testing a Tool ===== + // ===== EXAMPLE 4: AI SDK Version Comparison ===== + console.log('\n=== AI SDK Version Comparison ==='); + + // Create tools for v4 (legacy) + const v4Tools = createSimpleKnowledgeBaseTools( + availableResponse.knowledgeBases.slice(0, 1), // Just one for demo + client, + 'v4' + ); + + // Create tools for v5 (current) + const v5Tools = createSimpleKnowledgeBaseTools( + availableResponse.knowledgeBases.slice(0, 1), // Just one for demo + client, + 'v5' + ); + + const firstToolKey = Object.keys(v4Tools)[0]; + if (firstToolKey) { + console.log(`\nTool "${firstToolKey}" structure comparison:`); + console.log('v4 tool has "parameters":', 'parameters' in v4Tools[firstToolKey]); + console.log('v4 tool has "inputSchema":', 'inputSchema' in v4Tools[firstToolKey]); + console.log('v5 tool has "parameters":', 'parameters' in v5Tools[firstToolKey]); + console.log('v5 tool has "inputSchema":', 'inputSchema' in v5Tools[firstToolKey]); + } + + // ===== EXAMPLE 5: Testing a Tool ===== if (availableResponse.knowledgeBases.length > 0) { console.log('\n=== Testing a Knowledge Base Tool ==='); const firstKB = availableResponse.knowledgeBases[0]; @@ -111,7 +147,7 @@ async function exampleAISDKToolsUsage() { } } - // ===== EXAMPLE 5: Integration with Vercel AI SDK ===== + // ===== EXAMPLE 6: Integration with Vercel AI SDK ===== console.log('\n=== Vercel AI SDK Integration Example ==='); // This is how you would use it in a real Vercel AI SDK application @@ -142,7 +178,7 @@ async function exampleAISDKToolsUsage() { console.log('AI SDK configuration ready with tools:', Object.keys(aiSDKIntegration.tools)); - // ===== EXAMPLE 6: Dynamic Tool Loading ===== + // ===== EXAMPLE 7: Dynamic Tool Loading ===== console.log('\n=== Dynamic Tool Loading Example ==='); // Function to dynamically load tools based on external scope diff --git a/jest.config.ts b/jest.config.ts index 6397365..7f51284 100644 --- a/jest.config.ts +++ b/jest.config.ts @@ -23,5 +23,7 @@ module.exports = { 'html', 'json-summary' ], - setupFilesAfterEnv: ['/src/test/setup.ts'] + setupFilesAfterEnv: ['/src/test/setup.ts'], + testTimeout: 5000, // 5 second timeout for all tests + maxWorkers: '50%' // Use half the available CPU cores for parallel test execution }; \ No newline at end of file diff --git a/src/client/__tests__/http-client.test.ts b/src/client/__tests__/http-client.test.ts index a9c9005..5ff57dc 100644 --- a/src/client/__tests__/http-client.test.ts +++ b/src/client/__tests__/http-client.test.ts @@ -13,6 +13,12 @@ describe('VantigeHttpClient', () => { let mockAxiosInstance: any; beforeEach(() => { + // Mock setTimeout to make retry delays instant + jest.spyOn(global, 'setTimeout').mockImplementation((fn) => { + fn(); + return {} as any; + }); + mockAuth = { getAuthHeaders: jest.fn().mockReturnValue({ 'Authorization': 'Bearer test-token', @@ -39,8 +45,8 @@ describe('VantigeHttpClient', () => { httpClient = new VantigeHttpClient({ baseUrl: 'https://api.vantige.ai', - timeout: 30000, - retries: 3, + timeout: 1000, // Much shorter timeout for tests + retries: 2, // Fewer retries for faster tests auth: mockAuth, debug: false, }); @@ -48,6 +54,7 @@ describe('VantigeHttpClient', () => { afterEach(() => { jest.clearAllMocks(); + jest.restoreAllMocks(); }); describe('Constructor', () => { @@ -55,7 +62,7 @@ describe('VantigeHttpClient', () => { expect(httpClient).toBeInstanceOf(VantigeHttpClient); expect(mockedAxios.create).toHaveBeenCalledWith({ baseURL: 'https://api.vantige.ai', - timeout: 30000, + timeout: 1000, headers: { 'Authorization': 'Bearer test-token', 'Content-Type': 'application/json', @@ -70,6 +77,13 @@ describe('VantigeHttpClient', () => { }); describe('Error Handling', () => { + it('should pass through existing VantigeSDKError without wrapping', async () => { + const existing = new VantigeSDKError('Existing', VantigeErrorCode.NETWORK_ERROR); + mockAxiosInstance.get = jest.fn().mockRejectedValue(existing); + + await expect(httpClient.get('/test')).rejects.toBe(existing); + }); + it('should handle network timeout errors', async () => { const timeoutError = new Error('timeout of 30000ms exceeded') as AxiosError; timeoutError.code = 'ECONNABORTED'; @@ -78,7 +92,7 @@ describe('VantigeHttpClient', () => { mockAxiosInstance.get = jest.fn().mockRejectedValue(timeoutError); await expect(httpClient.get('/test')).rejects.toThrow(VantigeSDKError); - }, 10000); + }); it('should handle network errors without response', async () => { const networkError = new Error('Network Error') as AxiosError; @@ -88,7 +102,7 @@ describe('VantigeHttpClient', () => { mockAxiosInstance.get = jest.fn().mockRejectedValue(networkError); await expect(httpClient.get('/test')).rejects.toThrow(VantigeSDKError); - }, 10000); + }); it('should handle 401 Unauthorized errors', async () => { const unauthorizedError = { @@ -102,7 +116,7 @@ describe('VantigeHttpClient', () => { mockAxiosInstance.get = jest.fn().mockRejectedValue(unauthorizedError); await expect(httpClient.get('/test')).rejects.toThrow(VantigeSDKError); - }, 10000); + }); it('should handle 403 Forbidden errors', async () => { const forbiddenError = { @@ -116,7 +130,7 @@ describe('VantigeHttpClient', () => { mockAxiosInstance.get = jest.fn().mockRejectedValue(forbiddenError); await expect(httpClient.get('/test')).rejects.toThrow(VantigeSDKError); - }, 10000); + }); it('should handle 404 Not Found errors', async () => { const notFoundError = { @@ -130,7 +144,7 @@ describe('VantigeHttpClient', () => { mockAxiosInstance.get = jest.fn().mockRejectedValue(notFoundError); await expect(httpClient.get('/test')).rejects.toThrow(VantigeSDKError); - }, 10000); + }); it('should handle 422 Validation errors', async () => { const validationError = { @@ -144,7 +158,7 @@ describe('VantigeHttpClient', () => { mockAxiosInstance.get = jest.fn().mockRejectedValue(validationError); await expect(httpClient.get('/test')).rejects.toThrow(VantigeSDKError); - }, 10000); + }); it('should handle 429 Rate Limit errors', async () => { const rateLimitError = { @@ -158,7 +172,7 @@ describe('VantigeHttpClient', () => { mockAxiosInstance.get = jest.fn().mockRejectedValue(rateLimitError); await expect(httpClient.get('/test')).rejects.toThrow(VantigeSDKError); - }, 10000); + }); it('should handle 500 Internal Server errors', async () => { const serverError = { @@ -172,7 +186,7 @@ describe('VantigeHttpClient', () => { mockAxiosInstance.get = jest.fn().mockRejectedValue(serverError); await expect(httpClient.get('/test')).rejects.toThrow(VantigeSDKError); - }, 10000); + }); it('should handle 503 Service Unavailable errors', async () => { const serviceUnavailableError = { @@ -186,7 +200,7 @@ describe('VantigeHttpClient', () => { mockAxiosInstance.get = jest.fn().mockRejectedValue(serviceUnavailableError); await expect(httpClient.get('/test')).rejects.toThrow(VantigeSDKError); - }, 10000); + }); it('should handle unknown HTTP errors', async () => { const unknownError = { @@ -200,7 +214,7 @@ describe('VantigeHttpClient', () => { mockAxiosInstance.get = jest.fn().mockRejectedValue(unknownError); await expect(httpClient.get('/test')).rejects.toThrow(VantigeSDKError); - }, 10000); + }); it('should handle non-axios errors', async () => { const genericError = new Error('Generic error'); @@ -208,10 +222,25 @@ describe('VantigeHttpClient', () => { mockAxiosInstance.get = jest.fn().mockRejectedValue(genericError); await expect(httpClient.get('/test')).rejects.toThrow(VantigeSDKError); - }, 10000); + }); }); describe('Retry Logic', () => { + it('should throw after exhausting retries', async () => { + const networkError = new Error('Network Error') as AxiosError; + networkError.isAxiosError = true; + networkError.response = undefined; + + mockAxiosInstance.get = jest.fn() + .mockRejectedValueOnce(networkError) + .mockRejectedValueOnce(networkError) + .mockRejectedValueOnce(networkError) + .mockRejectedValueOnce(networkError); + + await expect(httpClient.get('/test')).rejects.toBeInstanceOf(VantigeSDKError); + expect(mockAxiosInstance.get).toHaveBeenCalled(); + }); + it('should retry on network errors', async () => { const networkError = new Error('Network Error') as AxiosError; networkError.isAxiosError = true; @@ -227,7 +256,7 @@ describe('VantigeHttpClient', () => { const result = await httpClient.get('/test'); expect(result).toEqual(successResponse); expect(mockAxiosInstance.get).toHaveBeenCalledTimes(3); - }, 10000); + }); it('should not retry on authentication errors', async () => { const authError = { @@ -247,7 +276,7 @@ describe('VantigeHttpClient', () => { await expect(httpClient.get('/test')).rejects.toThrow(VantigeSDKError); expect(mockAxiosInstance.get).toHaveBeenCalledTimes(1); - }, 10000); + }); it('should not retry on validation errors', async () => { const validationError = { @@ -267,7 +296,7 @@ describe('VantigeHttpClient', () => { await expect(httpClient.get('/test')).rejects.toThrow(VantigeSDKError); expect(mockAxiosInstance.get).toHaveBeenCalledTimes(1); - }, 10000); + }); it('should retry on rate limit errors with delay', async () => { const rateLimitError = { @@ -288,21 +317,32 @@ describe('VantigeHttpClient', () => { .mockRejectedValueOnce(rateLimitError) .mockResolvedValueOnce({ data: successResponse }); - // Mock setTimeout to avoid actual delays in tests - jest.spyOn(global, 'setTimeout').mockImplementation((fn) => { - fn(); - return {} as any; - }); - const result = await httpClient.get('/test'); expect(result).toEqual(successResponse); expect(mockAxiosInstance.get).toHaveBeenCalledTimes(2); - - jest.restoreAllMocks(); - }, 10000); + }); }); describe('HTTP Methods', () => { + it('should use interceptor success handler to return response unmodified', () => { + const handlers: any = {}; + mockAxiosInstance.interceptors.response.use = jest.fn((success: any, fail: any) => { + handlers.success = success; + handlers.fail = fail; + }); + + // Recreate client to register interceptor with captured handlers + httpClient = new VantigeHttpClient({ + baseUrl: 'https://api.vantige.ai', + timeout: 1000, + retries: 2, + auth: mockAuth, + }); + + const resp = { data: { ok: true } } as any; + const result = handlers.success(resp); + expect(result).toBe(resp); + }); it('should make GET requests', async () => { const mockResponse = { success: true }; mockAxiosInstance.get = jest.fn().mockResolvedValue({ data: mockResponse }); diff --git a/src/client/__tests__/vantige-client.test.ts b/src/client/__tests__/vantige-client.test.ts index 06d230d..539bd5f 100644 --- a/src/client/__tests__/vantige-client.test.ts +++ b/src/client/__tests__/vantige-client.test.ts @@ -309,5 +309,50 @@ describe('VantigeClient', () => { await expect(client.listAvailableCorpuses()).rejects.toThrow(VantigeSDKError); }); }); + + describe('query', () => { + it('should validate inputs and call correct endpoint', async () => { + const mockResponse = { + success: true, + corpusId: 'kb1', + query: 'hello', + retrieval_results: [] + }; + + jest.spyOn(client['httpClient'], 'post').mockResolvedValue(mockResponse as any); + + const result = await client.query('kb1', { query: 'hello', topK: 5 }); + expect(client['httpClient'].post).toHaveBeenCalledWith('/api/v1/knowledge-base/kb1/query', { query: 'hello', topK: 5 }); + expect(result).toEqual(mockResponse); + }); + + it('should throw on invalid inputs', async () => { + await expect(client.query('', { query: 'x' } as any)).rejects.toThrow(VantigeSDKError); + await expect(client.query('kb1', { query: '' } as any)).rejects.toThrow(VantigeSDKError); + await expect(client.query('kb1', { query: 'a'.repeat(1001) } as any)).rejects.toThrow(VantigeSDKError); + await expect(client.query('kb1', { query: 'ok', topK: 0 } as any)).rejects.toThrow(VantigeSDKError); + await expect(client.query('kb1', { query: 'ok', topK: 101 } as any)).rejects.toThrow(VantigeSDKError); + }); + + it('should throw when backend responds unsuccessfully', async () => { + jest.spyOn(client['httpClient'], 'post').mockResolvedValue({ success: false } as any); + await expect(client.query('kb1', { query: 'ok' })).rejects.toThrow(VantigeSDKError); + }); + }); + + describe('testConnection', () => { + it('should return success info when listing works', async () => { + jest.spyOn(client, 'listKnowledgeBases').mockResolvedValue({ success: true } as any); + const result = await client.testConnection(); + expect(result.success).toBe(true); + expect(result.environment).toBe('test'); + expect(typeof result.latency).toBe('number'); + }); + + it('should wrap errors into network error when listing fails', async () => { + jest.spyOn(client, 'listKnowledgeBases').mockRejectedValue(new Error('fail')); + await expect(client.testConnection()).rejects.toThrow(VantigeSDKError); + }); + }); }); }); \ No newline at end of file diff --git a/src/utils/__tests__/ai-sdk-tools.test.ts b/src/utils/__tests__/ai-sdk-tools.test.ts index 2126761..0ca129a 100644 --- a/src/utils/__tests__/ai-sdk-tools.test.ts +++ b/src/utils/__tests__/ai-sdk-tools.test.ts @@ -101,6 +101,24 @@ describe('AI SDK Tools', () => { expect(tools).toHaveProperty('user-guide-faq'); }); + it('should create tools with v4 schema format', () => { + const tools = createKnowledgeBaseTools(mockKnowledgeBases, mockClient, 'v4'); + const tool = tools['product-documentation']; + + expect(tool).toHaveProperty('parameters'); + expect(tool).not.toHaveProperty('inputSchema'); + expect(tool.parameters).toBeDefined(); + }); + + it('should create tools with v5 schema format by default', () => { + const tools = createKnowledgeBaseTools(mockKnowledgeBases, mockClient); + const tool = tools['product-documentation']; + + expect(tool).toHaveProperty('inputSchema'); + expect(tool).not.toHaveProperty('parameters'); + expect(tool.inputSchema).toBeDefined(); + }); + it('should create tools with correct structure', () => { const tools = createKnowledgeBaseTools(mockKnowledgeBases, mockClient); const tool = tools['product-documentation']; @@ -190,6 +208,24 @@ describe('AI SDK Tools', () => { expect(tools).toHaveProperty('product-documentation'); }); + it('should create tools with v4 schema format', () => { + const tools = createSimpleKnowledgeBaseTools(mockKnowledgeBases, mockClient, 'v4'); + const tool = tools['product-documentation']; + + expect(tool).toHaveProperty('parameters'); + expect(tool).not.toHaveProperty('inputSchema'); + expect(tool.parameters).toBeDefined(); + }); + + it('should create tools with v5 schema format by default', () => { + const tools = createSimpleKnowledgeBaseTools(mockKnowledgeBases, mockClient); + const tool = tools['product-documentation']; + + expect(tool).toHaveProperty('inputSchema'); + expect(tool).not.toHaveProperty('parameters'); + expect(tool.inputSchema).toBeDefined(); + }); + it('should execute with only query parameter', async () => { const mockQueryResponse: QueryResponse = { success: true, @@ -286,6 +322,59 @@ describe('AI SDK Tools', () => { const tool = tools['kb-kb1']; expect(tool.description).toBe('Query Product Documentation knowledge base'); }); + + it('should create tools with v4 schema format', () => { + const tools = createKnowledgeBaseToolsWithOptions( + mockKnowledgeBases, + mockClient, + { version: 'v4' } + ); + const tool = tools['product-documentation']; + + expect(tool).toHaveProperty('parameters'); + expect(tool).not.toHaveProperty('inputSchema'); + expect(tool.parameters).toBeDefined(); + }); + + it('should create simplified tools with v4 schema format', () => { + const tools = createKnowledgeBaseToolsWithOptions( + mockKnowledgeBases, + mockClient, + { simplified: true, version: 'v4' } + ); + const tool = tools['product-documentation']; + + expect(tool).toHaveProperty('parameters'); + expect(tool).not.toHaveProperty('inputSchema'); + expect(tool.parameters).toBeDefined(); + }); + + it('should execute simplified tools with v4 schema format', async () => { + const mockQueryResponse: QueryResponse = { + success: true, + corpusId: 'kb1', + query: 'test query', + retrieval_results: [], + }; + + mockClient.query.mockResolvedValue(mockQueryResponse); + + const tools = createKnowledgeBaseToolsWithOptions( + mockKnowledgeBases, + mockClient, + { simplified: true, version: 'v4' } + ); + const tool = tools['product-documentation']; + + const result = await tool.execute!({ + query: 'test query', + }); + + expect(mockClient.query).toHaveBeenCalledWith('kb1', { + query: 'test query', + }); + expect(result).toEqual(mockQueryResponse); + }); }); describe('Error handling', () => { diff --git a/src/utils/ai-sdk-tools.ts b/src/utils/ai-sdk-tools.ts index b3c0b82..4f655e9 100644 --- a/src/utils/ai-sdk-tools.ts +++ b/src/utils/ai-sdk-tools.ts @@ -21,7 +21,8 @@ export type MinimalToolExecuteFunction = ( export type AISDKToolLike = { description?: string; providerOptions?: Record; - inputSchema: unknown; // zod schema works here + inputSchema?: unknown; // zod schema works here (AI SDK v5+) + parameters?: unknown; // Legacy support for AI SDK v4 onInputStart?: (options: MinimalToolCallOptions) => void | Promise; onInputDelta?: ( options: { inputTextDelta: string } & MinimalToolCallOptions, @@ -59,11 +60,13 @@ function generateToolKey(knowledgeBase: AvailableKnowledgeBase): string { * * @param knowledgeBases - Array of available knowledge bases * @param client - VantigeClient instance for querying - * @returns Object of tools compatible with Vercel AI SDK + * @param version - AI SDK version to target ('v4' or 'v5', defaults to 'v5') + * @returns Object of tools compatible with the specified Vercel AI SDK version */ export function createKnowledgeBaseTools( knowledgeBases: AvailableKnowledgeBase[], client: VantigeClient, + version: 'v4' | 'v5' = 'v5', ): Record> { const tools: Record> = {}; @@ -71,40 +74,41 @@ export function createKnowledgeBaseTools( const toolKey = generateToolKey(kb); // Create the tool using zod schema and execute function - tools[toolKey] = { + const schema = z.object({ + query: z + .string() + .describe("The search query to find relevant information"), + topK: z + .number() + .min(1) + .max(100) + .optional() + .describe("Maximum number of results to return (1-100)"), + includeMetadata: z + .boolean() + .optional() + .describe("Whether to include metadata in results"), + useGeneration: z + .boolean() + .optional() + .describe("Whether to use AI generation for response"), + fieldMapping: z + .object({ + sourceUri: z + .string() + .optional() + .describe("Field name for source URI"), + sourceDisplayName: z + .string() + .optional() + .describe("Field name for source display name"), + }) + .optional() + .describe("Custom field mapping for results"), + }); + + const tool: AISDKToolLike = { description: kb.description || `Query the ${kb.name} knowledge base`, - inputSchema: z.object({ - query: z - .string() - .describe("The search query to find relevant information"), - topK: z - .number() - .min(1) - .max(100) - .optional() - .describe("Maximum number of results to return (1-100)"), - includeMetadata: z - .boolean() - .optional() - .describe("Whether to include metadata in results"), - useGeneration: z - .boolean() - .optional() - .describe("Whether to use AI generation for response"), - fieldMapping: z - .object({ - sourceUri: z - .string() - .optional() - .describe("Field name for source URI"), - sourceDisplayName: z - .string() - .optional() - .describe("Field name for source display name"), - }) - .optional() - .describe("Custom field mapping for results"), - }), execute: async ( params: { query: string; @@ -135,6 +139,15 @@ export function createKnowledgeBaseTools( return await client.query(kb.id, queryParams); }, }; + + // Set the appropriate schema property based on version + if (version === 'v4') { + tool.parameters = schema; + } else { + tool.inputSchema = schema; + } + + tools[toolKey] = tool; } return tools; @@ -146,11 +159,13 @@ export function createKnowledgeBaseTools( * * @param knowledgeBases - Array of available knowledge bases * @param client - VantigeClient instance for querying - * @returns Object of tools compatible with Vercel AI SDK + * @param version - AI SDK version to target ('v4' or 'v5', defaults to 'v5') + * @returns Object of tools compatible with the specified Vercel AI SDK version */ export function createSimpleKnowledgeBaseTools( knowledgeBases: AvailableKnowledgeBase[], client: VantigeClient, + version: 'v4' | 'v5' = 'v5', ): Record> { const tools: Record< string, @@ -161,13 +176,14 @@ export function createSimpleKnowledgeBaseTools( const toolKey = generateToolKey(kb); // Create the tool with simplified schema (only query required) - tools[toolKey] = { + const schema = z.object({ + query: z + .string() + .describe("The search query to find relevant information"), + }); + + const tool: AISDKToolLike<{ query: string }, QueryResponse> = { description: kb.description || `Query the ${kb.name} knowledge base`, - inputSchema: z.object({ - query: z - .string() - .describe("The search query to find relevant information"), - }), execute: async ( params: { query: string }, _options?: MinimalToolCallOptions, @@ -175,6 +191,15 @@ export function createSimpleKnowledgeBaseTools( return await client.query(kb.id, { query: params.query }); }, }; + + // Set the appropriate schema property based on version + if (version === 'v4') { + tool.parameters = schema; + } else { + tool.inputSchema = schema; + } + + tools[toolKey] = tool; } return tools; @@ -190,6 +215,12 @@ export interface CreateToolsOptions { */ simplified?: boolean; + /** + * AI SDK version to target ('v4' or 'v5') + * @default 'v5' + */ + version?: 'v4' | 'v5'; + /** * Custom function to generate tool keys from knowledge base names * @default toKebabCase @@ -218,6 +249,7 @@ export function createKnowledgeBaseToolsWithOptions( ): Record> { const { simplified = false, + version = 'v5', keyGenerator = generateToolKey, descriptionGenerator = (kb) => kb.description || `Query the ${kb.name} knowledge base`, @@ -230,13 +262,14 @@ export function createKnowledgeBaseToolsWithOptions( const description = descriptionGenerator(kb); if (simplified) { - tools[toolKey] = { + const schema = z.object({ + query: z + .string() + .describe("The search query to find relevant information"), + }); + + const tool: AISDKToolLike<{ query: string }, QueryResponse> = { description, - inputSchema: z.object({ - query: z - .string() - .describe("The search query to find relevant information"), - }), execute: async ( params: { query: string }, _options?: MinimalToolCallOptions, @@ -244,41 +277,51 @@ export function createKnowledgeBaseToolsWithOptions( return await client.query(kb.id, { query: params.query }); }, }; + + // Set the appropriate schema property based on version + if (version === 'v4') { + tool.parameters = schema; + } else { + tool.inputSchema = schema; + } + + tools[toolKey] = tool; } else { - tools[toolKey] = { + const schema = z.object({ + query: z + .string() + .describe("The search query to find relevant information"), + topK: z + .number() + .min(1) + .max(100) + .optional() + .describe("Maximum number of results to return (1-100)"), + includeMetadata: z + .boolean() + .optional() + .describe("Whether to include metadata in results"), + useGeneration: z + .boolean() + .optional() + .describe("Whether to use AI generation for response"), + fieldMapping: z + .object({ + sourceUri: z + .string() + .optional() + .describe("Field name for source URI"), + sourceDisplayName: z + .string() + .optional() + .describe("Field name for source display name"), + }) + .optional() + .describe("Custom field mapping for results"), + }); + + const tool: AISDKToolLike = { description, - inputSchema: z.object({ - query: z - .string() - .describe("The search query to find relevant information"), - topK: z - .number() - .min(1) - .max(100) - .optional() - .describe("Maximum number of results to return (1-100)"), - includeMetadata: z - .boolean() - .optional() - .describe("Whether to include metadata in results"), - useGeneration: z - .boolean() - .optional() - .describe("Whether to use AI generation for response"), - fieldMapping: z - .object({ - sourceUri: z - .string() - .optional() - .describe("Field name for source URI"), - sourceDisplayName: z - .string() - .optional() - .describe("Field name for source display name"), - }) - .optional() - .describe("Custom field mapping for results"), - }), execute: async ( params: { query: string; @@ -309,6 +352,15 @@ export function createKnowledgeBaseToolsWithOptions( return await client.query(kb.id, queryParams); }, }; + + // Set the appropriate schema property based on version + if (version === 'v4') { + tool.parameters = schema; + } else { + tool.inputSchema = schema; + } + + tools[toolKey] = tool; } }