@arizeai/openinference-instrumentation-openai 0.0.3 → 0.1.0

package/package.json

@@ -1,16 +1,20 @@
 {
   "name": "@arizeai/openinference-instrumentation-openai",
-  "version": "0.0.3",
+  "version": "0.1.0",
   "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.1.0"
   },
   "keywords": [],
   "files": [

package/src/instrumentation.ts

@@ -13,6 +13,7 @@ import {
   SpanKind,
   Attributes,
   SpanStatusCode,
+  Span,
 } from "@opentelemetry/api";
 import { VERSION } from "./version";
 import {
@@ -24,8 +25,17 @@ import {
   ChatCompletion,
   ChatCompletionChunk,
   ChatCompletionCreateParamsBase,
+  ChatCompletionMessage,
+  ChatCompletionMessageParam,
 } from "openai/resources/chat/completions";
+import { CompletionCreateParamsBase } from "openai/resources/completions";
 import { Stream } from "openai/streaming";
+import {
+  Completion,
+  CreateEmbeddingResponse,
+  EmbeddingCreateParams,
+} from "openai/resources";
+import { assertUnreachable } from "./typeUtils";
 
 const MODULE_NAME = "openai";
 
@@ -60,7 +70,8 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
     }
     // eslint-disable-next-line @typescript-eslint/no-this-alias
     const instrumentation: OpenAIInstrumentation = this;
-    type CompletionCreateType =
+
+    type ChatCompletionCreateType =
       typeof module.OpenAI.Chat.Completions.prototype.create;
 
     // Patch create chat completions
@@ -68,19 +79,17 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
       module.OpenAI.Chat.Completions.prototype,
       "create",
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      (original: CompletionCreateType): any => {
+      (original: ChatCompletionCreateType): any => {
         return function patchedCreate(
           this: unknown,
-          ...args: Parameters<
-            typeof module.OpenAI.Chat.Completions.prototype.create
-          >
+          ...args: Parameters<ChatCompletionCreateType>
         ) {
           const body = args[0];
           const { messages: _messages, ...invocationParameters } = body;
           const span = instrumentation.tracer.startSpan(
             `OpenAI Chat Completions`,
             {
-              kind: SpanKind.CLIENT,
+              kind: SpanKind.INTERNAL,
               attributes: {
                 [SemanticConventions.OPENINFERENCE_SPAN_KIND]:
                   OpenInferenceSpanKind.LLM,
@@ -95,7 +104,7 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
           );
           const execContext = trace.setSpan(context.active(), span);
           const execPromise = safeExecuteInTheMiddle<
-            ReturnType<CompletionCreateType>
+            ReturnType<ChatCompletionCreateType>
           >(
             () => {
               return context.with(execContext, () => {
@@ -106,18 +115,169 @@ 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();
               }
             },
           );
           const wrappedPromise = execPromise.then((result) => {
-            if (result) {
+            if (isChatCompletionResponse(result)) {
+              // Record the results
+              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]: result.model,
+                ...getChatCompletionLLMOutputMessagesAttributes(result),
+                ...getUsageAttributes(result),
+              });
+              span.setStatus({ code: SpanStatusCode.OK });
+              span.end();
+            } else {
+              // This is a streaming response
+              // handle the chunks and add them to the span
+              // First split the stream via tee
+              const [leftStream, rightStream] = result.tee();
+              consumeChatCompletionStreamChunks(rightStream, span);
+              result = leftStream;
+            }
+
+            return result;
+          });
+          return context.bind(execContext, wrappedPromise);
+        };
+      },
+    );
+
+    // Patch create completions
+    type CompletionsCreateType =
+      typeof module.OpenAI.Completions.prototype.create;
+
+    this._wrap(
+      module.OpenAI.Completions.prototype,
+      "create",
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      (original: CompletionsCreateType): any => {
+        return function patchedCreate(
+          this: unknown,
+          ...args: Parameters<CompletionsCreateType>
+        ) {
+          const body = args[0];
+          const { prompt: _prompt, ...invocationParameters } = body;
+          const span = instrumentation.tracer.startSpan(`OpenAI Completions`, {
+            kind: SpanKind.INTERNAL,
+            attributes: {
+              [SemanticConventions.OPENINFERENCE_SPAN_KIND]:
+                OpenInferenceSpanKind.LLM,
+              [SemanticConventions.LLM_MODEL_NAME]: body.model,
+              [SemanticConventions.LLM_INVOCATION_PARAMETERS]:
+                JSON.stringify(invocationParameters),
+              ...getCompletionInputValueAndMimeType(body),
+            },
+          });
+          const execContext = trace.setSpan(context.active(), span);
+          const execPromise = safeExecuteInTheMiddle<
+            ReturnType<CompletionsCreateType>
+          >(
+            () => {
+              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 (isCompletionResponse(result)) {
               // Record the results
               span.setAttributes({
                 [SemanticConventions.OUTPUT_VALUE]: JSON.stringify(result),
                 [SemanticConventions.OUTPUT_MIME_TYPE]: MimeType.JSON,
-                ...getLLMOutputMessagesAttributes(result),
+                // Override the model from the value sent by the server
+                [SemanticConventions.LLM_MODEL_NAME]: result.model,
+                ...getCompletionOutputValueAndMimeType(result),
                 ...getUsageAttributes(result),
               });
+              span.setStatus({ code: SpanStatusCode.OK });
+              span.end();
+            }
+            return result;
+          });
+          return context.bind(execContext, wrappedPromise);
+        };
+      },
+    );
+
+    // 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<EmbeddingsCreateType>
+        ) {
+          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();
@@ -127,6 +287,7 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
         };
       },
     );
