@arizeai/openinference-instrumentation-openai 3.0.0 → 3.1.1

package/src/instrumentation.ts

@@ -1,4 +1,4 @@
-import openai from "openai";
+import openai, { APIPromise } from "openai";
 import {
   InstrumentationBase,
   InstrumentationConfig,
@@ -188,9 +188,7 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
             },
           );
           const execContext = getExecContext(span);
-          const execPromise = safeExecuteInTheMiddle<
-            ReturnType<ChatCompletionCreateType>
-          >(
+          const execPromise = safeExecuteInTheMiddle(
             () => {
               return context.with(trace.setSpan(execContext, span), () => {
                 return original.apply(this, args);
@@ -209,7 +207,11 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
             },
           );
 
-          const wrappedPromise = execPromise.then((result) => {
+          const wrappedPromiseThen = (
+            result:
+              | Stream<openai.Chat.Completions.ChatCompletionChunk>
+              | openai.Chat.Completions.ChatCompletion,
+          ) => {
             if (isChatCompletionResponse(result)) {
               // Record the results
               span.setAttributes({
@@ -232,7 +234,11 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
             }
 
             return result;
-          });
+          };
+          const wrappedPromise = invokeMaybeAPIPromise(
+            execPromise,
+            wrappedPromiseThen,
+          );
           return context.bind(execContext, wrappedPromise);
         };
       },
@@ -271,9 +277,7 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
           );
           const execContext = getExecContext(span);
 
-          const execPromise = safeExecuteInTheMiddle<
-            ReturnType<CompletionsCreateType>
-          >(
+          const execPromise = safeExecuteInTheMiddle(
             () => {
               return context.with(trace.setSpan(execContext, span), () => {
                 return original.apply(this, args);
@@ -291,7 +295,9 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
               }
             },
           );
-          const wrappedPromise = execPromise.then((result) => {
+          const wrappedPromiseThen = (
+            result: Completion | Stream<Completion>,
+          ) => {
             if (isCompletionResponse(result)) {
               // Record the results
               span.setAttributes({
@@ -306,7 +312,11 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
               span.end();
             }
             return result;
-          });
+          };
+          const wrappedPromise = invokeMaybeAPIPromise(
+            execPromise,
+            wrappedPromiseThen,
+          );
           return context.bind(execContext, wrappedPromise);
         };
       },
@@ -343,9 +353,7 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
             },
           });
           const execContext = getExecContext(span);
-          const execPromise = safeExecuteInTheMiddle<
-            ReturnType<EmbeddingsCreateType>
-          >(
+          const execPromise = safeExecuteInTheMiddle(
             () => {
               return context.with(trace.setSpan(execContext, span), () => {
                 return original.apply(this, args);
@@ -363,7 +371,7 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
               }
             },
           );
-          const wrappedPromise = execPromise.then((result) => {
+          const wrappedPromiseThen = (result: CreateEmbeddingResponse) => {
             if (result) {
               // Record the results
               span.setAttributes({
@@ -374,7 +382,11 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
             span.setStatus({ code: SpanStatusCode.OK });
             span.end();
             return result;
-          });
+          };
+          const wrappedPromise = invokeMaybeAPIPromise(
+            execPromise,
+            wrappedPromiseThen,
+          );
           return context.bind(execContext, wrappedPromise);
         };
       },
@@ -416,9 +428,7 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
               },
             );
             const execContext = getExecContext(span);
-            const execPromise = safeExecuteInTheMiddle<
-              ReturnType<ResponsesCreateType>
-            >(
+            const execPromise = safeExecuteInTheMiddle(
               () => {
                 return context.with(trace.setSpan(execContext, span), () => {
                   return original.apply(this, args);
@@ -436,7 +446,9 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
                 }
               },
             );
-            const wrappedPromise = execPromise.then((result) => {
+            const wrappedPromiseThen = (
+              result: Stream<ResponseStreamEvent> | ResponseType,
+            ) => {
               const recordSpan = (result?: ResponseType) => {
                 if (!result) {
                   span.setStatus({ code: SpanStatusCode.ERROR });
@@ -469,7 +481,11 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
               }
 
               return result;
-            });
+            };
+            const wrappedPromise = invokeMaybeAPIPromise(
+              execPromise,
+              wrappedPromiseThen,
+            );
             return context.bind(execContext, wrappedPromise);
           };
         },
@@ -926,3 +942,46 @@ function getToolAndFunctionCallAttributesFromStreamChunk(
   }
   return attributes;
 }
+
+/**
+ * Type-guard that checks if the promise is an APIPromise.
+ *
+ * APIPromise is a class from openai that wraps promises with special behavior.
+ *
+ * @param promise - The promise to check
+ * @returns True if the promise is an APIPromise, false otherwise
+ */
+function isAPIPromise<T>(promise: unknown): promise is APIPromise<T> {
+  return promise instanceof APIPromise;
+}
+
+/**
+ * Invokes the thennable of a promise or an APIPromise.
+ *
+ * This is necessary to safely invoke promises returned by openai sdk methods.
+ *
+ * Without this wrapper, instrumentation will "consume" returned APIPromise instances,
+ * Making them unusable by other openai sdk methods, such as when completions.parse internally
+ * calls completions.create, expecting an APIPromise which we would otherwise consume by calling `then` on it.
+ *
+ * @param promise - The promise to invoke the thennable of
+ * @param then - The thennable to invoke
+ * @returns The promise with the thennable invoked
+ */
+function invokeMaybeAPIPromise<T>(
+  promise: T,
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  then: (value: any) => unknown,
+): T {
+  if (isAPIPromise<T>(promise)) {
+    return promise._thenUnwrap(then) as T;
+  } else if (promise instanceof Promise) {
+    return promise.then(then) as T;
+  } else {
+    // eslint-disable-next-line no-console
+    console.warn(
+      "Promise is not an APIPromise or a regular promise, cannot instrument.",
+    );
+    return promise;
+  }
+}

package/src/version.ts

@@ -1,2 +1,2 @@
 // this is an auto-generated file. see js/scripts/version-update.js
-export const VERSION = "3.0.0";
+export const VERSION = "3.1.1";