[MLOB-1560] LLMObs Span Processor

packages/dd-trace/src/exporters/agent/index.js

@@ -42,6 +42,12 @@ class AgentExporter {
   }
 
   export (spans) {
+    const { llmobs } = this._config
+    if (llmobs.enabled && llmobs.agentlessEnabled) {
+      log.debug('LLMObs agentless mode enabled. Not sending APM spans to the agent.')
+      return
+    }
+
     this._writer.append(spans)
 
     const { flushInterval } = this._config

packages/dd-trace/src/format.js

@@ -160,7 +160,9 @@ function extractTags (trace, span) {
           break
         }
       default: // eslint-disable-line no-fallthrough
-        addTag(trace.meta, trace.metrics, tag, tags[tag])
+        if (!tag.startsWith('_ml_obs')) { // don't add ml_obs-related tags
+          addTag(trace.meta, trace.metrics, tag, tags[tag])
+        }
     }
   }
   setSingleSpanIngestionTags(trace, context._spanSampling)

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

@@ -0,0 +1,147 @@
+'use strict'
+
+const {
+  SPAN_KIND,
+  MODEL_NAME,
+  MODEL_PROVIDER,
+  METADATA,
+  INPUT_MESSAGES,
+  INPUT_VALUE,
+  OUTPUT_MESSAGES,
+  INPUT_DOCUMENTS,
+  OUTPUT_DOCUMENTS,
+  OUTPUT_VALUE,
+  METRICS,
+  ML_APP,
+  TAGS,
+  PARENT_ID_KEY,
+  SESSION_ID,
+  NAME
+} = require('./constants')
+
+const {
+  ERROR_MESSAGE,
+  ERROR_TYPE,
+  ERROR_STACK
+} = require('../constants')
+
+const AgentlessWriter = require('./writers/spans/agentless')
+const AgentProxyWriter = require('./writers/spans/agentProxy')
+const { isLLMSpan } = require('./util')
+
+const tracerVersion = require('../../../../package.json').version
+
+class LLMObsSpanProcessor {
+  constructor (config) {
+    this._config = config
+    const { llmobs } = config
+
+    if (llmobs.enabled) {
+      const LLMObsSpanWriter = llmobs.agentlessEnabled ? AgentlessWriter : AgentProxyWriter
+      this._writer = new LLMObsSpanWriter(config)
+    }
+  }
+
+  process (span) {
+    if (!this._config.llmobs.enabled) return
+    if (!isLLMSpan(span)) return
+    const payload = this._process(span)
+
+    this._writer.append(payload)
+  }
+
+  _process (span) {
+    const tags = span.context()._tags
+    const spanKind = tags[SPAN_KIND]
+
+    const meta = { 'span.kind': spanKind, input: {}, output: {} }
+    const input = {}
+    const output = {}
+
+    if (['llm', 'embedding'].includes(spanKind) && tags[MODEL_NAME]) {
+      meta.model_name = tags[MODEL_NAME]
+      meta.model_provider = (tags[MODEL_PROVIDER] || 'custom').toLowerCase()
+    }
+    if (tags[METADATA]) {
+      meta.metadata = JSON.parse(tags[METADATA])
+    }
+    if (spanKind === 'llm' && tags[INPUT_MESSAGES]) {
+      input.messages = JSON.parse(tags[INPUT_MESSAGES])
+    }
+    if (tags[INPUT_VALUE]) {
+      input.value = tags[INPUT_VALUE]
+    }
+    if (spanKind === 'llm' && tags[OUTPUT_MESSAGES]) {
+      output.messages = JSON.parse(tags[OUTPUT_MESSAGES])
+    }
+    if (spanKind === 'embedding' && tags[INPUT_DOCUMENTS]) {
+      input.documents = JSON.parse(tags[INPUT_DOCUMENTS])
+    }
+    if (tags[OUTPUT_VALUE]) {
+      output.value = tags[OUTPUT_VALUE]
+    }
+    if (spanKind === 'retrieval' && tags[OUTPUT_DOCUMENTS]) {
+      output.documents = JSON.parse(tags[OUTPUT_DOCUMENTS])
+    }
+
+    const error = tags.error
+    if (error) {
+      meta[ERROR_MESSAGE] = tags[ERROR_MESSAGE] || error.message || error.code
+      meta[ERROR_TYPE] = tags[ERROR_TYPE] || error.name
+      meta[ERROR_STACK] = tags[ERROR_STACK] || error.stack
+    }
+
+    if (input) meta.input = input
+    if (output) meta.output = output
+
+    const metrics = JSON.parse(tags[METRICS] || '{}')
+
+    const mlApp = tags[ML_APP]
+    const sessionId = tags[SESSION_ID]
+    const parentId = tags[PARENT_ID_KEY]
+
+    const name = tags[NAME] || span._name
+
+    const llmObsSpanEvent = {
+      trace_id: span.context().toTraceId(true),
+      span_id: span.context().toSpanId(),
+      parent_id: parentId,
+      name,
+      tags: this._processTags(span, mlApp, sessionId, error),
+      start_ns: Math.round(span._startTime * 1e6),
+      duration: Math.round(span._duration * 1e6),
+      status: tags.error ? 'error' : 'ok',
+      meta,
+      metrics,
+      _dd: {
+        span_id: span.context().toSpanId(),
+        trace_id: span.context().toTraceId(true)
+      }
+    }
+
+    if (sessionId) llmObsSpanEvent.session_id = sessionId
+
+    return llmObsSpanEvent
+  }
+
+  _processTags (span, mlApp, sessionId, error) {
+    let tags = {
+      version: this._config.version,
+      env: this._config.env,
+      service: this._config.service,
+      source: 'integration',
+      ml_app: mlApp,
+      'dd-trace.version': tracerVersion,
+      error: Number(!!error) || 0,
+      language: 'javascript'
+    }
+    const errType = span.context()._tags[ERROR_TYPE] || error?.name
+    if (errType) tags.error_type = errType
+    if (sessionId) tags.session_id = sessionId
+    const existingTags = JSON.parse(span.context()._tags[TAGS] || '{}')
+    if (existingTags) tags = { ...tags, ...existingTags }
+    return Object.entries(tags).map(([key, value]) => `${key}:${value ?? ''}`)
+  }
+}
+
+module.exports = LLMObsSpanProcessor