+
     module.openInferencePatched = true;
     return module;
   }
@@ -136,71 +297,282 @@ 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");
   }
 }
 
 /**
- * Converts the body of the request to LLM input messages
+ * 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;
+}
+
+/**
+ * type-guard that checks if the response is a completion response
+ */
+function isCompletionResponse(
+  response: Stream<Completion> | Completion,
+): response is Completion {
+  return "choices" in response;
+}
+
+/**
+ * type-guard that checks if completion prompt attribute is an array of strings
+ */
+function isPromptStringArray(
+  prompt: CompletionCreateParamsBase["prompt"],
+): prompt is Array<string> {
+  return (
+    Array.isArray(prompt) && prompt.every((item) => typeof item === "string")
+  );
+}
+
+/**
+ * Converts the body of a chat completions request to LLM input messages
  */
 function getLLMInputMessagesAttributes(
   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,
-    );
+    const messageAttributes = getChatCompletionInputMessageAttributes(message);
+    const indexPrefix = `${SemanticConventions.LLM_INPUT_MESSAGES}.${index}.`;
+    // Flatten the attributes on the index prefix
+    for (const [key, value] of Object.entries(messageAttributes)) {
+      acc[`${indexPrefix}${key}`] = value;
+    }
     return acc;
   }, {} as Attributes);
 }
 
+function getChatCompletionInputMessageAttributes(
+  message: ChatCompletionMessageParam,
+): Attributes {
+  const role = message.role;
+  const attributes: Attributes = {
+    [SemanticConventions.MESSAGE_ROLE]: role,
+  };
+  // Add the content only if it is a string
+  if (typeof message.content === "string")
+    attributes[SemanticConventions.MESSAGE_CONTENT] = message.content;
+  switch (role) {
+    case "user":
+      // There's nothing to add for the user
+      break;
+    case "assistant":
+      if (message.tool_calls) {
+        message.tool_calls.forEach((toolCall, index) => {
+          // Make sure the tool call has a function
+          if (toolCall.function) {
+            const toolCallIndexPrefix = `${SemanticConventions.MESSAGE_TOOL_CALLS}.${index}.`;
+            attributes[
+              toolCallIndexPrefix + SemanticConventions.TOOL_CALL_FUNCTION_NAME
+            ] = toolCall.function.name;
+            attributes[
+              toolCallIndexPrefix +
+                SemanticConventions.TOOL_CALL_FUNCTION_ARGUMENTS_JSON
+            ] = toolCall.function.arguments;
+          }
+        });
+      }
+      break;
+    case "function":
+      attributes[SemanticConventions.MESSAGE_FUNCTION_CALL_NAME] = message.name;
+      break;
+    case "tool":
+      // There's nothing to add for the tool. There is a tool_id, but there are no
+      // semantic conventions for it
+      break;
+    case "system":
+      // There's nothing to add for the system. Content is captured above
+      break;
+    default:
+      assertUnreachable(role);
+      break;
+  }
+  return attributes;
+}
+
 /**
- * Get Usage attributes
+ * Converts the body of a completions request to input 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,
-      };
+function getCompletionInputValueAndMimeType(
+  body: CompletionCreateParamsBase,
+): Attributes {
+  if (typeof body.prompt === "string") {
+    return {
+      [SemanticConventions.INPUT_VALUE]: body.prompt,
+      [SemanticConventions.INPUT_MIME_TYPE]: MimeType.TEXT,
+    };
+  } else if (isPromptStringArray(body.prompt)) {
+    const prompt = body.prompt[0]; // Only single prompts are currently supported
+    if (prompt === undefined) {
+      return {};
     }
+    return {
+      [SemanticConventions.INPUT_VALUE]: prompt,
+      [SemanticConventions.INPUT_MIME_TYPE]: MimeType.TEXT,
+    };
   }
+  // Other cases in which the prompt is a token or array of tokens are currently unsupported
   return {};
 }
 
 /**
- * Converts the result to LLM output attributes
+ * Get usage attributes
  */
