From 6e7642c58826c52362052e1ffead5e5a6ffdfd12 Mon Sep 17 00:00:00 2001 From: Sam Brenner Date: Fri, 24 Jul 2026 15:34:21 -0400 Subject: [PATCH] add in llmobs tool definitions api --- .../trace/llmobs/domain/DDLLMObsSpan.java | 25 ++++++ .../llmobs/domain/DDLLMObsSpanTest.groovy | 83 +++++++++++++++++++ .../java/datadog/trace/api/llmobs/LLMObs.java | 48 +++++++++++ .../datadog/trace/api/llmobs/LLMObsSpan.java | 7 ++ .../trace/api/llmobs/noop/NoOpLLMObsSpan.java | 3 + .../datadog/trace/api/llmobs/LLMObsTest.java | 52 ++++++++++++ .../writer/ddintake/LLMObsSpanMapper.java | 35 ++++++++ .../writer/ddintake/LLMObsSpanMapperTest.java | 13 +-- 8 files changed, 261 insertions(+), 5 deletions(-) diff --git a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/DDLLMObsSpan.java b/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/DDLLMObsSpan.java index ac304698652..19a1c73346d 100644 --- a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/DDLLMObsSpan.java +++ b/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 toolDefinitions) { + if (finished || toolDefinitions == null || toolDefinitions.isEmpty()) { + return; + } + List 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 metadata) { if (finished) { diff --git a/dd-java-agent/agent-llmobs/src/test/groovy/datadog/trace/llmobs/domain/DDLLMObsSpanTest.groovy b/dd-java-agent/agent-llmobs/src/test/groovy/datadog/trace/llmobs/domain/DDLLMObsSpanTest.groovy index 34832914089..56ae1d42c17 100644 --- a/dd-java-agent/agent-llmobs/src/test/groovy/datadog/trace/llmobs/domain/DDLLMObsSpanTest.groovy +++ b/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() 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 index 629faa23f5a..25f3ff0a8ac 100644 --- 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 @@ -173,6 +173,54 @@ public Map getArguments() { } } + public static class ToolDefinition { + private String name; + private String description; + private Map 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 schema) { + return new ToolDefinition(name, description, schema, null); + } + + public static ToolDefinition from( + String name, String description, Map schema, String version) { + return new ToolDefinition(name, description, schema, version); + } + + private ToolDefinition( + String name, String description, Map 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 getSchema() { + return schema; + } + + public String getVersion() { + return version; + } + } + public static class ToolResult { private String name; private String type; 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 index e86c9717e24..fcf5ef93f82 100644 --- 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 @@ -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 toolDefinitions); + /** * Annotate the span with metadata * 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 index 17ea94101b4..e5d3146fe63 100644 --- 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 @@ -15,6 +15,9 @@ public void annotateIO(List inputData, List toolDefinitions) {} + @Override public void setMetadata(Map metadata) {} diff --git a/dd-trace-api/src/test/java/datadog/trace/api/llmobs/LLMObsTest.java b/dd-trace-api/src/test/java/datadog/trace/api/llmobs/LLMObsTest.java index 71f29740165..193a27af028 100644 --- a/dd-trace-api/src/test/java/datadog/trace/api/llmobs/LLMObsTest.java +++ b/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 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 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 args = new HashMap<>(); diff --git a/dd-trace-core/src/main/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapper.java b/dd-trace-core/src/main/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapper.java index 7849052b9d3..155f4a59406 100644 --- a/dd-trace-core/src/main/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapper.java +++ b/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 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()) { diff --git a/dd-trace-core/src/test/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapperTest.java b/dd-trace-core/src/test/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapperTest.java index af0b1476b89..1da6bf4018f 100644 --- a/dd-trace-core/src/test/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapperTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapperTest.java @@ -100,11 +100,11 @@ void testLLMObsSpanMapperSerialization() throws Exception { Map schema = new LinkedHashMap<>(); schema.put("type", "object"); schema.put("properties", properties); - Map 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> toolDefsResult = (List>) 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"));