dd-trace 5.116.0 → 5.118.0

package/packages/dd-trace/src/exporter.js

@@ -4,6 +4,7 @@ const fs = require('fs')
 const exporters = require('../../../ext/exporters')
 const { getEnvironmentVariable } = require('../../dd-trace/src/config/helper')
 const constants = require('./constants')
+const { isTrue } = require('./util')
 
 module.exports = function getExporter (name) {
   switch (name) {
@@ -20,20 +21,26 @@ module.exports = function getExporter (name) {
     case exporters.AGENT_PROXY:
       return require('./ci-visibility/exporters/agent-proxy')
     case exporters.CI_VALIDATION:
-      return require('./ci-visibility/exporters/ci-validation')
+      if (hasCiValidationEnvironment()) return require('./ci-visibility/exporters/ci-validation')
+      break
     case exporters.JEST_WORKER:
     case exporters.CUCUMBER_WORKER:
     case exporters.MOCHA_WORKER:
     case exporters.PLAYWRIGHT_WORKER:
     case exporters.VITEST_WORKER:
       return require('./ci-visibility/exporters/test-worker')
-    default: {
-      const inAWSLambda = getEnvironmentVariable('AWS_LAMBDA_FUNCTION_NAME') !== undefined
-      const usingAgent = inAWSLambda && (
-        fs.existsSync(constants.DATADOG_LAMBDA_EXTENSION_PATH) ||
-        fs.existsSync(constants.DATADOG_MINI_AGENT_PATH)
-      )
-      return inAWSLambda && !usingAgent ? require('./exporters/log') : require('./exporters/agent')
-    }
   }
+
+  const inAWSLambda = getEnvironmentVariable('AWS_LAMBDA_FUNCTION_NAME') !== undefined
+  const usingAgent = inAWSLambda && (
+    fs.existsSync(constants.DATADOG_LAMBDA_EXTENSION_PATH) ||
+    fs.existsSync(constants.DATADOG_MINI_AGENT_PATH)
+  )
+  return inAWSLambda && !usingAgent ? require('./exporters/log') : require('./exporters/agent')
+}
+
+function hasCiValidationEnvironment () {
+  return isTrue(getEnvironmentVariable('_DD_TEST_OPTIMIZATION_VALIDATION_MODE')) &&
+    getEnvironmentVariable('_DD_TEST_OPTIMIZATION_VALIDATION_MANIFEST_FILE') &&
+    getEnvironmentVariable('_DD_TEST_OPTIMIZATION_VALIDATION_OUTPUT_DIR')
 }

package/packages/dd-trace/src/llmobs/index.js

@@ -16,6 +16,8 @@ const {
   SAMPLING_DECISION,
   PROPAGATED_SAMPLE_RATE_KEY,
   PROPAGATED_SAMPLING_DECISION_KEY,
+  TRACE_ID,
+  PROPAGATED_TRACE_ID_KEY,
 } = require('./constants/tags')
 const { storage } = require('./storage')
 const telemetry = require('./telemetry')
@@ -25,6 +27,7 @@ const LLMObsTagger = require('./tagger')
 const LLMObsSpanWriter = require('./writers/spans')
 const { setAgentStrategy } = require('./writers/util')
 const { INCOMPATIBLE_INITIALIZATION } = require('./constants/text')
+const { llmObsTraceIdToWire } = require('./util')
 
 const spanFinishCh = channel('dd-trace:span:finish')
 const evalMetricAppendCh = channel('llmobs:eval-metric:append')
@@ -137,8 +140,12 @@ function handleLLMObsInjection ({ carrier }) {
     mlObsSpanTags?.[SESSION_ID] ??
     parentContext?._trace?.tags?.[SESSION_ID_TRACE_DEFAULT_KEY] ??
     parentContext?._trace?.tags?.[PROPAGATED_SESSION_ID_KEY]
+  const llmobsTraceId = mlObsSpanTags?.[TRACE_ID]
+  const propagatedTraceId = llmobsTraceId === undefined
+    ? parentContext?._trace?.tags?.[PROPAGATED_TRACE_ID_KEY]
+    : llmObsTraceIdToWire(llmobsTraceId)
 
-  if (!parentId && !mlApp && samplingDecision == null && !sessionId) return
+  if (!parentId && !mlApp && samplingDecision == null && !sessionId && !propagatedTraceId) return
 
   // `_injectTags` only writes `x-datadog-tags` when the trace has `_dd.p.*`
   // tags, so it may be undefined here — coalesce before appending.
@@ -149,6 +156,7 @@ function handleLLMObsInjection ({ carrier }) {
   if (sessionId) tags += `${tags ? ',' : ''}${PROPAGATED_SESSION_ID_KEY}=${sessionId}`
   if (sampleRate != null) tags += `${tags ? ',' : ''}${PROPAGATED_SAMPLE_RATE_KEY}=${sampleRate}`
   if (samplingDecision != null) tags += `${tags ? ',' : ''}${PROPAGATED_SAMPLING_DECISION_KEY}=${samplingDecision}`
+  if (propagatedTraceId != null) tags += `${tags ? ',' : ''}${PROPAGATED_TRACE_ID_KEY}=${propagatedTraceId}`
   if (tags !== existing) carrier['x-datadog-tags'] = tags
 }
 

package/packages/dd-trace/src/llmobs/plugins/ai/util.js

@@ -13,6 +13,8 @@ const MODEL_METADATA_KEYS = new Set([
 const VERCEL_AI_TELEMETRY_METADATA_PREFIX = 'ai.telemetry.metadata.'
 const VERCEL_AI_MODEL_METADATA_PREFIX = 'gen_ai.request.'
 const VERCEL_AI_GENERATION_METADATA_PREFIX = 'ai.settings.'
+const UNPARSABLE_TOOL_RESULT = '[Unparsable Tool Result]'
+const UNSUPPORTED_TOOL_RESULT = '[Unsupported Tool Result]'
 
 /**
  * @typedef {import('../../../opentracing/span')} Span
@@ -26,8 +28,31 @@ const VERCEL_AI_GENERATION_METADATA_PREFIX = 'ai.settings.'
 
 /**
  * @typedef {{
- *   type: string,
- *   value?: unknown
+ *   type: 'text',
+ *   text: string
+ * } | {
+ *   type: 'media' | 'file',
+ *   mediaType: string
+ * } | {
+ *   type: 'file-data' | 'file-url' | 'file-id' | 'file-reference'
+ * } | {
+ *   type: 'image-data' | 'image-url' | 'image-file-id' | 'image-file-reference'
+ * } | {
+ *   type: 'custom'
+ * }} ToolCallContentPart
+ *
+ * @typedef {{
+ *   type: 'text' | 'error-text',
+ *   value: string
+ * } | {
+ *   type: 'json' | 'error-json',
+ *   value: unknown
+ * } | {
+ *   type: 'content',
+ *   value: ToolCallContentPart[]
+ * } | {
+ *   type: 'execution-denied',
+ *   reason?: string
  * }} ToolCallOutput
  *
  * @typedef {{ output?: ToolCallOutput, result?: unknown } & Record<string, unknown>} ToolCallResultContent
@@ -304,32 +329,89 @@ function getToolNameFromTags (tags) {
 }
 
 /**
- * Get the content of a tool call result.
- * Version 5 of the ai sdk sets this tag as `content.output`, with a `
- * @param {{ output?: { type: string, value: unknown }, result?: unknown }} content
+ * @param {unknown} value
+ * @returns {string}
+ */
+function stringifyToolCallResult (value) {
+  return JSON.stringify(value) ?? UNPARSABLE_TOOL_RESULT
+}
+
+/**
+ * @param {ToolCallContentPart[]} value
+ * @returns {string}
+ */
+function formatToolCallContent (value) {
+  if (!Array.isArray(value)) return UNPARSABLE_TOOL_RESULT
+
+  let result = ''
+  for (const part of value) {
+    if (typeof part !== 'object' || part === null) return UNPARSABLE_TOOL_RESULT
+
+    const { type } = part
+    if (type === 'text') {
+      if (typeof part.text !== 'string') return UNPARSABLE_TOOL_RESULT
+      result += part.text
+    } else if (type === 'media' || type === 'file') {
+      const { mediaType } = part
+      if (typeof mediaType !== 'string') return UNPARSABLE_TOOL_RESULT
+      result += mediaType === 'image' || mediaType.startsWith('image/') ? '[Image]' : '[File]'
+    } else if (
+      type === 'file-data' ||
+      type === 'file-url' ||
+      type === 'file-id' ||
+      type === 'file-reference'
+    ) {
+      result += '[File]'
+    } else if (
+      type === 'image-data' ||
+      type === 'image-url' ||
+      type === 'image-file-id' ||
+      type === 'image-file-reference'
+    ) {
+      result += '[Image]'
+    } else if (type === 'custom') {
+      result += '[Custom Content]'
+    } else {
+      return UNPARSABLE_TOOL_RESULT
+    }
+  }
+
+  return result
+}
+
+/**
+ * @param {ToolCallResultContent | null | undefined} content
  * @returns {string}
  */
 function getToolCallResultContent (content) {
-  const { output, result } = content
-  if (output) {
-    if (output.type === 'text') {
-      return output.value
-    } else if (output.type === 'json') {
-      return JSON.stringify(output.value)
+  try {
+    if (typeof content !== 'object' || content === null) return UNPARSABLE_TOOL_RESULT
+
+    const { output, result } = content
+    if (output !== undefined) {
+      if (typeof output !== 'object' || output === null) return UNPARSABLE_TOOL_RESULT
+
+      const { type, value } = output
+      if (type === 'text' || type === 'error-text') {
+        return typeof value === 'string' ? value : UNPARSABLE_TOOL_RESULT
+      } else if (type === 'json' || type === 'error-json') {
+        return stringifyToolCallResult(value)
+      } else if (type === 'content') {
+        return formatToolCallContent(value)
+      } else if (type === 'execution-denied') {
+        const { reason } = output
+        if (reason === undefined) return '[Tool Execution Denied]'
+        return typeof reason === 'string' ? reason : UNPARSABLE_TOOL_RESULT
+      }
+      return UNPARSABLE_TOOL_RESULT
     }
-    return '[Unparsable Tool Result]'
-  } else if (result) {
-    if (typeof result === 'string') {
-      return result
+    if (result !== undefined) {
+      return typeof result === 'string' ? result : stringifyToolCallResult(result)
     }
 
-    try {
-      return JSON.stringify(result)
-    } catch {
-      return '[Unparsable Tool Result]'
-    }
-  } else {
-    return '[Unsupported Tool Result]'
+    return UNSUPPORTED_TOOL_RESULT
+  } catch {
+    return UNPARSABLE_TOOL_RESULT
   }
 }
 

package/packages/dd-trace/src/llmobs/plugins/openai/index.js

@@ -5,7 +5,6 @@ const {
   PROMPT_TRACKING_INSTRUMENTATION_METHOD,
   PROMPT_MULTIMODAL,
   INSTRUMENTATION_METHOD_AUTO,
-  UNKNOWN_MODEL_PROVIDER,
 } = require('../../constants/tags')
 const { audioMimeTypeFromFormat, formatAudioPart, safeJsonParse } = require('../../util')
 const { AUDIO_MIME_TYPES } = require('./constants')
@@ -15,6 +14,7 @@ const {
   extractTextFromContentItem,
   extractContentParts,
   hasMultimodalInputs,
+  getOpenAIModelProvider,
 } = require('./utils')
 
 const allowedParamKeys = new Set([
@@ -108,14 +108,10 @@ class OpenAiLLMObsPlugin extends LLMObsPlugin {
   }
 
   _getModelProviderAndClient (baseUrl = '') {
-    if (baseUrl.includes('azure')) {
-      return { modelProvider: 'azure_openai', client: 'AzureOpenAI' }
-    } else if (baseUrl.includes('openai')) {
-      return { modelProvider: 'openai', client: 'OpenAI' }
-    } else if (baseUrl.includes('deepseek')) {
-      return { modelProvider: 'deepseek', client: 'DeepSeek' }
-    }
-    return { modelProvider: UNKNOWN_MODEL_PROVIDER, client: 'OpenAI' }
+    const modelProvider = getOpenAIModelProvider(baseUrl)
+    if (modelProvider === 'azure_openai') return { modelProvider, client: 'AzureOpenAI' }
+    if (modelProvider === 'deepseek') return { modelProvider, client: 'DeepSeek' }
+    return { modelProvider, client: 'OpenAI' }
   }
 
   _extractMetrics (response) {

package/packages/dd-trace/src/llmobs/plugins/openai/utils.js

@@ -1,5 +1,6 @@
 'use strict'
 
+const { UNKNOWN_MODEL_PROVIDER } = require('../../constants/tags')
 const { audioMimeTypeFromFormat, formatAudioPart } = require('../../util')
 const {
   INPUT_TYPE_IMAGE,
@@ -160,10 +161,29 @@ function extractContentParts (parts) {
   return { content: extracted.join('\n'), audioParts }
 }
 
+/**
+ * Maps an OpenAI-compatible base URL to a model provider string. Covers
+ * OpenAI, Azure OpenAI, and DeepSeek; falls back to UNKNOWN_MODEL_PROVIDER
+ * for unrecognised hosts (e.g. local proxies or custom deployments).
+ *
+ * Shared with the openai-agents integration since both consume the same
+ * client baseURL convention.
+ *
+ * @param {string} baseUrl
+ * @returns {string}
+ */
+function getOpenAIModelProvider (baseUrl = '') {
+  if (baseUrl.includes('azure')) return 'azure_openai'
+  if (baseUrl.includes('deepseek')) return 'deepseek'
+  if (baseUrl.includes('openai')) return 'openai'
+  return UNKNOWN_MODEL_PROVIDER
+}
+
 module.exports = {
   extractChatTemplateFromInstructions,
   normalizePromptVariables,
   extractTextFromContentItem,
   extractContentParts,
   hasMultimodalInputs,
+  getOpenAIModelProvider,
 }

package/packages/dd-trace/src/llmobs/plugins/openai-agents/utils.js

@@ -0,0 +1,321 @@
+'use strict'
+
+const { extractContentParts, extractTextFromContentItem } = require('../openai/utils')
+const { safeJsonParse } = require('../../util')
+
+/**
+ * Flatten Responses and Chat Completions content parts without dropping
+ * provider-specific image or audio inputs.
+ *
+ * @param {Array<object>} parts
+ * @returns {{ content: string, audioParts: Array<{ mimeType: string, content: string }> }}
+ */
+function extractMessageContent (parts) {
+  const contentParts = []
+  const audioParts = []
+
+  for (const part of parts) {
+    if (!part) continue
+    const text = extractTextFromContentItem(part)
+    if (text) {
+      contentParts.push(text)
+      continue
+    }
+
+    const extracted = extractContentParts([part])
+    if (extracted.content) contentParts.push(extracted.content)
+    if (extracted.audioParts.length > 0) audioParts.push(...extracted.audioParts)
+  }
+
+  return { content: contentParts.join(''), audioParts }
+}
+
+/**
+ * Normalize OpenAI Chat Completions tool calls to the LLMObs message schema.
+ *
+ * @param {object} message
+ * @returns {Array<{ toolId?: string, name?: string, arguments: object, type?: string }>}
+ */
+function extractChatCompletionToolCalls (message) {
+  if (message.function_call) {
+    return [{
+      name: message.function_call.name,
+      arguments: safeJsonParse(message.function_call.arguments, {}),
+      type: 'function',
+    }]
+  }
+
+  const toolCalls = []
+  if (!Array.isArray(message.tool_calls)) return toolCalls
+
+  for (const toolCall of message.tool_calls) {
+    if (!toolCall) continue
+    toolCalls.push({
+      toolId: toolCall.id,
+      name: toolCall.function?.name,
+      arguments: safeJsonParse(toolCall.function?.arguments, {}),
+      type: toolCall.type,
+    })
+  }
+  return toolCalls
+}
+
+/**
+ * Normalize a Chat Completions message to the LLMObs message schema.
+ *
+ * @param {object} message
+ * @returns {object}
+ */
+function normalizeChatCompletionMessage (message) {
+  const normalized = {
+    role: message.role || 'assistant',
+    content: message.content ?? '',
+  }
+  if (message.audioParts?.length > 0) normalized.audioParts = message.audioParts
+  const toolCalls = extractChatCompletionToolCalls(message)
+  if (toolCalls.length > 0) normalized.toolCalls = toolCalls
+  if (message.tool_call_id) normalized.toolId = message.tool_call_id
+  return normalized
+}
+
+/**
+ * Extracts input messages for an LLM span. agents-openai stores only
+ * `request.input` on `spanData._input` (string or message-array), and the
+ * system instructions are echoed back on the response as `instructions`.
+ *
+ * @param {string|Array<unknown>} input - The raw `request.input` (`spanData._input`).
+ * @param {string} [instructions] - System instructions echoed on `response.instructions`.
+ * @returns {Array<{ role: string, content: string }>}
+ */
+function extractInputMessages (input, instructions) {
+  const messages = []
+
+  if (instructions) {
+    messages.push({ role: 'system', content: instructions })
+  }
+
+  if (typeof input === 'string') {
+    messages.push({ role: 'user', content: input })
+  } else if (Array.isArray(input)) {
+    for (const item of input) {
+      if (!item) continue
+      if (item.type === 'message' || item.role) {
+        const role = item.role
+        if (!role) continue
+
+        let content = ''
+        let audioParts
+        if (Array.isArray(item.content)) {
+          const extracted = extractMessageContent(item.content)
+          content = extracted.content
+          audioParts = extracted.audioParts
+        } else if (typeof item.content === 'string') {
+          content = item.content
+        }
+
+        const message = normalizeChatCompletionMessage({ ...item, content, role, audioParts })
+        if (content || message.audioParts || message.toolCalls || message.toolId) {
+          messages.push(message)
+        }
+      } else if (item.type === 'function_call') {
+        let args = item.arguments
+        if (typeof args === 'string') {
+          try {
+            args = JSON.parse(args)
+          } catch {
+            args = {}
+          }
+        }
+        messages.push({
+          role: 'assistant',
+          toolCalls: [{
+            toolId: item.call_id,
+            name: item.name,
+            arguments: args,
+            type: item.type,
+          }],
+        })
+      } else if (item.type === 'function_call_output') {
+        messages.push({
+          role: 'user',
+          toolResults: [{
+            toolId: item.call_id,
+            result: item.output,
+            name: item.name || '',
+            type: item.type,
+          }],
+        })
+      }
+    }
+  }
+
+  return messages.length > 0 ? messages : [{ role: 'user', content: '' }]
+}
+
+/**
+ * Extracts output messages from the model response.
+ *
+ * @param {{ output?: Array<unknown> }} result - The model response
+ * @returns {Array<{ role: string, content: string }>}
+ */
+function extractOutputMessages (result) {
+  const messages = []
+
+  if (result?.output) {
+    for (const item of result.output) {
+      if (!item) continue
+      if (item.type === 'message') {
+        let content = ''
+        if (Array.isArray(item.content)) {
+          const textParts = []
+          for (const contentItem of item.content) {
+            if (contentItem?.type === 'output_text' && contentItem.text) {
+              textParts.push(contentItem.text)
+            }
+          }
+          content = textParts.join('')
+        } else if (typeof item.content === 'string') {
+          content = item.content
+        }
+
+        messages.push({ role: item.role || 'assistant', content })
+      } else if (item.type === 'function_call') {
+        let args = item.arguments
+        if (typeof args === 'string') {
+          try {
+            args = JSON.parse(args)
+          } catch {
+            args = {}
+          }
+        }
+        messages.push({
+          role: 'assistant',
+          toolCalls: [{
+            toolId: item.call_id,
+            name: item.name,
+            arguments: args,
+            type: item.type,
+          }],
+        })
+      }
+    }
+  }
+
+  return messages.length > 0 ? messages : [{ content: '', role: '' }]
+}
+
+/**
+ * Extracts output messages from an OpenAI Chat Completions response.
+ *
+ * @param {Array<{ choices?: Array<{ message?: object }> }>} result - The model responses
+ * @returns {Array<object>}
+ */
+function extractGenerationOutputMessages (result) {
+  const messages = []
+
+  if (Array.isArray(result)) {
+    for (const response of result) {
+      if (!Array.isArray(response?.choices)) continue
+      for (const choice of response.choices) {
+        const message = choice?.message
+        if (!message) continue
+        messages.push(normalizeChatCompletionMessage(message))
+      }
+    }
+  }
+
+  return messages.length > 0 ? messages : [{ content: '', role: '' }]
+}
+
+/**
+ * Extracts token usage metrics from the model response. Returns `undefined`
+ * when there's nothing to tag, so callers can skip the tagger call without
+ * allocating an Object.keys array.
+ *
+ * @param {{ usage?: { inputTokens?: number, outputTokens?: number, totalTokens?: number,
+ *   outputTokensDetails?: { reasoningTokens?: number },
+ *   completion_tokens_details?: { reasoning_tokens?: number } } }} result
+ * @returns {{ inputTokens?: number, outputTokens?: number, totalTokens?: number,
+ *   reasoningOutputTokens?: number } | undefined}
+ */
+function extractMetrics (result) {
+  const usage = result?.usage
+  if (!usage) return
+
+  const inputTokens = usage.inputTokens ?? usage.input_tokens ?? usage.prompt_tokens
+  const outputTokens = usage.outputTokens ?? usage.output_tokens ?? usage.completion_tokens
+  const totalTokens = usage.totalTokens ?? usage.total_tokens
+  const reasoningTokens = usage.outputTokensDetails?.reasoningTokens ??
+    usage.output_tokens_details?.reasoning_tokens ??
+    usage.completion_tokens_details?.reasoning_tokens
+
+  if (inputTokens === undefined && outputTokens === undefined &&
+      totalTokens === undefined && !reasoningTokens) return
+
+  const metrics = {}
+  if (inputTokens !== undefined) metrics.inputTokens = inputTokens
+  if (outputTokens !== undefined) metrics.outputTokens = outputTokens
+  // Tagger maps `reasoningOutputTokens` → `reasoning_output_tokens` in the
+  // LLMObs span event. Skip when zero — emitting a zero just adds noise.
+  if (reasoningTokens) metrics.reasoningOutputTokens = reasoningTokens
+
+  if (totalTokens !== undefined) {
+    metrics.totalTokens = totalTokens
+  } else if (metrics.inputTokens !== undefined && metrics.outputTokens !== undefined) {
+    metrics.totalTokens = metrics.inputTokens + metrics.outputTokens
+  }
+
+  return metrics
+}
+
+// Fields the OpenAI Responses API echoes back from the request configuration.
+// agents-openai only stores `request.input` on the span — the user's
+// `modelSettings` aren't directly observable, so we read the response-echoed
+// values. Matches dd-trace-py's openai-agents integration (see
+// `OaiSpanAdapter.llmobs_metadata`); both ship without filtering OpenAI's
+// default values.
+const RESPONSE_METADATA_FIELDS = [
+  'temperature',
+  'max_output_tokens',
+  'top_p',
+  'tools',
+  'tool_choice',
+  'truncation',
+]
+
+/**
+ * Extracts metadata from the model response. Mirrors Python's
+ * `OaiSpanAdapter.llmobs_metadata` — emits all response-echoed configuration
+ * fields plus `text` when present. Returns `undefined` when nothing was
+ * captured, so callers can skip the tagger call without allocating.
+ *
+ * @param {object | undefined} response
+ * @returns {object | undefined}
+ */
+function extractMetadata (response) {
+  if (!response) return
+
+  let metadata
+  for (const field of RESPONSE_METADATA_FIELDS) {
+    const value = response[field]
+    if (value !== undefined && value !== null) {
+      metadata ??= {}
+      metadata[field] = value
+    }
+  }
+
+  if (response.text) {
+    metadata ??= {}
+    metadata.text = response.text
+  }
+
+  return metadata
+}
+
+module.exports = {
+  extractInputMessages,
+  extractOutputMessages,
+  extractGenerationOutputMessages,
+  extractMetrics,
+  extractMetadata,
+}

package/packages/dd-trace/src/llmobs/sdk.js

@@ -11,6 +11,7 @@ const {
   SPAN_KIND,
   OUTPUT_VALUE,
   INPUT_VALUE,
+  TRACE_ID,
 } = require('./constants/tags')
 const {
   getFunctionArguments,
@@ -345,7 +346,7 @@ class LLMObs extends NoopLLMObs {
     }
     try {
       return {
-        traceId: span.context().toTraceId(true),
+        traceId: LLMObsTagger.tagMap.get(span)[TRACE_ID],
         spanId: span.context().toSpanId(),
       }
     } catch {

package/packages/dd-trace/src/llmobs/span_processor.js

@@ -35,6 +35,7 @@ const {
   LLMOBS_SUBMITTED_TAG_KEY,
   SAMPLE_RATE,
   SAMPLING_DECISION,
+  TRACE_ID,
 } = require('./constants/tags')
 const { UNSERIALIZABLE_VALUE_TEXT } = require('./constants/text')
 const telemetry = require('./telemetry')
@@ -236,8 +237,11 @@ class LLMObsSpanProcessor {
       meta.input.prompt = prompt
     }
 
+    const apmTraceId = span.context().toTraceId(true)
+    const llmobsTraceId = mlObsTags[TRACE_ID] ?? apmTraceId
+
     const llmObsSpanEvent = {
-      trace_id: span.context().toTraceId(true),
+      trace_id: llmobsTraceId,
       span_id: span.context().toSpanId(),
       parent_id: parentId,
       name,
@@ -249,9 +253,10 @@ class LLMObsSpanProcessor {
       metrics,
       _dd: {
         span_id: span.context().toSpanId(),
-        trace_id: span.context().toTraceId(true),
+        trace_id: apmTraceId,
         sample_rate: mlObsTags[SAMPLE_RATE],
         sampling_decision: mlObsTags[SAMPLING_DECISION],
+        apm_trace_id: apmTraceId,
       },
     }