wip tests + cleanup

Author: Nayeem Kamal

communication/src/main/java/datadog/communication/BackendApiFactory.java

@@ -72,7 +72,7 @@ private HttpUrl getAgentlessUrl(Intake intake) {
 
   public enum Intake {
     API("api", "v2", Config::isCiVisibilityAgentlessEnabled, Config::getCiVisibilityAgentlessUrl),
-    LLMOBS_API("api", "v2", Config::isLlmObsAgentlessEnabled, Config::getLlMObsAgentlessUrl),
+    LLMOBS_API("api", "v2", Config::isLlmObsAgentlessEnabled, Config::getLlmObsAgentlessUrl),
     LOGS(
         "http-intake.logs",
         "v2",

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

@@ -29,17 +29,15 @@ public static void start(Instrumentation inst, SharedCommunicationObjects sco) {
     sco.createRemaining(config);
 
     LLMObsInternal.setLLMObsSpanFactory(
-        new LLMObsManualSpanFactory(
-            config.getLlmObsMlApp(), config.getServiceName()));
+        new LLMObsManualSpanFactory(config.getLlmObsMlApp(), config.getServiceName()));
   }
 
   private static class LLMObsManualSpanFactory implements LLMObs.LLMObsSpanFactory {
 
     private final String serviceName;
     private final String defaultMLApp;
 
-    public LLMObsManualSpanFactory(
-        String defaultMLApp, String serviceName) {
+    public LLMObsManualSpanFactory(String defaultMLApp, String serviceName) {
       this.defaultMLApp = defaultMLApp;
       this.serviceName = serviceName;
     }

dd-trace-core/build.gradle

@@ -68,7 +68,6 @@ dependencies {
   implementation project(':components:json')
   implementation project(':utils:container-utils')
   implementation project(':utils:socket-utils')
-
   // for span exception debugging
   compileOnly project(':dd-java-agent:agent-debugger:debugger-bootstrap')
 

dd-trace-core/src/main/java/datadog/trace/common/writer/DDAgentWriter.java

@@ -151,7 +151,6 @@ public DDAgentWriter build() {
       }
 
       final DDAgentMapperDiscovery mapperDiscovery = new DDAgentMapperDiscovery(featureDiscovery);
-
       final PayloadDispatcher dispatcher =
           new PayloadDispatcherImpl(mapperDiscovery, agentApi, healthMetrics, monitoring);
       final TraceProcessingWorker traceProcessingWorker =

dd-trace-core/src/main/java/datadog/trace/common/writer/WriterFactory.java

@@ -82,7 +82,7 @@ public static Writer createWriter(
 
     // The AgentWriter doesn't support the CI Visibility protocol. If CI Visibility is
     // enabled, check if we can use the IntakeWriter instead.
-    if (DD_AGENT_WRITER_TYPE.equals(configuredType) && (config.isCiVisibilityEnabled())) {
+    if (DD_AGENT_WRITER_TYPE.equals(configuredType) && config.isCiVisibilityEnabled()) {
       if (featuresDiscovery.supportsEvpProxy() || config.isCiVisibilityAgentlessEnabled()) {
         configuredType = DD_INTAKE_WRITER_TYPE;
       } else {
@@ -192,7 +192,7 @@ private static RemoteApi createDDIntakeRemoteApi(
           .build();
     } else {
       HttpUrl hostUrl = null;
-      String llmObsAgentlessUrl = config.getLlMObsAgentlessUrl();
+      String llmObsAgentlessUrl = config.getLlmObsAgentlessUrl();
 
       if (config.getCiVisibilityAgentlessUrl() != null) {
         hostUrl = HttpUrl.get(config.getCiVisibilityAgentlessUrl());

dd-trace-core/src/main/java/datadog/trace/common/writer/ddagent/DDAgentApi.java

@@ -90,7 +90,6 @@ public void addResponseListener(final RemoteResponseListener listener) {
   public Response sendSerializedTraces(final Payload payload) {
     final int sizeInBytes = payload.sizeInBytes();
     String tracesEndpoint = featuresDiscovery.getTraceEndpoint();
-
     if (null == tracesEndpoint) {
       featuresDiscovery.discoverIfOutdated();
       tracesEndpoint = featuresDiscovery.getTraceEndpoint();

dd-trace-core/src/main/java/datadog/trace/common/writer/ddintake/DDIntakeApi.java

@@ -131,7 +131,6 @@ public Response sendSerializedTraces(Payload payload) {
             .post(payload.toRequest())
             .tag(OkHttpUtils.CustomListener.class, telemetryListener)
             .build();
-
     totalTraces += payload.traceCount();
     receivedTraces += payload.traceCount();
 

dd-trace-core/src/main/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapper.java

@@ -243,7 +243,7 @@ public void accept(Metadata metadata) {
         String tagKey = tag.getKey();
         if (tagKey.startsWith(LLMOBS_METRIC_PREFIX) && tag.getValue() instanceof Number) {
           writable.writeString(tagKey.substring(LLMOBS_METRIC_PREFIX.length()), null);
-          writable.writeDouble((double) tag.getValue());
+          writable.writeObjectString(tag.getValue(), null);
         }
       }
 

dd-trace-core/src/test/groovy/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapperTest.groovy

@@ -0,0 +1,226 @@
+package datadog.trace.llmobs.writer.ddintake
+
+import datadog.trace.common.writer.ListWriter
+import datadog.communication.serialization.ByteBufferConsumer
+import datadog.communication.serialization.FlushingBuffer
+import datadog.communication.serialization.msgpack.MsgPackWriter
+import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes
+import datadog.trace.bootstrap.instrumentation.api.Tags
+import datadog.trace.api.llmobs.LLMObsTags
+import datadog.trace.api.DDTags
+import datadog.trace.core.DDSpan
+import datadog.trace.core.test.DDCoreSpecification
+import org.msgpack.core.MessagePack
+import org.msgpack.core.MessageUnpacker
+
+import java.nio.ByteBuffer
+
+class LLMObsSpanMapperTest extends DDCoreSpecification {
+
+  def "test LLMObs span mapper"() {
+    setup:
+    def tracer = tracerBuilder().writer(new ListWriter()).build()
+    DDSpan span = (DDSpan) tracer.buildSpan("llm-operation")
+      .withServiceName("my-llm-service")
+      .withSpanType(InternalSpanTypes.LLMOBS)
+      .start()
+
+    // Add LLM-specific tags with proper prefixes
+    span.setTag("_ml_obs_tag.span.kind", Tags.LLMOBS_WORKFLOW_SPAN_KIND)
+    span.setTag("_ml_obs_tag." + LLMObsTags.MODEL_NAME, "gpt-4")
+    span.setTag("_ml_obs_tag." + LLMObsTags.MODEL_PROVIDER, "openai")
+    span.setTag("_ml_obs_tag.input", "What is the weather?")
+    span.setTag("_ml_obs_tag.output", "It's sunny today.")
+    span.setTag("_ml_obs_tag.custom_tag", "test-value")
+    span.setTag("_ml_obs_metric.input_tokens", 10.0)
+    span.setTag("_ml_obs_metric.output_tokens", 5.0)
+    span.setTag("_ml_obs_metric.total_cost", 0.005)
+
+    // Add some metadata
+    Map<String, Object> metadata = [
+      "temperature": 0.7,
+      "max_tokens": 100
+    ]
+    span.setTag("_ml_obs_tag." + LLMObsTags.METADATA, metadata)
+
+    def trace = [span]
+
+    when:
+    LLMObsSpanMapper spanMapper = new LLMObsSpanMapper()
+    CapturingByteBufferConsumer sink = new CapturingByteBufferConsumer()
+    MsgPackWriter packer = new MsgPackWriter(new FlushingBuffer(1024, sink))
+    packer.format(trace, spanMapper)
+    packer.flush()
+
+    then:
+    sink.captured != null
+
+    MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(sink.captured)
+
+    int topLevelMapSize = unpacker.unpackMapHeader()
+    topLevelMapSize == 3
+
+    Map<String, Object> topLevel = [:]
+    for (int i = 0; i < topLevelMapSize; i++) {
+      String key = unpacker.unpackString()
+      if (key == "event_type") {
+        topLevel[key] = unpacker.unpackString()
+      } else if (key == "_dd.stage") {
+        topLevel[key] = unpacker.unpackString()
+      } else if (key == "spans") {
+        int spansArraySize = unpacker.unpackArrayHeader()
+        topLevel[key] = spansArraySize
+        unpacker.skipValue() // TODO: add check for span data
+      }
+    }
+
+    topLevel["event_type"] == "span"
+    topLevel["_dd.stage"] == "raw"
+    topLevel["spans"] == 1 // Should have 1 span
+
+    cleanup:
+    tracer.close()
+  }
+
+  def "test non-LLMObs span is filtered out"() {
+    setup:
+    def tracer = tracerBuilder().writer(new ListWriter()).build()
+    DDSpan regularSpan = (DDSpan) tracer.buildSpan("regular-operation")
+      .withServiceName("my-service")
+      .start()
+
+    def trace = [regularSpan]
+
+    when:
+    LLMObsSpanMapper spanMapper = new LLMObsSpanMapper()
+    CapturingByteBufferConsumer sink = new CapturingByteBufferConsumer()
+    MsgPackWriter packer = new MsgPackWriter(new FlushingBuffer(1024, sink))
+    packer.format(trace, spanMapper)
+    packer.flush()
+
+    then:
+    sink.captured != null
+
+    MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(sink.captured)
+
+    int topLevelMapSize = unpacker.unpackMapHeader()
+    topLevelMapSize == 3
+
+    Map<String, Object> topLevel = [:]
+    for (int i = 0; i < topLevelMapSize; i++) {
+      String key = unpacker.unpackString()
+      if (key == "event_type") {
+        topLevel[key] = unpacker.unpackString()
+      } else if (key == "_dd.stage") {
+        topLevel[key] = unpacker.unpackString()
+      } else if (key == "spans") {
+        int spansArraySize = unpacker.unpackArrayHeader()
+        topLevel[key] = spansArraySize
+        // Since array is empty, no need to skip anything
+      }
+    }
+
+    // Verify that no spans are included since regular span is filtered out
+    topLevel["spans"] == 0
+
+    cleanup:
+    tracer.close()
+  }
+
+  def "test LLM span with error"() {
+    setup:
+    def tracer = tracerBuilder().writer(new ListWriter()).build()
+    DDSpan span = (DDSpan) tracer.buildSpan("llm-operation")
+      .withServiceName("my-llm-service")
+      .withSpanType(InternalSpanTypes.LLMOBS)
+      .start()
+
+    span.setTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND)
+    span.setError(true)
+    span.setTag(DDTags.ERROR_MSG, "API rate limit exceeded")
+    span.setTag(DDTags.ERROR_TYPE, "RateLimitError")
+    span.setTag(DDTags.ERROR_STACK, "java.lang.RuntimeException: API rate limit exceeded")
+
+    def trace = [span]
+
+    when:
+    LLMObsSpanMapper spanMapper = new LLMObsSpanMapper()
+    CapturingByteBufferConsumer sink = new CapturingByteBufferConsumer()
+    MsgPackWriter packer = new MsgPackWriter(new FlushingBuffer(1024, sink))
+    packer.format(trace, spanMapper)
+    packer.flush()
+
+    then:
+    sink.captured != null
+
+    MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(sink.captured)
+
+    int topLevelMapSize = unpacker.unpackMapHeader()
+    topLevelMapSize == 3
+
+    Map<String, Object> topLevel = [:]
+    for (int i = 0; i < topLevelMapSize; i++) {
+      String key = unpacker.unpackString()
+      if (key == "event_type") {
+        topLevel[key] = unpacker.unpackString()
+      } else if (key == "_dd.stage") {
+        topLevel[key] = unpacker.unpackString()
+      } else if (key == "spans") {
+        int spansArraySize = unpacker.unpackArrayHeader()
+        topLevel[key] = spansArraySize
+
+        // Parse the spans array to check error information
+        for (int spanIndex = 0; spanIndex < spansArraySize; spanIndex++) {
+          int spanMapSize = unpacker.unpackMapHeader()
+          spanMapSize == 11
+
+          Map<String, Object> spanData = [:]
+          for (int fieldIndex = 0; fieldIndex < spanMapSize; fieldIndex++) {
+            String fieldKey = unpacker.unpackString()
+            if (fieldKey == "error") {
+              spanData[fieldKey] = unpacker.unpackInt()
+            } else if (fieldKey == "status") {
+              spanData[fieldKey] = unpacker.unpackString()
+            } else if (fieldKey == "meta") {
+              int metaMapSize = unpacker.unpackMapHeader()
+              Map<String, String> metaMap = [:]
+              for (int metaIndex = 0; metaIndex < metaMapSize; metaIndex++) {
+                String metaKey = unpacker.unpackString()
+                String metaValue = unpacker.unpackString()
+                metaMap[metaKey] = metaValue
+              }
+              spanData[fieldKey] = metaMap
+            } else {
+              // Skip other fields
+              unpacker.skipValue()
+            }
+          }
+
+          // Verify error information
+          spanData["error"] == 1
+          spanData["status"] == "error"
+
+          Map<String, String> meta = (Map<String, String>) spanData["meta"]
+          meta[DDTags.ERROR_MSG] == "API rate limit exceeded"
+          meta[DDTags.ERROR_TYPE] == "RateLimitError"
+          meta[DDTags.ERROR_STACK] == "java.lang.RuntimeException: API rate limit exceeded"
+        }
+      }
+    }
+
+    topLevel["spans"] == 1
+
+    cleanup:
+    tracer.close()
+  }
+
+  static class CapturingByteBufferConsumer implements ByteBufferConsumer {
+
+    ByteBuffer captured
+
+    @Override
+    void accept(int messageCount, ByteBuffer buffer) {
+      captured = buffer
+    }
+  }
+}

internal-api/src/main/java/datadog/trace/api/Config.java

@@ -2898,7 +2898,7 @@ public boolean isLlmObsAgentlessEnabled() {
     return llmObsAgentlessEnabled;
   }
 
-  public String getLlMObsAgentlessUrl() {
+  public String getLlmObsAgentlessUrl() {
     return llmObsAgentlessUrl;
   }