Skip to content

Add in LLM Observability Prompt Tracking API #12161

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sabrenner wants to merge 4 commits into
base: master
Choose a base branch
from
Open

Add in LLM Observability Prompt Tracking API #12161

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
21b2f2e
prompt tracking via manual annotatePrompt method with prompt builder
sabrenner Aug 6, 2026
866782e
codex review comments
sabrenner Aug 7, 2026
b387155
Merge branch 'master' into sabrenner/llmobs-manual-prompt-tracking
sabrenner Aug 7, 2026
dbe9f51
Merge branch 'master' of github.com:DataDog/dd-trace-java into sabren…
sabrenner Aug 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
92 changes: 91 additions & 1 deletion dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/DDLLMObsSpan.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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";
Expand All @@ -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;

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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<String, Object> 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());
Comment thread
sabrenner marked this conversation as resolved.
annotatedPrompt.put(
QUERY_VARIABLE_KEYS,
prompt.getQueryVariables() == null
? Collections.singletonList("question")
: prompt.getQueryVariables());
Comment thread
sabrenner marked this conversation as resolved.

span.setTag(INPUT_PROMPT, annotatedPrompt);
span.setTag(PROMPT_TRACKING_INSTRUMENTATION_METHOD, INSTRUMENTATION_METHOD_ANNOTATED);
}

private static Map<String, Object> copyStringKeyedMap(Map<?, ?> source) {
Map<String, Object> 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<String, Object> prompt, String key, Object value) {
if (value != null) {
prompt.put(key, value);
}
}

private static List<Map<String, String>> toChatTemplate(List<LLMObs.LLMMessage> messages) {
List<Map<String, String>> 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<String, String> templateMessage = new LinkedHashMap<>();
templateMessage.put("role", message.getRole());
templateMessage.put("content", message.getContent());
chatTemplate.add(templateMessage);
}
return chatTemplate;
}

@Override
public void setToolDefinitions(List<LLMObs.ToolDefinition> toolDefinitions) {
if (finished || toolDefinitions == null || toolDefinitions.isEmpty()) {
Expand Down Expand Up @@ -269,7 +357,9 @@ public void setMetadata(Map<String, Object> metadata) {
}

if (value instanceof Map) {
((Map) value).putAll(metadata);
Map<String, Object> mergedMetadata = copyStringKeyedMap((Map<?, ?>) value);
mergedMetadata.putAll(metadata);
span.setTag(METADATA, mergedMetadata);
} else {
LOGGER.debug(
"unexpected instance type for metadata {}, overwriting for now",
Expand Down
196 changes: 196 additions & 0 deletions ...va-agent/agent-llmobs/src/test/groovy/datadog/trace/llmobs/domain/DDLLMObsSpanTest.groovy
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"() {
Expand Down Expand Up @@ -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")
Expand Down
Loading