diff --git a/packages/dd-trace/src/llmobs/plugins/ai/index.js b/packages/dd-trace/src/llmobs/plugins/ai/index.js index 0d4c9dc8091..e732ce18ab9 100644 --- a/packages/dd-trace/src/llmobs/plugins/ai/index.js +++ b/packages/dd-trace/src/llmobs/plugins/ai/index.js @@ -82,7 +82,9 @@ class VercelAILLMObsPlugin extends BaseLLMObsPlugin { * @param {string} toolDescription * @returns {string | undefined} */ - findToolName (toolDescription) { + findToolName (toolName, toolDescription) { + if (Number.isNaN(Number.parseInt(toolName))) return toolName + for (const availableTool of this.#availableTools) { const description = availableTool.description if (description === toolDescription && availableTool.id) { @@ -260,16 +262,17 @@ class VercelAILLMObsPlugin extends BaseLLMObsPlugin { const formattedToolCalls = [] for (const toolCall of outputMessageToolCalls) { - const toolCallArgs = getJsonStringValue(toolCall.args, {}) + const toolArgs = toolCall.args ?? toolCall.input + const toolCallArgs = typeof toolArgs === 'string' ? getJsonStringValue(toolArgs, {}) : toolArgs const toolDescription = toolsForModel?.find(tool => toolCall.toolName === tool.name)?.description - const name = this.findToolName(toolDescription) + const name = this.findToolName(toolCall.toolName, toolDescription) this.#toolCallIdsToName[toolCall.toolCallId] = name formattedToolCalls.push({ arguments: toolCallArgs, name, toolId: toolCall.toolCallId, - type: 'function' + type: toolCall.toolCallType ?? 'function' }) } @@ -317,10 +320,10 @@ class VercelAILLMObsPlugin extends BaseLLMObsPlugin { finalContent += part.text ?? part.data } else if (type === 'tool-call') { const toolDescription = toolsForModel?.find(tool => part.toolName === tool.name)?.description - const name = this.findToolName(toolDescription) + const name = this.findToolName(part.toolName, toolDescription) toolCalls.push({ - arguments: part.args, + arguments: part.args ?? part.input, name, toolId: part.toolCallId, type: 'function' diff --git a/packages/dd-trace/test/llmobs/plugins/ai/index.spec.js b/packages/dd-trace/test/llmobs/plugins/ai/index.spec.js index 318deca46d4..4bd6278d117 100644 --- a/packages/dd-trace/test/llmobs/plugins/ai/index.spec.js +++ b/packages/dd-trace/test/llmobs/plugins/ai/index.spec.js @@ -346,8 +346,7 @@ describe('Plugin', () => { }) }) - // TODO(sabrenner): Fix this test for v5.0.0 - tool "input" instead of "arguments" - it.skip('creates a span for a tool call', async () => { // eslint-disable-line mocha/no-pending-tests + it('creates a span for a tool call', async () => { let tools let additionalOptions = {} const toolSchema = ai.jsonSchema({ @@ -504,8 +503,7 @@ describe('Plugin', () => { }) }) - // TODO(sabrenner): Fix this test for v5.0.0 - tool "input" instead of "arguments" & parsing, streaming - it.skip('created a span for a tool call from a stream', async () => { // eslint-disable-line mocha/no-pending-tests + it('created a span for a tool call from a stream', async () => { let tools let additionalOptions = {} const toolSchema = ai.jsonSchema({ @@ -629,15 +627,13 @@ describe('Plugin', () => { span: apmSpans[2], parentId: llmobsSpans[0].span_id, /** - * MOCK_STRING used as the stream implementation for ai does not finish the initial llm spans + * Before ai@4.0.2, the stream implementation did not finish the initial llm spans * first to associate the tool call id with the tool itself (by matching descriptions). * - * Usually, this would mean the tool call name is 'toolCall'. - * - * However, because we used mocked responses, the second time this test is called, the tool call - * will have the name 'weather' instead. We just assert that the name exists and is a string to simplify. + * Usually, this would mean the tool call name is 'toolCall'. This is a limitation with the older library + * versions. In v5+, this is resolved as the tool name is not its index in the tools array, but its actual name. */ - name: MOCK_STRING, + name: semifies(realVersion, NODE_MAJOR < 22 ? '<=4.0.2' : '<4.0.2') ? 'toolCall' : 'weather', spanKind: 'tool', inputValue: JSON.stringify({ location: 'Tokyo' }), outputValue: JSON.stringify({ location: 'Tokyo', temperature: 72 }),