fix(openai): streaming tool_call + logging multiple tool_call (#463)

guillaumefontan · nirga · web-flow · commit 5d5de09e6d32 · 2024-10-16T09:47:49.000+02:00
Co-authored-by: Nir Gazit &lt;nirga@users.noreply.github.com&gt;
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -33,6 +33,10 @@ jobs:
         image: qdrant/qdrant
         ports:
           - 6333:6333
+      chroma:
+        image: chromadb/chroma
+        ports:
+          - 8000:8000
     permissions:
       contents: "read"
       id-token: "write"
@@ -47,8 +51,6 @@ jobs:
           cache: "npm"
           node-version-file: ".nvmrc"
 
-      - run: pip3 install chromadb
-
       - uses: nrwl/nx-set-shas@v3
       - run: npm ci
       - name: Build
diff --git a/packages/instrumentation-chromadb/tests/instrumentation.test.ts b/packages/instrumentation-chromadb/tests/instrumentation.test.ts
@@ -26,8 +26,6 @@ import {
 import * as chromadb from "chromadb";
 import * as assert from "assert";
 
-import { exec, ChildProcess } from "child_process";
-
 const memoryExporter = new InMemorySpanExporter();
 
 describe("Test ChromaDB instrumentation", function () {
@@ -36,13 +34,8 @@ describe("Test ChromaDB instrumentation", function () {
   let contextManager: AsyncHooksContextManager;
   let chromaDbClient: chromadb.ChromaClient;
   let collection: chromadb.Collection;
-  let chromaRun: ChildProcess;
 
   this.beforeAll(async () => {
-    // Run ChromaDB instance on different terminal instance
-    chromaRun = exec("/bin/sh");
-    chromaRun.stdin?.write("chroma run --path .\n");
-
     chromaDbClient = new chromadb.ChromaClient();
 
     // Wait for ChromaDB to spin up
@@ -92,14 +85,6 @@ describe("Test ChromaDB instrumentation", function () {
     context.disable();
   });
 
-  after(() => {
-    // Terminate the Chroma client process after tests
-    if (chromaRun.pid) {
-      // DEBT: chromaRun.kill() is not killing spawned process
-      process.exit(0);
-    }
-  });
-
   it("should set span attributes for Query", async () => {
     const input: chromadb.QueryParams = {
       nResults: 2,
diff --git a/packages/instrumentation-openai/recordings/Test-OpenAI-instrumentation_1770406427/should-set-function_call-attributes-in-span-for-stream-completion-when-multiple-tools-cal_2278551756/recording.har b/packages/instrumentation-openai/recordings/Test-OpenAI-instrumentation_1770406427/should-set-function_call-attributes-in-span-for-stream-completion-when-multiple-tools-cal_2278551756/recording.har
@@ -0,0 +1,231 @@
+{
+  "log": {
+    "_recordingName": "Test OpenAI instrumentation/should set function_call attributes in span for stream completion when multiple tools called",
+    "creator": {
+      "comment": "persister:fs",
+      "name": "Polly.JS",
+      "version": "6.0.6"
+    },
+    "entries": [
+      {
+        "_id": "4e7f1c094bf406a3975b24c0c8857547",
+        "_order": 0,
+        "cache": {},
+        "request": {
+          "bodySize": 1516,
+          "cookies": [],
+          "headers": [
+            {
+              "_fromType": "array",
+              "name": "content-length",
+              "value": "1516"
+            },
+            {
+              "_fromType": "array",
+              "name": "accept",
+              "value": "application/json"
+            },
+            {
+              "_fromType": "array",
+              "name": "content-type",
+              "value": "application/json"
+            },
+            {
+              "_fromType": "array",
+              "name": "user-agent",
+              "value": "OpenAI/JS 4.57.0"
+            },
+            {
+              "_fromType": "array",
+              "name": "x-stainless-lang",
+              "value": "js"
+            },
+            {
+              "_fromType": "array",
+              "name": "x-stainless-package-version",
+              "value": "4.57.0"
+            },
+            {
+              "_fromType": "array",
+              "name": "x-stainless-os",
+              "value": "MacOS"
+            },
+            {
+              "_fromType": "array",
+              "name": "x-stainless-arch",
+              "value": "arm64"
+            },
+            {
+              "_fromType": "array",
+              "name": "x-stainless-runtime",
+              "value": "node"
+            },
+            {
+              "_fromType": "array",
+              "name": "x-stainless-runtime-version",
+              "value": "v22.1.0"
+            },
+            {
+              "_fromType": "array",
+              "name": "accept-encoding",
+              "value": "gzip,deflate"
+            },
+            {
+              "name": "host",
+              "value": "api.openai.com"
+            }
+          ],
+          "headersSize": 470,
+          "httpVersion": "HTTP/1.1",
+          "method": "POST",
+          "postData": {
+            "mimeType": "application/json",
+            "params": [],
+            "text": "{\n  \"model\": \"gpt-4o-mini\",\n  \"messages\": [\n    {\n      \"role\": \"user\",\n      \"content\": \"What's the weather today in Boston and what will the weather be tomorrow in Chicago?\"\n    }\n  ],\n  \"stream\": true,\n  \"tools\": [\n    {\n      \"type\": \"function\",\n      \"function\": {\n        \"name\": \"get_current_weather\",\n        \"description\": \"Get the current weather in a given location\",\n        \"parameters\": {\n          \"type\": \"object\",\n          \"properties\": {\n            \"location\": {\n              \"type\": \"string\",\n              \"description\": \"The city and state, e.g. San Francisco, CA\"\n            },\n            \"unit\": {\n              \"type\": \"string\",\n              \"enum\": [\n                \"celsius\",\n                \"fahrenheit\"\n              ]\n            }\n          },\n          \"required\": [\n            \"location\"\n          ]\n        }\n      }\n    },\n    {\n      \"type\": \"function\",\n      \"function\": {\n        \"name\": \"get_tomorrow_weather\",\n        \"description\": \"Get tomorrow's weather in a given location\",\n        \"parameters\": {\n          \"type\": \"object\",\n          \"properties\": {\n            \"location\": {\n              \"type\": \"string\",\n              \"description\": \"The city and state, e.g. San Francisco, CA\"\n            },\n            \"unit\": {\n              \"type\": \"string\",\n              \"enum\": [\n                \"celsius\",\n                \"fahrenheit\"\n              ]\n            }\n          },\n          \"required\": [\n            \"location\"\n          ]\n        }\n      }\n    }\n  ]\n}"
+          },
+          "queryString": [],
+          "url": "https://api.openai.com/v1/chat/completions"
+        },
+        "response": {
+          "bodySize": 5057,
+          "content": {
+            "mimeType": "text/event-stream; charset=utf-8",
+            "size": 5057,
+            "text": "data: {\"id\":\"chatcmpl-AFzJUI0700ButGSyUudPGR9IqmGJi\",\"object\":\"chat.completion.chunk\",\"created\":1728373972,\"model\":\"gpt-4o-mini-2024-07-18\",\"system_fingerprint\":\"fp_f85bea6784\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":null},\"logprobs\":null,\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-AFzJUI0700ButGSyUudPGR9IqmGJi\",\"object\":\"chat.completion.chunk\",\"created\":1728373972,\"model\":\"gpt-4o-mini-2024-07-18\",\"system_fingerprint\":\"fp_f85bea6784\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_5LSV532rbNrhEKj3XddxLXKb\",\"type\":\"function\",\"function\":{\"name\":\"get_current_weather\",\"arguments\":\"\"}}]},\"logprobs\":null,\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-AFzJUI0700ButGSyUudPGR9IqmGJi\",\"object\":\"chat.completion.chunk\",\"created\":1728373972,\"model\":\"gpt-4o-mini-2024-07-18\",\"system_fingerprint\":\"fp_f85bea6784\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"{\\\"lo\"}}]},\"logprobs\":null,\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-AFzJUI0700ButGSyUudPGR9IqmGJi\",\"object\":\"chat.completion.chunk\",\"created\":1728373972,\"model\":\"gpt-4o-mini-2024-07-18\",\"system_fingerprint\":\"fp_f85bea6784\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"catio\"}}]},\"logprobs\":null,\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-AFzJUI0700ButGSyUudPGR9IqmGJi\",\"object\":\"chat.completion.chunk\",\"created\":1728373972,\"model\":\"gpt-4o-mini-2024-07-18\",\"system_fingerprint\":\"fp_f85bea6784\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"n\\\": \\\"B\"}}]},\"logprobs\":null,\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-AFzJUI0700ButGSyUudPGR9IqmGJi\",\"object\":\"chat.completion.chunk\",\"created\":1728373972,\"model\":\"gpt-4o-mini-2024-07-18\",\"system_fingerprint\":\"fp_f85bea6784\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"osto\"}}]},\"logprobs\":null,\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-AFzJUI0700ButGSyUudPGR9IqmGJi\",\"object\":\"chat.completion.chunk\",\"created\":1728373972,\"model\":\"gpt-4o-mini-2024-07-18\",\"system_fingerprint\":\"fp_f85bea6784\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"n, MA\"}}]},\"logprobs\":null,\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-AFzJUI0700ButGSyUudPGR9IqmGJi\",\"object\":\"chat.completion.chunk\",\"created\":1728373972,\"model\":\"gpt-4o-mini-2024-07-18\",\"system_fingerprint\":\"fp_f85bea6784\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"\\\"}\"}}]},\"logprobs\":null,\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-AFzJUI0700ButGSyUudPGR9IqmGJi\",\"object\":\"chat.completion.chunk\",\"created\":1728373972,\"model\":\"gpt-4o-mini-2024-07-18\",\"system_fingerprint\":\"fp_f85bea6784\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":1,\"id\":\"call_8rndUDfjLo174WcwdN7aT7mP\",\"type\":\"function\",\"function\":{\"name\":\"get_tomorrow_weather\",\"arguments\":\"\"}}]},\"logprobs\":null,\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-AFzJUI0700ButGSyUudPGR9IqmGJi\",\"object\":\"chat.completion.chunk\",\"created\":1728373972,\"model\":\"gpt-4o-mini-2024-07-18\",\"system_fingerprint\":\"fp_f85bea6784\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":1,\"function\":{\"arguments\":\"{\\\"lo\"}}]},\"logprobs\":null,\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-AFzJUI0700ButGSyUudPGR9IqmGJi\",\"object\":\"chat.completion.chunk\",\"created\":1728373972,\"model\":\"gpt-4o-mini-2024-07-18\",\"system_fingerprint\":\"fp_f85bea6784\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":1,\"function\":{\"arguments\":\"catio\"}}]},\"logprobs\":null,\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-AFzJUI0700ButGSyUudPGR9IqmGJi\",\"object\":\"chat.completion.chunk\",\"created\":1728373972,\"model\":\"gpt-4o-mini-2024-07-18\",\"system_fingerprint\":\"fp_f85bea6784\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":1,\"function\":{\"arguments\":\"n\\\": \\\"C\"}}]},\"logprobs\":null,\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-AFzJUI0700ButGSyUudPGR9IqmGJi\",\"object\":\"chat.completion.chunk\",\"created\":1728373972,\"model\":\"gpt-4o-mini-2024-07-18\",\"system_fingerprint\":\"fp_f85bea6784\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":1,\"function\":{\"arguments\":\"hica\"}}]},\"logprobs\":null,\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-AFzJUI0700ButGSyUudPGR9IqmGJi\",\"object\":\"chat.completion.chunk\",\"created\":1728373972,\"model\":\"gpt-4o-mini-2024-07-18\",\"system_fingerprint\":\"fp_f85bea6784\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":1,\"function\":{\"arguments\":\"go, I\"}}]},\"logprobs\":null,\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-AFzJUI0700ButGSyUudPGR9IqmGJi\",\"object\":\"chat.completion.chunk\",\"created\":1728373972,\"model\":\"gpt-4o-mini-2024-07-18\",\"system_fingerprint\":\"fp_f85bea6784\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":1,\"function\":{\"arguments\":\"L\\\"}\"}}]},\"logprobs\":null,\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-AFzJUI0700ButGSyUudPGR9IqmGJi\",\"object\":\"chat.completion.chunk\",\"created\":1728373972,\"model\":\"gpt-4o-mini-2024-07-18\",\"system_fingerprint\":\"fp_f85bea6784\",\"choices\":[{\"index\":0,\"delta\":{},\"logprobs\":null,\"finish_reason\":\"tool_calls\"}]}\n\ndata: [DONE]\n\n"
+          },
+          "cookies": [
+            {
+              "domain": ".api.openai.com",
+              "expires": "2024-10-08T08:22:52.000Z",
+              "httpOnly": true,
+              "name": "__cf_bm",
+              "path": "/",
+              "sameSite": "None",
+              "secure": true,
+              "value": "7VLzjD6bk2vl3dQ1buulluBXUcZWRHRf3FMpMkNfDe4-1728373972-1.0.1.1-M5XcPulIjIntr0YG3dvI4hnwTDHLaXKHuLZd_JUTOH40kvNyQnhrPvm2PB12Tf7rR.E_D7F80cPaEEHKxli0SA"
+            },
+            {
+              "domain": ".api.openai.com",
+              "httpOnly": true,
+              "name": "_cfuvid",
+              "path": "/",
+              "sameSite": "None",
+              "secure": true,
+              "value": "M4JyAwhhEaoDpYpUXniLhAW_YlQUIJCcsDEx1ZsVvTs-1728373972937-0.0.1.1-604800000"
+            }
+          ],
+          "headers": [
+            {
+              "name": "date",
+              "value": "Tue, 08 Oct 2024 07:52:52 GMT"
+            },
+            {
+              "name": "content-type",
+              "value": "text/event-stream; charset=utf-8"
+            },
+            {
+              "name": "transfer-encoding",
+              "value": "chunked"
+            },
+            {
+              "name": "connection",
+              "value": "keep-alive"
+            },
+            {
+              "name": "access-control-expose-headers",
+              "value": "X-Request-ID"
+            },
+            {
+              "name": "openai-organization",
+              "value": "trubrics"
+            },
+            {
+              "name": "openai-processing-ms",
+              "value": "625"
+            },
+            {
+              "name": "openai-version",
+              "value": "2020-10-01"
+            },
+            {
+              "name": "strict-transport-security",
+              "value": "max-age=31536000; includeSubDomains; preload"
+            },
+            {
+              "name": "x-ratelimit-limit-requests",
+              "value": "10000"
+            },
+            {
+              "name": "x-ratelimit-limit-tokens",
+              "value": "10000000"
+            },
+            {
+              "name": "x-ratelimit-remaining-requests",
+              "value": "9999"
+            },
+            {
+              "name": "x-ratelimit-remaining-tokens",
+              "value": "9999961"
+            },
+            {
+              "name": "x-ratelimit-reset-requests",
+              "value": "6ms"
+            },
+            {
+              "name": "x-ratelimit-reset-tokens",
+              "value": "0s"
+            },
+            {
+              "name": "x-request-id",
+              "value": "req_a89738719c34e11ea192e73a5497d00a"
+            },
+            {
+              "name": "cf-cache-status",
+              "value": "DYNAMIC"
+            },
+            {
+              "_fromType": "array",
+              "name": "set-cookie",
+              "value": "__cf_bm=7VLzjD6bk2vl3dQ1buulluBXUcZWRHRf3FMpMkNfDe4-1728373972-1.0.1.1-M5XcPulIjIntr0YG3dvI4hnwTDHLaXKHuLZd_JUTOH40kvNyQnhrPvm2PB12Tf7rR.E_D7F80cPaEEHKxli0SA; path=/; expires=Tue, 08-Oct-24 08:22:52 GMT; domain=.api.openai.com; HttpOnly; Secure; SameSite=None"
+            },
+            {
+              "_fromType": "array",
+              "name": "set-cookie",
+              "value": "_cfuvid=M4JyAwhhEaoDpYpUXniLhAW_YlQUIJCcsDEx1ZsVvTs-1728373972937-0.0.1.1-604800000; path=/; domain=.api.openai.com; HttpOnly; Secure; SameSite=None"
+            },
+            {
+              "name": "x-content-type-options",
+              "value": "nosniff"
+            },
+            {
+              "name": "server",
+              "value": "cloudflare"
+            },
+            {
+              "name": "cf-ray",
+              "value": "8cf48dcd0a786fdb-CDG"
+            }
+          ],
+          "headersSize": 1150,
+          "httpVersion": "HTTP/1.1",
+          "redirectURL": "",
+          "status": 200,
+          "statusText": "OK"
+        },
+        "startedDateTime": "2024-10-08T07:52:51.703Z",
+        "time": 1328,
+        "timings": {
+          "blocked": -1,
+          "connect": -1,
+          "dns": -1,
+          "receive": 0,
+          "send": 0,
+          "ssl": -1,
+          "wait": 1328
+        }
+      }
+    ],
+    "pages": [],
+    "version": "1.2"
+  }
+}
diff --git a/packages/instrumentation-openai/src/instrumentation.ts b/packages/instrumentation-openai/src/instrumentation.ts
@@ -356,7 +356,12 @@ export class OpenAIInstrumentation extends InstrumentationBase {
             index: 0,
             logprobs: null,
             finish_reason: "stop",
-            message: { role: "assistant", content: "", refusal: null },
+            message: {
+              role: "assistant",
+              content: "",
+              refusal: null,
+              tool_calls: [],
+            },
           },
         ],
         object: "chat.completion",
@@ -388,24 +393,39 @@ export class OpenAIInstrumentation extends InstrumentationBase {
             arguments: chunk.choices[0].delta.function_call.arguments,
           };
         }
-        if (chunk.choices[0]?.delta.tool_calls) {
-          // I needed to re-build the object so that Typescript will understand that arguments are not null.
-          result.choices[0].message.tool_calls = [];
-          for (const toolCall of chunk.choices[0].delta.tool_calls) {
-            if (
-              toolCall.id &&
-              toolCall.type &&
-              toolCall.function?.name &&
-              toolCall.function?.arguments
-            ) {
-              result.choices[0].message.tool_calls.push({
-                id: toolCall.id,
-                type: toolCall.type,
-                function: {
-                  name: toolCall.function.name,
-                  arguments: toolCall.function.arguments,
-                },
-              });
+        for (const toolCall of chunk.choices[0]?.delta?.tool_calls ?? []) {
+          if (
+            (result.choices[0].message.tool_calls?.length ?? 0) <
+            toolCall.index + 1
+          ) {
+            result.choices[0].message.tool_calls?.push({
+              function: {
+                name: "",
+                arguments: "",
+              },
+              id: "",
+              type: "function",
+            });
+          }
+
+          if (result.choices[0].message.tool_calls) {
+            if (toolCall.id) {
+              result.choices[0].message.tool_calls[toolCall.index].id +=
+                toolCall.id;
+            }
+            if (toolCall.type) {
+              result.choices[0].message.tool_calls[toolCall.index].type +=
+                toolCall.type;
+            }
+            if (toolCall.function?.name) {
+              result.choices[0].message.tool_calls[
+                toolCall.index
+              ].function.name += toolCall.function.name;
+            }
+            if (toolCall.function?.arguments) {
+              result.choices[0].message.tool_calls[
+                toolCall.index
+              ].function.arguments += toolCall.function.arguments;
             }
           }
         }
@@ -621,14 +641,17 @@ export class OpenAIInstrumentation extends InstrumentationBase {
                 choice.message.function_call.arguments,
               );
             }
-            if (choice.message.tool_calls) {
+            for (const [
+              toolIndex,
+              toolCall,
+            ] of choice?.message?.tool_calls?.entries() || []) {
               span.setAttribute(
-                `${SpanAttributes.LLM_COMPLETIONS}.${index}.function_call.name`,
-                choice.message.tool_calls[0].function.name,
+                `${SpanAttributes.LLM_COMPLETIONS}.${index}.tool_calls.${toolIndex}.name`,
+                toolCall.function.name,
               );
               span.setAttribute(
-                `${SpanAttributes.LLM_COMPLETIONS}.${index}.function_call.arguments`,
-                choice.message.tool_calls[0].function.arguments,
+                `${SpanAttributes.LLM_COMPLETIONS}.${index}.tool_calls.${toolIndex}.arguments`,
+                toolCall.function.arguments,
               );
             }
           });
@@ -722,6 +745,10 @@ export class OpenAIInstrumentation extends InstrumentationBase {
   private _encodingCache = new Map<string, Tiktoken>();
 
   private tokenCountFromString(text: string, model: string) {
+    if (!text) {
+      return 0;
+    }
+
     let encoding = this._encodingCache.get(model);
 
     if (!encoding) {
diff --git a/packages/instrumentation-openai/test/instrumentation.test.ts b/packages/instrumentation-openai/test/instrumentation.test.ts
