@arizeai/openinference-instrumentation-openai 0.0.4 → 0.1.1

package/src/instrumentation.ts

@@ -13,6 +13,7 @@ import {
   SpanKind,
   Attributes,
   SpanStatusCode,
+  Span,
 } from "@opentelemetry/api";
 import { VERSION } from "./version";
 import {
@@ -24,12 +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, isString } from "./typeUtils";
 
 const MODULE_NAME = "openai";
 
@@ -64,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
@@ -72,12 +79,10 @@ 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;
@@ -99,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, () => {
@@ -119,23 +124,95 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
             },
           );
           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]: isChatCompletionResponse(
-                  result,
-                )
-                  ? result.model
-                  : body.model,
-                ...getLLMOutputMessagesAttributes(result),
+                [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,
+                // 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();
             }
-            span.setStatus({ code: SpanStatusCode.OK });
-            span.end();
             return result;
           });
           return context.bind(execContext, wrappedPromise);
@@ -153,11 +230,11 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
       (original: EmbeddingsCreateType): any => {
         return function patchedEmbeddingCreate(
           this: unknown,
-          ...args: Parameters<typeof module.OpenAI.Embeddings.prototype.create>
+          ...args: Parameters<EmbeddingsCreateType>
         ) {
           const body = args[0];
           const { input } = body;
-          const isStringInput = typeof input == "string";
+          const isStringInput = typeof input === "string";
           const span = instrumentation.tracer.startSpan(`OpenAI Embeddings`, {
             kind: SpanKind.INTERNAL,
             attributes: {
@@ -234,78 +311,215 @@ function isChatCompletionResponse(
 }
 
 /**
- * Converts the body of the request to LLM input messages
+ * 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 {};
-    }
-    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);
+  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 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") {
+  if (typeof request.input === "string") {
     return {
       [`${SemanticConventions.EMBEDDING_EMBEDDINGS}.0.${SemanticConventions.EMBEDDING_TEXT}`]:
         request.input,
@@ -313,11 +527,11 @@ function getEmbeddingTextAttributes(
   } else if (
     Array.isArray(request.input) &&
     request.input.length > 0 &&
-    typeof request.input[0] == "string"
+    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;
+      const indexPrefix = `${SemanticConventions.EMBEDDING_EMBEDDINGS}.${index}.`;
+      acc[`${indexPrefix}${SemanticConventions.EMBEDDING_TEXT}`] = input;
       return acc;
     }, {} as Attributes);
   }
@@ -332,9 +546,98 @@ 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}`] =
+    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 = "";
+  // Tool and function call attributes can also arrive in the stream
+  // NB: the tools and function calls arrive in partial diffs
+  // So the final tool and function calls need to be aggregated
+  // across chunks
+  const toolAndFunctionCallAttributes: Attributes = {};
+  // The first message is for the assistant response so we start at 1
+  for await (const chunk of stream) {
+    if (chunk.choices.length <= 0) {
+      continue;
+    }
+    const choice = chunk.choices[0];
+    if (choice.delta.content) {
+      streamResponse += choice.delta.content;
+    }
+    // Accumulate the tool and function call attributes
+    const toolAndFunctionCallAttributesDiff =
+      getToolAndFunctionCallAttributesFromStreamChunk(chunk);
+    for (const [key, value] of Object.entries(
+      toolAndFunctionCallAttributesDiff,
+    )) {
+      if (isString(toolAndFunctionCallAttributes[key]) && isString(value)) {
+        toolAndFunctionCallAttributes[key] += value;
+      } else if (isString(value)) {
+        toolAndFunctionCallAttributes[key] = value;
+      }
+    }
+  }
+  const messageIndexPrefix = `${SemanticConventions.LLM_OUTPUT_MESSAGES}.0.`;
+
+  // Append the attributes to the span as a message
+  const attributes: Attributes = {
+    [SemanticConventions.OUTPUT_VALUE]: streamResponse,
+    [SemanticConventions.OUTPUT_MIME_TYPE]: MimeType.TEXT,
+    [`${messageIndexPrefix}${SemanticConventions.MESSAGE_CONTENT}`]:
+      streamResponse,
+    [`${messageIndexPrefix}${SemanticConventions.MESSAGE_ROLE}`]: "assistant",
+  };
+  // Add the tool and function call attributes
+  for (const [key, value] of Object.entries(toolAndFunctionCallAttributes)) {
+    attributes[`${messageIndexPrefix}${key}`] = value;
+  }
+  span.setAttributes(attributes);
+  span.end();
+}
+
+/**
+ * Extracts the semantic attributes from the stream chunk for tool_calls and function_calls
+ */
+function getToolAndFunctionCallAttributesFromStreamChunk(
+  chunk: ChatCompletionChunk,
+): Attributes {
+  if (chunk.choices.length <= 0) {
+    return {};
+  }
+  const choice = chunk.choices[0];
+  const attributes: Attributes = {};
+  if (choice.delta.tool_calls) {
+    choice.delta.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 (choice.delta.function_call) {
+    attributes[SemanticConventions.MESSAGE_FUNCTION_CALL_NAME] =
+      choice.delta.function_call.name;
+    attributes[SemanticConventions.MESSAGE_FUNCTION_CALL_ARGUMENTS_JSON] =
+      choice.delta.function_call.arguments;
+  }
+  return attributes;
+}

package/src/typeUtils.ts

@@ -0,0 +1,13 @@
+/**
+ * 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");
+}
+
+export function isString(value: unknown): value is string {
+  return typeof value === "string";
+}

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