-function getLLMOutputMessagesAttributes(
-  response: Stream<ChatCompletionChunk> | ChatCompletion,
+function getUsageAttributes(
+  completion: ChatCompletion | Completion,
 ): Attributes {
-  // 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 {};
+  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 chat completion result to LLM output attributes
+ */
+function getChatCompletionLLMOutputMessagesAttributes(
+  chatCompletion: ChatCompletion,
+): Attributes {
+  // Right now support just the first choice
+  const choice = chatCompletion.choices[0];
+  if (!choice) {
+    return {};
+  }
+  return [choice.message].reduce((acc, message, index) => {
+    const indexPrefix = `${SemanticConventions.LLM_OUTPUT_MESSAGES}.${index}.`;
+    const messageAttributes = getChatCompletionOutputMessageAttributes(message);
+    // Flatten the attributes on the index prefix
+    for (const [key, value] of Object.entries(messageAttributes)) {
+      acc[`${indexPrefix}${key}`] = value;
     }
-    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);
+}
+
+function getChatCompletionOutputMessageAttributes(
+  message: ChatCompletionMessage,
+): Attributes {
+  const role = message.role;
+  const attributes: Attributes = {
+    [SemanticConventions.MESSAGE_ROLE]: role,
+  };
+  if (message.content) {
+    attributes[SemanticConventions.MESSAGE_CONTENT] = message.content;
+  }
+  if (message.tool_calls) {
+    message.tool_calls.forEach((toolCall, index) => {
+      const toolCallIndexPrefix = `${SemanticConventions.MESSAGE_TOOL_CALLS}.${index}.`;
+      // Double check that the tool call has a function
+      // NB: OpenAI only supports tool calls with functions right now but this may change
+      if (toolCall.function) {
+        attributes[
+          toolCallIndexPrefix + SemanticConventions.TOOL_CALL_FUNCTION_NAME
+        ] = toolCall.function.name;
+        attributes[
+          toolCallIndexPrefix +
+            SemanticConventions.TOOL_CALL_FUNCTION_ARGUMENTS_JSON
+        ] = toolCall.function.arguments;
+      }
+    });
+  }
+  if (message.function_call) {
+    attributes[SemanticConventions.MESSAGE_FUNCTION_CALL_NAME] =
+      message.function_call.name;
+    attributes[SemanticConventions.MESSAGE_FUNCTION_CALL_ARGUMENTS_JSON] =
+      message.function_call.arguments;
+  }
+  return attributes;
+}
+
+/**
+ * Converts the completion result to output attributes
+ */
+function getCompletionOutputValueAndMimeType(
+  completion: Completion,
+): Attributes {
+  // Right now support just the first choice
+  const choice = completion.choices[0];
+  if (!choice) {
+    return {};
+  }
+  return {
+    [SemanticConventions.OUTPUT_VALUE]: String(choice.text),
+    [SemanticConventions.OUTPUT_MIME_TYPE]: MimeType.TEXT,
+  };
+}
+
+/**
+ * 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 indexPrefix = `${SemanticConventions.EMBEDDING_EMBEDDINGS}.${index}.`;
+      acc[`${indexPrefix}${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 indexPrefix = `${SemanticConventions.EMBEDDING_EMBEDDINGS}.${index}.`;
+    acc[`${indexPrefix}${SemanticConventions.EMBEDDING_VECTOR}`] =
+      embedding.embedding;
+    return acc;
+  }, {} as Attributes);
+}
+
+/**
+ * Consumes the stream chunks and adds them to the span
+ */
+async function consumeChatCompletionStreamChunks(
+  stream: Stream<ChatCompletionChunk>,
+  span: Span,
+) {
+  let streamResponse = "";
+  for await (const chunk of stream) {
+    if (chunk.choices.length > 0 && chunk.choices[0].delta.content) {
+      streamResponse += chunk.choices[0].delta.content;
+    }
+  }
+  span.setAttributes({
+    [SemanticConventions.OUTPUT_VALUE]: streamResponse,
+    [SemanticConventions.OUTPUT_MIME_TYPE]: MimeType.TEXT,
+    [`${SemanticConventions.LLM_OUTPUT_MESSAGES}.0.${SemanticConventions.MESSAGE_CONTENT}`]:
+      streamResponse,
+    [`${SemanticConventions.LLM_OUTPUT_MESSAGES}.0.${SemanticConventions.MESSAGE_ROLE}`]:
+      "assistant",
+  });
+  span.end();
+}

package/src/typeUtils.ts

@@ -0,0 +1,9 @@
+/**
+ * Utility function that uses the type system to check if a switch statement is exhaustive.
+ * If the switch statement is not exhaustive, there will be a type error caught in typescript
+ *
+ * See https://stackoverflow.com/questions/39419170/how-do-i-check-that-a-switch-block-is-exhaustive-in-typescript for more details.
+ */
+export function assertUnreachable(_: never): never {
+  throw new Error("Unreachable");
+}

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.1.0";