@arizeai/openinference-instrumentation-openai 0.0.1-rc8 → 0.0.2

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@arizeai/openinference-instrumentation-openai",
-  "version": "0.0.1-rc8",
+  "version": "0.0.2",
   "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.11"
   },
   "keywords": [],
   "files": [

package/src/instrumentation.ts

@@ -10,17 +10,22 @@ import {
     diag,
     context,
     trace,
-    isSpanContextValid,
-    Span,
     SpanKind,
     Attributes,
+    SpanStatusCode,
 } from "@opentelemetry/api";
 import { VERSION } from "./version";
 import {
     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";
 
@@ -42,18 +47,22 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
         );
         return module;
     }
+    /**
+     * Patches the OpenAI module
+     */
     private patch(
         module: typeof openai & { openInferencePatched?: boolean },
         moduleVersion?: string
     ) {
-        console.log(`Applying patch for ${MODULE_NAME}@${moduleVersion}`);
-        const plugin = this;
+        diag.debug(`Applying patch for ${MODULE_NAME}@${moduleVersion}`);
         if (module?.openInferencePatched) {
             return module;
         }
         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",
@@ -65,6 +74,7 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
                     >
                 ) {
                     const body = args[0];
+                    const { messages, ...invocationParameters } = body;
                     const span = instrumentation.tracer.startSpan(
                         `OpenAI Chat Completions`,
                         {
@@ -72,6 +82,14 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
                             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),
                             },
                         }
@@ -86,11 +104,25 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
                             });
                         },
                         (error) => {
-                            // TODO handler the error
+                            // Push the error to the span
+                            if (error) {
+                                span.recordException(error);
+                            }
                         }
                     );
                     const wrappedPromise = execPromise.then((result) => {
-                        // TODO set span attributes
+                        if (result) {
+                            // Record the results
+                            span.setAttributes({
+                                [SemanticConventions.OUTPUT_VALUE]:
+                                    JSON.stringify(result),
+                                [SemanticConventions.OUTPUT_MIME_TYPE]:
+                                    MimeType.JSON,
+                                ...getLLMOutputMessagesAttributes(result),
+                                ...getUsageAttributes(result),
+                            });
+                        }
+                        span.setStatus({ code: SpanStatusCode.OK });
                         span.end();
                         return result;
                     });
@@ -101,13 +133,17 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
         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
@@ -123,3 +159,51 @@ function getLLMInputMessagesAttributes(
         return acc;
     }, {} as Attributes);
 }
+
+/**
+ * Get Usage attributes
+ */
+function getUsageAttributes(
+    response: Stream<ChatCompletionChunk> | ChatCompletion
+) {
+    if (response.hasOwnProperty("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 {
+    // Handle chat completion
+    if (response.hasOwnProperty("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-rc7";
+export const VERSION = "0.0.2";