From 1e94d01b18d40606770f507bfeaa00e0d7fb93e1 Mon Sep 17 00:00:00 2001 From: Sam Brenner Date: Wed, 29 Jul 2026 13:26:54 -0400 Subject: [PATCH 1/3] add in llmobs user span processor --- .../datadog/trace/llmobs/LLMObsSystem.java | 6 +- .../trace/llmobs/domain/LLMObsInternal.java | 13 -- dd-trace-api/build.gradle.kts | 2 + .../java/datadog/trace/api/llmobs/LLMObs.java | 27 +++ .../trace/api/llmobs/LLMObsSpanData.java | 59 +++++ .../trace/api/llmobs/LLMObsSpanProcessor.java | 20 ++ .../datadog/trace/api/llmobs/LLMObsTest.java | 26 +++ .../ddintake/LLMObsSpanDataAdapter.java | 219 ++++++++++++++++++ .../writer/ddintake/LLMObsSpanMapper.java | 76 ++++-- .../writer/ddintake/LLMObsSpanMapperTest.java | 179 ++++++++++++++ .../trace/api/llmobs/LLMObsInternal.java | 24 ++ .../api/telemetry/LLMObsMetricCollector.java | 20 ++ .../LLMObsMetricCollectorTest.groovy | 13 +- 13 files changed, 652 insertions(+), 32 deletions(-) delete mode 100644 dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/LLMObsInternal.java create mode 100644 dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObsSpanData.java create mode 100644 dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObsSpanProcessor.java create mode 100644 dd-trace-core/src/main/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanDataAdapter.java create mode 100644 internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsInternal.java diff --git a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java b/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java index a57dd858b45..914fac8e124 100644 --- a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java +++ b/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java @@ -4,12 +4,12 @@ import datadog.trace.api.Config; import datadog.trace.api.WellKnownTags; import datadog.trace.api.llmobs.LLMObs; +import datadog.trace.api.llmobs.LLMObsInternal; import datadog.trace.api.llmobs.LLMObsSpan; import datadog.trace.api.llmobs.LLMObsTags; import datadog.trace.bootstrap.instrumentation.api.Tags; import datadog.trace.llmobs.domain.DDLLMObsSpan; import datadog.trace.llmobs.domain.LLMObsEval; -import datadog.trace.llmobs.domain.LLMObsInternal; import java.lang.instrument.Instrumentation; import java.util.Map; import java.util.concurrent.TimeUnit; @@ -34,9 +34,9 @@ public static void start(Instrumentation inst, SharedCommunicationObjects sco) { String mlApp = config.getLlmObsMlApp(); WellKnownTags wellKnownTags = config.getWellKnownTags(); - LLMObsInternal.setLLMObsSpanFactory(new LLMObsManualSpanFactory(mlApp, wellKnownTags)); + LLMObsInternal.setSpanFactory(new LLMObsManualSpanFactory(mlApp, wellKnownTags)); - LLMObsInternal.setLLMObsEvalProcessor(new LLMObsCustomEvalProcessor(mlApp, sco, config)); + LLMObsInternal.setEvalProcessor(new LLMObsCustomEvalProcessor(mlApp, sco, config)); } private static class LLMObsCustomEvalProcessor implements LLMObs.LLMObsEvalProcessor { diff --git a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/LLMObsInternal.java b/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/LLMObsInternal.java deleted file mode 100644 index 85e1482b412..00000000000 --- a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/LLMObsInternal.java +++ /dev/null @@ -1,13 +0,0 @@ -package datadog.trace.llmobs.domain; - -import datadog.trace.api.llmobs.LLMObs; - -public class LLMObsInternal extends LLMObs { - public static void setLLMObsSpanFactory(final LLMObsSpanFactory factory) { - LLMObs.SPAN_FACTORY = factory; - } - - public static void setLLMObsEvalProcessor(final LLMObsEvalProcessor evalProcessor) { - LLMObs.EVAL_PROCESSOR = evalProcessor; - } -} diff --git a/dd-trace-api/build.gradle.kts b/dd-trace-api/build.gradle.kts index 70a003348c9..d144ac2c87c 100644 --- a/dd-trace-api/build.gradle.kts +++ b/dd-trace-api/build.gradle.kts @@ -56,6 +56,8 @@ extra["excludedClassesCoverage"] = listOf( "datadog.trace.api.llmobs.LLMObs.ToolCall", "datadog.trace.api.llmobs.LLMObs.ToolResult", "datadog.trace.api.llmobs.LLMObsSpan", + "datadog.trace.api.llmobs.LLMObsSpanData", + "datadog.trace.api.llmobs.LLMObsSpanProcessor", "datadog.trace.api.llmobs.noop.NoOpLLMObsSpan", "datadog.trace.api.llmobs.noop.NoOpLLMObsSpanFactory", "datadog.trace.api.llmobs.noop.NoOpLLMObsEvalProcessor", 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..720b840f4de 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 @@ -4,6 +4,7 @@ import datadog.trace.api.llmobs.noop.NoOpLLMObsSpanFactory; import java.util.List; import java.util.Map; +import java.util.Objects; import javax.annotation.Nullable; public class LLMObs { @@ -11,6 +12,7 @@ protected LLMObs() {} protected static LLMObsSpanFactory SPAN_FACTORY = NoOpLLMObsSpanFactory.INSTANCE; protected static LLMObsEvalProcessor EVAL_PROCESSOR = NoOpLLMObsEvalProcessor.INSTANCE; + @Nullable protected static volatile LLMObsSpanProcessor SPAN_PROCESSOR; public static LLMObsSpan startLLMSpan( String spanName, @@ -60,6 +62,31 @@ public static LLMObsSpan startRetrievalSpan( return SPAN_FACTORY.startRetrievalSpan(spanName, mlApp, sessionId); } + /** + * Registers a processor to be called for each LLM Observability span before it is sent. + * + *

The processor can modify the span input and output, or return {@code null} to omit the span + * from LLM Observability. Only one processor can be registered at a time. + * + * @param processor the processor to register + * @throws NullPointerException if {@code processor} is {@code null} + * @throws IllegalStateException if a processor is already registered + */ + public static synchronized void registerProcessor(LLMObsSpanProcessor processor) { + Objects.requireNonNull(processor, "processor"); + if (SPAN_PROCESSOR != null) { + throw new IllegalStateException( + "An LLM Observability span processor is already registered. " + + "Deregister it before registering another."); + } + SPAN_PROCESSOR = processor; + } + + /** Deregisters the current LLM Observability span processor, if one is registered. */ + public static synchronized void deregisterProcessor() { + SPAN_PROCESSOR = null; + } + public static void SubmitEvaluation( LLMObsSpan llmObsSpan, String label, String categoricalValue, Map tags) { EVAL_PROCESSOR.SubmitEvaluation(llmObsSpan, label, categoricalValue, tags); diff --git a/dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObsSpanData.java b/dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObsSpanData.java new file mode 100644 index 00000000000..f76e8208a84 --- /dev/null +++ b/dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObsSpanData.java @@ -0,0 +1,59 @@ +package datadog.trace.api.llmobs; + +import java.util.List; +import javax.annotation.Nullable; + +/** + * Mutable view of an LLM Observability span passed to a registered {@link LLMObsSpanProcessor}. + * + *

Changes to the input and output are applied immediately before the span is sent to LLM + * Observability. + */ +public interface LLMObsSpanData { + + /** + * Gets the LLM Observability span kind. + * + * @return the span kind + */ + String getKind(); + + /** + * Gets the input content associated with the span. + * + * @return the input represented as messages + */ + List getInput(); + + /** + * Replaces the input content associated with the span. + * + * @param input the new input represented as messages + * @throws NullPointerException if {@code input} is {@code null} + */ + void setInput(List input); + + /** + * Gets the output content associated with the span. + * + * @return the output represented as messages + */ + List getOutput(); + + /** + * Replaces the output content associated with the span. + * + * @param output the new output represented as messages + * @throws NullPointerException if {@code output} is {@code null} + */ + void setOutput(List output); + + /** + * Gets an LLM Observability tag from the span. + * + * @param key the unprefixed tag name + * @return the tag value, or {@code null} when the tag is not present + */ + @Nullable + String getTag(String key); +} diff --git a/dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObsSpanProcessor.java b/dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObsSpanProcessor.java new file mode 100644 index 00000000000..2e035a3c9ac --- /dev/null +++ b/dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObsSpanProcessor.java @@ -0,0 +1,20 @@ +package datadog.trace.api.llmobs; + +import javax.annotation.Nullable; + +/** Processes LLM Observability spans before they are sent. */ +@FunctionalInterface +public interface LLMObsSpanProcessor { + + /** + * Processes an LLM Observability span. + * + *

The processor may mutate and return {@code span}, or return {@code null} to omit the span + * from LLM Observability. + * + * @param span the span being processed + * @return the span to send, or {@code null} to omit it + */ + @Nullable + LLMObsSpanData process(LLMObsSpanData span); +} 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..212702c628e 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 @@ -5,6 +5,7 @@ import static org.junit.jupiter.api.Assertions.assertNotSame; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -27,23 +28,48 @@ class LLMObsTest { private static Object originalSpanFactory; private static Object originalEvalProcessor; + private static Object originalSpanProcessor; @BeforeAll static void setupSpec() throws Exception { originalSpanFactory = getStaticField("SPAN_FACTORY"); originalEvalProcessor = getStaticField("EVAL_PROCESSOR"); + originalSpanProcessor = getStaticField("SPAN_PROCESSOR"); } @AfterAll static void cleanupSpec() throws Exception { setStaticField("SPAN_FACTORY", originalSpanFactory); setStaticField("EVAL_PROCESSOR", originalEvalProcessor); + setStaticField("SPAN_PROCESSOR", originalSpanProcessor); } @AfterEach void cleanup() throws Exception { setStaticField("SPAN_FACTORY", NoOpLLMObsSpanFactory.INSTANCE); setStaticField("EVAL_PROCESSOR", NoOpLLMObsEvalProcessor.INSTANCE); + LLMObs.deregisterProcessor(); + } + + @Test + void testRegisterAndDeregisterProcessor() throws Exception { + LLMObsSpanData span = mock(LLMObsSpanData.class); + LLMObsSpanProcessor processor = registeredSpan -> registeredSpan; + + LLMObs.registerProcessor(processor); + + assertSame(processor, getStaticField("SPAN_PROCESSOR")); + assertSame(span, processor.process(span)); + assertThrows(IllegalStateException.class, () -> LLMObs.registerProcessor(processor)); + + LLMObs.deregisterProcessor(); + + assertNull(getStaticField("SPAN_PROCESSOR")); + } + + @Test + void testRegisterNullProcessor() { + assertThrows(NullPointerException.class, () -> LLMObs.registerProcessor(null)); } @Test diff --git a/dd-trace-core/src/main/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanDataAdapter.java b/dd-trace-core/src/main/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanDataAdapter.java new file mode 100644 index 00000000000..0c5c2bfaaf8 --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanDataAdapter.java @@ -0,0 +1,219 @@ +package datadog.trace.llmobs.writer.ddintake; + +import static java.util.Objects.requireNonNull; + +import datadog.trace.api.DDTags; +import datadog.trace.api.llmobs.LLMObs; +import datadog.trace.api.llmobs.LLMObsSpanData; +import datadog.trace.bootstrap.instrumentation.api.Tags; +import datadog.trace.core.CoreSpan; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** Adapts the internal tag representation of an LLM Observability span to the public API. */ +final class LLMObsSpanDataAdapter implements LLMObsSpanData { + private static final String LLMOBS_TAG_PREFIX = "_ml_obs_tag."; + private static final String INPUT_TAG = LLMOBS_TAG_PREFIX + "input"; + private static final String OUTPUT_TAG = LLMOBS_TAG_PREFIX + "output"; + private static final String SPAN_KIND_TAG = LLMOBS_TAG_PREFIX + Tags.SPAN_KIND; + + private enum IOType { + NONE, + MESSAGES, + DOCUMENTS, + VALUE + } + + private final CoreSpan span; + private final String kind; + private final Object originalInput; + private final Object originalOutput; + private final IOType inputType; + private final IOType outputType; + private List input; + private List output; + + LLMObsSpanDataAdapter(CoreSpan span) { + this.span = span; + Object rawKind = span.getTag(SPAN_KIND_TAG); + kind = rawKind == null ? "unknown" : String.valueOf(rawKind); + originalInput = span.getTag(INPUT_TAG); + originalOutput = span.getTag(OUTPUT_TAG); + inputType = ioType(kind, originalInput, true); + outputType = ioType(kind, originalOutput, false); + input = asMessages(originalInput, inputType); + output = asMessages(originalOutput, outputType); + } + + @Override + public String getKind() { + return kind; + } + + @Override + public List getInput() { + return input; + } + + @Override + public void setInput(List input) { + this.input = copyMessages(requireNonNull(input, "input")); + } + + @Override + public List getOutput() { + return output; + } + + @Override + public void setOutput(List output) { + this.output = copyMessages(requireNonNull(output, "output")); + } + + @Override + public String getTag(String key) { + Object value = span.getTag(LLMOBS_TAG_PREFIX + key); + if (value == null && "error".equals(key)) { + return String.valueOf(span.getError()); + } + if (value == null && "error_type".equals(key)) { + value = span.getTag(DDTags.ERROR_TYPE); + } + return value == null ? null : String.valueOf(value); + } + + void apply(LLMObsSpanData processedSpan) { + applyIO( + span, + INPUT_TAG, + originalInput, + inputType, + copyMessages(requireNonNull(processedSpan.getInput(), "processed input"))); + applyIO( + span, + OUTPUT_TAG, + originalOutput, + outputType, + copyMessages(requireNonNull(processedSpan.getOutput(), "processed output"))); + } + + private static IOType ioType(String kind, Object value, boolean input) { + if (value == null) { + return IOType.NONE; + } + Object unwrapped = unwrapMessages(value); + if (Tags.LLMOBS_LLM_SPAN_KIND.equals(kind)) { + return unwrapped instanceof List && allMessages((List) unwrapped) + ? IOType.MESSAGES + : IOType.NONE; + } + if (input + && Tags.LLMOBS_EMBEDDING_SPAN_KIND.equals(kind) + && value instanceof List + && allDocuments((List) value)) { + return IOType.DOCUMENTS; + } + return IOType.VALUE; + } + + private static List asMessages(Object value, IOType type) { + if (type == IOType.NONE) { + return new ArrayList<>(); + } + if (type == IOType.MESSAGES) { + return copyMessages((List) unwrapMessages(value)); + } + if (type == IOType.DOCUMENTS) { + List messages = new ArrayList<>(((List) value).size()); + for (Object valueElement : (List) value) { + LLMObs.Document document = (LLMObs.Document) valueElement; + messages.add(LLMObs.LLMMessage.from("", document.getText())); + } + return messages; + } + List messages = new ArrayList<>(1); + messages.add(LLMObs.LLMMessage.from("", String.valueOf(value))); + return messages; + } + + private static List copyMessages(List values) { + List messages = new ArrayList<>(values.size()); + for (Object value : values) { + if (!(value instanceof LLMObs.LLMMessage)) { + throw new IllegalArgumentException( + "LLM Observability input and output must contain messages"); + } + messages.add((LLMObs.LLMMessage) value); + } + return messages; + } + + private static Object unwrapMessages(Object value) { + if (value instanceof Map) { + return ((Map) value).get("messages"); + } + return value; + } + + private static boolean allMessages(List values) { + for (Object value : values) { + if (!(value instanceof LLMObs.LLMMessage)) { + return false; + } + } + return true; + } + + private static boolean allDocuments(List values) { + for (Object value : values) { + if (!(value instanceof LLMObs.Document)) { + return false; + } + } + return true; + } + + private static void applyIO( + CoreSpan span, + String tag, + Object originalValue, + IOType type, + List messages) { + if (type == IOType.NONE) { + return; + } + if (messages.isEmpty()) { + if (type == IOType.MESSAGES && originalValue instanceof Map) { + Map updatedValue = new HashMap<>((Map) originalValue); + updatedValue.remove("messages"); + if (updatedValue.isEmpty()) { + span.removeTag(tag); + } else { + span.setTag(tag, updatedValue); + } + } else { + span.removeTag(tag); + } + return; + } + if (type == IOType.MESSAGES) { + if (originalValue instanceof Map) { + Map updatedValue = new HashMap<>((Map) originalValue); + updatedValue.put("messages", messages); + span.setTag(tag, updatedValue); + } else { + span.setTag(tag, messages); + } + } else if (type == IOType.DOCUMENTS) { + List documents = new ArrayList<>(messages.size()); + for (LLMObs.LLMMessage message : messages) { + documents.add(LLMObs.Document.from(message.getContent())); + } + span.setTag(tag, documents); + } else { + span.setTag(tag, messages.get(0).getContent()); + } + } +} 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..22c7e86e676 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 @@ -8,7 +8,11 @@ import datadog.trace.api.DDTags; import datadog.trace.api.intake.TrackType; import datadog.trace.api.llmobs.LLMObs; +import datadog.trace.api.llmobs.LLMObsInternal; +import datadog.trace.api.llmobs.LLMObsSpanData; +import datadog.trace.api.llmobs.LLMObsSpanProcessor; import datadog.trace.api.llmobs.LLMObsTags; +import datadog.trace.api.telemetry.LLMObsMetricCollector; import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes; import datadog.trace.bootstrap.instrumentation.api.Tags; import datadog.trace.common.writer.Payload; @@ -24,6 +28,7 @@ import java.util.Collections; import java.util.HashMap; import java.util.HashSet; +import java.util.IdentityHashMap; import java.util.List; import java.util.Map; import java.util.Set; @@ -93,9 +98,11 @@ public class LLMObsSpanMapper implements RemoteMapper { LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID; private final MetaWriter metaWriter = new MetaWriter(); + private final Set> droppedSpans = Collections.newSetFromMap(new IdentityHashMap<>()); private final int size; private final ByteBuffer header; + private List> pendingTrace; private int spansWritten; public LLMObsSpanMapper() { @@ -123,12 +130,27 @@ public void map(List> trace, Writable writable) { List> llmobsSpans = trace.stream().filter(LLMObsSpanMapper::isLLMObsSpan).collect(Collectors.toList()); + boolean retry = trace == pendingTrace; + if (!retry) { + pendingTrace = null; + droppedSpans.clear(); + } if (llmobsSpans.isEmpty()) { // do nothing if no llmobs spans in the trace return; } + if (!retry) { + prepareSpans(llmobsSpans); + pendingTrace = trace; + } + + int writtenSpans = 0; for (CoreSpan span : llmobsSpans) { + if (droppedSpans.contains(span)) { + continue; + } + // Read session_id off the span before opening the map so we can size it correctly. // We deliberately do NOT remove the tag (unlike parent_id) — the session_id: // entry must remain in the tags[] array to match dd-trace-py and dd-trace-js behavior. @@ -151,7 +173,6 @@ public void map(List> trace, Writable writable) { // 3 writable.writeUTF8(PARENT_ID); writable.writeString(span.getTag(PARENT_ID_TAG_INTERNAL_FULL), null); - span.removeTag(PARENT_ID_TAG_INTERNAL_FULL); // 4 writable.writeUTF8(NAME); @@ -188,11 +209,42 @@ public void map(List> trace, Writable writable) { /* 10 (metrics), 11 (tags), 12 meta — shift down 1 if session_id absent */ span.processTagsAndBaggage(metaWriter.withWritable(writable, getErrorsMap(span))); + writtenSpans++; } // Increase only after all spans have been written. This way, if it rolls back because of a // buffer overflow, the counter won't be skewed. - spansWritten += llmobsSpans.size(); + spansWritten += writtenSpans; + pendingTrace = null; + droppedSpans.clear(); + } + + private void prepareSpans(List> spans) { + LLMObsSpanProcessor processor = LLMObsInternal.getSpanProcessor(); + for (CoreSpan span : spans) { + boolean dropped = false; + if (processor != null) { + boolean processorError = false; + try { + LLMObsSpanDataAdapter adapter = new LLMObsSpanDataAdapter(span); + LLMObsSpanData result = processor.process(adapter); + if (result == null) { + dropped = true; + } else { + adapter.apply(result); + } + } catch (RuntimeException error) { + processorError = true; + dropped = true; + LOGGER.warn("Error in LLM Observability span processor, dropping span", error); + } finally { + LLMObsMetricCollector.get().recordUserProcessorCalled(processorError); + } + } + if (dropped) { + droppedSpans.add(span); + } + } } private CharSequence llmObsSpanName(CoreSpan span) { @@ -282,21 +334,12 @@ public void accept(Metadata metadata) { tagsToRemapToMeta.put(key, tag.getValue()); } else if (key.startsWith(LLMOBS_METRIC_PREFIX) && tag.getValue() instanceof Number) { ++metricsSize; - } else if (key.startsWith(LLMOBS_TAG_PREFIX)) { - if (key.startsWith(LLMOBS_TAG_PREFIX)) { - key = key.substring(LLMOBS_TAG_PREFIX.length()); - } - if (TAGS_FOR_REMAPPING.contains(key)) { - tagsToRemapToMeta.put(key, tag.getValue()); - } else { - ++tagsSize; - } + } else if (key.startsWith(LLMOBS_TAG_PREFIX) && !key.equals(PARENT_ID_TAG_INTERNAL_FULL)) { + ++tagsSize; } } - if (!spanKind.equals("unknown")) { - metadata.getTags().remove(SPAN_KIND_TAG_KEY); - } else { + if (spanKind.equals("unknown")) { LOGGER.warn("missing span kind"); } @@ -318,7 +361,10 @@ public void accept(Metadata metadata) { for (Map.Entry tag : metadata.getTags().entrySet()) { String key = tag.getKey(); Object value = tag.getValue(); - if (!tagsToRemapToMeta.containsKey(key) && key.startsWith(LLMOBS_TAG_PREFIX)) { + if (!tagsToRemapToMeta.containsKey(key) + && key.startsWith(LLMOBS_TAG_PREFIX) + && !key.equals(SPAN_KIND_TAG_KEY) + && !key.equals(PARENT_ID_TAG_INTERNAL_FULL)) { writable.writeObject(key.substring(LLMOBS_TAG_PREFIX.length()) + ":" + value, null); } } 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..b79118ac122 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 @@ -3,6 +3,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import com.fasterxml.jackson.databind.ObjectMapper; @@ -11,6 +12,7 @@ import datadog.communication.serialization.msgpack.MsgPackWriter; import datadog.trace.api.DDTags; import datadog.trace.api.llmobs.LLMObs; +import datadog.trace.api.telemetry.LLMObsMetricCollector; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes; import datadog.trace.bootstrap.instrumentation.api.Tags; @@ -28,6 +30,7 @@ import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.Test; import org.msgpack.jackson.dataformat.MessagePackFactory; @@ -381,6 +384,180 @@ void testLLMObsSpanMapperOmitsTopLevelSessionIdWhenNotSet() throws Exception { tracer.close(); } + @Test + void testLLMObsSpanProcessorModifiesInputAndOutput() throws Exception { + LLMObs.registerProcessor( + span -> { + assertEquals(Tags.LLMOBS_LLM_SPAN_KIND, span.getKind()); + assertEquals("true", span.getTag("redact")); + assertEquals("secret input", span.getInput().get(0).getContent()); + span.setInput(Collections.singletonList(LLMObs.LLMMessage.from("user", "[REDACTED]"))); + span.setOutput(Collections.emptyList()); + return span; + }); + + try { + CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build(); + Map originalInput = new LinkedHashMap<>(); + originalInput.put( + "messages", Collections.singletonList(LLMObs.LLMMessage.from("user", "secret input"))); + originalInput.put("prompt", Collections.singletonMap("id", "prompt-id")); + AgentSpan llmSpan = + tracer + .buildSpan("datadog", "processed") + .withTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND) + .withTag("_ml_obs_tag.input", originalInput) + .withTag( + "_ml_obs_tag.output", + Collections.singletonList(LLMObs.LLMMessage.from("assistant", "secret output"))) + .withTag("_ml_obs_tag.redact", true) + .start(); + llmSpan.setSpanType(InternalSpanTypes.LLMOBS); + llmSpan.finish(); + + List> spans = + serialize(Collections.singletonList((DDSpan) llmSpan), new LLMObsSpanMapper()); + Map meta = (Map) spans.get(0).get("meta"); + Map input = (Map) meta.get("input"); + List> messages = (List>) input.get("messages"); + + assertEquals("[REDACTED]", messages.get(0).get("content")); + assertEquals(Collections.singletonMap("id", "prompt-id"), input.get("prompt")); + assertFalse(meta.containsKey("output")); + tracer.close(); + } finally { + LLMObs.deregisterProcessor(); + } + } + + @Test + void testLLMObsSpanProcessorCanDropSpan() throws Exception { + LLMObs.registerProcessor(span -> "true".equals(span.getTag("drop")) ? null : span); + + try { + CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build(); + AgentSpan dropped = newLlmObsSpan(tracer, "dropped", true); + AgentSpan retained = newLlmObsSpan(tracer, "retained", false); + + List> spans = + serialize(Arrays.asList((DDSpan) dropped, (DDSpan) retained), new LLMObsSpanMapper()); + + assertEquals(1, spans.size()); + assertEquals("retained", spans.get(0).get("name")); + tracer.close(); + } finally { + LLMObs.deregisterProcessor(); + } + } + + @Test + void testLLMObsSpanProcessorExceptionDropsSpan() throws Exception { + LLMObs.registerProcessor( + span -> { + if ("true".equals(span.getTag("drop"))) { + throw new IllegalStateException("processor failure"); + } + return span; + }); + + try { + CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build(); + AgentSpan dropped = newLlmObsSpan(tracer, "dropped", true); + AgentSpan retained = newLlmObsSpan(tracer, "retained", false); + + List> spans = + serialize(Arrays.asList((DDSpan) dropped, (DDSpan) retained), new LLMObsSpanMapper()); + + assertEquals(1, spans.size()); + assertEquals("retained", spans.get(0).get("name")); + tracer.close(); + } finally { + LLMObs.deregisterProcessor(); + } + } + + @Test + void testLLMObsSpanProcessorRunsOnceWhenSerializationRetries() { + AtomicInteger calls = new AtomicInteger(); + LLMObsMetricCollector.get().drain(); + LLMObs.registerProcessor( + span -> { + calls.incrementAndGet(); + return span; + }); + + try { + CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build(); + AgentSpan first = newLlmObsSpan(tracer, "first", false); + AgentSpan second = newLlmObsSpan(tracer, "second", false); + String largeInput = String.join("", Collections.nCopies(600, "x")); + first.setTag( + "_ml_obs_tag.input", + Collections.singletonList(LLMObs.LLMMessage.from("user", largeInput))); + second.setTag( + "_ml_obs_tag.input", + Collections.singletonList(LLMObs.LLMMessage.from("user", largeInput))); + + LLMObsSpanMapper mapper = new LLMObsSpanMapper(); + CapturingByteBufferConsumer sink = new CapturingByteBufferConsumer(); + MsgPackWriter packer = new MsgPackWriter(new FlushingBuffer(1024, sink)); + + assertTrue(packer.format(Collections.singletonList((DDSpan) first), mapper)); + assertTrue(packer.format(Collections.singletonList((DDSpan) second), mapper)); + assertEquals(1, sink.accepts); + assertEquals(2, calls.get()); + assertEquals( + 2, + LLMObsMetricCollector.get().drain().stream() + .filter( + metric -> + LLMObsMetricCollector.USER_PROCESSOR_CALLED_METRIC.equals(metric.metricName)) + .count()); + tracer.close(); + } finally { + LLMObs.deregisterProcessor(); + LLMObsMetricCollector.get().drain(); + } + } + + @Test + void testLLMObsSpanProcessorInputAndOutputRejectNull() { + CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build(); + LLMObsSpanDataAdapter adapter = + new LLMObsSpanDataAdapter((DDSpan) newLlmObsSpan(tracer, "processed", false)); + + assertThrows(NullPointerException.class, () -> adapter.setInput(null)); + assertThrows(NullPointerException.class, () -> adapter.setOutput(null)); + tracer.close(); + } + + private static AgentSpan newLlmObsSpan(CoreTracer tracer, String name, boolean drop) { + AgentSpan span = + tracer + .buildSpan("datadog", name) + .withTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND) + .withTag("_ml_obs_tag.drop", drop) + .start(); + span.setSpanType(InternalSpanTypes.LLMOBS); + span.finish(); + return span; + } + + private static List> serialize(List trace, LLMObsSpanMapper mapper) + throws Exception { + CapturingByteBufferConsumer sink = new CapturingByteBufferConsumer(); + MsgPackWriter packer = new MsgPackWriter(new FlushingBuffer(16 * 1024, sink)); + + packer.format(trace, mapper); + packer.flush(); + + assertNotNull(sink.captured); + datadog.trace.common.writer.Payload payload = mapper.newPayload(); + payload.withBody(trace.size(), sink.captured); + Map result = objectMapper.readValue(writeTo(payload), Map.class); + return (List>) result.get("spans"); + } + private static byte[] writeTo(datadog.trace.common.writer.Payload payload) throws IOException { ByteArrayOutputStream channel = new ByteArrayOutputStream(); payload.writeTo( @@ -407,10 +584,12 @@ public void close() throws IOException {} static class CapturingByteBufferConsumer implements ByteBufferConsumer { ByteBuffer captured; + int accepts; @Override public void accept(int messageCount, ByteBuffer buffer) { captured = buffer; + accepts++; } } } diff --git a/internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsInternal.java b/internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsInternal.java new file mode 100644 index 00000000000..a1b2ee169d1 --- /dev/null +++ b/internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsInternal.java @@ -0,0 +1,24 @@ +package datadog.trace.api.llmobs; + +import javax.annotation.Nullable; + +/** Internal bridge to LLM Observability API state. */ +public final class LLMObsInternal extends LLMObs { + private LLMObsInternal() {} + + /** Sets the LLM Observability span factory. */ + public static void setSpanFactory(LLMObsSpanFactory factory) { + SPAN_FACTORY = factory; + } + + /** Sets the LLM Observability evaluation processor. */ + public static void setEvalProcessor(LLMObsEvalProcessor evalProcessor) { + EVAL_PROCESSOR = evalProcessor; + } + + /** Returns the registered user span processor, if any. */ + @Nullable + public static LLMObsSpanProcessor getSpanProcessor() { + return SPAN_PROCESSOR; + } +} diff --git a/internal-api/src/main/java/datadog/trace/api/telemetry/LLMObsMetricCollector.java b/internal-api/src/main/java/datadog/trace/api/telemetry/LLMObsMetricCollector.java index f43d92cb741..bf0f3de2e52 100644 --- a/internal-api/src/main/java/datadog/trace/api/telemetry/LLMObsMetricCollector.java +++ b/internal-api/src/main/java/datadog/trace/api/telemetry/LLMObsMetricCollector.java @@ -24,6 +24,7 @@ public static LLMObsMetricCollector get() { } public static final String SPAN_FINISHED_METRIC = "span.finished"; + public static final String USER_PROCESSOR_CALLED_METRIC = "user_processor_called"; public static final String COUNT_METRIC_TYPE = "count"; private static final String IS_ROOT_SPAN_TRUE = "is_root_span:1"; @@ -81,6 +82,25 @@ public void recordSpanFinished( } } + /** + * Records that a user-provided LLM Observability span processor was called. + * + * @param error whether the processor failed + */ + public void recordUserProcessorCalled(boolean error) { + LLMObsMetric metric = + new LLMObsMetric( + METRIC_NAMESPACE, + true, + USER_PROCESSOR_CALLED_METRIC, + COUNT_METRIC_TYPE, + 1L, + Collections.singletonList(error ? ERROR_TRUE : ERROR_FALSE)); + if (!metricsQueue.offer(metric)) { + log.debug("Unable to add telemetry metric {}", USER_PROCESSOR_CALLED_METRIC); + } + } + @Override public void prepareMetrics() { // metrics are added directly via recordSpanFinished; no additional preparation needed diff --git a/internal-api/src/test/groovy/datadog/trace/api/telemetry/LLMObsMetricCollectorTest.groovy b/internal-api/src/test/groovy/datadog/trace/api/telemetry/LLMObsMetricCollectorTest.groovy index e3927c03900..6a7595fce49 100644 --- a/internal-api/src/test/groovy/datadog/trace/api/telemetry/LLMObsMetricCollectorTest.groovy +++ b/internal-api/src/test/groovy/datadog/trace/api/telemetry/LLMObsMetricCollectorTest.groovy @@ -71,5 +71,16 @@ class LLMObsMetricCollectorTest extends DDSpecification { 'has_session_id:0' ].toSet() } -} + def "record and drain user processor called metrics"() { + when: + collector.recordUserProcessorCalled(false) + collector.recordUserProcessorCalled(true) + def metrics = collector.drain() + + then: + metrics.size() == 2 + metrics*.metricName == ['user_processor_called', 'user_processor_called'] + metrics*.tags == [['error:0'], ['error:1']] + } +} From 9266b4581be945f56919492c8d00a501c42631e7 Mon Sep 17 00:00:00 2001 From: Sam Brenner Date: Thu, 30 Jul 2026 12:36:56 -0400 Subject: [PATCH 2/3] llmobsspandataadapter test coverage --- .../ddintake/LLMObsSpanDataAdapterTest.java | 104 ++++++++++++++++++ 1 file changed, 104 insertions(+) create mode 100644 dd-trace-core/src/test/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanDataAdapterTest.java diff --git a/dd-trace-core/src/test/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanDataAdapterTest.java b/dd-trace-core/src/test/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanDataAdapterTest.java new file mode 100644 index 00000000000..5e447a16c67 --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanDataAdapterTest.java @@ -0,0 +1,104 @@ +package datadog.trace.llmobs.writer.ddintake; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import datadog.trace.api.DDTags; +import datadog.trace.api.llmobs.LLMObs; +import datadog.trace.bootstrap.instrumentation.api.Tags; +import datadog.trace.core.CoreSpan; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; + +class LLMObsSpanDataAdapterTest { + private static final String INPUT_TAG = "_ml_obs_tag.input"; + private static final String OUTPUT_TAG = "_ml_obs_tag.output"; + private static final String SPAN_KIND_TAG = "_ml_obs_tag.span.kind"; + + @Test + void convertsAndAppliesEmbeddingDocumentsAndValueOutput() { + CoreSpan span = mock(CoreSpan.class); + when(span.getTag(SPAN_KIND_TAG)).thenReturn(Tags.LLMOBS_EMBEDDING_SPAN_KIND); + when(span.getTag(INPUT_TAG)) + .thenReturn(Collections.singletonList(LLMObs.Document.from("original document"))); + when(span.getTag(OUTPUT_TAG)).thenReturn("original output"); + + LLMObsSpanDataAdapter adapter = new LLMObsSpanDataAdapter(span); + + assertEquals(Tags.LLMOBS_EMBEDDING_SPAN_KIND, adapter.getKind()); + assertEquals("original document", adapter.getInput().get(0).getContent()); + assertEquals("original output", adapter.getOutput().get(0).getContent()); + + adapter.setInput(Collections.singletonList(LLMObs.LLMMessage.from("", "processed document"))); + adapter.setOutput(Collections.singletonList(LLMObs.LLMMessage.from("", "processed output"))); + adapter.apply(adapter); + + ArgumentCaptor inputCaptor = ArgumentCaptor.forClass(Object.class); + verify(span).setTag(eq(INPUT_TAG), inputCaptor.capture()); + List documents = (List) inputCaptor.getValue(); + assertEquals("processed document", ((LLMObs.Document) documents.get(0)).getText()); + verify(span).setTag(OUTPUT_TAG, "processed output"); + } + + @Test + void removesEmptyMessageInputAndOutput() { + CoreSpan span = mock(CoreSpan.class); + Map input = new LinkedHashMap<>(); + input.put("messages", Collections.singletonList(LLMObs.LLMMessage.from("user", "input"))); + when(span.getTag(SPAN_KIND_TAG)).thenReturn(Tags.LLMOBS_LLM_SPAN_KIND); + when(span.getTag(INPUT_TAG)).thenReturn(input); + when(span.getTag(OUTPUT_TAG)) + .thenReturn(Collections.singletonList(LLMObs.LLMMessage.from("assistant", "output"))); + + LLMObsSpanDataAdapter adapter = new LLMObsSpanDataAdapter(span); + adapter.setInput(Collections.emptyList()); + adapter.setOutput(Collections.emptyList()); + adapter.apply(adapter); + + verify(span).removeTag(INPUT_TAG); + verify(span).removeTag(OUTPUT_TAG); + } + + @Test + @SuppressWarnings({"rawtypes", "unchecked"}) + void rejectsInvalidMessagesAndIgnoresMalformedLlmIo() { + CoreSpan span = mock(CoreSpan.class); + List invalidMessages = Collections.singletonList("not a message"); + when(span.getTag(SPAN_KIND_TAG)).thenReturn(Tags.LLMOBS_LLM_SPAN_KIND); + when(span.getTag(INPUT_TAG)).thenReturn(invalidMessages); + when(span.getTag(OUTPUT_TAG)).thenReturn(Collections.singletonMap("messages", invalidMessages)); + + LLMObsSpanDataAdapter adapter = new LLMObsSpanDataAdapter(span); + + assertEquals(Collections.emptyList(), adapter.getInput()); + assertEquals(Collections.emptyList(), adapter.getOutput()); + assertThrows(IllegalArgumentException.class, () -> adapter.setInput(invalidMessages)); + assertThrows(IllegalArgumentException.class, () -> adapter.setOutput(invalidMessages)); + adapter.apply(adapter); + } + + @Test + void readsPublicAndErrorTags() { + CoreSpan span = mock(CoreSpan.class); + when(span.getTag("_ml_obs_tag.custom")).thenReturn(123); + when(span.getError()).thenReturn(1); + when(span.getTag(DDTags.ERROR_TYPE)).thenReturn("java.lang.IllegalStateException"); + + LLMObsSpanDataAdapter adapter = new LLMObsSpanDataAdapter(span); + + assertEquals("unknown", adapter.getKind()); + assertEquals("123", adapter.getTag("custom")); + assertEquals("1", adapter.getTag("error")); + assertEquals("java.lang.IllegalStateException", adapter.getTag("error_type")); + assertNull(adapter.getTag("missing")); + } +} From 113074bd0c8f7f7265412e516fda57e88ced1d21 Mon Sep 17 00:00:00 2001 From: Sam Brenner Date: Thu, 30 Jul 2026 14:51:18 -0400 Subject: [PATCH 3/3] add llmobsinternal to internal-api build.gradle.kts --- internal-api/build.gradle.kts | 2 ++ 1 file changed, 2 insertions(+) diff --git a/internal-api/build.gradle.kts b/internal-api/build.gradle.kts index 189c01ff1fd..d415d4e6fc8 100644 --- a/internal-api/build.gradle.kts +++ b/internal-api/build.gradle.kts @@ -153,6 +153,8 @@ extra["excludedClassesCoverage"] = listOf( "datadog.trace.api.civisibility.CiVisibilityWellKnownTags", "datadog.trace.api.civisibility.InstrumentationBridge", "datadog.trace.api.civisibility.InstrumentationTestBridge", + // Internal cross-module bridge + "datadog.trace.api.llmobs.LLMObsInternal", // POJO "datadog.trace.api.git.GitInfo", "datadog.trace.api.git.GitInfoProvider",