diff --git a/js/src/wrappers/ai-sdk/ai-sdk.test.ts b/js/src/wrappers/ai-sdk/ai-sdk.test.ts index 027450690..732bd0ae3 100644 --- a/js/src/wrappers/ai-sdk/ai-sdk.test.ts +++ b/js/src/wrappers/ai-sdk/ai-sdk.test.ts @@ -1774,6 +1774,103 @@ describe("ai sdk client unit tests", TEST_SUITE_OPTIONS, () => { }); }); + test("metadata.tools is in ToolFunctionDefinition format", async () => { + // Test that metadata.tools follows the ToolFunctionDefinition schema: + // { type: "function", function: { name, description?, parameters?, strict? } } + expect(await backgroundLogger.drain()).toHaveLength(0); + + const calculatorTool = { + description: "Perform a calculation", + inputSchema: z.object({ + operation: z.enum(["add", "subtract"]), + a: z.number(), + b: z.number(), + }), + execute: async (args: { operation: string; a: number; b: number }) => { + return args.operation === "add" + ? String(args.a + args.b) + : String(args.a - args.b); + }, + }; + + const model = openai(TEST_MODEL); + + const result = await wrappedAI.generateText({ + model, + tools: { + calculator: calculatorTool, + }, + toolChoice: "required", + prompt: "Use the calculator to add 5 and 3", + stopWhen: ai.stepCountIs(1), + }); + + assert.ok(result); + + const spans = await backgroundLogger.drain(); + + // Find the generateText root span + const generateTextSpan = spans.find( + (s: any) => s.span_attributes?.name === "generateText", + ) as any; + + expect(generateTextSpan).toBeDefined(); + expect(generateTextSpan.metadata.tools).toBeDefined(); + expect(Array.isArray(generateTextSpan.metadata.tools)).toBe(true); + expect(generateTextSpan.metadata.tools).toHaveLength(1); + + // Verify ToolFunctionDefinition format + const toolDef = generateTextSpan.metadata.tools[0]; + expect(toolDef.type).toBe("function"); + expect(toolDef.function).toBeDefined(); + expect(toolDef.function.name).toBe("calculator"); + expect(toolDef.function.description).toBe("Perform a calculation"); + expect(toolDef.function.parameters).toBeDefined(); + expect(toolDef.function.parameters.type).toBe("object"); + expect(toolDef.function.parameters.properties).toBeDefined(); + expect(toolDef.function.parameters.properties.operation).toBeDefined(); + expect(toolDef.function.parameters.properties.a).toBeDefined(); + expect(toolDef.function.parameters.properties.b).toBeDefined(); + + // Find doGenerate child span - should also have tools in metadata + const doGenerateSpan = spans.find( + (s: any) => s.span_attributes?.name === "doGenerate", + ) as any; + + expect(doGenerateSpan).toBeDefined(); + expect(doGenerateSpan.metadata.tools).toBeDefined(); + expect(Array.isArray(doGenerateSpan.metadata.tools)).toBe(true); + + // Verify child span also has ToolFunctionDefinition format + const childToolDef = doGenerateSpan.metadata.tools[0]; + expect(childToolDef.type).toBe("function"); + expect(childToolDef.function.name).toBe("calculator"); + }); + + test("metadata.tools is not set when no tools are used", async () => { + expect(await backgroundLogger.drain()).toHaveLength(0); + + const model = openai(TEST_MODEL); + + const result = await wrappedAI.generateText({ + model, + prompt: "Say hello", + maxOutputTokens: 50, + }); + + assert.ok(result); + + const spans = await backgroundLogger.drain(); + + const generateTextSpan = spans.find( + (s: any) => s.span_attributes?.name === "generateText", + ) as any; + + expect(generateTextSpan).toBeDefined(); + // metadata.tools should not be set (not even an empty array) + expect(generateTextSpan.metadata.tools).toBeUndefined(); + }); + test("ai sdk string model ID resolution with per-step spans", async () => { expect(await backgroundLogger.drain()).toHaveLength(0); diff --git a/js/src/wrappers/ai-sdk/ai-sdk.ts b/js/src/wrappers/ai-sdk/ai-sdk.ts index f6fcf3a6e..56a77d5df 100644 --- a/js/src/wrappers/ai-sdk/ai-sdk.ts +++ b/js/src/wrappers/ai-sdk/ai-sdk.ts @@ -242,6 +242,9 @@ const makeGenerateTextWrapper = ( ...spanInfoMetadata, model, ...(provider ? { provider } : {}), + ...(params.tools + ? { tools: transformToolsForMetadata(params.tools) } + : {}), braintrust: { integration_name: "ai-sdk", sdk_language: "typescript", @@ -341,6 +344,9 @@ const wrapModel = (model: any, ai?: any): any => { metadata: { model: modelId, ...(provider ? { provider } : {}), + ...(options.tools + ? { tools: transformToolsForMetadata(options.tools) } + : {}), braintrust: { integration_name: "ai-sdk", sdk_language: "typescript", @@ -368,6 +374,9 @@ const wrapModel = (model: any, ai?: any): any => { metadata: { model: modelId, ...(provider ? { provider } : {}), + ...(options.tools + ? { tools: transformToolsForMetadata(options.tools) } + : {}), braintrust: { integration_name: "ai-sdk", sdk_language: "typescript", @@ -574,6 +583,9 @@ const wrapGenerateObject = ( ...spanInfoMetadata, model, ...(provider ? { provider } : {}), + ...(params.tools + ? { tools: transformToolsForMetadata(params.tools) } + : {}), braintrust: { integration_name: "ai-sdk", sdk_language: "typescript", @@ -624,6 +636,9 @@ const makeStreamTextWrapper = ( ...spanInfoMetadata, model, ...(provider ? { provider } : {}), + ...(params.tools + ? { tools: transformToolsForMetadata(params.tools) } + : {}), braintrust: { integration_name: "ai-sdk", sdk_language: "typescript", @@ -794,6 +809,9 @@ const wrapStreamObject = ( ...spanInfoMetadata, model, ...(provider ? { provider } : {}), + ...(params.tools + ? { tools: transformToolsForMetadata(params.tools) } + : {}), braintrust: { integration_name: "ai-sdk", sdk_language: "typescript", @@ -1247,6 +1265,67 @@ const processTool = (tool: any): any => { return processed; }; +/** + * Transforms tools into ToolFunctionDefinition format for metadata.tools + * Output format: Array<{ type: "function", function: { name, description?, parameters?, strict? } }> + */ +const transformToolsForMetadata = ( + tools: any, +): Array<{ + type: "function"; + function: { + name: string; + description?: string; + parameters?: Record; + strict?: boolean | null; + }; +}> => { + if (!tools || typeof tools !== "object") return []; + + const result: Array<{ + type: "function"; + function: { + name: string; + description?: string; + parameters?: Record; + strict?: boolean | null; + }; + }> = []; + + // Handle both object format { toolName: toolDef } and array format [toolDef] + const entries = Array.isArray(tools) + ? tools.map((tool, i) => [tool?.name || `tool_${i}`, tool]) + : Object.entries(tools); + + for (const [name, tool] of entries) { + if (!tool || typeof tool !== "object") continue; + + // Get parameters from inputSchema or parameters field + let parameters: Record | undefined; + if (isZodSchema(tool.inputSchema)) { + parameters = serializeZodSchema(tool.inputSchema); + } else if (tool.inputSchema && typeof tool.inputSchema === "object") { + parameters = tool.inputSchema; + } else if (isZodSchema(tool.parameters)) { + parameters = serializeZodSchema(tool.parameters); + } else if (tool.parameters && typeof tool.parameters === "object") { + parameters = tool.parameters; + } + + result.push({ + type: "function", + function: { + name: String(name), + ...(tool.description ? { description: String(tool.description) } : {}), + ...(parameters ? { parameters } : {}), + ...(tool.strict !== undefined ? { strict: tool.strict } : {}), + }, + }); + } + + return result; +}; + /** * Detects if an object is an AI SDK Output object (from Output.object() or Output.text()) * Output objects have a responseFormat property (function, object, or Promise).