add llm message class to support llm spans

DataDog · Mar 7, 2025 · c1f63dc · c1f63dc
1 parent cb95fbb
commit c1f63dc
Show file tree

Hide file tree

Showing 4 changed files with 84 additions and 8 deletions.
diff --git a/dd-trace-api/build.gradle b/dd-trace-api/build.gradle
@@ -32,6 +32,8 @@ excludedClassesCoverage += [
   'datadog.trace.api.profiling.ProfilingContext',
   'datadog.trace.api.profiling.ProfilingContextAttribute.NoOp',
   'datadog.trace.api.llmobs.LLMObs',
+  'datadog.trace.api.llmobs.LLMObs.LLMMessage',
+  'datadog.trace.api.llmobs.LLMObs.ToolCall',
   'datadog.trace.api.llmobs.LLMObsSpan',
   'datadog.trace.api.llmobs.noop.NoOpLLMObsSpan',
   'datadog.trace.api.llmobs.noop.NoOpLLMObsSpanFactory',

diff --git a/dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObs.java b/dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObs.java
@@ -1,9 +1,13 @@
 package datadog.trace.api.llmobs;
 
 import datadog.trace.api.llmobs.noop.NoOpLLMObsSpanFactory;
+import java.util.List;
+import java.util.Map;
 import javax.annotation.Nullable;
 
 public class LLMObs {
+  private LLMObs() {}
+
   private static LLMObsSpanFactory SPAN_FACTORY = NoOpLLMObsSpanFactory.INSTANCE;
 
   public static LLMObsSpan startLLMSpan(
@@ -57,4 +61,76 @@ LLMObsSpan startLLMSpan(
     LLMObsSpan startWorkflowSpan(
         String spanName, @Nullable String mlApp, @Nullable String sessionID);
   }
+
+  public static class ToolCall {
+    private String name;
+    private String type;
+    private String toolID;
+    private Map<String, Object> arguments;
+
+    public static ToolCall from(
+        String name, String type, String toolID, Map<String, Object> arguments) {
+      return new ToolCall(name, type, toolID, arguments);
+    }
+
+    private ToolCall(String name, String type, String toolID, Map<String, Object> arguments) {
+      this.name = name;
+      this.type = type;
+      this.toolID = toolID;
+      this.arguments = arguments;
+    }
+
+    public String getName() {
+      return name;
+    }
+
+    public String getType() {
+      return type;
+    }
+
+    public String getToolID() {
+      return toolID;
+    }
+
+    public Map<String, Object> getArguments() {
+      return arguments;
+    }
+  }
+
+  public static class LLMMessage {
+    private String role;
+    private String content;
+    private List<ToolCall> toolCalls;
+
+    public static LLMMessage from(String role, String content, List<ToolCall> toolCalls) {
+      return new LLMMessage(role, content, toolCalls);
+    }
+
+    public static LLMMessage from(String role, String content) {
+      return new LLMMessage(role, content);
+    }
+
+    private LLMMessage(String role, String content, List<ToolCall> toolCalls) {
+      this.role = role;
+      this.content = content;
+      this.toolCalls = toolCalls;
+    }
+
+    private LLMMessage(String role, String content) {
+      this.role = role;
+      this.content = content;
+    }
+
+    public String getRole() {
+      return role;
+    }
+
+    public String getContent() {
+      return content;
+    }
+
+    public List<ToolCall> getToolCalls() {
+      return toolCalls;
+    }
+  }
 }
diff --git a/dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObsSpan.java b/dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObsSpan.java
@@ -7,14 +7,12 @@
 public interface LLMObsSpan {
 
   /**
-   * Annotate the span with inputs and outputs
+   * Annotate the span with inputs and outputs for LLM spans
    *
-   * @param inputData The input data of the span in the form of a list, for example a list of input
-   *     messages
-   * @param outputData The output data of the span in the form of a list, for example a list of
-   *     output messages
+   * @param inputMessages The input messages of the span in the form of a list
+   * @param outputMessages The output messages of the span in the form of a list
    */
-  void annotateIO(List<Map<String, Object>> inputData, List<Map<String, Object>> outputData);
+  void annotateIO(List<LLMObs.LLMMessage> inputMessages, List<LLMObs.LLMMessage> outputMessages);
 
   /**
    * Annotate the span with inputs and outputs

diff --git a/dd-trace-api/src/main/java/datadog/trace/api/llmobs/noop/NoOpLLMObsSpan.java b/dd-trace-api/src/main/java/datadog/trace/api/llmobs/noop/NoOpLLMObsSpan.java
@@ -1,5 +1,6 @@
 package datadog.trace.api.llmobs.noop;
 
+import datadog.trace.api.llmobs.LLMObs;
 import datadog.trace.api.llmobs.LLMObsSpan;
 import java.util.List;
 import java.util.Map;
@@ -8,8 +9,7 @@ public class NoOpLLMObsSpan implements LLMObsSpan {
   public static final LLMObsSpan INSTANCE = new NoOpLLMObsSpan();
 
   @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) {}
 
   @Override
   public void annotateIO(String inputData, String outputData) {}