@arizeai/openinference-instrumentation-openai 0.0.3 → 0.0.4

package/package.json

@@ -1,16 +1,20 @@
 {
   "name": "@arizeai/openinference-instrumentation-openai",
-  "version": "0.0.3",
+  "version": "0.0.4",
   "description": "OpenInference instrumentation for OpenAI",
   "private": false,
   "main": "dist/src/index.js",
   "module": "dist/esm/index.js",
   "esnext": "dist/esnext/index.js",
   "types": "dist/src/index.d.ts",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/Arize-ai/openinference.git"
+  },
   "dependencies": {
     "@opentelemetry/api": "^1.7.0",
     "@opentelemetry/instrumentation": "^0.46.0",
-    "@arizeai/openinference-semantic-conventions": "0.0.12"
+    "@arizeai/openinference-semantic-conventions": "0.0.13"
   },
   "keywords": [],
   "files": [

package/src/instrumentation.ts

@@ -26,6 +26,10 @@ import {
   ChatCompletionCreateParamsBase,
 } from "openai/resources/chat/completions";
 import { Stream } from "openai/streaming";
+import {
+  CreateEmbeddingResponse,
+  EmbeddingCreateParams,
+} from "openai/resources";
 
 const MODULE_NAME = "openai";
 
@@ -80,7 +84,7 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
           const span = instrumentation.tracer.startSpan(
             `OpenAI Chat Completions`,
             {
-              kind: SpanKind.CLIENT,
+              kind: SpanKind.INTERNAL,
               attributes: {
                 [SemanticConventions.OPENINFERENCE_SPAN_KIND]:
                   OpenInferenceSpanKind.LLM,
@@ -106,6 +110,11 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
               // Push the error to the span
               if (error) {
                 span.recordException(error);
+                span.setStatus({
+                  code: SpanStatusCode.ERROR,
+                  message: error.message,
+                });
+                span.end();
               }
             },
           );
@@ -115,6 +124,12 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
               span.setAttributes({
                 [SemanticConventions.OUTPUT_VALUE]: JSON.stringify(result),
                 [SemanticConventions.OUTPUT_MIME_TYPE]: MimeType.JSON,
+                // Override the model from the value sent by the server
+                [SemanticConventions.LLM_MODEL_NAME]: isChatCompletionResponse(
+                  result,
+                )
+                  ? result.model
+                  : body.model,
                 ...getLLMOutputMessagesAttributes(result),
                 ...getUsageAttributes(result),
               });
@@ -127,6 +142,75 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
         };
       },
     );
+
+    // Patch embeddings
+    type EmbeddingsCreateType =
+      typeof module.OpenAI.Embeddings.prototype.create;
+    this._wrap(
+      module.OpenAI.Embeddings.prototype,
+      "create",
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      (original: EmbeddingsCreateType): any => {
+        return function patchedEmbeddingCreate(
+          this: unknown,
+          ...args: Parameters<typeof module.OpenAI.Embeddings.prototype.create>
+        ) {
+          const body = args[0];
+          const { input } = body;
+          const isStringInput = typeof input == "string";
+          const span = instrumentation.tracer.startSpan(`OpenAI Embeddings`, {
+            kind: SpanKind.INTERNAL,
+            attributes: {
+              [SemanticConventions.OPENINFERENCE_SPAN_KIND]:
+                OpenInferenceSpanKind.EMBEDDING,
+              [SemanticConventions.EMBEDDING_MODEL_NAME]: body.model,
+              [SemanticConventions.INPUT_VALUE]: isStringInput
+                ? input
+                : JSON.stringify(input),
+              [SemanticConventions.INPUT_MIME_TYPE]: isStringInput
+                ? MimeType.TEXT
+                : MimeType.JSON,
+              ...getEmbeddingTextAttributes(body),
+            },
+          });
+          const execContext = trace.setSpan(context.active(), span);
+          const execPromise = safeExecuteInTheMiddle<
+            ReturnType<EmbeddingsCreateType>
+          >(
+            () => {
+              return context.with(execContext, () => {
+                return original.apply(this, args);
+              });
+            },
+            (error) => {
+              // Push the error to the span
+              if (error) {
+                span.recordException(error);
+                span.setStatus({
+                  code: SpanStatusCode.ERROR,
+                  message: error.message,
+                });
+                span.end();
+              }
+            },
+          );
+          const wrappedPromise = execPromise.then((result) => {
+            if (result) {
+              // Record the results
+              span.setAttributes({
+                // Do not record the output data as it can be large
+                ...getEmbeddingEmbeddingsAttributes(result),
+              });
+            }
+            span.setStatus({ code: SpanStatusCode.OK });
+            span.end();
+            return result;
+          });
+          return context.bind(execContext, wrappedPromise);
+        };
+      },
+    );
+
     module.openInferencePatched = true;
     return module;
   }
@@ -136,9 +220,19 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
   private unpatch(moduleExports: typeof openai, moduleVersion?: string) {
     diag.debug(`Removing patch for ${MODULE_NAME}@${moduleVersion}`);
     this._unwrap(moduleExports.OpenAI.Chat.Completions.prototype, "create");
+    this._unwrap(moduleExports.OpenAI.Embeddings.prototype, "create");
   }
 }
 
+/**
+ * type-guard that checks if the response is a chat completion response
+ */
+function isChatCompletionResponse(
+  response: Stream<ChatCompletionChunk> | ChatCompletion,
+): response is ChatCompletion {
+  return "choices" in response;
+}
+
 /**
  * Converts the body of the request to LLM input messages
  */
@@ -204,3 +298,43 @@ function getLLMOutputMessagesAttributes(
   }
   return {};
 }
+
+/**
+ * Converts the embedding result payload to embedding attributes
+ */
+function getEmbeddingTextAttributes(
+  request: EmbeddingCreateParams,
+): Attributes {
+  if (typeof request.input == "string") {
+    return {
+      [`${SemanticConventions.EMBEDDING_EMBEDDINGS}.0.${SemanticConventions.EMBEDDING_TEXT}`]:
+        request.input,
+    };
+  } else if (
+    Array.isArray(request.input) &&
+    request.input.length > 0 &&
+    typeof request.input[0] == "string"
+  ) {
+    return request.input.reduce((acc, input, index) => {
+      const index_prefix = `${SemanticConventions.EMBEDDING_EMBEDDINGS}.${index}`;
+      acc[`${index_prefix}.${SemanticConventions.EMBEDDING_TEXT}`] = input;
+      return acc;
+    }, {} as Attributes);
+  }
+  // Ignore other cases where input is a number or an array of numbers
+  return {};
+}
+
+/**
+ * Converts the embedding result payload to embedding attributes
+ */
+function getEmbeddingEmbeddingsAttributes(
+  response: CreateEmbeddingResponse,
+): Attributes {
+  return response.data.reduce((acc, embedding, index) => {
+    const index_prefix = `${SemanticConventions.EMBEDDING_EMBEDDINGS}.${index}`;
+    acc[`${index_prefix}.${SemanticConventions.EMBEDDING_VECTOR}`] =
+      embedding.embedding;
+    return acc;
+  }, {} as Attributes);
+}

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.2";
+export const VERSION = "0.0.4";