feat(llmobs): add user llm observability span processor to modify spans (#6329)

Author: sabrenner

docs/test.ts

@@ -603,6 +603,27 @@ llmobs.enable({
 // manually disable
 llmobs.disable()
 
+// register a processor
+llmobs.registerProcessor((llmobsSpan) => {
+  const drop = llmobsSpan.getTag('drop')
+  if (drop) {
+    return null
+  }
+
+  const redactInput = llmobsSpan.getTag('redactInput')
+  if (redactInput) {
+    llmobsSpan.input = llmobsSpan.input.map(input => {
+      return {
+        ...input,
+      }
+    })
+  }
+
+  return llmobsSpan
+})
+
+llmobs.deregisterProcessor()
+
 // trace block of code
 llmobs.trace({ name: 'name', kind: 'llm' }, () => {})
 llmobs.trace({ kind: 'llm', name: 'myLLM', modelName: 'myModel', modelProvider: 'myProvider' }, () => {})

index.d.ts

@@ -2566,6 +2566,25 @@ declare namespace tracer {
       annotate (options: llmobs.AnnotationOptions): void
       annotate (span: tracer.Span | undefined, options: llmobs.AnnotationOptions): void
 
+      /**
+       * Register a processor to be called on each LLMObs span.
+       *
+       * This can be used to modify the span before it is sent to LLMObs. For example, you can modify the input/output.
+       * You can also return `null` to omit the span entirely from being sent to LLM Observability.
+       *
+       * Otherwise, if the return value from the processor is not an instance of `LLMObservabilitySpan`, the span will be dropped.
+       *
+       * To deregister the processor, call `llmobs.deregisterProcessor()`
+       * @param processor A function that will be called for each span.
+       * @throws {Error} If a processor is already registered.
+       */
+      registerProcessor (processor: ((span: LLMObservabilitySpan) => LLMObservabilitySpan | null)): void
+
+      /**
+       * Deregister a processor.
+       */
+      deregisterProcessor (): void
+
       /**
        * Submits a custom evaluation metric for a given span ID and trace ID.
        * @param spanContext The span context of the span to submit the evaluation metric for.
@@ -2579,6 +2598,25 @@ declare namespace tracer {
       flush (): void
     }
 
+    interface LLMObservabilitySpan {
+      /**
+       * The input content associated with the span.
+       */
+      input: { content: string, role?: string }[]
+
+      /**
+       * The output content associated with the span.
+       */
+      output: { content: string, role?: string }[]
+
+      /**
+       * Get a tag from the span.
+       * @param key The key of the tag to get.
+       * @returns The value of the tag, or `undefined` if the tag does not exist.
+       */
+      getTag (key: string): string | undefined
+    }
+
     interface EvaluationOptions {
       /**
        * The name of the evaluation metric

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

@@ -16,6 +16,7 @@ const spanProcessCh = channel('dd-trace:span:process')
 const evalMetricAppendCh = channel('llmobs:eval-metric:append')
 const flushCh = channel('llmobs:writers:flush')
 const injectCh = channel('dd-trace:span:inject')
+const registerUserSpanProcessorCh = channel('llmobs:register-processor')
 
 const LLMObsEvalMetricsWriter = require('./writers/evaluations')
 const LLMObsTagger = require('./tagger')
@@ -56,6 +57,7 @@ function enable (config) {
 
   evalMetricAppendCh.subscribe(handleEvalMetricAppend)
   flushCh.subscribe(handleFlush)
+  registerUserSpanProcessorCh.subscribe(handleRegisterProcessor)
 
   // span processing
   spanProcessor = new LLMObsSpanProcessor(config)
@@ -86,6 +88,7 @@ function disable () {
   if (flushCh.hasSubscribers) flushCh.unsubscribe(handleFlush)
   if (spanProcessCh.hasSubscribers) spanProcessCh.unsubscribe(handleSpanProcess)
   if (injectCh.hasSubscribers) injectCh.unsubscribe(handleLLMObsParentIdInjection)
+  if (registerUserSpanProcessorCh.hasSubscribers) registerUserSpanProcessorCh.unsubscribe(handleRegisterProcessor)
 
   spanWriter?.destroy()
   evalWriter?.destroy()
@@ -126,6 +129,10 @@ function handleFlush () {
   telemetry.recordUserFlush(err)
 }
 
+function handleRegisterProcessor (userSpanProcessor) {
+  spanProcessor.setUserSpanProcessor(userSpanProcessor)
+}
+
 function handleSpanProcess (data) {
   spanProcessor.process(data)
 }

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

@@ -23,9 +23,16 @@ const LLMObsTagger = require('./tagger')
 const { channel } = require('dc-polyfill')
 const evalMetricAppendCh = channel('llmobs:eval-metric:append')
 const flushCh = channel('llmobs:writers:flush')
+const registerUserSpanProcessorCh = channel('llmobs:register-processor')
 const NoopLLMObs = require('./noop')
 
 class LLMObs extends NoopLLMObs {
+  /**
+   * flag representing if a user span processor has been registered
+   * @type {boolean}
+   */
+  #hasUserSpanProcessor = false
+
   constructor (tracer, llmobsModule, config) {
     super(tracer)
 
@@ -309,6 +316,27 @@ class LLMObs extends NoopLLMObs {
     }
   }
 
+  registerProcessor (processor) {
+    if (!this.enabled) return
+
+    if (this.#hasUserSpanProcessor) {
+      throw new Error(
+        '[LLMObs] Only one user span processor can be registered. ' +
+        'To register a new processor, deregister the existing processor first using `llmobs.deregisterProcessor()`.'
+      )
+    }
+
+    this.#hasUserSpanProcessor = true
+    registerUserSpanProcessorCh.publish(processor)
+  }
+
+  deregisterProcessor () {
+    if (!this.enabled) return
+
+    this.#hasUserSpanProcessor = false
+    registerUserSpanProcessorCh.publish(null)
+  }
+
   submitEvaluation (llmobsSpanContext, options = {}) {
     if (!this.enabled) return