@arizeai/openinference-instrumentation-openai 0.0.1-rc9 → 0.0.3

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@arizeai/openinference-instrumentation-openai",
-  "version": "0.0.1-rc9",
+  "version": "0.0.3",
   "description": "OpenInference instrumentation for OpenAI",
   "private": false,
   "main": "dist/src/index.js",
@@ -10,7 +10,7 @@
   "dependencies": {
     "@opentelemetry/api": "^1.7.0",
     "@opentelemetry/instrumentation": "^0.46.0",
-    "@arizeai/openinference-semantic-conventions": "0.0.8"
+    "@arizeai/openinference-semantic-conventions": "0.0.12"
   },
   "keywords": [],
   "files": [
@@ -30,6 +30,7 @@
     "prebuild": "rimraf dist & pnpm run version:update",
     "build": "tsc --build tsconfig.json tsconfig.esm.json tsconfig.esnext.json",
     "version:update": "../../scripts/version-update.js",
+    "type:check": "tsc --noEmit",
     "test": "jest ."
   }
 }

package/src/instrumentation.ts

@@ -1,125 +1,206 @@
 import * as openai from "openai";
 import {
-    InstrumentationBase,
-    InstrumentationConfig,
-    InstrumentationModuleDefinition,
-    InstrumentationNodeModuleDefinition,
-    safeExecuteInTheMiddle,
+  InstrumentationBase,
+  InstrumentationConfig,
+  InstrumentationModuleDefinition,
+  InstrumentationNodeModuleDefinition,
+  safeExecuteInTheMiddle,
 } from "@opentelemetry/instrumentation";
 import {
-    diag,
-    context,
-    trace,
-    isSpanContextValid,
-    Span,
-    SpanKind,
-    Attributes,
+  diag,
+  context,
+  trace,
+  SpanKind,
+  Attributes,
+  SpanStatusCode,
 } from "@opentelemetry/api";
 import { VERSION } from "./version";
 import {
-    SemanticConventions,
-    OpenInferenceSpanKind,
+  SemanticConventions,
+  OpenInferenceSpanKind,
+  MimeType,
 } from "@arizeai/openinference-semantic-conventions";
-import { ChatCompletionCreateParamsBase } from "openai/resources/chat/completions";
+import {
+  ChatCompletion,
+  ChatCompletionChunk,
+  ChatCompletionCreateParamsBase,
+} from "openai/resources/chat/completions";
+import { Stream } from "openai/streaming";
 
 const MODULE_NAME = "openai";
 
 export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
-    constructor(config?: InstrumentationConfig) {
-        super(
-            "@arizeai/openinference-instrumentation-openai",
-            VERSION,
-            Object.assign({}, config)
-        );
-    }
+  constructor(config?: InstrumentationConfig) {
+    super(
+      "@arizeai/openinference-instrumentation-openai",
+      VERSION,
+      Object.assign({}, config),
+    );
+  }
 
