feat: Parse tool call ids as per semantic conventions

Arize-ai · Nov 7, 2024 · 8174c9b · 8174c9b
1 parent 33e1f53
commit 8174c9b
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 6 deletions.
diff --git a/js/packages/openinference-instrumentation-openai/src/instrumentation.ts b/js/packages/openinference-instrumentation-openai/src/instrumentation.ts
@@ -503,15 +503,21 @@ function getChatCompletionInputMessageAttributes(
     case "assistant":
       if (message.tool_calls) {
         message.tool_calls.forEach((toolCall, index) => {
+          const toolCallIndexPrefix = `${SemanticConventions.MESSAGE_TOOL_CALLS}.${index}.`;
+
+          // Add the tool call id if it exists
+          if (toolCall.id) {
+            attributes[
+              `${toolCallIndexPrefix}${SemanticConventions.TOOL_CALL_ID}`
+            ] = toolCall.id;
+          }
           // Make sure the tool call has a function
           if (toolCall.function) {
-            const toolCallIndexPrefix = `${SemanticConventions.MESSAGE_TOOL_CALLS}.${index}.`;
             attributes[
-              toolCallIndexPrefix + SemanticConventions.TOOL_CALL_FUNCTION_NAME
+              `${toolCallIndexPrefix}${SemanticConventions.TOOL_CALL_FUNCTION_NAME}`
             ] = toolCall.function.name;
             attributes[
-              toolCallIndexPrefix +
-                SemanticConventions.TOOL_CALL_FUNCTION_ARGUMENTS_JSON
+              `${toolCallIndexPrefix}${SemanticConventions.TOOL_CALL_FUNCTION_ARGUMENTS_JSON}`
             ] = toolCall.function.arguments;
           }
         });
@@ -521,8 +527,10 @@ function getChatCompletionInputMessageAttributes(
       attributes[SemanticConventions.MESSAGE_FUNCTION_CALL_NAME] = message.name;
       break;
     case "tool":
-      // There's nothing to add for the tool. There is a tool_id, but there are no
-      // semantic conventions for it
+      if (message.tool_call_id) {
+        attributes[`${SemanticConventions.MESSAGE_TOOL_CALL_ID}`] =
+          message.tool_call_id;
+      }
       break;
     case "system":
       // There's nothing to add for the system. Content is captured above
@@ -613,6 +621,12 @@ function getChatCompletionOutputMessageAttributes(
   if (message.tool_calls) {
     message.tool_calls.forEach((toolCall, index) => {
       const toolCallIndexPrefix = `${SemanticConventions.MESSAGE_TOOL_CALLS}.${index}.`;
+      // Add the tool call id if it exists
+      if (toolCall.id) {
+        attributes[
+          `${toolCallIndexPrefix}${SemanticConventions.TOOL_CALL_ID}`
+        ] = toolCall.id;
+      }
       // Double check that the tool call has a function
       // NB: OpenAI only supports tool calls with functions right now but this may change
       if (toolCall.function) {
@@ -759,6 +773,12 @@ function getToolAndFunctionCallAttributesFromStreamChunk(
   if (choice.delta.tool_calls) {
     choice.delta.tool_calls.forEach((toolCall, index) => {
       const toolCallIndexPrefix = `${SemanticConventions.MESSAGE_TOOL_CALLS}.${index}.`;
+      // Add the tool call id if it exists
+      if (toolCall.id) {
+        attributes[
+          `${toolCallIndexPrefix}${SemanticConventions.TOOL_CALL_ID}`
+        ] = toolCall.id;
+      }
       // Double check that the tool call has a function
       // NB: OpenAI only supports tool calls with functions right now but this may change
       if (toolCall.function) {

diff --git a/js/packages/openinference-instrumentation-openai/test/openai.test.ts b/js/packages/openinference-instrumentation-openai/test/openai.test.ts
@@ -398,6 +398,7 @@ describe("OpenAIInstrumentation", () => {
   "llm.output_messages.0.message.role": "assistant",
   "llm.output_messages.0.message.tool_calls.0.tool_call.function.arguments": "{}",
   "llm.output_messages.0.message.tool_calls.0.tool_call.function.name": "getCurrentLocation",
+  "llm.output_messages.0.message.tool_calls.0.tool_call.id": "call_5ERYvu4iTGSvDlcDQjDP3g3J",
   "llm.provider": "openai",
   "llm.system": "openai",
   "llm.token_count.completion": 7,
@@ -420,15 +421,18 @@ describe("OpenAIInstrumentation", () => {
   "llm.input_messages.1.message.role": "assistant",
   "llm.input_messages.1.message.tool_calls.0.tool_call.function.arguments": "{}",
   "llm.input_messages.1.message.tool_calls.0.tool_call.function.name": "getCurrentLocation",
+  "llm.input_messages.1.message.tool_calls.0.tool_call.id": "call_5ERYvu4iTGSvDlcDQjDP3g3J",
   "llm.input_messages.2.message.content": "Boston",
   "llm.input_messages.2.message.role": "tool",
+  "llm.input_messages.2.message.tool_call_id": "call_5ERYvu4iTGSvDlcDQjDP3g3J",
   "llm.invocation_parameters": "{"model":"gpt-3.5-turbo","tools":[{"type":"function","function":{"name":"getCurrentLocation","parameters":{"type":"object","properties":{}},"description":"Get the current location of the user."}},{"type":"function","function":{"name":"getWeather","parameters":{"type":"object","properties":{"location":{"type":"string"}}},"description":"Get the weather for a location."}}],"tool_choice":"auto","stream":false}",
   "llm.model_name": "gpt-3.5-turbo-0613",
   "llm.output_messages.0.message.role": "assistant",
   "llm.output_messages.0.message.tool_calls.0.tool_call.function.arguments": "{
   "location": "Boston"
 }",
   "llm.output_messages.0.message.tool_calls.0.tool_call.function.name": "getWeather",
+  "llm.output_messages.0.message.tool_calls.0.tool_call.id": "call_0LCdYLkdRUt3rV3dawoIFHBf",
   "llm.provider": "openai",
   "llm.system": "openai",
   "llm.token_count.completion": 15,
@@ -451,15 +455,19 @@ describe("OpenAIInstrumentation", () => {
   "llm.input_messages.1.message.role": "assistant",
   "llm.input_messages.1.message.tool_calls.0.tool_call.function.arguments": "{}",
   "llm.input_messages.1.message.tool_calls.0.tool_call.function.name": "getCurrentLocation",
+  "llm.input_messages.1.message.tool_calls.0.tool_call.id": "call_5ERYvu4iTGSvDlcDQjDP3g3J",
   "llm.input_messages.2.message.content": "Boston",
   "llm.input_messages.2.message.role": "tool",
+  "llm.input_messages.2.message.tool_call_id": "call_5ERYvu4iTGSvDlcDQjDP3g3J",
   "llm.input_messages.3.message.role": "assistant",
   "llm.input_messages.3.message.tool_calls.0.tool_call.function.arguments": "{
   "location": "Boston"
 }",
   "llm.input_messages.3.message.tool_calls.0.tool_call.function.name": "getWeather",
+  "llm.input_messages.3.message.tool_calls.0.tool_call.id": "call_0LCdYLkdRUt3rV3dawoIFHBf",
   "llm.input_messages.4.message.content": "{"temperature":52,"precipitation":"rainy"}",
   "llm.input_messages.4.message.role": "tool",
+  "llm.input_messages.4.message.tool_call_id": "call_0LCdYLkdRUt3rV3dawoIFHBf",
   "llm.invocation_parameters": "{"model":"gpt-3.5-turbo","tools":[{"type":"function","function":{"name":"getCurrentLocation","parameters":{"type":"object","properties":{}},"description":"Get the current location of the user."}},{"type":"function","function":{"name":"getWeather","parameters":{"type":"object","properties":{"location":{"type":"string"}}},"description":"Get the weather for a location."}}],"tool_choice":"auto","stream":false}",
   "llm.model_name": "gpt-3.5-turbo-0613",
   "llm.output_messages.0.message.content": "The weather in Boston this week is expected to be rainy with a temperature of 52 degrees.",
@@ -598,6 +606,7 @@ describe("OpenAIInstrumentation", () => {
   "llm.output_messages.0.message.role": "assistant",
   "llm.output_messages.0.message.tool_calls.0.tool_call.function.arguments": "{}",
   "llm.output_messages.0.message.tool_calls.0.tool_call.function.name": "getWeather",
+  "llm.output_messages.0.message.tool_calls.0.tool_call.id": "call_PGkcUg2u6vYrCpTn0e9ofykY",
   "llm.provider": "openai",
   "llm.system": "openai",
   "llm.tools.0.tool.json_schema": "{"type":"function","function":{"name":"getCurrentLocation","parameters":{"type":"object","properties":{}},"description":"Get the current location of the user."}}",