fix(openai): structured output promise unwrapping exception (#474)

Author: galkleinman

package-lock.json

@@ -14,14 +14,7 @@
  * limitations under the License.
  */
 import type * as openai from "openai";
-import {
-  context,
-  trace,
-  Span,
-  Attributes,
-  SpanKind,
-  SpanStatusCode,
-} from "@opentelemetry/api";
+import { context, trace, Span, Attributes, SpanKind } from "@opentelemetry/api";
 import {
   InstrumentationBase,
   InstrumentationModuleDefinition,
@@ -46,6 +39,7 @@ import type {
 import type { Stream } from "openai/streaming";
 import { version } from "../package.json";
 import { encodingForModel, TiktokenModel, Tiktoken } from "js-tiktoken";
+import { APIPromise } from "openai/core";
 
 export class OpenAIInstrumentation extends InstrumentationBase {
   protected declare _config: OpenAIInstrumentationConfig;
@@ -338,13 +332,13 @@ export class OpenAIInstrumentation extends InstrumentationBase {
         span: Span;
         type: "chat";
         params: ChatCompletionCreateParamsStreaming;
-        promise: Promise<Stream<ChatCompletionChunk>>;
+        promise: APIPromise<Stream<ChatCompletionChunk>>;
       }
     | {
         span: Span;
         params: CompletionCreateParamsStreaming;
         type: "completion";
-        promise: Promise<Stream<Completion>>;
+        promise: APIPromise<Stream<Completion>>;
       }) {
     if (type === "chat") {
       const result: ChatCompletion = {
@@ -532,63 +526,49 @@ export class OpenAIInstrumentation extends InstrumentationBase {
     type: "chat" | "completion",
     version: "v3" | "v4",
     span: Span,
-    promise: Promise<T>,
-  ): Promise<T> {
-    return promise
-      .then((result) => {
-        return new Promise<T>((resolve) => {
-          if (version === "v3") {
-            if (type === "chat") {
-              this._addLogProbsEvent(
-                span,
-                ((result as any).data as ChatCompletion).choices[0].logprobs,
-              );
-              this._endSpan({
-                type,
-                span,
-                result: (result as any).data as ChatCompletion,
-              });
-            } else {
-              this._addLogProbsEvent(
-                span,
-                ((result as any).data as Completion).choices[0].logprobs,
-              );
-              this._endSpan({
-                type,
-                span,
-                result: (result as any).data as Completion,
-              });
-            }
-          } else {
-            if (type === "chat") {
-              this._addLogProbsEvent(
-                span,
-                (result as ChatCompletion).choices[0].logprobs,
-              );
-              this._endSpan({ type, span, result: result as ChatCompletion });
-            } else {
-              this._addLogProbsEvent(
-                span,
-                (result as Completion).choices[0].logprobs,
-              );
-              this._endSpan({ type, span, result: result as Completion });
-            }
-          }
-          resolve(result);
-        });
-      })
-      .catch((error: Error) => {
-        return new Promise<T>((_, reject) => {
-          span.setStatus({
-            code: SpanStatusCode.ERROR,
-            message: error.message,
+    promise: APIPromise<T>,
+  ): APIPromise<T> {
+    return promise._thenUnwrap((result) => {
+      if (version === "v3") {
+        if (type === "chat") {
+          this._addLogProbsEvent(
+            span,
+            ((result as any).data as ChatCompletion).choices[0].logprobs,
+          );
+          this._endSpan({
+            type,
+            span,
+            result: (result as any).data as ChatCompletion,
+          });
+        } else {
+          this._addLogProbsEvent(
+            span,
+            ((result as any).data as Completion).choices[0].logprobs,
+          );
+          this._endSpan({
+            type,
+            span,
+            result: (result as any).data as Completion,
           });
-          span.recordException(error);
-          span.end();
+        }
+      } else {
+        if (type === "chat") {
+          this._addLogProbsEvent(
+            span,
+            (result as ChatCompletion).choices[0].logprobs,
+          );
+          this._endSpan({ type, span, result: result as ChatCompletion });
+        } else {
+          this._addLogProbsEvent(
+            span,
+            (result as Completion).choices[0].logprobs,
+          );
+          this._endSpan({ type, span, result: result as Completion });
+        }
+      }
 
-          reject(error);
-        });
-      });
+      return result;
+    });
   }
 
   private _endSpan({

packages/instrumentation-openai/src/instrumentation.ts

@@ -24,6 +24,7 @@
     "run:llamaindex": "npm run build && node dist/src/sample_llamaindex.js",
     "run:pinecone": "npm run build && node dist/src/sample_pinecone.js",
     "run:langchain": "npm run build && node dist/src/sample_langchain.js",
+    "run:sample_structured_output": "npm run build && node dist/src/sample_structured_output.js",
     "lint": "eslint . --ext .ts",
     "lint:fix": "eslint . --ext .ts --fix"
   },
@@ -46,7 +47,8 @@
     "cohere-ai": "^7.7.5",
     "langchain": "^0.2.12",
     "llamaindex": "^0.5.20",
-    "openai": "^4.57.0"
+    "openai": "^4.57.0",
+    "zod": "^3.23.8"
   },
   "private": true,
   "gitHead": "ef1e70d6037f7b5c061056ef2be16e3f55f02ed5"

packages/sample-app/package.json

@@ -0,0 +1,37 @@
+import * as traceloop from "@traceloop/node-server-sdk";
+import OpenAI from "openai";
+import { zodResponseFormat } from "openai/helpers/zod";
+import { z } from "zod";
+
+traceloop.initialize({
+  appName: "sample_same_structured_output",
+  apiKey: process.env.TRACELOOP_API_KEY,
+  disableBatch: true,
+});
+const openai = new OpenAI();
+
+const CalendarEvent = z.object({
+  name: z.string(),
+  date: z.string(),
+  participants: z.array(z.string()),
+});
+
+async function create_event() {
+  const completion = await openai.beta.chat.completions.parse({
+    model: "gpt-4o",
+    messages: [
+      { role: "system", content: "Extract the event information." },
+      {
+        role: "user",
+        content: "Alice and Bob are going to a science fair on Friday.",
+      },
+    ],
+    response_format: zodResponseFormat(CalendarEvent, "event"),
+  });
+
+  console.log(completion.choices[0].message.parsed);
+
+  return completion.choices[0].message.parsed;
+}
+
+create_event();