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 86913f110c7..b5283b73dad 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 @@ -18,6 +18,7 @@ import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import javax.annotation.Nonnull; @@ -36,9 +37,16 @@ public class DDLLMObsSpan implements LLMObsSpan { // internal tags to be prefixed private static final String INPUT = LLMOBS_TAG_PREFIX + "input"; private static final String OUTPUT = LLMOBS_TAG_PREFIX + "output"; + private static final String INPUT_PROMPT = LLMOBS_TAG_PREFIX + "input_prompt"; 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 PROMPT_TRACKING_INSTRUMENTATION_METHOD = + LLMOBS_TAG_PREFIX + "prompt_tracking_instrumentation_method"; + private static final String INSTRUMENTATION_METHOD_ANNOTATED = "annotated"; + private static final String DEFAULT_PROMPT_NAME = "unnamed-prompt"; + private static final String CONTEXT_VARIABLE_KEYS = "_dd_context_variable_keys"; + private static final String QUERY_VARIABLE_KEYS = "_dd_query_variable_keys"; private static final String PARENT_ID_TAG_INTERNAL = "parent_id"; private static final String SERVICE = LLMOBS_TAG_PREFIX + "service"; @@ -52,6 +60,7 @@ public class DDLLMObsSpan implements LLMObsSpan { private final AgentSpan span; private final String spanKind; + private final String mlApp; private final ContextScope scope; private final boolean hasSessionId; @@ -90,6 +99,7 @@ public DDLLMObsSpan( span.setTag(SPAN_KIND, kind); spanKind = kind; + this.mlApp = mlApp; span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.ML_APP, mlApp); // Resolve effective parent_id and session_id from the LLMObs context, both gated on // trace-id consistency. A stale context from a different trace (e.g. async boundary @@ -234,6 +244,84 @@ public void annotateIO(String inputData, String outputData) { } } + @Override + public void annotatePrompt(LLMObs.Prompt prompt) { + if (finished || prompt == null) { + return; + } + if (!Tags.LLMOBS_LLM_SPAN_KIND.equals(spanKind)) { + LOGGER.warn( + "dropping prompt on non-LLM span kind, annotating prompts is only supported for LLM span kinds"); + return; + } + + Map annotatedPrompt = new LinkedHashMap<>(); + Object currentPrompt = span.getTag(INPUT_PROMPT); + if (currentPrompt instanceof Map) { + annotatedPrompt.putAll(copyStringKeyedMap((Map) currentPrompt)); + } + if (prompt.getId() != null && !prompt.getId().isEmpty()) { + annotatedPrompt.put("id", prompt.getId()); + } + if (!annotatedPrompt.containsKey("id")) { + annotatedPrompt.put("id", mlApp + "_" + DEFAULT_PROMPT_NAME); + } + putIfPresent(annotatedPrompt, "version", prompt.getVersion()); + putIfPresent(annotatedPrompt, "variables", prompt.getVariables()); + if (prompt.getTemplate() != null) { + annotatedPrompt.remove("chat_template"); + annotatedPrompt.put("template", prompt.getTemplate()); + } else if (prompt.getChatTemplate() != null && !prompt.getChatTemplate().isEmpty()) { + annotatedPrompt.remove("template"); + annotatedPrompt.put("chat_template", toChatTemplate(prompt.getChatTemplate())); + } + putIfPresent(annotatedPrompt, "tags", prompt.getTags()); + annotatedPrompt.put( + CONTEXT_VARIABLE_KEYS, + prompt.getContextVariables() == null + ? Collections.singletonList("context") + : prompt.getContextVariables()); + annotatedPrompt.put( + QUERY_VARIABLE_KEYS, + prompt.getQueryVariables() == null + ? Collections.singletonList("question") + : prompt.getQueryVariables()); + + span.setTag(INPUT_PROMPT, annotatedPrompt); + span.setTag(PROMPT_TRACKING_INSTRUMENTATION_METHOD, INSTRUMENTATION_METHOD_ANNOTATED); + } + + private static Map copyStringKeyedMap(Map source) { + Map copy = new LinkedHashMap<>(); + for (Map.Entry entry : source.entrySet()) { + if (entry.getKey() instanceof String) { + copy.put((String) entry.getKey(), entry.getValue()); + } + } + return copy; + } + + private static void putIfPresent(Map prompt, String key, Object value) { + if (value != null) { + prompt.put(key, value); + } + } + + private static List> toChatTemplate(List messages) { + List> chatTemplate = new ArrayList<>(messages.size()); + for (LLMObs.LLMMessage message : messages) { + if (message == null || message.getRole() == null || message.getContent() == null) { + LOGGER.warn("prompt chat template messages must define both role and content; skipping"); + continue; + } + Map templateMessage = new LinkedHashMap<>(); + templateMessage.put("role", message.getRole()); + templateMessage.put("content", message.getContent()); + chatTemplate.add(templateMessage); + } + return chatTemplate; + } + @Override public void setToolDefinitions(List toolDefinitions) { if (finished || toolDefinitions == null || toolDefinitions.isEmpty()) { @@ -269,7 +357,9 @@ public void setMetadata(Map metadata) { } if (value instanceof Map) { - ((Map) value).putAll(metadata); + Map mergedMetadata = copyStringKeyedMap((Map) value); + mergedMetadata.putAll(metadata); + span.setTag(METADATA, mergedMetadata); } else { LOGGER.debug( "unexpected instance type for metadata {}, overwriting for now", 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 56ae1d42c17..bef904409a2 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 @@ -59,9 +59,12 @@ class DDLLMObsSpanTest extends DDSpecification{ // internal tags to be prefixed private static final String INPUT = LLMOBS_TAG_PREFIX + "input" + private static final String INPUT_PROMPT = LLMOBS_TAG_PREFIX + "input_prompt" 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 + private static final String PROMPT_TRACKING_INSTRUMENTATION_METHOD = + LLMOBS_TAG_PREFIX + "prompt_tracking_instrumentation_method" def "test span simple"() { @@ -338,6 +341,199 @@ class DDLLMObsSpanTest extends DDSpecification{ DDTraceApiInfo.VERSION == innerSpan.getTag(LLMOBS_TAG_PREFIX + "ddtrace.version") } + def "test llm span with prompt"() { + setup: + def test = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "test-span") + def messages = [LLMObs.LLMMessage.from("user", "What is the weather in Paris?")] + def prompt = LLMObs.Prompt.builder() + .id("weather-prompt") + .version("1.0.0") + .template("What is the weather in {{city}}?") + .variables([city: "Paris"]) + .tags([team: "weather"]) + .contextVariables(["forecast"]) + .queryVariables(["city"]) + .build() + test.annotateIO(messages, null) + + when: + test.annotatePrompt(prompt) + + then: + def innerSpan = (AgentSpan)test.span + innerSpan.getTag(INPUT) == messages + innerSpan.getTag(INPUT_PROMPT) == [ + id: "weather-prompt", + version: "1.0.0", + template: "What is the weather in {{city}}?", + variables: [city: "Paris"], + tags: [team: "weather"], + _dd_context_variable_keys: ["forecast"], + _dd_query_variable_keys: ["city"] + ] + innerSpan.getTag(PROMPT_TRACKING_INSTRUMENTATION_METHOD) == "annotated" + } + + def "test llm span with chat prompt"() { + setup: + def test = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "test-span") + def chatTemplate = [ + LLMObs.LLMMessage.from("system", "You are a weather assistant."), + LLMObs.LLMMessage.from("user", "What is the weather in {{city}}?") + ] + def prompt = LLMObs.Prompt.builder() + .id("weather-prompt") + .template(chatTemplate) + .build() + + when: + test.annotatePrompt(prompt) + + then: + def innerSpan = (AgentSpan)test.span + innerSpan.getTag(INPUT_PROMPT) == [ + id: "weather-prompt", + chat_template: [ + [role: "system", content: "You are a weather assistant."], + [role: "user", content: "What is the weather in {{city}}?"] + ], + _dd_context_variable_keys: ["context"], + _dd_query_variable_keys: ["question"] + ] + } + + def "input messages added after a prompt preserve the prompt"() { + setup: + def test = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "test-span") + def prompt = LLMObs.Prompt.builder().id("weather-prompt").build() + def messages = [LLMObs.LLMMessage.from("user", "What is the weather?")] + test.annotatePrompt(prompt) + + when: + test.annotateIO(messages, null) + + then: + def innerSpan = (AgentSpan)test.span + innerSpan.getTag(INPUT_PROMPT).id == "weather-prompt" + innerSpan.getTag(INPUT) == messages + } + + def "prompt annotations merge with prior prompt attributes"() { + setup: + def test = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "test-span") + test.annotatePrompt( + LLMObs.Prompt.builder() + .id("weather-prompt") + .template("Weather in {{city}}") + .variables([city: "Paris"]) + .build()) + + when: + test.annotatePrompt( + LLMObs.Prompt.builder() + .version("2.0.0") + .tags([team: "weather"]) + .build()) + + then: + def prompt = ((AgentSpan)test.span).getTag(INPUT_PROMPT) + prompt.id == "weather-prompt" + prompt.version == "2.0.0" + prompt.template == "Weather in {{city}}" + prompt.variables == [city: "Paris"] + prompt.tags == [team: "weather"] + } + + def "later prompt template replaces the prior template representation"() { + setup: + def test = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "test-span") + test.annotatePrompt(LLMObs.Prompt.builder().template("Weather in {{city}}").build()) + + when: + test.annotatePrompt(LLMObs.Prompt.builder().template([LLMObs.LLMMessage.from("user", "Weather in {{city}}")]).build()) + + then: + def chatPrompt = ((AgentSpan)test.span).getTag(INPUT_PROMPT) + chatPrompt.template == null + chatPrompt.chat_template == [[role: "user", content: "Weather in {{city}}"]] + + when: + test.annotatePrompt(LLMObs.Prompt.builder().template("Forecast for {{city}}").build()) + + then: + def textPrompt = ((AgentSpan)test.span).getTag(INPUT_PROMPT) + textPrompt.template == "Forecast for {{city}}" + textPrompt.chat_template == null + } + + def "prompt without an id uses the ml app default"() { + setup: + def test = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "test-span") + + when: + test.annotatePrompt(LLMObs.Prompt.builder().template("Hello").build()) + + then: + ((AgentSpan)test.span).getTag(INPUT_PROMPT).id == "test-ml-app_unnamed-prompt" + } + + def "empty prompt id is treated as missing"() { + setup: + def test = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "test-span") + + when: + test.annotatePrompt(LLMObs.Prompt.builder().id("").build()) + + then: + ((AgentSpan)test.span).getTag(INPUT_PROMPT).id == "test-ml-app_unnamed-prompt" + + when: + test.annotatePrompt(LLMObs.Prompt.builder().id("weather-prompt").build()) + test.annotatePrompt(LLMObs.Prompt.builder().id("").version("2.0.0").build()) + + then: + def prompt = ((AgentSpan)test.span).getTag(INPUT_PROMPT) + prompt.id == "weather-prompt" + prompt.version == "2.0.0" + } + + def "prompts are ignored on non-LLM spans"() { + setup: + def test = llmObsSpan(spanKind, "test-span") + + when: + test.annotatePrompt(LLMObs.Prompt.builder().id("weather-prompt").build()) + + then: + def innerSpan = (AgentSpan)test.span + innerSpan.getTag(INPUT) == null + innerSpan.getTag(INPUT_PROMPT) == null + innerSpan.getTag(PROMPT_TRACKING_INSTRUMENTATION_METHOD) == null + + where: + spanKind << [ + Tags.LLMOBS_AGENT_SPAN_KIND, + Tags.LLMOBS_TOOL_SPAN_KIND, + Tags.LLMOBS_TASK_SPAN_KIND, + Tags.LLMOBS_WORKFLOW_SPAN_KIND, + Tags.LLMOBS_EMBEDDING_SPAN_KIND, + Tags.LLMOBS_RETRIEVAL_SPAN_KIND + ] + } + + def "prompts cannot be annotated after the span finishes"() { + setup: + def test = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "test-span") + test.annotatePrompt(LLMObs.Prompt.builder().id("first").build()) + test.finish() + + when: + test.annotatePrompt(LLMObs.Prompt.builder().id("second").build()) + + then: + ((AgentSpan)test.span).getTag(INPUT_PROMPT).id == "first" + } + def "test llm span with tool definitions"() { setup: def test = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "test-span") 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 522e871e98c..86ef5ccb07d 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 @@ -2,6 +2,9 @@ import datadog.trace.api.llmobs.noop.NoOpLLMObsEvalProcessor; import datadog.trace.api.llmobs.noop.NoOpLLMObsSpanFactory; +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import javax.annotation.Nullable; @@ -138,6 +141,132 @@ void SubmitEvaluation( Map tags); } + /** A prompt template and its associated attributes for an LLM call. */ + public static final class Prompt { + private final String id; + private final String version; + private final String template; + private final List chatTemplate; + private final Map variables; + private final Map tags; + private final List contextVariables; + private final List queryVariables; + + public static Builder builder() { + return new Builder(); + } + + private Prompt(Builder builder) { + this.id = builder.id; + this.version = builder.version; + this.template = builder.template; + this.chatTemplate = immutableList(builder.chatTemplate); + this.variables = immutableMap(builder.variables); + this.tags = immutableMap(builder.tags); + this.contextVariables = immutableList(builder.contextVariables); + this.queryVariables = immutableList(builder.queryVariables); + } + + public String getId() { + return id; + } + + public String getVersion() { + return version; + } + + public String getTemplate() { + return template; + } + + public List getChatTemplate() { + return chatTemplate; + } + + public Map getVariables() { + return variables; + } + + public Map getTags() { + return tags; + } + + public List getContextVariables() { + return contextVariables; + } + + public List getQueryVariables() { + return queryVariables; + } + + private static List immutableList(List values) { + return values == null ? null : Collections.unmodifiableList(new ArrayList<>(values)); + } + + private static Map immutableMap(Map values) { + return values == null ? null : Collections.unmodifiableMap(new LinkedHashMap<>(values)); + } + + public static final class Builder { + private String id; + private String version; + private String template; + private List chatTemplate; + private Map variables; + private Map tags; + private List contextVariables; + private List queryVariables; + + private Builder() {} + + public Builder id(String id) { + this.id = id; + return this; + } + + public Builder version(String version) { + this.version = version; + return this; + } + + public Builder template(String template) { + this.template = template; + this.chatTemplate = null; + return this; + } + + public Builder template(List chatTemplate) { + this.template = null; + this.chatTemplate = chatTemplate; + return this; + } + + public Builder variables(Map variables) { + this.variables = variables; + return this; + } + + public Builder tags(Map tags) { + this.tags = tags; + return this; + } + + public Builder contextVariables(List contextVariables) { + this.contextVariables = contextVariables; + return this; + } + + public Builder queryVariables(List queryVariables) { + this.queryVariables = queryVariables; + return this; + } + + public Prompt build() { + return new Prompt(this); + } + } + } + public static class ToolCall { 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 817cfc96e85..efaf13123b7 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 @@ -43,6 +43,15 @@ default void annotateRetrievalIO(String inputData, List outputD */ void annotateIO(String inputData, String outputData); + /** + * Annotate an LLM span with the prompt used for the LLM call. + * + *

This annotation is ignored for non-LLM spans. + * + * @param prompt The prompt used for the LLM call + */ + default void annotatePrompt(LLMObs.Prompt prompt) {} + /** * Annotate the span with the definitions of tools available to the LLM. * 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 e5d3146fe63..bcde38ccfae 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) {} 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 40bac293d60..0fdbac58993 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 @@ -129,6 +129,59 @@ void testSetToolDefinitionsIsCompatibilityPreservingDefaultMethod() throws Excep assertTrue(LLMObsSpan.class.getMethod("setToolDefinitions", List.class).isDefault()); } + @Test + void testAnnotatePromptIsCompatibilityPreservingDefaultMethod() throws Exception { + assertTrue(LLMObsSpan.class.getMethod("annotatePrompt", LLMObs.Prompt.class).isDefault()); + } + + @Test + void testPromptBuilderWithTextTemplate() { + Map variables = new HashMap<>(); + variables.put("city", "Paris"); + Map tags = new HashMap<>(); + tags.put("team", "weather"); + List contextVariables = Arrays.asList("forecast", "history"); + List queryVariables = Collections.singletonList("city"); + + LLMObs.Prompt prompt = + LLMObs.Prompt.builder() + .id("weather-prompt") + .version("1.0.0") + .template("What is the weather in {{city}}?") + .variables(variables) + .tags(tags) + .contextVariables(contextVariables) + .queryVariables(queryVariables) + .build(); + + assertEquals("weather-prompt", prompt.getId()); + assertEquals("1.0.0", prompt.getVersion()); + assertEquals("What is the weather in {{city}}?", prompt.getTemplate()); + assertNull(prompt.getChatTemplate()); + assertEquals(variables, prompt.getVariables()); + assertEquals(tags, prompt.getTags()); + assertEquals(contextVariables, prompt.getContextVariables()); + assertEquals(queryVariables, prompt.getQueryVariables()); + assertNotSame(variables, prompt.getVariables()); + assertNotSame(tags, prompt.getTags()); + assertNotSame(contextVariables, prompt.getContextVariables()); + assertNotSame(queryVariables, prompt.getQueryVariables()); + } + + @Test + void testPromptBuilderWithChatTemplate() { + List chatTemplate = + Arrays.asList( + LLMObs.LLMMessage.from("system", "You are a weather assistant."), + LLMObs.LLMMessage.from("user", "What is the weather in {{city}}?")); + + LLMObs.Prompt prompt = LLMObs.Prompt.builder().template(chatTemplate).build(); + + assertNull(prompt.getTemplate()); + assertEquals(chatTemplate, prompt.getChatTemplate()); + assertNotSame(chatTemplate, prompt.getChatTemplate()); + } + @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 d55935e9835..4a3b3e581c6 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 @@ -42,6 +42,7 @@ public class LLMObsSpanMapper implements RemoteMapper { // internal tags to be prefixed private static final String INPUT = "input"; + private static final String INPUT_PROMPT = "input_prompt"; private static final String OUTPUT = "output"; private static final String SPAN_KIND_TAG_KEY = LLMOBS_TAG_PREFIX + Tags.SPAN_KIND; @@ -67,6 +68,7 @@ public class LLMObsSpanMapper implements RemoteMapper { private static final byte[] META = "meta".getBytes(StandardCharsets.UTF_8); private static final byte[] METADATA = "metadata".getBytes(StandardCharsets.UTF_8); + private static final byte[] PROMPT = "prompt".getBytes(StandardCharsets.UTF_8); private static final byte[] SPAN_KIND = "span.kind".getBytes(StandardCharsets.UTF_8); private static final byte[] SPANS = "spans".getBytes(StandardCharsets.UTF_8); private static final byte[] METRICS = "metrics".getBytes(StandardCharsets.UTF_8); @@ -256,6 +258,7 @@ private static final class MetaWriter implements MetadataConsumer { new HashSet<>( Arrays.asList( LLMOBS_TAG_PREFIX + INPUT, + LLMOBS_TAG_PREFIX + INPUT_PROMPT, LLMOBS_TAG_PREFIX + OUTPUT, LLMOBS_TAG_PREFIX + LLMObsTags.MODEL_NAME, LLMOBS_TAG_PREFIX + LLMObsTags.MODEL_PROVIDER, @@ -300,6 +303,20 @@ public void accept(Metadata metadata) { LOGGER.warn("missing span kind"); } + String inputTag = LLMOBS_TAG_PREFIX + INPUT; + String inputPromptTag = LLMOBS_TAG_PREFIX + INPUT_PROMPT; + boolean hasInput = tagsToRemapToMeta.containsKey(inputTag); + boolean hasInputPrompt = tagsToRemapToMeta.containsKey(inputPromptTag); + Object inputPrompt = null; + if (hasInputPrompt) { + if (spanKind.equals(Tags.LLMOBS_LLM_SPAN_KIND)) { + inputPrompt = tagsToRemapToMeta.get(inputPromptTag); + } else { + LOGGER.warn( + "dropping prompt on non-LLM span kind, annotating prompts is only supported for LLM span kinds"); + } + } + // write metrics (9) writable.writeUTF8(METRICS); writable.startMap(metricsSize); @@ -325,7 +342,11 @@ public void accept(Metadata metadata) { // write meta (11) int metaSize = - tagsToRemapToMeta.size() + 1 + (null != errorInfo && !errorInfo.isEmpty() ? 1 : 0); + tagsToRemapToMeta.size() + - (hasInputPrompt ? 1 : 0) + + (inputPrompt != null && !hasInput ? 1 : 0) + + 1 + + (null != errorInfo && !errorInfo.isEmpty() ? 1 : 0); writable.writeUTF8(META); writable.startMap(metaSize); writable.writeUTF8(SPAN_KIND); @@ -363,11 +384,9 @@ public void accept(Metadata metadata) { if (spanKind.equals(Tags.LLMOBS_LLM_SPAN_KIND)) { writable.writeString(key, null); if (val instanceof List) { - writable.startMap(1); - writable.writeString("messages", null); - writeLlmMessages((List) val); + writeLlmMessagesField((List) val, key.equals(INPUT) ? inputPrompt : null); } else if (key.equals(INPUT) && val instanceof Map) { - writeLlmInputMap((Map) val); + writeLlmInputMap((Map) val, inputPrompt); } else { LOGGER.warn( "unexpectedly found incorrect type for LLM span IO {}, expecting list", @@ -413,6 +432,9 @@ public void accept(Metadata metadata) { writable.writeString("value", null); writable.writeObject(val, null); } + } else if (key.equals(INPUT_PROMPT)) { + // Serialized as meta.input.prompt above, or after this loop when no input is present. + continue; } else if (key.equals(LLMObsTags.TOOL_DEFINITIONS) && val instanceof List) { writable.writeString(key, null); writeToolDefinitions((List) val); @@ -429,6 +451,13 @@ public void accept(Metadata metadata) { writable.writeObject(val, null); } } + + if (inputPrompt != null && !hasInput) { + writable.writeString(INPUT, null); + writable.startMap(1); + writable.writeUTF8(PROMPT); + writable.writeObject(inputPrompt, null); + } } private void writeToolDefinitions(List toolDefinitions) { @@ -475,8 +504,19 @@ private static boolean isDocumentList(Object value) { return true; } - private void writeLlmInputMap(Map inputMap) { - writable.startMap(inputMap.size()); + private void writeLlmMessagesField(List messages, Object inputPrompt) { + writable.startMap(inputPrompt == null ? 1 : 2); + writable.writeString("messages", null); + writeLlmMessages(messages); + if (inputPrompt != null) { + writable.writeUTF8(PROMPT); + writable.writeObject(inputPrompt, null); + } + } + + private void writeLlmInputMap(Map inputMap, Object inputPrompt) { + boolean addInputPrompt = inputPrompt != null && !inputMap.containsKey("prompt"); + writable.startMap(inputMap.size() + (addInputPrompt ? 1 : 0)); for (Map.Entry entry : inputMap.entrySet()) { String inputKey = String.valueOf(entry.getKey()); Object inputValue = entry.getValue(); @@ -487,6 +527,10 @@ private void writeLlmInputMap(Map inputMap) { writable.writeObject(inputValue, null); } } + if (addInputPrompt) { + writable.writeUTF8(PROMPT); + writable.writeObject(inputPrompt, null); + } } private void writeLlmMessages(List messages) { 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 a07ca4b6210..f683e64ab91 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 @@ -84,10 +84,8 @@ void testLLMObsSpanMapperSerialization() throws Exception { prompt.put("variables", Collections.singletonMap("city", "San Francisco")); prompt.put("chat_template", Collections.singletonList(chatTemplateEntry)); - Map inputMap = new LinkedHashMap<>(); - inputMap.put("messages", inputMessages); - inputMap.put("prompt", prompt); - llmSpan.setTag("_ml_obs_tag.input", inputMap); + llmSpan.setTag("_ml_obs_tag.input", inputMessages); + llmSpan.setTag("_ml_obs_tag.input_prompt", prompt); llmSpan.setTag("_ml_obs_tag.output", outputMessages); Map metadataMap = new LinkedHashMap<>(); @@ -197,6 +195,7 @@ void testLLMObsSpanMapperSerialization() throws Exception { assertTrue(meta.containsKey("output")); Map outputResult = (Map) meta.get("output"); assertTrue(outputResult.containsKey("messages")); + assertFalse(outputResult.containsKey("prompt")); List> outputMsgs = (List>) outputResult.get("messages"); assertTrue(outputMsgs.get(0).containsKey("content")); assertEquals("I'll help you check the weather.", outputMsgs.get(0).get("content")); @@ -222,6 +221,66 @@ void testLLMObsSpanMapperSerialization() throws Exception { List tags = (List) spanData.get("tags"); assertTrue(tags.contains("language:jvm")); assertTrue(tags.contains("session_id:abc-123-session")); + assertFalse(tags.stream().anyMatch(tag -> tag.startsWith("input_prompt:"))); + + tracer.close(); + } + + @Test + void testLLMObsSpanMapperSerializesPromptWithoutInputMessages() throws Exception { + LLMObsSpanMapper mapper = new LLMObsSpanMapper(); + CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build(); + + Map prompt = new LinkedHashMap<>(); + prompt.put("id", "prompt_123"); + prompt.put("template", "Hello {{name}}"); + prompt.put("variables", Collections.singletonMap("name", "Sam")); + + AgentSpan llmSpan = + tracer + .buildSpan("datadog", "chat-completion") + .withTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND) + .withTag("_ml_obs_tag.input_prompt", prompt) + .start(); + llmSpan.setSpanType(InternalSpanTypes.LLMOBS); + llmSpan.finish(); + + Map spanData = serializeSingleSpan(mapper, llmSpan); + Map meta = (Map) spanData.get("meta"); + Map input = (Map) meta.get("input"); + + assertEquals(Collections.singletonMap("prompt", prompt), input); + List tags = (List) spanData.get("tags"); + assertFalse(tags.stream().anyMatch(tag -> tag.startsWith("input_prompt:"))); + + tracer.close(); + } + + @Test + void testLLMObsSpanMapperPreservesNestedPromptInputCompatibility() throws Exception { + LLMObsSpanMapper mapper = new LLMObsSpanMapper(); + CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build(); + + Map prompt = Collections.singletonMap("id", "legacy_prompt"); + Map input = new LinkedHashMap<>(); + input.put("messages", Collections.singletonList(LLMObs.LLMMessage.from("user", "Hello"))); + input.put("prompt", prompt); + + AgentSpan llmSpan = + tracer + .buildSpan("datadog", "chat-completion") + .withTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND) + .withTag("_ml_obs_tag.input", input) + .start(); + llmSpan.setSpanType(InternalSpanTypes.LLMOBS); + llmSpan.finish(); + + Map spanData = serializeSingleSpan(mapper, llmSpan); + Map meta = (Map) spanData.get("meta"); + Map serializedInput = (Map) meta.get("input"); + + assertEquals(prompt, serializedInput.get("prompt")); + assertTrue(serializedInput.containsKey("messages")); tracer.close(); } @@ -510,6 +569,21 @@ public void close() throws IOException {} return channel.toByteArray(); } + private static Map serializeSingleSpan(LLMObsSpanMapper mapper, AgentSpan span) + throws IOException { + CapturingByteBufferConsumer sink = new CapturingByteBufferConsumer(); + MsgPackWriter packer = new MsgPackWriter(new FlushingBuffer(16 * 1024, sink)); + packer.format(Collections.singletonList((DDSpan) span), mapper); + packer.flush(); + + assertNotNull(sink.captured); + datadog.trace.common.writer.Payload payload = mapper.newPayload(); + payload.withBody(1, sink.captured); + Map result = objectMapper.readValue(writeTo(payload), Map.class); + List> spans = (List>) result.get("spans"); + return spans.get(0); + } + static class CapturingByteBufferConsumer implements ByteBufferConsumer { ByteBuffer captured;