fix(ai): fix streamed tool call parsing, tool call arguments in v5 (#6813)

* fix tool call arguments parsing for v5

* fix tool call streaming and re-enable test

* update assertion

* fix remaining tests

Author: sabrenner

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'

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 [email protected], 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 }),