support llm messages with tool calls

Author: gary-huang

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/Agent.java

@@ -103,9 +103,8 @@ private enum AgentFeature {
     SPAN_ORIGIN(TraceInstrumentationConfig.CODE_ORIGIN_FOR_SPANS_ENABLED, false),
     DATA_JOBS(GeneralConfig.DATA_JOBS_ENABLED, false),
     AGENTLESS_LOG_SUBMISSION(GeneralConfig.AGENTLESS_LOG_SUBMISSION_ENABLED, false),
-    LLMOBS(propertyNameToSystemPropertyName(LlmObsConfig.LLMOBS_ENABLED), false),
-    LLMOBS_AGENTLESS(
-        propertyNameToSystemPropertyName(LlmObsConfig.LLMOBS_AGENTLESS_ENABLED), false);
+    LLMOBS(LlmObsConfig.LLMOBS_ENABLED, false),
+    LLMOBS_AGENTLESS(LlmObsConfig.LLMOBS_AGENTLESS_ENABLED, false);
 
     private final String configKey;
     private final String systemProp;

dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/DDLLMObsSpan.java

@@ -1,34 +1,22 @@
 package datadog.trace.llmobs.domain;
 
 import datadog.trace.api.DDSpanTypes;
+import datadog.trace.api.llmobs.LLMObs;
 import datadog.trace.api.llmobs.LLMObsSpan;
 import datadog.trace.api.llmobs.LLMObsTags;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
-import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import javax.annotation.Nonnull;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 public class DDLLMObsSpan implements LLMObsSpan {
-
-  private enum State {
-    VALID,
-    INVALID_IO_MESSAGE_KEY
-  }
-
-  private static final String MESSAGE_KEY_ROLE = "role";
-  private static final String MESSAGE_KEY_CONTENT = "content";
-
-  private static final Set<String> VALID_MESSAGE_KEYS =
-      new HashSet<>(Arrays.asList(MESSAGE_KEY_ROLE, MESSAGE_KEY_CONTENT));
+  private static final String LLM_MESSAGE_UNKNOWN_ROLE = "unknown";
 
   // Well known tags for LLM obs will be prefixed with _ml_obs_(tags|metrics).
   // Prefix for tags
@@ -92,35 +80,15 @@ public String toString() {
         + this.span.getTag(SPAN_KIND);
   }
 
-  private static State validateIOMessages(List<Map<String, Object>> messages) {
-    for (Map<String, Object> message : messages) {
-      for (String key : message.keySet()) {
-        if (!VALID_MESSAGE_KEYS.contains(key)) {
-          return State.INVALID_IO_MESSAGE_KEY;
-        }
-      }
-    }
-    return State.VALID;
-  }
-
   @Override
-  public void annotateIO(
-      List<Map<String, Object>> inputData, List<Map<String, Object>> outputData) {
+  public void annotateIO(List<LLMObs.LLMMessage> inputData, List<LLMObs.LLMMessage> outputData) {
     if (finished) {
       return;
     }
     if (inputData != null && !inputData.isEmpty()) {
-      State inputState = validateIOMessages(inputData);
-      if (validateIOMessages(inputData) != State.VALID) {
-        LOGGER.debug("malformed/unexpected input message, state={}", inputState);
-      }
       this.span.setTag(INPUT, inputData);
     }
     if (outputData != null && !outputData.isEmpty()) {
-      State outputState = validateIOMessages(outputData);
-      if (validateIOMessages(outputData) != State.VALID) {
-        LOGGER.debug("malformed/unexpected output message, state={}", outputState);
-      }
       this.span.setTag(OUTPUT, outputData);
     }
   }
@@ -130,24 +98,31 @@ public void annotateIO(String inputData, String outputData) {
     if (finished) {
       return;
     }
+    boolean wrongSpanKind = false;
     if (inputData != null && !inputData.isEmpty()) {
       if (Tags.LLMOBS_LLM_SPAN_KIND.equals(this.spanKind)) {
+        wrongSpanKind = true;
         annotateIO(
-            Collections.singletonList(Collections.singletonMap(MESSAGE_KEY_CONTENT, inputData)),
+            Collections.singletonList(LLMObs.LLMMessage.from(LLM_MESSAGE_UNKNOWN_ROLE, inputData)),
             null);
       } else {
         this.span.setTag(INPUT, inputData);
       }
     }
     if (outputData != null && !outputData.isEmpty()) {
       if (Tags.LLMOBS_LLM_SPAN_KIND.equals(this.spanKind)) {
+        wrongSpanKind = true;
         annotateIO(
             null,
-            Collections.singletonList(Collections.singletonMap(MESSAGE_KEY_CONTENT, outputData)));
+            Collections.singletonList(LLMObs.LLMMessage.from(LLM_MESSAGE_UNKNOWN_ROLE, outputData)));
       } else {
         this.span.setTag(OUTPUT, outputData);
       }
     }
+    if (wrongSpanKind) {
+      LOGGER.warn(
+          "the span being annotated is an LLM span, it is recommended to use the overload with List<LLMObs.LLMMessage> as arguments");
+    }
   }
 
   @Override

dd-java-agent/agent-llmobs/src/test/groovy/datadog/trace/llmobs/domain/DDLLMObsSpanTest.groovy

@@ -3,6 +3,7 @@ package datadog.trace.llmobs.domain
 import datadog.trace.agent.tooling.TracerInstaller
 import datadog.trace.api.DDTags
 import datadog.trace.api.IdGenerationStrategy
+import datadog.trace.api.llmobs.LLMObs
 import datadog.trace.api.llmobs.LLMObsSpan
 import datadog.trace.api.llmobs.LLMObsTags
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan
@@ -199,11 +200,65 @@ class DDLLMObsSpanTest  extends DDSpecification{
     assert Tags.LLMOBS_LLM_SPAN_KIND.equals(innerSpan.getTag(LLMOBS_TAG_PREFIX + "span.kind"))
 
     assert null == innerSpan.getTag("input")
-    def expectedInput = Arrays.asList(Maps.of("content", input))
-    assert expectedInput.equals(innerSpan.getTag(INPUT))
+    def spanInput = innerSpan.getTag(INPUT)
+    assert spanInput instanceof List
+    assert ((List)spanInput).size() == 1
+    assert spanInput.get(0) instanceof LLMObs.LLMMessage
+    def expectedInputMsg = LLMObs.LLMMessage.from("unknown", input)
+    assert expectedInputMsg.getContent().equals(input)
+    assert expectedInputMsg.getRole().equals("unknown")
+    assert expectedInputMsg.getToolCalls().equals(null)
+
     assert null == innerSpan.getTag("output")
-    def expectedOutput = Arrays.asList(Maps.of("content", output))
-    assert expectedOutput.equals(innerSpan.getTag(OUTPUT))
+    def spanOutput = innerSpan.getTag(OUTPUT)
+    assert spanOutput instanceof List
+    assert ((List)spanOutput).size() == 1
+    assert spanOutput.get(0) instanceof LLMObs.LLMMessage
+    def expectedOutputMsg = LLMObs.LLMMessage.from("unknown", output)
+    assert expectedOutputMsg.getContent().equals(output)
+    assert expectedOutputMsg.getRole().equals("unknown")
+    assert expectedOutputMsg.getToolCalls().equals(null)
+  }
+
+  def "test llm span with messages"() {
+    setup:
+    def test = givenALLMObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "test-span")
+
+    when:
+    def inputMsg =  LLMObs.LLMMessage.from("user", "input")
+    def outputMsg = LLMObs.LLMMessage.from("assistant", "output", Arrays.asList(LLMObs.ToolCall.from("weather-tool", "function", "6176241000", Maps.of("location", "paris"))))
+    // initial set
+    test.annotateIO(Arrays.asList(inputMsg), Arrays.asList(outputMsg))
+
+    then:
+    def innerSpan = (AgentSpan)test.span
+    assert Tags.LLMOBS_LLM_SPAN_KIND.equals(innerSpan.getTag(LLMOBS_TAG_PREFIX + "span.kind"))
+
+    assert null == innerSpan.getTag("input")
+    def spanInput = innerSpan.getTag(INPUT)
+    assert spanInput instanceof List
+    assert ((List)spanInput).size() == 1
+    def spanInputMsg = spanInput.get(0)
+    assert spanInputMsg instanceof LLMObs.LLMMessage
+    assert spanInputMsg.getContent().equals(inputMsg.getContent())
+    assert spanInputMsg.getRole().equals("user")
+    assert spanInputMsg.getToolCalls().equals(null)
+
+    assert null == innerSpan.getTag("output")
+    def spanOutput = innerSpan.getTag(OUTPUT)
+    assert spanOutput instanceof List
+    assert ((List)spanOutput).size() == 1
+    def spanOutputMsg = spanOutput.get(0)
+    assert spanOutputMsg instanceof LLMObs.LLMMessage
+    assert spanOutputMsg.getContent().equals(outputMsg.getContent())
+    assert spanOutputMsg.getRole().equals("assistant")
+    assert spanOutputMsg.getToolCalls().size() == 1
+    def toolCall = spanOutputMsg.getToolCalls().get(0)
+    assert toolCall.getName().equals("weather-tool")
+    assert toolCall.getType().equals("function")
+    assert toolCall.getToolID().equals("6176241000")
+    def expectedToolArgs = Maps.of("location", "paris")
+    assert toolCall.getArguments().equals(expectedToolArgs)
   }
 
   private LLMObsSpan givenALLMObsSpan(String kind, name){