@arizeai/openinference-instrumentation-openai 0.6.0 → 1.1.0

package/src/instrumentation.ts

@@ -40,7 +40,11 @@ import {
 import { assertUnreachable, isString } from "./typeUtils";
 import { isTracingSuppressed } from "@opentelemetry/core";
 
-import { safelyJSONStringify } from "@arizeai/openinference-core";
+import {
+  OITracer,
+  safelyJSONStringify,
+  TraceConfigOptions,
+} from "@arizeai/openinference-core";
 
 const MODULE_NAME = "openai";
 
@@ -74,13 +78,34 @@ function getExecContext(span: Span) {
   }
   return execContext;
 }
+/**
+ * An auto instrumentation class for OpenAI that creates {@link https://github.com/Arize-ai/openinference/blob/main/spec/semantic_conventions.md|OpenInference} Compliant spans for the OpenAI API
+ * @param instrumentationConfig The config for the instrumentation @see {@link InstrumentationConfig}
+ * @param traceConfig The OpenInference trace configuration. Can be used to mask or redact sensitive information on spans. @see {@link TraceConfigOptions}
+ */
 export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
-  constructor(config?: InstrumentationConfig) {
+  private oiTracer: OITracer;
+  constructor({
+    instrumentationConfig,
+    traceConfig,
+  }: {
+    /**
+     * The config for the instrumentation
+     * @see {@link InstrumentationConfig}
+     */
+    instrumentationConfig?: InstrumentationConfig;
+    /**
+     * The OpenInference trace configuration. Can be used to mask or redact sensitive information on spans.
+     * @see {@link TraceConfigOptions}
+     */
+    traceConfig?: TraceConfigOptions;
+  } = {}) {
     super(
       "@arizeai/openinference-instrumentation-openai",
       VERSION,
-      Object.assign({}, config),
+      Object.assign({}, instrumentationConfig),
     );
+    this.oiTracer = new OITracer({ tracer: this.tracer, traceConfig });
   }
 
   protected init(): InstrumentationModuleDefinition<typeof openai> {
@@ -131,7 +156,7 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
         ) {
           const body = args[0];
           const { messages: _messages, ...invocationParameters } = body;
-          const span = instrumentation.tracer.startSpan(
+          const span = instrumentation.oiTracer.startSpan(
             `OpenAI Chat Completions`,
             {
               kind: SpanKind.INTERNAL,
@@ -150,13 +175,12 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
               },
             },
           );
-
           const execContext = getExecContext(span);
           const execPromise = safeExecuteInTheMiddle<
             ReturnType<ChatCompletionCreateType>
           >(
             () => {
-              return context.with(execContext, () => {
+              return context.with(trace.setSpan(execContext, span), () => {
                 return original.apply(this, args);
               });
             },
@@ -216,26 +240,29 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
         ) {
           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),
-              [SemanticConventions.LLM_SYSTEM]: LLMSystem.OPENAI,
-              [SemanticConventions.LLM_PROVIDER]: LLMProvider.OPENAI,
-              ...getCompletionInputValueAndMimeType(body),
+          const span = instrumentation.oiTracer.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),
+                [SemanticConventions.LLM_SYSTEM]: LLMSystem.OPENAI,
+                [SemanticConventions.LLM_PROVIDER]: LLMProvider.OPENAI,
+                ...getCompletionInputValueAndMimeType(body),
+              },
             },
-          });
+          );
           const execContext = getExecContext(span);
 
           const execPromise = safeExecuteInTheMiddle<
             ReturnType<CompletionsCreateType>
           >(
             () => {
-              return context.with(execContext, () => {
+              return context.with(trace.setSpan(execContext, span), () => {
                 return original.apply(this, args);
               });
             },
@@ -287,7 +314,7 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
           const body = args[0];
           const { input } = body;
           const isStringInput = typeof input === "string";
-          const span = instrumentation.tracer.startSpan(`OpenAI Embeddings`, {
+          const span = instrumentation.oiTracer.startSpan(`OpenAI Embeddings`, {
             kind: SpanKind.INTERNAL,
             attributes: {
               [SemanticConventions.OPENINFERENCE_SPAN_KIND]:
@@ -307,7 +334,7 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
             ReturnType<EmbeddingsCreateType>
           >(
             () => {
-              return context.with(execContext, () => {
+              return context.with(trace.setSpan(execContext, span), () => {
                 return original.apply(this, args);
               });
             },
@@ -345,7 +372,7 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
       // This can fail if the module is made immutable via the runtime or bundler
       module.openInferencePatched = true;
     } catch (e) {
-      diag.warn(`Failed to set ${MODULE_NAME} patched flag on the module`, e);
+      diag.debug(`Failed to set ${MODULE_NAME} patched flag on the module`, e);
     }
 
     return module;
@@ -475,15 +502,21 @@ function getChatCompletionInputMessageAttributes(
     case "assistant":
       if (message.tool_calls) {
         message.tool_calls.forEach((toolCall, index) => {
+          const toolCallIndexPrefix = `${SemanticConventions.MESSAGE_TOOL_CALLS}.${index}.`;
+
+          // Add the tool call id if it exists
+          if (toolCall.id) {
+            attributes[
+              `${toolCallIndexPrefix}${SemanticConventions.TOOL_CALL_ID}`
+            ] = toolCall.id;
+          }
           // 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
+              `${toolCallIndexPrefix}${SemanticConventions.TOOL_CALL_FUNCTION_NAME}`
             ] = toolCall.function.name;
             attributes[
-              toolCallIndexPrefix +
-                SemanticConventions.TOOL_CALL_FUNCTION_ARGUMENTS_JSON
+              `${toolCallIndexPrefix}${SemanticConventions.TOOL_CALL_FUNCTION_ARGUMENTS_JSON}`
             ] = toolCall.function.arguments;
           }
         });
@@ -493,8 +526,10 @@ function getChatCompletionInputMessageAttributes(
       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
+      if (message.tool_call_id) {
+        attributes[`${SemanticConventions.MESSAGE_TOOL_CALL_ID}`] =
+          message.tool_call_id;
+      }
       break;
     case "system":
       // There's nothing to add for the system. Content is captured above
@@ -585,6 +620,12 @@ function getChatCompletionOutputMessageAttributes(
   if (message.tool_calls) {
     message.tool_calls.forEach((toolCall, index) => {
       const toolCallIndexPrefix = `${SemanticConventions.MESSAGE_TOOL_CALLS}.${index}.`;
+      // Add the tool call id if it exists
+      if (toolCall.id) {
+        attributes[
+          `${toolCallIndexPrefix}${SemanticConventions.TOOL_CALL_ID}`
+        ] = toolCall.id;
+      }
       // 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) {
@@ -731,6 +772,12 @@ function getToolAndFunctionCallAttributesFromStreamChunk(
   if (choice.delta.tool_calls) {
     choice.delta.tool_calls.forEach((toolCall, index) => {
       const toolCallIndexPrefix = `${SemanticConventions.MESSAGE_TOOL_CALLS}.${index}.`;
+      // Add the tool call id if it exists
+      if (toolCall.id) {
+        attributes[
+          `${toolCallIndexPrefix}${SemanticConventions.TOOL_CALL_ID}`
+        ] = toolCall.id;
+      }
       // 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) {

package/src/version.ts

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