diff --git a/communication/src/main/java/datadog/communication/serialization/Mapper.java b/communication/src/main/java/datadog/communication/serialization/Mapper.java index f0c07681ccf..71c38641abc 100644 --- a/communication/src/main/java/datadog/communication/serialization/Mapper.java +++ b/communication/src/main/java/datadog/communication/serialization/Mapper.java @@ -4,5 +4,9 @@ public interface Mapper { void map(T data, Writable packer); + default void map(T data, Writable packer, boolean retry) { + map(data, packer); + } + default void reset() {} } diff --git a/communication/src/main/java/datadog/communication/serialization/msgpack/MsgPackWriter.java b/communication/src/main/java/datadog/communication/serialization/msgpack/MsgPackWriter.java index 8c96226dc0f..4fc9e8f967a 100644 --- a/communication/src/main/java/datadog/communication/serialization/msgpack/MsgPackWriter.java +++ b/communication/src/main/java/datadog/communication/serialization/msgpack/MsgPackWriter.java @@ -91,7 +91,7 @@ public boolean format(T message, Mapper mapper) { if (buffer.flush()) { try { mapper.reset(); - mapper.map(message, this); + mapper.map(message, this, true); buffer.mark(); return true; } catch (BufferOverflowException fatal) { 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 25f3ff0a8ac..dff642d1752 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..5abfad20940 --- /dev/null +++ b/dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObsSpanData.java @@ -0,0 +1,61 @@ +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. + * + *

Other input metadata, including prompt tracking information, is preserved. + * + * @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 40bac293d60..1f80ebd2e81 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.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; @@ -28,23 +29,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..f6aaa3d457f --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanDataAdapter.java @@ -0,0 +1,257 @@ +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 LLMObs.LLMMessage[] NO_MESSAGES = new LLMObs.LLMMessage[0]; + 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 final LLMObs.LLMMessage[] initialInput; + private final LLMObs.LLMMessage[] initialOutput; + private List input; + private List output; + private boolean inputModified; + private boolean outputModified; + + 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); + initialInput = snapshot(input); + initialOutput = snapshot(output); + } + + @Override + public String getKind() { + return kind; + } + + @Override + public List getInput() { + return input; + } + + @Override + public void setInput(List input) { + this.input = new ArrayList<>(requireNonNull(input, "input")); + inputModified = true; + } + + @Override + public List getOutput() { + return output; + } + + @Override + public void setOutput(List output) { + this.output = new ArrayList<>(requireNonNull(output, "output")); + outputModified = true; + } + + @Override + public String getTag(String key) { + Object value = span.getTag(LLMOBS_TAG_PREFIX + key); + if (value == null && "error".equals(key)) { + int error = span.getError(); + value = error == 0 ? null : error; + } + if (value == null && "error_type".equals(key)) { + value = span.getTag(DDTags.ERROR_TYPE); + } + return value == null ? null : String.valueOf(value); + } + + void apply(LLMObsSpanData processedSpan) { + if (processedSpan != this || inputModified || wasModified(input, initialInput)) { + applyIO( + span, + INPUT_TAG, + originalInput, + inputType, + new ArrayList<>(requireNonNull(processedSpan.getInput(), "processed input"))); + } + if (processedSpan != this || outputModified || wasModified(output, initialOutput)) { + applyIO( + span, + OUTPUT_TAG, + originalOutput, + outputType, + new ArrayList<>(requireNonNull(processedSpan.getOutput(), "processed output"))); + } + } + + private static IOType ioType(String kind, Object value, boolean input) { + if (value == null) { + if (Tags.LLMOBS_LLM_SPAN_KIND.equals(kind)) { + return IOType.MESSAGES; + } + if ((input && Tags.LLMOBS_EMBEDDING_SPAN_KIND.equals(kind)) + || (!input && Tags.LLMOBS_RETRIEVAL_SPAN_KIND.equals(kind))) { + return IOType.DOCUMENTS; + } + return IOType.VALUE; + } + Object unwrapped = unwrapMessages(value); + if (Tags.LLMOBS_LLM_SPAN_KIND.equals(kind)) { + if (input && value instanceof Map && !((Map) value).containsKey("messages")) { + return IOType.MESSAGES; + } + return unwrapped instanceof List && allMessages((List) unwrapped) + ? IOType.MESSAGES + : IOType.NONE; + } + if (((input && Tags.LLMOBS_EMBEDDING_SPAN_KIND.equals(kind)) + || (!input && Tags.LLMOBS_RETRIEVAL_SPAN_KIND.equals(kind))) + && value instanceof List + && allDocuments((List) value)) { + return IOType.DOCUMENTS; + } + return IOType.VALUE; + } + + @SuppressWarnings("unchecked") + private static List asMessages(Object value, IOType type) { + if (value == null || type == IOType.NONE) { + return new ArrayList<>(); + } + if (type == IOType.MESSAGES) { + Object messages = unwrapMessages(value); + return messages == null + ? new ArrayList<>() + : new ArrayList<>((List) messages); + } + 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 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 LLMObs.LLMMessage[] snapshot(List messages) { + return messages.isEmpty() + ? NO_MESSAGES + : messages.toArray(new LLMObs.LLMMessage[messages.size()]); + } + + private static boolean wasModified( + List messages, LLMObs.LLMMessage[] initialMessages) { + if (messages.size() != initialMessages.length) { + return true; + } + for (int i = 0; i < initialMessages.length; i++) { + if (messages.get(i) != initialMessages[i]) { + return true; + } + } + return false; + } + + private static void applyIO( + CoreSpan span, + String tag, + Object originalValue, + IOType type, + List messages) { + if (type == IOType.NONE) { + return; + } + if (messages.isEmpty()) { + if (originalValue == null) { + return; + } + 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 155f4a59406..0ff43d7389d 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; @@ -20,6 +24,7 @@ import java.nio.ByteBuffer; import java.nio.channels.WritableByteChannel; import java.nio.charset.StandardCharsets; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; @@ -34,6 +39,8 @@ public class LLMObsSpanMapper implements RemoteMapper { + private static final boolean[] NO_DROPPED_SPANS = new boolean[0]; + // Well known tags for LLM obs will be prefixed with _ml_obs_(tags|metrics). // Prefix for tags private static final String LLMOBS_TAG_PREFIX = "_ml_obs_tag."; @@ -96,6 +103,7 @@ public class LLMObsSpanMapper implements RemoteMapper { private final int size; private final ByteBuffer header; + private boolean[] pendingDroppedSpans; private int spansWritten; public LLMObsSpanMapper() { @@ -120,6 +128,15 @@ private LLMObsSpanMapper(int size) { @Override public void map(List> trace, Writable writable) { + this.map(trace, writable, false); + } + + @Override + public void map(List> trace, Writable writable, boolean retry) { + if (!retry) { + pendingDroppedSpans = null; + } + List> llmobsSpans = trace.stream().filter(LLMObsSpanMapper::isLLMObsSpan).collect(Collectors.toList()); @@ -128,6 +145,8 @@ public void map(List> trace, Writable writable) { return; } + llmobsSpans = processSpans(llmobsSpans, retry); + for (CoreSpan span : llmobsSpans) { // 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: @@ -193,6 +212,62 @@ public void map(List> trace, Writable writable) { // 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(); + pendingDroppedSpans = null; + } + + private List> processSpans( + List> spans, boolean retry) { + if (retry && pendingDroppedSpans != null) { + boolean[] droppedSpans = pendingDroppedSpans; + pendingDroppedSpans = null; + if (droppedSpans.length == 0) { + return spans; + } + + List> processedSpans = new ArrayList<>(spans.size()); + for (int i = 0; i < spans.size(); i++) { + if (!droppedSpans[i]) { + processedSpans.add(spans.get(i)); + } + } + return processedSpans; + } + + LLMObsSpanProcessor processor = LLMObsInternal.getSpanProcessor(); + if (processor == null) { + return spans; + } + + pendingDroppedSpans = NO_DROPPED_SPANS; + List> processedSpans = new ArrayList<>(spans.size()); + for (int i = 0; i < spans.size(); i++) { + CoreSpan span = spans.get(i); + boolean processorError = false; + try { + LLMObsSpanDataAdapter adapter = new LLMObsSpanDataAdapter(span); + LLMObsSpanData result = processor.process(adapter); + if (result != null) { + adapter.apply(result); + processedSpans.add(span); + } else { + markDroppedSpan(i, spans.size()); + } + } catch (RuntimeException error) { + processorError = true; + markDroppedSpan(i, spans.size()); + LOGGER.warn("Error in LLM Observability span processor, dropping span", error); + } finally { + LLMObsMetricCollector.get().recordUserProcessorCalled(processorError); + } + } + return processedSpans; + } + + private void markDroppedSpan(int index, int spanCount) { + if (pendingDroppedSpans.length == 0) { + pendingDroppedSpans = new boolean[spanCount]; + } + pendingDroppedSpans[index] = true; } private CharSequence llmObsSpanName(CoreSpan span) { @@ -371,13 +446,9 @@ public void accept(Metadata metadata) { val.getClass().getName()); continue; } - } else if (spanKind.equals(Tags.LLMOBS_EMBEDDING_SPAN_KIND) && key.equals(INPUT)) { - if (!(val instanceof List)) { - LOGGER.warn( - "unexpectedly found incorrect type for embedding span input {}, expecting list", - val.getClass().getName()); - continue; - } + } else if (((spanKind.equals(Tags.LLMOBS_EMBEDDING_SPAN_KIND) && key.equals(INPUT)) + || (spanKind.equals(Tags.LLMOBS_RETRIEVAL_SPAN_KIND) && key.equals(OUTPUT))) + && val instanceof List) { writable.writeString(key, null); writable.startMap(1); List documents = (List) val; 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..0f6f8102c56 --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanDataAdapterTest.java @@ -0,0 +1,182 @@ +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.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.clearInvocations; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +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 convertsAndAppliesRetrievalOutputDocuments() { + CoreSpan span = mock(CoreSpan.class); + when(span.getTag(SPAN_KIND_TAG)).thenReturn(Tags.LLMOBS_RETRIEVAL_SPAN_KIND); + when(span.getTag(OUTPUT_TAG)) + .thenReturn(Collections.singletonList(LLMObs.Document.from("original document"))); + + LLMObsSpanDataAdapter adapter = new LLMObsSpanDataAdapter(span); + + assertEquals("original document", adapter.getOutput().get(0).getContent()); + + adapter.setOutput(Collections.singletonList(LLMObs.LLMMessage.from("", "processed document"))); + adapter.apply(adapter); + + ArgumentCaptor outputCaptor = ArgumentCaptor.forClass(Object.class); + verify(span).setTag(eq(OUTPUT_TAG), outputCaptor.capture()); + List documents = (List) outputCaptor.getValue(); + assertEquals("processed document", ((LLMObs.Document) documents.get(0)).getText()); + } + + @Test + void appliesInPlaceClearToMessageInputAndOutput() { + 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.getInput().clear(); + adapter.getOutput().clear(); + adapter.apply(adapter); + + verify(span).removeTag(INPUT_TAG); + verify(span).removeTag(OUTPUT_TAG); + } + + @Test + void ignoresMalformedLlmIo() { + 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()); + adapter.apply(adapter); + } + + @Test + void addsMissingInputAndOutput() { + CoreSpan span = mock(CoreSpan.class); + when(span.getTag(SPAN_KIND_TAG)).thenReturn(Tags.LLMOBS_LLM_SPAN_KIND); + + LLMObsSpanDataAdapter adapter = new LLMObsSpanDataAdapter(span); + List input = + Collections.singletonList(LLMObs.LLMMessage.from("user", "input")); + List output = + Collections.singletonList(LLMObs.LLMMessage.from("assistant", "output")); + adapter.setInput(input); + adapter.setOutput(output); + adapter.apply(adapter); + + verify(span).setTag(INPUT_TAG, input); + verify(span).setTag(OUTPUT_TAG, output); + } + + @Test + void addsMessagesToPromptOnlyInput() { + CoreSpan span = mock(CoreSpan.class); + Map prompt = Collections.singletonMap("id", "prompt-id"); + when(span.getTag(SPAN_KIND_TAG)).thenReturn(Tags.LLMOBS_LLM_SPAN_KIND); + when(span.getTag(INPUT_TAG)).thenReturn(Collections.singletonMap("prompt", prompt)); + + LLMObsSpanDataAdapter adapter = new LLMObsSpanDataAdapter(span); + adapter.setInput(Collections.singletonList(LLMObs.LLMMessage.from("user", "input"))); + adapter.apply(adapter); + + ArgumentCaptor inputCaptor = ArgumentCaptor.forClass(Object.class); + verify(span).setTag(eq(INPUT_TAG), inputCaptor.capture()); + Map input = (Map) inputCaptor.getValue(); + assertEquals(prompt, input.get("prompt")); + assertEquals(adapter.getInput(), input.get("messages")); + } + + @Test + void preservesIoWhenProcessorDoesNotModifyIt() { + CoreSpan span = mock(CoreSpan.class); + Map input = Collections.singletonMap("key", "input"); + List output = Collections.singletonList("output"); + when(span.getTag(INPUT_TAG)).thenReturn(input); + when(span.getTag(OUTPUT_TAG)).thenReturn(output); + + LLMObsSpanDataAdapter adapter = new LLMObsSpanDataAdapter(span); + clearInvocations(span); + adapter.apply(adapter); + + verifyNoInteractions(span); + } + + @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")); + } + + @Test + void returnsNullForAbsentErrorTag() { + CoreSpan span = mock(CoreSpan.class); + + LLMObsSpanDataAdapter adapter = new LLMObsSpanDataAdapter(span); + + assertNull(adapter.getTag("error")); + } +} 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 1da6bf4018f..97fcadb3a14 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; @@ -384,6 +387,270 @@ 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 testLLMObsSpanProcessorAddsMissingInputAndOutput() throws Exception { + LLMObs.registerProcessor( + span -> { + span.setInput(Collections.singletonList(LLMObs.LLMMessage.from("user", "added input"))); + span.setOutput( + Collections.singletonList(LLMObs.LLMMessage.from("assistant", "added output"))); + return span; + }); + + try { + CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build(); + AgentSpan llmSpan = newLlmObsSpan(tracer, "processed", false); + + List> spans = + serialize(Collections.singletonList((DDSpan) llmSpan), new LLMObsSpanMapper()); + Map meta = (Map) spans.get(0).get("meta"); + Map input = (Map) meta.get("input"); + Map output = (Map) meta.get("output"); + List> inputMessages = (List>) input.get("messages"); + List> outputMessages = (List>) output.get("messages"); + + assertEquals("added input", inputMessages.get(0).get("content")); + assertEquals("added output", outputMessages.get(0).get("content")); + tracer.close(); + } finally { + LLMObs.deregisterProcessor(); + } + } + + @Test + void testLLMObsSpanProcessorModifiesRetrievalOutputDocuments() throws Exception { + LLMObs.registerProcessor( + span -> { + span.setOutput( + Collections.singletonList(LLMObs.LLMMessage.from("", "processed document"))); + return span; + }); + + try { + CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build(); + AgentSpan retrievalSpan = + tracer + .buildSpan("datadog", "retrieval") + .withTag("_ml_obs_tag.span.kind", Tags.LLMOBS_RETRIEVAL_SPAN_KIND) + .withTag( + "_ml_obs_tag.output", + Collections.singletonList(LLMObs.Document.from("original document"))) + .start(); + retrievalSpan.setSpanType(InternalSpanTypes.LLMOBS); + retrievalSpan.finish(); + + List> spans = + serialize(Collections.singletonList((DDSpan) retrievalSpan), new LLMObsSpanMapper()); + Map meta = (Map) spans.get(0).get("meta"); + Map output = (Map) meta.get("output"); + List> documents = (List>) output.get("documents"); + + assertEquals("processed document", documents.get(0).get("text")); + tracer.close(); + } finally { + LLMObs.deregisterProcessor(); + } + } + + @Test + void testLLMObsSpanMapperPreservesStringRetrievalOutput() throws Exception { + CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build(); + AgentSpan retrievalSpan = + tracer + .buildSpan("datadog", "retrieval") + .withTag("_ml_obs_tag.span.kind", Tags.LLMOBS_RETRIEVAL_SPAN_KIND) + .withTag("_ml_obs_tag.output", "retrieval output") + .start(); + retrievalSpan.setSpanType(InternalSpanTypes.LLMOBS); + retrievalSpan.finish(); + + List> spans = + serialize(Collections.singletonList((DDSpan) retrievalSpan), new LLMObsSpanMapper()); + Map meta = (Map) spans.get(0).get("meta"); + Map output = (Map) meta.get("output"); + + assertEquals("retrieval output", output.get("value")); + tracer.close(); + } + + @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 "true".equals(span.getTag("drop")) ? null : span; + }); + + try { + CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build(); + AgentSpan first = newLlmObsSpan(tracer, "first", false); + AgentSpan dropped = newLlmObsSpan(tracer, "dropped", true); + AgentSpan retained = newLlmObsSpan(tracer, "retained", false); + String largeInput = String.join("", Collections.nCopies(600, "x")); + first.setTag( + "_ml_obs_tag.input", + Collections.singletonList(LLMObs.LLMMessage.from("user", largeInput))); + dropped.setTag( + "_ml_obs_tag.input", + Collections.singletonList(LLMObs.LLMMessage.from("user", largeInput))); + retained.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(Arrays.asList((DDSpan) dropped, (DDSpan) retained), mapper)); + assertEquals(1, sink.accepts); + assertEquals(3, calls.get()); + assertEquals( + 3, + 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( @@ -410,10 +677,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/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", 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/java/datadog/trace/api/telemetry/LLMObsMetricCollectorUserProcessorTest.java b/internal-api/src/test/java/datadog/trace/api/telemetry/LLMObsMetricCollectorUserProcessorTest.java new file mode 100644 index 00000000000..a5c2daff6a3 --- /dev/null +++ b/internal-api/src/test/java/datadog/trace/api/telemetry/LLMObsMetricCollectorUserProcessorTest.java @@ -0,0 +1,34 @@ +package datadog.trace.api.telemetry; + +import static java.util.Collections.singletonList; +import static org.junit.jupiter.api.Assertions.assertEquals; + +import java.util.ArrayList; +import java.util.List; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class LLMObsMetricCollectorUserProcessorTest { + private final LLMObsMetricCollector collector = LLMObsMetricCollector.get(); + + @BeforeEach + void clearMetrics() { + collector.drain(); + } + + @Test + void recordsUserProcessorCalledMetrics() { + collector.recordUserProcessorCalled(false); + collector.recordUserProcessorCalled(true); + + List metrics = new ArrayList<>(collector.drain()); + + assertEquals(2, metrics.size()); + assertEquals(LLMObsMetricCollector.USER_PROCESSOR_CALLED_METRIC, metrics.get(0).metricName); + assertEquals(LLMObsMetricCollector.COUNT_METRIC_TYPE, metrics.get(0).type); + assertEquals(singletonList("error:0"), metrics.get(0).tags); + assertEquals(LLMObsMetricCollector.USER_PROCESSOR_CALLED_METRIC, metrics.get(1).metricName); + assertEquals(LLMObsMetricCollector.COUNT_METRIC_TYPE, metrics.get(1).type); + assertEquals(singletonList("error:1"), metrics.get(1).tags); + } +}