@arizeai/openinference-instrumentation-openai 0.1.4 → 0.3.0

package/src/instrumentation.ts

@@ -40,6 +40,19 @@ import { isTracingSuppressed } from "@opentelemetry/core";
 
 const MODULE_NAME = "openai";
 
+/**
+ * Flag to check if the openai module has been patched
+ * Note: This is a fallback in case the module is made immutable (e.x. Deno, webpack, etc.)
+ */
+let _isOpenInferencePatched = false;
+
+/**
+ * function to check if instrumentation is enabled / disabled
+ */
+export function isPatched() {
+  return _isOpenInferencePatched;
+}
+
 /**
  * Resolves the execution context for the current span
  * If tracing is suppressed, the span is dropped and the current context is returned
@@ -75,6 +88,16 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
     );
     return module;
   }
+
+  /**
+   * Manually instruments the OpenAI module. This is needed when the module is not loaded via require (commonjs)
+   * @param {openai} module
+   */
+  manuallyInstrument(module: typeof openai) {
+    diag.debug(`Manually instrumenting ${MODULE_NAME}`);
+    this.patch(module);
+  }
+
   /**
    * Patches the OpenAI module
    */
@@ -83,7 +106,7 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
     moduleVersion?: string,
   ) {
     diag.debug(`Applying patch for ${MODULE_NAME}@${moduleVersion}`);
-    if (module?.openInferencePatched) {
+    if (module?.openInferencePatched || _isOpenInferencePatched) {
       return module;
     }
     // eslint-disable-next-line @typescript-eslint/no-this-alias
@@ -308,16 +331,35 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
       },
     );
 
-    module.openInferencePatched = true;
+    _isOpenInferencePatched = true;
+    try {
+      // This can fail if the module is made immutable via the runtime or bundler
+      module.openInferencePatched = true;
+    } catch (e) {
+      diag.warn(`Failed to set ${MODULE_NAME} patched flag on the module`, e);
+    }
+
     return module;
   }
   /**
    * Un-patches the OpenAI module's chat completions API
    */
-  private unpatch(moduleExports: typeof openai, moduleVersion?: string) {
+  private unpatch(
+    moduleExports: typeof openai & { openInferencePatched?: boolean },
+    moduleVersion?: string,
+  ) {
     diag.debug(`Removing patch for ${MODULE_NAME}@${moduleVersion}`);
     this._unwrap(moduleExports.OpenAI.Chat.Completions.prototype, "create");
+    this._unwrap(moduleExports.OpenAI.Completions.prototype, "create");
     this._unwrap(moduleExports.OpenAI.Embeddings.prototype, "create");
+
+    _isOpenInferencePatched = false;
+    try {
+      // This can fail if the module is made immutable via the runtime or bundler
+      moduleExports.openInferencePatched = false;
+    } catch (e) {
+      diag.warn(`Failed to unset ${MODULE_NAME} patched flag on the module`, e);
+    }
   }
 }
 

package/src/version.ts

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