packages/dd-trace/src/llmobs/util.js

@@ -1,5 +1,7 @@
 'use strict'
 
+const { SPAN_TYPE } = require('../../../../ext/tags')
+
 function encodeUnicode (str) {
   if (!str) return str
   return str.split('').map(char => {
@@ -11,6 +13,11 @@ function encodeUnicode (str) {
   }).join('')
 }
 
+function isLLMSpan (span) {
+  return ['llm', 'openai'].includes(span?.context()._tags[SPAN_TYPE])
+}
+
 module.exports = {
-  encodeUnicode
+  encodeUnicode,
+  isLLMSpan
 }

packages/dd-trace/src/span_processor.js

@@ -6,6 +6,7 @@ const SpanSampler = require('./span_sampler')
 const GitMetadataTagger = require('./git_metadata_tagger')
 
 const { SpanStatsProcessor } = require('./span_stats')
+const LLMObsSpanProcessor = require('./llmobs/span_processor')
 
 const startedSpans = new WeakSet()
 const finishedSpans = new WeakSet()
@@ -20,6 +21,7 @@ class SpanProcessor {
     this._stats = new SpanStatsProcessor(config)
     this._spanSampler = new SpanSampler(config.sampler)
     this._gitMetadataTagger = new GitMetadataTagger(config)
+    this._llmobs = new LLMObsSpanProcessor(config)
   }
 
   process (span) {
@@ -42,6 +44,7 @@ class SpanProcessor {
 
       for (const span of started) {
         if (span._duration !== undefined) {
+          this._llmobs.process(span)
           const formattedSpan = format(span)
           this._stats.onSpanFinished(formattedSpan)
           formatted.push(formattedSpan)

packages/dd-trace/test/exporters/agent/exporter.spec.js

@@ -15,11 +15,13 @@ describe('Exporter', () => {
   let writer
   let prioritySampler
   let span
+  let llmobs
 
   beforeEach(() => {
     url = 'www.example.com'
     flushInterval = 1000
     span = {}
+    llmobs = {}
     writer = {
       append: sinon.spy(),
       flush: sinon.spy(),
@@ -35,7 +37,7 @@ describe('Exporter', () => {
 
   it('should pass computed stats header through to writer', () => {
     const stats = { enabled: true }
-    exporter = new Exporter({ url, flushInterval, stats }, prioritySampler)
+    exporter = new Exporter({ url, flushInterval, stats, llmobs }, prioritySampler)
     expect(Writer).to.have.been.calledWithMatch({
       headers: {
         'Datadog-Client-Computed-Stats': 'yes'
@@ -46,7 +48,7 @@ describe('Exporter', () => {
   it('should pass computed stats header through to writer if standalone appsec is enabled', () => {
     const stats = { enabled: false }
     const appsec = { standalone: { enabled: true } }
-    exporter = new Exporter({ url, flushInterval, stats, appsec }, prioritySampler)
+    exporter = new Exporter({ url, flushInterval, stats, appsec, llmobs }, prioritySampler)
 
     expect(Writer).to.have.been.calledWithMatch({
       headers: {
@@ -57,27 +59,27 @@ describe('Exporter', () => {
 
   it('should support IPv6', () => {
     const stats = { enabled: true }
-    exporter = new Exporter({ hostname: '::1', flushInterval, stats }, prioritySampler)
+    exporter = new Exporter({ hostname: '::1', flushInterval, stats, llmobs }, prioritySampler)
     expect(Writer).to.have.been.calledWithMatch({
       url: new URL('http://[::1]')
     })
   })
 
   describe('when interval is set to a positive number', () => {
     beforeEach(() => {
-      exporter = new Exporter({ url, flushInterval }, prioritySampler)
+      exporter = new Exporter({ url, flushInterval, llmobs }, prioritySampler)
     })
 
     it('should not flush if export has not been called', (done) => {
-      exporter = new Exporter({ url, flushInterval }, prioritySampler)
+      exporter = new Exporter({ url, flushInterval, llmobs }, prioritySampler)
       setTimeout(() => {
         expect(writer.flush).not.to.have.been.called
         done()
       }, flushInterval + 100)
     })
 
     it('should flush after the configured interval if a payload has been exported', (done) => {
-      exporter = new Exporter({ url, flushInterval }, prioritySampler)
+      exporter = new Exporter({ url, flushInterval, llmobs }, prioritySampler)
       exporter.export([{}])
       setTimeout(() => {
         expect(writer.flush).to.have.been.called
@@ -101,7 +103,7 @@ describe('Exporter', () => {
 
   describe('when interval is set to 0', () => {
     beforeEach(() => {
-      exporter = new Exporter({ url, flushInterval: 0 })
+      exporter = new Exporter({ url, flushInterval: 0, llmobs })
     })
 
     it('should flush right away when interval is set to 0', () => {
@@ -112,7 +114,7 @@ describe('Exporter', () => {
 
   describe('setUrl', () => {
     beforeEach(() => {
-      exporter = new Exporter({ url })
+      exporter = new Exporter({ url, llmobs })
     })
 
     it('should set the URL on self and writer', () => {
@@ -122,4 +124,22 @@ describe('Exporter', () => {
       expect(writer.setUrl).to.have.been.calledWith(url)
     })
   })
+
+  describe('with llmobs agentless enabled', () => {
+    beforeEach(() => {
+      exporter = new Exporter({
+        url,
+        llmobs: {
+          enabled: true,
+          agentlessEnabled: true
+        }
+      })
+    })
+
+    it('does not write the span', () => {
+      exporter.export([span])
+
+      expect(writer.append).to.not.have.been.called
+    })
+  })
 })

packages/dd-trace/test/format.spec.js

@@ -577,5 +577,15 @@ describe('format', () => {
 
       expect(trace.metrics).to.have.property('_dd1.sr.eausr', 1)
     })
+
+    it('should skip formatting temporary _ml_obs.* tags', () => {
+      spanContext._tags['_ml_obs.meta.span.kind'] = 'llm'
+      spanContext._tags['_ml_obs.meta.ml_app'] = 'test'
+
+      trace = format(span)
+
+      expect(trace.meta['_ml_obs.meta.span.kind']).to.be.undefined
+      expect(trace.meta['_ml_obs.meta.ml_app']).to.be.undefined
+    })
   })
 })