simplify

Author: nicohrubec

packages/core/src/tracing/openai/index.ts

@@ -12,7 +12,6 @@ import {
   GEN_AI_OPERATION_NAME_ATTRIBUTE,
   GEN_AI_REQUEST_AVAILABLE_TOOLS_ATTRIBUTE,
   GEN_AI_REQUEST_MODEL_ATTRIBUTE,
-  GEN_AI_RESPONSE_TEXT_ATTRIBUTE,
   GEN_AI_SYSTEM_ATTRIBUTE,
   GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE,
 } from '../ai/gen-ai-attributes';
@@ -26,18 +25,8 @@ import {
 } from '../ai/utils';
 import { OPENAI_METHOD_REGISTRY } from './constants';
 import { instrumentStream } from './streaming';
-import type { ChatCompletionChunk, OpenAiOptions, OpenAiResponse, OpenAIStream, ResponseStreamingEvent } from './types';
-import {
-  addChatCompletionAttributes,
-  addConversationAttributes,
-  addEmbeddingsAttributes,
-  addResponsesApiAttributes,
-  extractRequestParameters,
-  isChatCompletionResponse,
-  isConversationResponse,
-  isEmbeddingsResponse,
-  isResponsesApiResponse,
-} from './utils';
+import type { ChatCompletionChunk, OpenAiOptions, OpenAIStream, ResponseStreamingEvent } from './types';
+import { addResponseAttributes, extractRequestParameters } from './utils';
 
 /**
  * Extract available tools from request parameters
@@ -88,33 +77,6 @@ function extractRequestAttributes(args: unknown[], operationName: string): Recor
   return attributes;
 }
 
-/**
- * Add response attributes to spans
- * This supports Chat Completion, Responses API, Embeddings, and Conversations API responses
- */
-function addResponseAttributes(span: Span, result: unknown, recordOutputs?: boolean): void {
-  if (!result || typeof result !== 'object') return;
-
-  const response = result as OpenAiResponse;
-
-  if (isChatCompletionResponse(response)) {
-    addChatCompletionAttributes(span, response, recordOutputs);
-    if (recordOutputs && response.choices?.length) {
-      const responseTexts = response.choices.map(choice => choice.message?.content || '');
-      span.setAttributes({ [GEN_AI_RESPONSE_TEXT_ATTRIBUTE]: JSON.stringify(responseTexts) });
-    }
-  } else if (isResponsesApiResponse(response)) {
-    addResponsesApiAttributes(span, response, recordOutputs);
-    if (recordOutputs && response.output_text) {
-      span.setAttributes({ [GEN_AI_RESPONSE_TEXT_ATTRIBUTE]: response.output_text });
-    }
-  } else if (isEmbeddingsResponse(response)) {
-    addEmbeddingsAttributes(span, response);
-  } else if (isConversationResponse(response)) {
-    addConversationAttributes(span, response);
-  }
-}
-
 // Extract and record AI request inputs, if present. This is intentionally separate from response attributes.
 function addRequestAttributes(span: Span, params: Record<string, unknown>, operationName: string): void {
   // Store embeddings input on a separate attribute and do not truncate it

packages/core/src/tracing/openai/streaming.ts

@@ -3,9 +3,14 @@ import { SPAN_STATUS_ERROR } from '../../tracing';
 import type { Span } from '../../types-hoist/span';
 import {
   GEN_AI_RESPONSE_FINISH_REASONS_ATTRIBUTE,
+  GEN_AI_RESPONSE_ID_ATTRIBUTE,
+  GEN_AI_RESPONSE_MODEL_ATTRIBUTE,
   GEN_AI_RESPONSE_STREAMING_ATTRIBUTE,
   GEN_AI_RESPONSE_TEXT_ATTRIBUTE,
   GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE,
+  GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE,
+  GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE,
+  GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE,
 } from '../ai/gen-ai-attributes';
 import { RESPONSE_EVENT_TYPES } from './constants';
 import type {
@@ -15,12 +20,7 @@ import type {
   ResponseFunctionCall,
   ResponseStreamingEvent,
 } from './types';
-import {
-  isChatCompletionChunk,
-  isResponsesApiStreamEvent,
-  setCommonResponseAttributes,
-  setTokenUsageAttributes,
-} from './utils';
+import { isChatCompletionChunk, isResponsesApiStreamEvent } from './utils';
 
 /**
  * State object used to accumulate information from a stream of OpenAI events/chunks.
@@ -36,8 +36,6 @@ interface StreamingState {
   responseId: string;
   /** The model name. */
   responseModel: string;
-  /** The timestamp of the response. */
-  responseTimestamp: number;
   /** Number of prompt/input tokens used. */
   promptTokens: number | undefined;
   /** Number of completion/output tokens used. */
@@ -99,7 +97,6 @@ function processChatCompletionToolCalls(toolCalls: ChatCompletionToolCall[], sta
 function processChatCompletionChunk(chunk: ChatCompletionChunk, state: StreamingState, recordOutputs: boolean): void {
   state.responseId = chunk.id ?? state.responseId;
   state.responseModel = chunk.model ?? state.responseModel;
-  state.responseTimestamp = chunk.created ?? state.responseTimestamp;
 
   if (chunk.usage) {
     // For stream responses, the input tokens remain constant across all events in the stream.
@@ -183,7 +180,6 @@ function processResponsesApiEvent(
     const { response } = event as { response: OpenAIResponseObject };
     state.responseId = response.id ?? state.responseId;
     state.responseModel = response.model ?? state.responseModel;
-    state.responseTimestamp = response.created_at ?? state.responseTimestamp;
 
     if (response.usage) {
       // For stream responses, the input tokens remain constant across all events in the stream.
@@ -227,7 +223,6 @@ export async function* instrumentStream<T>(
     finishReasons: [],
     responseId: '',
     responseModel: '',
-    responseTimestamp: 0,
     promptTokens: undefined,
     completionTokens: undefined,
     totalTokens: undefined,
@@ -245,35 +240,31 @@ export async function* instrumentStream<T>(
       yield event;
     }
   } finally {
-    setCommonResponseAttributes(span, state.responseId, state.responseModel);
-    setTokenUsageAttributes(span, state.promptTokens, state.completionTokens, state.totalTokens);
-
-    span.setAttributes({
+    const attrs: Record<string, string | number | boolean> = {
+      [GEN_AI_RESPONSE_ID_ATTRIBUTE]: state.responseId,
+      [GEN_AI_RESPONSE_MODEL_ATTRIBUTE]: state.responseModel,
       [GEN_AI_RESPONSE_STREAMING_ATTRIBUTE]: true,
-    });
+    };
+
+    if (state.promptTokens !== undefined) attrs[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE] = state.promptTokens;
+    if (state.completionTokens !== undefined) attrs[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE] = state.completionTokens;
+    if (state.totalTokens !== undefined) attrs[GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE] = state.totalTokens;
 
     if (state.finishReasons.length) {
-      span.setAttributes({
-        [GEN_AI_RESPONSE_FINISH_REASONS_ATTRIBUTE]: JSON.stringify(state.finishReasons),
-      });
+      attrs[GEN_AI_RESPONSE_FINISH_REASONS_ATTRIBUTE] = JSON.stringify(state.finishReasons);
     }
 
     if (recordOutputs && state.responseTexts.length) {
-      span.setAttributes({
-        [GEN_AI_RESPONSE_TEXT_ATTRIBUTE]: state.responseTexts.join(''),
-      });
+      attrs[GEN_AI_RESPONSE_TEXT_ATTRIBUTE] = state.responseTexts.join('');
     }
 
-    // Set tool calls attribute if any were accumulated
     const chatCompletionToolCallsArray = Object.values(state.chatCompletionToolCalls);
     const allToolCalls = [...chatCompletionToolCallsArray, ...state.responsesApiToolCalls];
-
     if (allToolCalls.length > 0) {
-      span.setAttributes({
-        [GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE]: JSON.stringify(allToolCalls),
-      });
+      attrs[GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE] = JSON.stringify(allToolCalls);
     }
 
+    span.setAttributes(attrs);
     span.end();
   }
 }