@mondaydotcomorg/atp-vercel-sdk 0.20.3 → 0.20.4

package/src/tools.ts

@@ -1,82 +1,113 @@
 import { z } from 'zod';
 import { VercelAIATPClient } from './client.js';
 import type { CreateATPToolsOptions, ATPToolsResult, StreamingToolsOptions } from './types.js';
-import { createToolsFromATPClient, ToolNames } from '@mondaydotcomorg/atp-client';
+import {
+	ToolNames,
+	executeCodeInputSchema,
+	exploreApiInputSchema,
+	searchApiInputSchema,
+	fetchAllApisInputSchema,
+} from '@mondaydotcomorg/atp-client';
 import { ExecutionStatus } from '@mondaydotcomorg/atp-protocol';
 import { tool } from 'ai';
 import { createVercelEventHandler, type UIMessageStreamWriter } from './event-adapter.js';
 
+const TOOL_SCHEMAS = {
+	[ToolNames.EXECUTE_CODE]: executeCodeInputSchema.pick({ code: true }),
+	[ToolNames.EXPLORE_API]: exploreApiInputSchema,
+	[ToolNames.SEARCH_API]: searchApiInputSchema,
+	[ToolNames.FETCH_ALL_APIS]: fetchAllApisInputSchema,
+} as const;
+
 export async function createATPTools(options: CreateATPToolsOptions): Promise<ATPToolsResult> {
 	const { defaultExecutionConfig, ...clientOptions } = options;
 
 	const client = new VercelAIATPClient(clientOptions);
 	await client.connect();
 
-	const atpTools = createToolsFromATPClient(client.getUnderlyingClient());
-
+	const underlyingClient = client.getUnderlyingClient();
 	const vercelTools: Record<string, any> = {};
 
-	for (const atpTool of atpTools) {
-		if (atpTool.name === ToolNames.EXECUTE_CODE) {
-			vercelTools.atp_execute_code = tool({
-				description:
-					atpTool.description ||
-					'Execute TypeScript code in ATP sandbox with access to runtime APIs (atp.llm.*, atp.embedding.*, atp.approval.*)',
-				parameters: z.object({
-					code: z.string().describe('TypeScript code to execute in the ATP sandbox'),
-				}),
-				execute: async ({ code }: { code: string }) => {
-					try {
-						const result = await client.execute(code, defaultExecutionConfig);
-
-						if (result.status === ExecutionStatus.COMPLETED) {
-							return {
-								success: true,
-								result: result.result,
-								stats: result.stats,
-							};
-						} else if (result.status === ExecutionStatus.FAILED) {
-							return {
-								success: false,
-								error: result.error,
-								stats: result.stats,
-							};
-						} else {
-							return {
-								success: false,
-								error: 'Execution in unexpected state: ' + result.status,
-							};
-						}
-					} catch (error: any) {
-						return {
-							success: false,
-							error: error.message || 'Unknown error',
-						};
-					}
-				},
-			});
-		} else {
-			const toolName = `atp_${atpTool.name}`;
-			vercelTools[toolName] = tool({
-				description: atpTool.description || '',
-				parameters: atpTool.zodSchema || z.object({}),
-				execute: async (input: any) => {
-					try {
-						const result = await atpTool.func(input);
-						return {
-							success: true,
-							result,
-						};
-					} catch (error: any) {
-						return {
-							success: false,
-							error: error.message,
-						};
-					}
-				},
-			});
-		}
-	}
+	vercelTools.atp_execute_code = tool({
+		description:
+			'Execute TypeScript code in ATP sandbox with access to runtime APIs (atp.llm.*, atp.embedding.*, atp.approval.*)',
+		parameters: TOOL_SCHEMAS[ToolNames.EXECUTE_CODE],
+		execute: async ({ code }: { code: string }) => {
+			try {
+				const result = await client.execute(code, defaultExecutionConfig);
+
+				if (result.status === ExecutionStatus.COMPLETED) {
+					return {
+						success: true,
+						result: result.result,
+						stats: result.stats,
+					};
+				} else if (result.status === ExecutionStatus.FAILED) {
+					return {
+						success: false,
+						error: result.error,
+						stats: result.stats,
+					};
+				} else {
+					return {
+						success: false,
+						error: 'Execution in unexpected state: ' + result.status,
+					};
+				}
+			} catch (error: any) {
+				return {
+					success: false,
+					error: error.message || 'Unknown error',
+				};
+			}
+		},
+	});
+
+	vercelTools.atp_explore_api = tool({
+		description:
+			'Explore APIs using filesystem-like navigation. Navigate through directories to discover available functions.',
+		parameters: TOOL_SCHEMAS[ToolNames.EXPLORE_API],
+		execute: async ({ path }: { path: string }) => {
+			try {
+				const result = await underlyingClient.exploreAPI(path);
+				return {
+					success: true,
+					result,
+				};
+			} catch (error: any) {
+				return {
+					success: false,
+					error: error.message,
+				};
+			}
+		},
+	});
+
+	vercelTools.atp_search_api = tool({
+		description:
+			'Search for APIs by keyword. Provide search term as string like "add", "math", "user", etc.',
+		parameters: TOOL_SCHEMAS[ToolNames.SEARCH_API],
+		execute: async ({ query }: { query: string }) => {
+			try {
+				const results = await underlyingClient.searchAPI(query);
+				return {
+					success: true,
+					results: results.map((r: any) => ({
+						apiGroup: r.apiGroup,
+						functionName: r.functionName,
+						description: r.description,
+						signature: r.signature,
+					})),
+					count: results.length,
+				};
+			} catch (error: any) {
+				return {
+					success: false,
+					error: error.message,
+				};
+			}
+		},
+	});
 
 	vercelTools.atp_get_type_definitions = tool({
 		description:
@@ -146,9 +177,7 @@ export async function createATPStreamingTools(
 	vercelTools.atp_execute_code = tool({
 		description:
 			'Execute TypeScript code in ATP sandbox with streaming events for thinking, tool execution, and text output',
-		parameters: z.object({
-			code: z.string().describe('TypeScript code to execute in the ATP sandbox'),
-		}),
+		parameters: TOOL_SCHEMAS[ToolNames.EXECUTE_CODE],
 		execute: async ({ code }: { code: string }) => {
 			try {
 				const result = await underlyingClient.executeStream(
@@ -184,6 +213,52 @@ export async function createATPStreamingTools(
 		},
 	});
 
+	vercelTools.atp_explore_api = tool({
+		description:
+			'Explore APIs using filesystem-like navigation. Navigate through directories to discover available functions.',
+		parameters: TOOL_SCHEMAS[ToolNames.EXPLORE_API],
+		execute: async ({ path }: { path: string }) => {
+			try {
+				const result = await underlyingClient.exploreAPI(path);
+				return {
+					success: true,
+					result,
+				};
+			} catch (error: any) {
+				return {
+					success: false,
+					error: error.message,
+				};
+			}
+		},
+	});
+
+	vercelTools.atp_search_api = tool({
+		description:
+			'Search for APIs by keyword. Provide search term as string like "add", "math", "user", etc.',
+		parameters: TOOL_SCHEMAS[ToolNames.SEARCH_API],
+		execute: async ({ query }: { query: string }) => {
+			try {
+				const results = await underlyingClient.searchAPI(query);
+				return {
+					success: true,
+					results: results.map((r: any) => ({
+						apiGroup: r.apiGroup,
+						functionName: r.functionName,
+						description: r.description,
+						signature: r.signature,
+					})),
+					count: results.length,
+				};
+			} catch (error: any) {
+				return {
+					success: false,
+					error: error.message,
+				};
+			}
+		},
+	});
+
 	vercelTools.atp_get_type_definitions = tool({
 		description:
 			'Get TypeScript type definitions for ATP runtime APIs to understand available functions',