Add in LLMObs Tool Definitions API

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

@@ -15,6 +15,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -37,6 +38,7 @@ public class DDLLMObsSpan implements LLMObsSpan {
   private static final String OUTPUT = LLMOBS_TAG_PREFIX + "output";
   private static final String SPAN_KIND = LLMOBS_TAG_PREFIX + Tags.SPAN_KIND;
   private static final String METADATA = LLMOBS_TAG_PREFIX + LLMObsTags.METADATA;
+  private static final String TOOL_DEFINITIONS = LLMOBS_TAG_PREFIX + LLMObsTags.TOOL_DEFINITIONS;
   private static final String PARENT_ID_TAG_INTERNAL = "parent_id";
 
   private static final String SERVICE = LLMOBS_TAG_PREFIX + "service";
@@ -186,6 +188,29 @@ public void annotateIO(String inputData, String outputData) {
     }
   }
 
+  @Override
+  public void setToolDefinitions(List<LLMObs.ToolDefinition> toolDefinitions) {
+    if (finished || toolDefinitions == null || toolDefinitions.isEmpty()) {
+      return;
+    }
+    List<LLMObs.ToolDefinition> validToolDefinitions = new ArrayList<>(toolDefinitions.size());
+    for (int i = 0; i < toolDefinitions.size(); i++) {
+      LLMObs.ToolDefinition toolDefinition = toolDefinitions.get(i);
+      if (toolDefinition == null) {
+        LOGGER.warn("tool definition at index {} is null; skipping", i);
+        continue;
+      }
+      if (toolDefinition.getName() == null || toolDefinition.getName().isEmpty()) {
+        LOGGER.warn("tool definition at index {} must have a non-empty name; skipping", i);
+        continue;
+      }
+      validToolDefinitions.add(toolDefinition);
+    }
+    if (!validToolDefinitions.isEmpty()) {
+      span.setTag(TOOL_DEFINITIONS, validToolDefinitions);
+    }
+  }
+
   @Override
   public void setMetadata(Map<String, Object> metadata) {
     if (finished) {

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

@@ -61,6 +61,7 @@ class DDLLMObsSpanTest  extends DDSpecification{
   private static final String INPUT = LLMOBS_TAG_PREFIX + "input"
   private static final String OUTPUT = LLMOBS_TAG_PREFIX + "output"
   private static final String METADATA = LLMOBS_TAG_PREFIX + LLMObsTags.METADATA
+  private static final String TOOL_DEFINITIONS = LLMOBS_TAG_PREFIX + LLMObsTags.TOOL_DEFINITIONS
 
 
   def "test span simple"() {
@@ -337,6 +338,88 @@ class DDLLMObsSpanTest  extends DDSpecification{
     DDTraceApiInfo.VERSION == innerSpan.getTag(LLMOBS_TAG_PREFIX + "ddtrace.version")
   }
 
+  def "test llm span with tool definitions"() {
+    setup:
+    def test = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "test-span")
+    def schema = Maps.of(
+      "type", "object",
+      "properties", Maps.of("location", Maps.of("type", "string")))
+    def toolDefinition =
+      LLMObs.ToolDefinition.from("get_weather", "Get the weather by location", schema, "1.2.3")
+
+    when:
+    test.setToolDefinitions(Arrays.asList(
+      toolDefinition,
+      LLMObs.ToolDefinition.from("get_time"),
+      LLMObs.ToolDefinition.from(""),
+      LLMObs.ToolDefinition.from(null),
+      null))
+
+    then:
+    def innerSpan = (AgentSpan)test.span
+    def toolDefinitions = innerSpan.getTag(TOOL_DEFINITIONS)
+    toolDefinitions instanceof List
+    toolDefinitions.size() == 2
+    toolDefinitions.get(0).is(toolDefinition)
+    toolDefinitions.get(1).name == "get_time"
+    toolDefinitions.get(1).description == null
+    toolDefinitions.get(1).schema == null
+    toolDefinitions.get(1).version == null
+  }
+
+  def "tool definitions overwrite prior annotations and ignore invalid definitions"() {
+    setup:
+    def test = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "test-span")
+    test.setToolDefinitions(Arrays.asList(LLMObs.ToolDefinition.from("first")))
+
+    when:
+    test.setToolDefinitions(Arrays.asList(
+      LLMObs.ToolDefinition.from("second")))
+
+    then:
+    def innerSpan = (AgentSpan)test.span
+    innerSpan.getTag(TOOL_DEFINITIONS)*.name == ["second"]
+
+    when:
+    test.setToolDefinitions(Arrays.asList(
+      LLMObs.ToolDefinition.from(""),
+      LLMObs.ToolDefinition.from(null),
+      null))
+
+    then:
+    innerSpan.getTag(TOOL_DEFINITIONS)*.name == ["second"]
+  }
+
+  def "null and empty tool definition lists do not overwrite prior annotations"() {
+    setup:
+    def test = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "test-span")
+    test.setToolDefinitions(Arrays.asList(LLMObs.ToolDefinition.from("first")))
+
+    when:
+    test.setToolDefinitions(toolDefinitions)
+
+    then:
+    def innerSpan = (AgentSpan)test.span
+    innerSpan.getTag(TOOL_DEFINITIONS)*.name == ["first"]
+
+    where:
+    toolDefinitions << [null, Collections.emptyList()]
+  }
+
+  def "tool definitions cannot be annotated after the span finishes"() {
+    setup:
+    def test = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "test-span")
+    test.setToolDefinitions(Arrays.asList(LLMObs.ToolDefinition.from("first")))
+    test.finish()
+
+    when:
+    test.setToolDefinitions(Arrays.asList(LLMObs.ToolDefinition.from("second")))
+
+    then:
+    def innerSpan = (AgentSpan)test.span
+    innerSpan.getTag(TOOL_DEFINITIONS)*.name == ["first"]
+  }
+
   def "finish records span.finished telemetry when LLMObs enabled"() {
     setup:
     LLMObsMetricCollector collector = LLMObsMetricCollector.get()

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

@@ -173,6 +173,54 @@ public Map<String, Object> getArguments() {
     }
   }
 