-    protected init(): InstrumentationModuleDefinition<typeof openai> {
-        const module = new InstrumentationNodeModuleDefinition<typeof openai>(
-            "openai",
-            ["^4.0.0"],
-            this.patch.bind(this),
-            this.unpatch.bind(this)
-        );
-        return module;
+  protected init(): InstrumentationModuleDefinition<typeof openai> {
+    const module = new InstrumentationNodeModuleDefinition<typeof openai>(
+      "openai",
+      ["^4.0.0"],
+      this.patch.bind(this),
+      this.unpatch.bind(this),
+    );
+    return module;
+  }
+  /**
+   * Patches the OpenAI module
+   */
+  private patch(
+    module: typeof openai & { openInferencePatched?: boolean },
+    moduleVersion?: string,
+  ) {
+    diag.debug(`Applying patch for ${MODULE_NAME}@${moduleVersion}`);
+    if (module?.openInferencePatched) {
+      return module;
     }
-    private patch(
-        module: typeof openai & { openInferencePatched?: boolean },
-        moduleVersion?: string
-    ) {
-        console.log(`Applying patch for ${MODULE_NAME}@${moduleVersion}`);
-        const plugin = this;
-        if (module?.openInferencePatched) {
-            return module;
-        }
-        const instrumentation: OpenAIInstrumentation = this;
-        type CompletionCreateType =
-            typeof module.OpenAI.Chat.Completions.prototype.create;
-        this._wrap(
-            module.OpenAI.Chat.Completions.prototype,
-            "create",
-            (original: CompletionCreateType): any => {
-                return function patchedCreate(
-                    this: unknown,
-                    ...args: Parameters<
-                        typeof module.OpenAI.Chat.Completions.prototype.create
-                    >
-                ) {
-                    const body = args[0];
-                    const span = instrumentation.tracer.startSpan(
-                        `OpenAI Chat Completions`,
-                        {
-                            kind: SpanKind.CLIENT,
-                            attributes: {
-                                [SemanticConventions.OPENINFERENCE_SPAN_KIND]:
-                                    OpenInferenceSpanKind.LLM,
-                                ...getLLMInputMessagesAttributes(body),
-                            },
-                        }
-                    );
-                    const execContext = trace.setSpan(context.active(), span);
-                    const execPromise = safeExecuteInTheMiddle<
-                        ReturnType<CompletionCreateType>
-                    >(
-                        () => {
-                            return context.with(execContext, () => {
-                                return original.apply(this, args);
-                            });
-                        },
-                        (error) => {
-                            // TODO handler the error
-                        }
-                    );
-                    const wrappedPromise = execPromise.then((result) => {
-                        // TODO set span attributes
-                        span.end();
-                        return result;
-                    });
-                    return context.bind(execContext, wrappedPromise);
-                };
+    // eslint-disable-next-line @typescript-eslint/no-this-alias
+    const instrumentation: OpenAIInstrumentation = this;
+    type CompletionCreateType =
+      typeof module.OpenAI.Chat.Completions.prototype.create;
+
+    // Patch create chat completions
+    this._wrap(
+      module.OpenAI.Chat.Completions.prototype,
+      "create",
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      (original: CompletionCreateType): any => {
+        return function patchedCreate(
+          this: unknown,
+          ...args: Parameters<
+            typeof module.OpenAI.Chat.Completions.prototype.create
+          >
+        ) {
+          const body = args[0];
+          const { messages: _messages, ...invocationParameters } = body;
+          const span = instrumentation.tracer.startSpan(
+            `OpenAI Chat Completions`,
+            {
+              kind: SpanKind.CLIENT,
+              attributes: {
+                [SemanticConventions.OPENINFERENCE_SPAN_KIND]:
+                  OpenInferenceSpanKind.LLM,
+                [SemanticConventions.LLM_MODEL_NAME]: body.model,
+                [SemanticConventions.INPUT_VALUE]: JSON.stringify(body),
+                [SemanticConventions.INPUT_MIME_TYPE]: MimeType.JSON,
+                [SemanticConventions.LLM_INVOCATION_PARAMETERS]:
+                  JSON.stringify(invocationParameters),
+                ...getLLMInputMessagesAttributes(body),
+              },
+            },
+          );
+          const execContext = trace.setSpan(context.active(), span);
+          const execPromise = safeExecuteInTheMiddle<
+            ReturnType<CompletionCreateType>
+          >(
+            () => {
+              return context.with(execContext, () => {
+                return original.apply(this, args);
+              });
+            },
+            (error) => {
+              // Push the error to the span
+              if (error) {
+                span.recordException(error);
+              }
+            },
+          );
+          const wrappedPromise = execPromise.then((result) => {
+            if (result) {
+              // Record the results
+              span.setAttributes({
+                [SemanticConventions.OUTPUT_VALUE]: JSON.stringify(result),
+                [SemanticConventions.OUTPUT_MIME_TYPE]: MimeType.JSON,
+                ...getLLMOutputMessagesAttributes(result),
+                ...getUsageAttributes(result),
+              });
             }
-        );
-        module.openInferencePatched = true;
-        return module;
-    }
-    private unpatch(moduleExports: typeof openai, moduleVersion?: string) {
-        diag.debug(`Removing patch for ${MODULE_NAME}@${moduleVersion}`);
-    }
+            span.setStatus({ code: SpanStatusCode.OK });
+            span.end();
+            return result;
+          });
+          return context.bind(execContext, wrappedPromise);
+        };
+      },
+    );
+    module.openInferencePatched = true;
+    return module;
+  }
+  /**
+   * Un-patches the OpenAI module's chat completions API
+   */
+  private unpatch(moduleExports: typeof openai, moduleVersion?: string) {
+    diag.debug(`Removing patch for ${MODULE_NAME}@${moduleVersion}`);
+    this._unwrap(moduleExports.OpenAI.Chat.Completions.prototype, "create");
+  }
 }
 
 /**
- * Converts the body input
+ * Converts the body of the request to LLM input messages
  */
 function getLLMInputMessagesAttributes(
-    body: ChatCompletionCreateParamsBase
+  body: ChatCompletionCreateParamsBase,
+): Attributes {
+  return body.messages.reduce((acc, message, index) => {
+    const index_prefix = `${SemanticConventions.LLM_INPUT_MESSAGES}.${index}`;
+    acc[`${index_prefix}.${SemanticConventions.MESSAGE_CONTENT}`] = String(
+      message.content,
+    );
+    acc[`${index_prefix}.${SemanticConventions.MESSAGE_ROLE}`] = String(
+      message.role,
+    );
+    return acc;
+  }, {} as Attributes);
+}
+
+/**
+ * Get Usage attributes
+ */
+function getUsageAttributes(
+  response: Stream<ChatCompletionChunk> | ChatCompletion,
+) {
+  if (Object.prototype.hasOwnProperty.call(response, "usage")) {
+    const completion = response as ChatCompletion;
+    if (completion.usage) {
+      return {
+        [SemanticConventions.LLM_TOKEN_COUNT_COMPLETION]:
+          completion.usage.completion_tokens,
+        [SemanticConventions.LLM_TOKEN_COUNT_PROMPT]:
+          completion.usage.prompt_tokens,
+        [SemanticConventions.LLM_TOKEN_COUNT_TOTAL]:
+          completion.usage.total_tokens,
+      };
+    }
+  }
+  return {};
+}
+
+/**
+ * Converts the result to LLM output attributes
+ */
+function getLLMOutputMessagesAttributes(
+  response: Stream<ChatCompletionChunk> | ChatCompletion,
 ): Attributes {
-    return body.messages.reduce((acc, message, index) => {
-        const index_prefix = `${SemanticConventions.LLM_INPUT_MESSAGES}.${index}`;
-        acc[`${index_prefix}.${SemanticConventions.MESSAGE_CONTENT}`] = String(
-            message.content
-        );
-        acc[`${index_prefix}.${SemanticConventions.MESSAGE_ROLE}`] = String(
-            message.role
-        );
-        return acc;
+  // Handle chat completion
+  if (Object.prototype.hasOwnProperty.call(response, "choices")) {
+    const completion = response as ChatCompletion;
+    // Right now support just the first choice
+    const choice = completion.choices[0];
+    if (!choice) {
+      return {};
+    }
+    return [choice.message].reduce((acc, message, index) => {
+      const index_prefix = `${SemanticConventions.LLM_OUTPUT_MESSAGES}.${index}`;
+      acc[`${index_prefix}.${SemanticConventions.MESSAGE_CONTENT}`] = String(
+        message.content,
+      );
+      acc[`${index_prefix}.${SemanticConventions.MESSAGE_ROLE}`] = message.role;
+      return acc;
     }, {} as Attributes);
+  }
+  return {};
 }

package/src/version.ts

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