Add in LLM Observability Span Processor API

dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java

@@ -4,12 +4,12 @@
 import datadog.trace.api.Config;
 import datadog.trace.api.WellKnownTags;
 import datadog.trace.api.llmobs.LLMObs;
+import datadog.trace.api.llmobs.LLMObsInternal;
 import datadog.trace.api.llmobs.LLMObsSpan;
 import datadog.trace.api.llmobs.LLMObsTags;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
 import datadog.trace.llmobs.domain.DDLLMObsSpan;
 import datadog.trace.llmobs.domain.LLMObsEval;
-import datadog.trace.llmobs.domain.LLMObsInternal;
 import java.lang.instrument.Instrumentation;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
@@ -34,9 +34,9 @@ public static void start(Instrumentation inst, SharedCommunicationObjects sco) {
 
     String mlApp = config.getLlmObsMlApp();
     WellKnownTags wellKnownTags = config.getWellKnownTags();
-    LLMObsInternal.setLLMObsSpanFactory(new LLMObsManualSpanFactory(mlApp, wellKnownTags));
+    LLMObsInternal.setSpanFactory(new LLMObsManualSpanFactory(mlApp, wellKnownTags));
 
-    LLMObsInternal.setLLMObsEvalProcessor(new LLMObsCustomEvalProcessor(mlApp, sco, config));
+    LLMObsInternal.setEvalProcessor(new LLMObsCustomEvalProcessor(mlApp, sco, config));
   }
 
   private static class LLMObsCustomEvalProcessor implements LLMObs.LLMObsEvalProcessor {

dd-trace-api/build.gradle.kts

@@ -56,6 +56,8 @@ extra["excludedClassesCoverage"] = listOf(
   "datadog.trace.api.llmobs.LLMObs.ToolCall",
   "datadog.trace.api.llmobs.LLMObs.ToolResult",
   "datadog.trace.api.llmobs.LLMObsSpan",
+  "datadog.trace.api.llmobs.LLMObsSpanData",
+  "datadog.trace.api.llmobs.LLMObsSpanProcessor",
   "datadog.trace.api.llmobs.noop.NoOpLLMObsSpan",
   "datadog.trace.api.llmobs.noop.NoOpLLMObsSpanFactory",
   "datadog.trace.api.llmobs.noop.NoOpLLMObsEvalProcessor",

dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObs.java

@@ -4,13 +4,15 @@
 import datadog.trace.api.llmobs.noop.NoOpLLMObsSpanFactory;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import javax.annotation.Nullable;
 
 public class LLMObs {
   protected LLMObs() {}
 
   protected static LLMObsSpanFactory SPAN_FACTORY = NoOpLLMObsSpanFactory.INSTANCE;
   protected static LLMObsEvalProcessor EVAL_PROCESSOR = NoOpLLMObsEvalProcessor.INSTANCE;
+  @Nullable protected static volatile LLMObsSpanProcessor SPAN_PROCESSOR;
 
   public static LLMObsSpan startLLMSpan(
       String spanName,
@@ -60,6 +62,31 @@ public static LLMObsSpan startRetrievalSpan(
     return SPAN_FACTORY.startRetrievalSpan(spanName, mlApp, sessionId);
   }
 
+  /**
+   * Registers a processor to be called for each LLM Observability span before it is sent.
+   *
+   * <p>The processor can modify the span input and output, or return {@code null} to omit the span
+   * from LLM Observability. Only one processor can be registered at a time.
+   *
+   * @param processor the processor to register
+   * @throws NullPointerException if {@code processor} is {@code null}
+   * @throws IllegalStateException if a processor is already registered
+   */
+  public static synchronized void registerProcessor(LLMObsSpanProcessor processor) {
+    Objects.requireNonNull(processor, "processor");
+    if (SPAN_PROCESSOR != null) {
+      throw new IllegalStateException(
+          "An LLM Observability span processor is already registered. "
+              + "Deregister it before registering another.");
+    }
+    SPAN_PROCESSOR = processor;
+  }
+
+  /** Deregisters the current LLM Observability span processor, if one is registered. */
+  public static synchronized void deregisterProcessor() {
+    SPAN_PROCESSOR = null;
+  }
+
   public static void SubmitEvaluation(
       LLMObsSpan llmObsSpan, String label, String categoricalValue, Map<String, Object> tags) {
     EVAL_PROCESSOR.SubmitEvaluation(llmObsSpan, label, categoricalValue, tags);

dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObsSpanData.java

@@ -0,0 +1,59 @@
+package datadog.trace.api.llmobs;
+
+import java.util.List;
+import javax.annotation.Nullable;
+
+/**
+ * Mutable view of an LLM Observability span passed to a registered {@link LLMObsSpanProcessor}.
+ *
+ * <p>Changes to the input and output are applied immediately before the span is sent to LLM
+ * Observability.
+ */
+public interface LLMObsSpanData {
+
+  /**
+   * Gets the LLM Observability span kind.
+   *
+   * @return the span kind
+   */
+  String getKind();
+
+  /**
+   * Gets the input content associated with the span.
+   *
+   * @return the input represented as messages
+   */
+  List<LLMObs.LLMMessage> getInput();
+
+  /**
+   * Replaces the input content associated with the span.
+   *
+   * @param input the new input represented as messages
+   * @throws NullPointerException if {@code input} is {@code null}
+   */
+  void setInput(List<LLMObs.LLMMessage> input);
+
+  /**
+   * Gets the output content associated with the span.
+   *
+   * @return the output represented as messages
+   */
+  List<LLMObs.LLMMessage> getOutput();
+
+  /**
+   * Replaces the output content associated with the span.
+   *
+   * @param output the new output represented as messages
+   * @throws NullPointerException if {@code output} is {@code null}
+   */
+  void setOutput(List<LLMObs.LLMMessage> output);
+
+  /**
+   * Gets an LLM Observability tag from the span.
+   *
+   * @param key the unprefixed tag name
+   * @return the tag value, or {@code null} when the tag is not present
+   */
+  @Nullable
+  String getTag(String key);
+}

dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObsSpanProcessor.java

@@ -0,0 +1,20 @@
+package datadog.trace.api.llmobs;
+
+import javax.annotation.Nullable;
+
+/** Processes LLM Observability spans before they are sent. */
+@FunctionalInterface
+public interface LLMObsSpanProcessor {
+
+  /**
+   * Processes an LLM Observability span.
+   *
+   * <p>The processor may mutate and return {@code span}, or return {@code null} to omit the span
+   * from LLM Observability.
+   *
+   * @param span the span being processed
+   * @return the span to send, or {@code null} to omit it
+   */
+  @Nullable
+  LLMObsSpanData process(LLMObsSpanData span);
+}

dd-trace-api/src/test/java/datadog/trace/api/llmobs/LLMObsTest.java

@@ -5,6 +5,7 @@
 import static org.junit.jupiter.api.Assertions.assertNotSame;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
@@ -28,23 +29,48 @@ class LLMObsTest {
 
   private static Object originalSpanFactory;
   private static Object originalEvalProcessor;
+  private static Object originalSpanProcessor;
 
   @BeforeAll
   static void setupSpec() throws Exception {
     originalSpanFactory = getStaticField("SPAN_FACTORY");
     originalEvalProcessor = getStaticField("EVAL_PROCESSOR");
+    originalSpanProcessor = getStaticField("SPAN_PROCESSOR");
   }
 
   @AfterAll
   static void cleanupSpec() throws Exception {
     setStaticField("SPAN_FACTORY", originalSpanFactory);
     setStaticField("EVAL_PROCESSOR", originalEvalProcessor);
+    setStaticField("SPAN_PROCESSOR", originalSpanProcessor);
   }
 
   @AfterEach
   void cleanup() throws Exception {
     setStaticField("SPAN_FACTORY", NoOpLLMObsSpanFactory.INSTANCE);
     setStaticField("EVAL_PROCESSOR", NoOpLLMObsEvalProcessor.INSTANCE);
+    LLMObs.deregisterProcessor();
+  }
+
+  @Test
+  void testRegisterAndDeregisterProcessor() throws Exception {
+    LLMObsSpanData span = mock(LLMObsSpanData.class);
+    LLMObsSpanProcessor processor = registeredSpan -> registeredSpan;
+
+    LLMObs.registerProcessor(processor);
+
+    assertSame(processor, getStaticField("SPAN_PROCESSOR"));
+    assertSame(span, processor.process(span));
+    assertThrows(IllegalStateException.class, () -> LLMObs.registerProcessor(processor));
+
+    LLMObs.deregisterProcessor();
+
+    assertNull(getStaticField("SPAN_PROCESSOR"));
+  }
+
+  @Test
+  void testRegisterNullProcessor() {
+    assertThrows(NullPointerException.class, () -> LLMObs.registerProcessor(null));
   }
 
   @Test