+  public static class ToolDefinition {
+    private String name;
+    private String description;
+    private Map<String, Object> schema;
+    private String version;
+
+    public static ToolDefinition from(String name) {
+      return new ToolDefinition(name, null, null, null);
+    }
+
+    public static ToolDefinition from(String name, String description) {
+      return new ToolDefinition(name, description, null, null);
+    }
+
+    public static ToolDefinition from(String name, String description, Map<String, Object> schema) {
+      return new ToolDefinition(name, description, schema, null);
+    }
+
+    public static ToolDefinition from(
+        String name, String description, Map<String, Object> schema, String version) {
+      return new ToolDefinition(name, description, schema, version);
+    }
+
+    private ToolDefinition(
+        String name, String description, Map<String, Object> schema, String version) {
+      this.name = name;
+      this.description = description;
+      this.schema = schema;
+      this.version = version;
+    }
+
+    public String getName() {
+      return name;
+    }
+
+    public String getDescription() {
+      return description;
+    }
+
+    public Map<String, Object> getSchema() {
+      return schema;
+    }
+
+    public String getVersion() {
+      return version;
+    }
+  }
+
   public static class ToolResult {
     private String name;
     private String type;

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

@@ -23,6 +23,13 @@ public interface LLMObsSpan {
    */
   void annotateIO(String inputData, String outputData);
 
+  /**
+   * Annotate the span with the definitions of tools available to the LLM.
+   *
+   * @param toolDefinitions The tool definitions supplied to the LLM
+   */
+  void setToolDefinitions(List<LLMObs.ToolDefinition> toolDefinitions);
+
   /**
    * Annotate the span with metadata
    *

dd-trace-api/src/main/java/datadog/trace/api/llmobs/noop/NoOpLLMObsSpan.java

@@ -15,6 +15,9 @@ public void annotateIO(List<LLMObs.LLMMessage> inputData, List<LLMObs.LLMMessage
   @Override
   public void annotateIO(String inputData, String outputData) {}
 
+  @Override
+  public void setToolDefinitions(List<LLMObs.ToolDefinition> toolDefinitions) {}
+
   @Override
   public void setMetadata(Map<String, Object> metadata) {}
 

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

@@ -71,6 +71,58 @@ void testToolCallWithNullArguments() {
     assertNull(toolCall.getArguments());
   }
 
+  @Test
+  void testToolDefinitionFromName() {
+    LLMObs.ToolDefinition toolDefinition = LLMObs.ToolDefinition.from("get_weather");
+
+    assertEquals("get_weather", toolDefinition.getName());
+    assertNull(toolDefinition.getDescription());
+    assertNull(toolDefinition.getSchema());
+    assertNull(toolDefinition.getVersion());
+  }
+
+  @Test
+  void testToolDefinitionFromNameAndDescription() {
+    LLMObs.ToolDefinition toolDefinition =
+        LLMObs.ToolDefinition.from("get_weather", "Get the weather by location");
+
+    assertEquals("get_weather", toolDefinition.getName());
+    assertEquals("Get the weather by location", toolDefinition.getDescription());
+    assertNull(toolDefinition.getSchema());
+    assertNull(toolDefinition.getVersion());
+  }
+
+  @Test
+  void testToolDefinitionFromNameDescriptionAndSchema() {
+    Map<String, Object> schema = new HashMap<>();
+    schema.put("type", "object");
+    schema.put(
+        "properties",
+        Collections.singletonMap("location", Collections.singletonMap("type", "string")));
+
+    LLMObs.ToolDefinition toolDefinition =
+        LLMObs.ToolDefinition.from("get_weather", "Get the weather by location", schema);
+
+    assertEquals("get_weather", toolDefinition.getName());
+    assertEquals("Get the weather by location", toolDefinition.getDescription());
+    assertEquals(schema, toolDefinition.getSchema());
+    assertNull(toolDefinition.getVersion());
+  }
+
+  @Test
+  void testToolDefinitionFromAllFields() {
+    Map<String, Object> schema = new HashMap<>();
+    schema.put("type", "object");
+
+    LLMObs.ToolDefinition toolDefinition =
+        LLMObs.ToolDefinition.from("get_weather", "Get the weather by location", schema, "1.2.3");
+
+    assertEquals("get_weather", toolDefinition.getName());
+    assertEquals("Get the weather by location", toolDefinition.getDescription());
+    assertEquals(schema, toolDefinition.getSchema());
+    assertEquals("1.2.3", toolDefinition.getVersion());
+  }
+
   @Test
   void testLLMMessageCreationWithToolCalls() {
     Map<String, Object> args = new HashMap<>();

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

@@ -394,6 +394,9 @@ public void accept(Metadata metadata) {
             writable.writeString("value", null);
             writable.writeObject(val, null);
           }
+        } else if (key.equals(LLMObsTags.TOOL_DEFINITIONS) && val instanceof List) {
+          writable.writeString(key, null);
+          writeToolDefinitions((List<?>) val);
         } else if (key.equals(LLMObsTags.METADATA) && val instanceof Map) {
           Map<String, Object> metadataMap = (Map) val;
           writable.writeUTF8(METADATA);
@@ -409,6 +412,38 @@ public void accept(Metadata metadata) {
       }
     }
 
+    private void writeToolDefinitions(List<?> toolDefinitions) {
+      writable.startArray(toolDefinitions.size());
+      for (Object toolDefinitionObject : toolDefinitions) {
+        if (!(toolDefinitionObject instanceof LLMObs.ToolDefinition)) {
+          writable.writeObject(toolDefinitionObject, null);
+          continue;
+        }
+
+        LLMObs.ToolDefinition toolDefinition = (LLMObs.ToolDefinition) toolDefinitionObject;
+        int mapSize = 1;
+        if (toolDefinition.getDescription() != null) mapSize++;
+        if (toolDefinition.getSchema() != null) mapSize++;
+        if (toolDefinition.getVersion() != null) mapSize++;
+
+        writable.startMap(mapSize);
+        writable.writeString("name", null);
+        writable.writeString(toolDefinition.getName(), null);
+        if (toolDefinition.getDescription() != null) {
+          writable.writeString("description", null);
+          writable.writeString(toolDefinition.getDescription(), null);
+        }
+        if (toolDefinition.getSchema() != null) {
+          writable.writeString("schema", null);
+          writable.writeObject(toolDefinition.getSchema(), null);
+        }
+        if (toolDefinition.getVersion() != null) {
+          writable.writeString("version", null);
+          writable.writeString(toolDefinition.getVersion(), null);
+        }
+      }
+    }
+
     private void writeLlmInputMap(Map<?, ?> inputMap) {
       writable.startMap(inputMap.size());
       for (Map.Entry<?, ?> entry : inputMap.entrySet()) {

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

@@ -100,11 +100,11 @@ void testLLMObsSpanMapperSerialization() throws Exception {
     Map<String, Object> schema = new LinkedHashMap<>();
     schema.put("type", "object");
     schema.put("properties", properties);
-    Map<String, Object> toolDef = new LinkedHashMap<>();
-    toolDef.put("name", "get_weather");
-    toolDef.put("description", "Get weather by city");
-    toolDef.put("schema", schema);
-    llmSpan.setTag("_ml_obs_tag.tool_definitions", Collections.singletonList(toolDef));
+    LLMObs.ToolDefinition toolDefinition =
+        LLMObs.ToolDefinition.from("get_weather", "Get weather by city", schema, "1.2.3");
+    llmSpan.setTag(
+        "_ml_obs_tag.tool_definitions",
+        Arrays.asList(toolDefinition, LLMObs.ToolDefinition.from("get_time")));
 
     llmSpan.setError(true);
     llmSpan.setTag(DDTags.ERROR_MSG, "boom");
@@ -204,9 +204,12 @@ void testLLMObsSpanMapperSerialization() throws Exception {
     assertEquals("assistant", outputMsgs.get(0).get("role"));
     List<Map<String, Object>> toolDefsResult =
         (List<Map<String, Object>>) meta.get("tool_definitions");
+    assertEquals(2, toolDefsResult.size());
     assertEquals("get_weather", toolDefsResult.get(0).get("name"));
     assertEquals("Get weather by city", toolDefsResult.get(0).get("description"));
     assertEquals(schema, toolDefsResult.get(0).get("schema"));
+    assertEquals("1.2.3", toolDefsResult.get(0).get("version"));
+    assertEquals(Collections.singletonMap("name", "get_time"), toolDefsResult.get(1));
     assertTrue(meta.containsKey("metadata"));
 
     assertTrue(spanData.containsKey("metrics"));