diff --git a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/DDLLMObsSpan.java b/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/DDLLMObsSpan.java index 473d118cdc3..3aaac795186 100644 --- a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/DDLLMObsSpan.java +++ b/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/DDLLMObsSpan.java @@ -6,6 +6,7 @@ import datadog.trace.api.DDTraceApiInfo; import datadog.trace.api.DDTraceId; import datadog.trace.api.WellKnownTags; +import datadog.trace.api.llmobs.GenAiApmTags; import datadog.trace.api.llmobs.LLMObs; import datadog.trace.api.llmobs.LLMObsContext; import datadog.trace.api.llmobs.LLMObsSampler; @@ -691,6 +692,12 @@ public void finish() { if (finished) { return; } + // While the span is still mutable, and isolated so a failure cannot cost the event. + try { + GenAiApmTags.apply(span); + } catch (Throwable t) { + LOGGER.debug("failed to set gen_ai APM tags", t); + } span.finish(); if (standaloneApmScope != null) { standaloneApmScope.close(); diff --git a/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanGenAiApmTagsTest.java b/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanGenAiApmTagsTest.java new file mode 100644 index 00000000000..7b4720699ae --- /dev/null +++ b/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanGenAiApmTagsTest.java @@ -0,0 +1,125 @@ +package datadog.trace.llmobs.domain; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; + +import datadog.trace.agent.tooling.TracerInstaller; +import datadog.trace.api.WellKnownTags; +import datadog.trace.api.llmobs.GenAiApmTags; +import datadog.trace.api.llmobs.LLMObsTags; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.Tags; +import datadog.trace.core.CoreTracer; +import java.lang.reflect.Field; +import java.util.HashMap; +import java.util.Map; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +/** Covers the {@code gen_ai.*} attributes a manual LLM Observability span emits at finish. */ +class DDLLMObsSpanGenAiApmTagsTest { + private static final Field SPAN_FIELD; + + private static CoreTracer tracer; + + static { + try { + SPAN_FIELD = DDLLMObsSpan.class.getDeclaredField("span"); + SPAN_FIELD.setAccessible(true); + } catch (ReflectiveOperationException error) { + throw new ExceptionInInitializerError(error); + } + } + + @BeforeAll + static void installTracer() { + tracer = CoreTracer.builder().build(); + TracerInstaller.forceInstallGlobalTracer(tracer); + } + + @AfterAll + static void closeTracer() { + TracerInstaller.forceInstallGlobalTracer(null); + tracer.close(); + } + + @Test + void llmSpanEmitsEveryScalarAndTokenUsage() { + DDLLMObsSpan llm = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "chat", "sess-1"); + llm.setTag(LLMObsTags.MODEL_NAME, "gpt-4"); + llm.setTag(LLMObsTags.MODEL_PROVIDER, "OpenAI"); + Map metrics = new HashMap<>(); + metrics.put("input_tokens", 10); + metrics.put("output_tokens", 20); + metrics.put("total_tokens", 30); + metrics.put("cache_read_input_tokens", 4); + metrics.put("cache_write_input_tokens", 5); + metrics.put("reasoning_output_tokens", 6); + llm.setMetrics(metrics); + llm.finish(); + + AgentSpan span = spanOf(llm); + assertEquals(Tags.LLMOBS_LLM_SPAN_KIND, span.getTag(GenAiApmTags.OPERATION_NAME)); + assertEquals("gpt-4", span.getTag(GenAiApmTags.REQUEST_MODEL)); + assertEquals("openai", span.getTag(GenAiApmTags.PROVIDER_NAME)); + assertEquals("test-ml-app", span.getTag(GenAiApmTags.APPLICATION_NAME)); + assertEquals("sess-1", span.getTag(GenAiApmTags.CONVERSATION_ID)); + assertEquals(10.0, span.getTag(GenAiApmTags.USAGE_INPUT_TOKENS)); + assertEquals(20.0, span.getTag(GenAiApmTags.USAGE_OUTPUT_TOKENS)); + assertEquals(30.0, span.getTag(GenAiApmTags.USAGE_TOTAL_TOKENS)); + assertEquals(4.0, span.getTag(GenAiApmTags.USAGE_CACHE_READ_INPUT_TOKENS)); + assertEquals(5.0, span.getTag(GenAiApmTags.USAGE_CACHE_WRITE_INPUT_TOKENS)); + assertEquals(6.0, span.getTag(GenAiApmTags.USAGE_REASONING_OUTPUT_TOKENS)); + } + + @Test + void modelBackedSpanWithoutModelFallsBackToCustom() { + DDLLMObsSpan embedding = llmObsSpan(Tags.LLMOBS_EMBEDDING_SPAN_KIND, "embed", null); + embedding.finish(); + + AgentSpan span = spanOf(embedding); + assertEquals("custom", span.getTag(GenAiApmTags.REQUEST_MODEL)); + assertEquals("custom", span.getTag(GenAiApmTags.PROVIDER_NAME)); + assertNull(span.getTag(GenAiApmTags.CONVERSATION_ID)); + } + + @Test + void nonModelBackedSpanEmitsNoModelFieldsOrTokenUsage() { + DDLLMObsSpan workflow = llmObsSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "flow", null); + workflow.setMetric("input_tokens", 10); + workflow.finish(); + + AgentSpan span = spanOf(workflow); + assertEquals(Tags.LLMOBS_WORKFLOW_SPAN_KIND, span.getTag(GenAiApmTags.OPERATION_NAME)); + assertNull(span.getTag(GenAiApmTags.REQUEST_MODEL)); + assertNull(span.getTag(GenAiApmTags.PROVIDER_NAME)); + assertNull(span.getTag(GenAiApmTags.USAGE_INPUT_TOKENS)); + } + + @Test + void nonModelBackedSpanKeepsExplicitModelFields() { + DDLLMObsSpan agent = llmObsSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "agent", null); + agent.setTag(LLMObsTags.MODEL_NAME, "gpt-4"); + agent.setTag(LLMObsTags.MODEL_PROVIDER, "OpenAI"); + agent.finish(); + + AgentSpan span = spanOf(agent); + assertEquals("gpt-4", span.getTag(GenAiApmTags.REQUEST_MODEL)); + assertEquals("openai", span.getTag(GenAiApmTags.PROVIDER_NAME)); + } + + private static DDLLMObsSpan llmObsSpan(String kind, String name, String sessionId) { + WellKnownTags tags = + new WellKnownTags("runtime-id", "hostname", "test", "service", "version", "java"); + return new DDLLMObsSpan(kind, name, "test-ml-app", sessionId, "service", tags); + } + + private static AgentSpan spanOf(DDLLMObsSpan llmObsSpan) { + try { + return (AgentSpan) SPAN_FIELD.get(llmObsSpan); + } catch (IllegalAccessException error) { + throw new AssertionError(error); + } + } +} diff --git a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/OpenAiDecorator.java b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/OpenAiDecorator.java index 796490bdda5..a31cf5c95b3 100644 --- a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/OpenAiDecorator.java +++ b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/OpenAiDecorator.java @@ -1,5 +1,6 @@ package datadog.trace.instrumentation.openai_java; +import static datadog.trace.api.llmobs.GenAiApmTags.stringTag; import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; import com.openai.core.ClientOptions; @@ -9,6 +10,7 @@ import datadog.trace.api.DDTags; import datadog.trace.api.DDTraceApiInfo; import datadog.trace.api.WellKnownTags; +import datadog.trace.api.llmobs.GenAiApmTags; import datadog.trace.api.llmobs.LLMObsContext; import datadog.trace.api.llmobs.LLMObsSampler; import datadog.trace.api.telemetry.LLMObsMetricCollector; @@ -16,6 +18,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes; +import datadog.trace.bootstrap.instrumentation.api.Tags; import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString; import datadog.trace.bootstrap.instrumentation.decorator.ClientDecorator; import java.util.List; @@ -37,6 +40,8 @@ public class OpenAiDecorator extends ClientDecorator { private static final String TOKENS_LIMIT_METRIC = METRIC_PREFIX + "tokens.limit"; private static final String TOKENS_REMAINING_METRIC = METRIC_PREFIX + "tokens.remaining"; + private static final String EMBEDDINGS_ENDPOINT = "/v1/embeddings"; + private static final String HEADER_PREFIX = "x-ratelimit-"; private static final String LIMIT_REQUESTS_HEADER = HEADER_PREFIX + "limit-requests"; private static final String REMAINING_REQUESTS_HEADER = HEADER_PREFIX + "remaining-requests"; @@ -179,6 +184,8 @@ protected void doBeforeFinish(@Nonnull Context context) { span.setTag(CommonTags.ERROR, span.isError() ? 1 : 0); span.setTag(CommonTags.ERROR_TYPE, span.getTag(DDTags.ERROR_TYPE)); + GenAiApmTags.apply(span); + Object spanKindTag = span.getTag(CommonTags.SPAN_KIND); if (spanKindTag != null) { String spanKind = spanKindTag.toString(); @@ -186,10 +193,30 @@ protected void doBeforeFinish(@Nonnull Context context) { LLMObsMetricCollector.get() .recordSpanFinished(INTEGRATION, spanKind, isRootSpan, true, span.isError(), false); } + } else if (span != null) { + // Tracing still runs with LLM Observability off, where these remain resolvable. + GenAiApmTags.apply( + span, operationName(span), requestedModel(span), Config.get().getLlmObsMlApp()); } super.doBeforeFinish(context); } + private static String operationName(AgentSpan span) { + String endpoint = stringTag(span, CommonTags.OPENAI_REQUEST_ENDPOINT); + if (endpoint == null) { + return null; + } + return EMBEDDINGS_ENDPOINT.equals(endpoint) + ? Tags.LLMOBS_EMBEDDING_SPAN_KIND + : Tags.LLMOBS_LLM_SPAN_KIND; + } + + /** The response model resolves aliases the request used, so it wins. */ + private static String requestedModel(AgentSpan span) { + String model = stringTag(span, CommonTags.OPENAI_RESPONSE_MODEL); + return model != null ? model : stringTag(span, CommonTags.OPENAI_REQUEST_MODEL); + } + public void withHttpResponse(AgentSpan span, Headers headers) { if (!llmObsEnabled) { return; diff --git a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/ChatCompletionServiceTest.groovy b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/ChatCompletionServiceTest.groovy index 0a8553e7930..37ba7178988 100644 --- a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/ChatCompletionServiceTest.groovy +++ b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/ChatCompletionServiceTest.groovy @@ -401,6 +401,17 @@ class ChatCompletionServiceTest extends OpenAiTest { "_ml_obs_tag.sampling_decision" "1" "_ml_obs_tag.sample_rate" "1" "_ml_obs_tag.ml_app" String + "gen_ai.operation.name" "llm" + "gen_ai.request.model" String + "gen_ai.provider.name" "openai" + "gen_ai.application.name" String + if (!isStreaming) { + // streamed completions missing usage data + "gen_ai.usage.input_tokens" Double + "gen_ai.usage.output_tokens" Double + "gen_ai.usage.total_tokens" Double + "gen_ai.usage.cache_read_input_tokens" Double + } "_ml_obs_tag.service" String "$CommonTags.DDTRACE_VERSION" String if (expectToolDefinitions) { diff --git a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/CompletionServiceTest.groovy b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/CompletionServiceTest.groovy index 50529ee02a6..8723549ab00 100644 --- a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/CompletionServiceTest.groovy +++ b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/CompletionServiceTest.groovy @@ -183,6 +183,13 @@ class CompletionServiceTest extends OpenAiTest { "_ml_obs_tag.sampling_decision" "1" "_ml_obs_tag.sample_rate" "1" "_ml_obs_tag.ml_app" String + "gen_ai.operation.name" "llm" + "gen_ai.request.model" String + "gen_ai.provider.name" "openai" + "gen_ai.application.name" String + "gen_ai.usage.input_tokens" Double + "gen_ai.usage.output_tokens" Double + "gen_ai.usage.total_tokens" Double "_ml_obs_tag.service" String "$CommonTags.DDTRACE_VERSION" String "$CommonTags.SOURCE" "integration" diff --git a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/EmbeddingServiceTest.groovy b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/EmbeddingServiceTest.groovy index 41e341284f3..0d718ee5e3d 100644 --- a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/EmbeddingServiceTest.groovy +++ b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/EmbeddingServiceTest.groovy @@ -76,6 +76,12 @@ class EmbeddingServiceTest extends OpenAiTest { "_ml_obs_tag.sampling_decision" "1" "_ml_obs_tag.sample_rate" "1" "_ml_obs_tag.ml_app" String + "gen_ai.operation.name" "embedding" + "gen_ai.request.model" "text-embedding-ada-002-v2" + "gen_ai.provider.name" "openai" + "gen_ai.application.name" String + "gen_ai.usage.input_tokens" Double + "gen_ai.usage.total_tokens" Double "_ml_obs_tag.service" String "$CommonTags.DDTRACE_VERSION" String "$CommonTags.SOURCE" "integration" diff --git a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/ResponseServiceTest.groovy b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/ResponseServiceTest.groovy index b22389a4d41..c455bffdd96 100644 --- a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/ResponseServiceTest.groovy +++ b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/ResponseServiceTest.groovy @@ -478,6 +478,15 @@ class ResponseServiceTest extends OpenAiTest { "_ml_obs_tag.sampling_decision" "1" "_ml_obs_tag.sample_rate" "1" "_ml_obs_tag.ml_app" String + "gen_ai.operation.name" "llm" + "gen_ai.request.model" String + "gen_ai.provider.name" "openai" + "gen_ai.application.name" String + "gen_ai.usage.input_tokens" Double + "gen_ai.usage.output_tokens" Double + "gen_ai.usage.total_tokens" Double + "gen_ai.usage.reasoning_output_tokens" Double + "gen_ai.usage.cache_read_input_tokens" Double "$CommonTags.INTEGRATION" "openai" "_ml_obs_tag.service" String "$CommonTags.DDTRACE_VERSION" String diff --git a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/java/datadog/trace/instrumentation/openai_java/LlmObsContextPropagationForkedTest.java b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/java/datadog/trace/instrumentation/openai_java/LlmObsContextPropagationForkedTest.java index 66d4f6aa285..16830ca0d2a 100644 --- a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/java/datadog/trace/instrumentation/openai_java/LlmObsContextPropagationForkedTest.java +++ b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/java/datadog/trace/instrumentation/openai_java/LlmObsContextPropagationForkedTest.java @@ -9,6 +9,8 @@ import com.openai.credential.BearerTokenCredential; import com.openai.models.ChatModel; import com.openai.models.chat.completions.ChatCompletionCreateParams; +import com.openai.models.embeddings.EmbeddingCreateParams; +import com.openai.models.embeddings.EmbeddingModel; import com.sun.net.httpserver.HttpServer; import datadog.context.ContextScope; import datadog.trace.agent.test.AbstractInstrumentationTest; @@ -76,6 +78,13 @@ protected static ChatCompletionCreateParams buildMinimalChatParams() { .build(); } + protected static EmbeddingCreateParams buildMinimalEmbeddingParams() { + return EmbeddingCreateParams.builder() + .model(EmbeddingModel.TEXT_EMBEDDING_ADA_002) + .input("") + .build(); + } + protected static DDSpan findSpanByOperationName(List> traces, String operationName) { return traces.stream() .flatMap(List::stream) @@ -318,3 +327,56 @@ void parentlessOpenAiRequestSpanIsDroppedAtZeroSampleRate() throws Exception { assertEquals("0", openAiSpan.getTag("_ml_obs_tag.sample_rate")); } } + +/** + * Verifies the gen_ai.* attributes an openai.request span carries with LLM Observability disabled: + * operation, model, provider and application, but never token usage or conversation id. + */ +@WithConfig(key = "llmobs.enabled", value = "false") +class LlmObsDisabledForkedTest extends AbstractLlmObsOpenAiForkedTest { + + @Test + void chatCompletionEmitsTheGenAiAttributesAvailableWithoutLlmObs() throws Exception { + try { + openAiClient.chat().completions().create(buildMinimalChatParams()); + } catch (Exception ignored) { + // The mock server returns no body, so the SDK may throw while parsing the response. + } + + writer.waitForTraces(1); + DDSpan openAiSpan = findSpanByOperationName(writer, "openai.request"); + assertNotNull(openAiSpan, "openai.request span should have been created"); + + assertEquals("llm", openAiSpan.getTag("gen_ai.operation.name")); + // The mock returns no body, so the request model stands in for the absent response model. + assertEquals( + openAiSpan.getTag("openai.request.model"), openAiSpan.getTag("gen_ai.request.model")); + assertEquals("openai", openAiSpan.getTag("gen_ai.provider.name")); + assertNotNull(openAiSpan.getTag("gen_ai.application.name")); + + assertNull(openAiSpan.getTag("gen_ai.conversation.id")); + assertNull(openAiSpan.getTag("gen_ai.usage.input_tokens")); + assertNull(openAiSpan.getTag("gen_ai.usage.output_tokens")); + assertNull(openAiSpan.getTag("gen_ai.usage.total_tokens")); + + assertNull(openAiSpan.getTag("_ml_obs_tag.span.kind")); + } + + @Test + void embeddingMapsToTheEmbeddingOperation() throws Exception { + try { + openAiClient.embeddings().create(buildMinimalEmbeddingParams()); + } catch (Exception ignored) { + // The mock server returns no body, so the SDK may throw while parsing the response. + } + + writer.waitForTraces(1); + DDSpan openAiSpan = findSpanByOperationName(writer, "openai.request"); + assertNotNull(openAiSpan, "openai.request span should have been created"); + + assertEquals("embedding", openAiSpan.getTag("gen_ai.operation.name")); + assertEquals( + openAiSpan.getTag("openai.request.model"), openAiSpan.getTag("gen_ai.request.model")); + assertEquals("openai", openAiSpan.getTag("gen_ai.provider.name")); + } +} diff --git a/internal-api/src/main/java/datadog/trace/api/llmobs/GenAiApmTags.java b/internal-api/src/main/java/datadog/trace/api/llmobs/GenAiApmTags.java new file mode 100644 index 00000000000..1324ec98a7a --- /dev/null +++ b/internal-api/src/main/java/datadog/trace/api/llmobs/GenAiApmTags.java @@ -0,0 +1,117 @@ +package datadog.trace.api.llmobs; + +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.Tags; +import java.util.Locale; + +/** + * Emits the scalar {@code gen_ai.*} attributes of an LLM Observability span onto the APM span, so + * model, provider, application, conversation and token usage are searchable in APM. Message bodies + * stay off the APM span and keep coming from the LLM Observability track. + */ +public final class GenAiApmTags { + public static final String OPERATION_NAME = "gen_ai.operation.name"; + public static final String REQUEST_MODEL = "gen_ai.request.model"; + public static final String PROVIDER_NAME = "gen_ai.provider.name"; + public static final String APPLICATION_NAME = "gen_ai.application.name"; + public static final String CONVERSATION_ID = "gen_ai.conversation.id"; + + public static final String USAGE_INPUT_TOKENS = "gen_ai.usage.input_tokens"; + public static final String USAGE_OUTPUT_TOKENS = "gen_ai.usage.output_tokens"; + public static final String USAGE_TOTAL_TOKENS = "gen_ai.usage.total_tokens"; + public static final String USAGE_CACHE_READ_INPUT_TOKENS = "gen_ai.usage.cache_read_input_tokens"; + public static final String USAGE_CACHE_WRITE_INPUT_TOKENS = + "gen_ai.usage.cache_write_input_tokens"; + public static final String USAGE_REASONING_OUTPUT_TOKENS = "gen_ai.usage.reasoning_output_tokens"; + + private static final String LLMOBS_TAG_PREFIX = "_ml_obs_tag."; + private static final String LLMOBS_METRIC_PREFIX = "_ml_obs_metric."; + + private static final String SPAN_KIND_TAG = LLMOBS_TAG_PREFIX + Tags.SPAN_KIND; + private static final String MODEL_NAME_TAG = LLMOBS_TAG_PREFIX + LLMObsTags.MODEL_NAME; + private static final String MODEL_PROVIDER_TAG = LLMOBS_TAG_PREFIX + LLMObsTags.MODEL_PROVIDER; + private static final String ML_APP_TAG = LLMOBS_TAG_PREFIX + LLMObsTags.ML_APP; + private static final String SESSION_ID_TAG = LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID; + + /** Matches the fallback the LLM Observability event uses. */ + private static final String DEFAULT_MODEL = "custom"; + + /** LLM Observability metric name paired with the {@code gen_ai.usage.*} key it maps to. */ + private static final String[][] TOKEN_METRICS = { + {LLMOBS_METRIC_PREFIX + "input_tokens", USAGE_INPUT_TOKENS}, + {LLMOBS_METRIC_PREFIX + "output_tokens", USAGE_OUTPUT_TOKENS}, + {LLMOBS_METRIC_PREFIX + "total_tokens", USAGE_TOTAL_TOKENS}, + {LLMOBS_METRIC_PREFIX + "cache_read_input_tokens", USAGE_CACHE_READ_INPUT_TOKENS}, + {LLMOBS_METRIC_PREFIX + "cache_write_input_tokens", USAGE_CACHE_WRITE_INPUT_TOKENS}, + {LLMOBS_METRIC_PREFIX + "reasoning_output_tokens", USAGE_REASONING_OUTPUT_TOKENS}, + }; + + public static void apply(AgentSpan span) { + apply(span, null, null, null); + } + + /** + * Writes the attributes onto a span that is not yet finished, reading them from its {@code + * _ml_obs_tag.} / {@code _ml_obs_metric.} tags. The arguments take precedence over those tags and + * cover instrumentation that traces with LLM Observability disabled, where they are not all set. + * No-op for a span with no resolvable operation. + */ + public static void apply(AgentSpan span, String operationName, String modelName, String mlApp) { + if (span == null) { + return; + } + String operation = firstNonEmpty(operationName, stringTag(span, SPAN_KIND_TAG)); + if (operation == null) { + return; + } + span.setTag(OPERATION_NAME, operation); + + boolean modelBacked = + Tags.LLMOBS_LLM_SPAN_KIND.equals(operation) + || Tags.LLMOBS_EMBEDDING_SPAN_KIND.equals(operation); + + String model = firstNonEmpty(modelName, stringTag(span, MODEL_NAME_TAG)); + if (model != null || modelBacked) { + span.setTag(REQUEST_MODEL, model == null ? DEFAULT_MODEL : model); + } + String provider = stringTag(span, MODEL_PROVIDER_TAG); + if (provider != null || modelBacked) { + span.setTag( + PROVIDER_NAME, (provider == null ? DEFAULT_MODEL : provider).toLowerCase(Locale.ROOT)); + } + String application = firstNonEmpty(mlApp, stringTag(span, ML_APP_TAG)); + if (application != null) { + span.setTag(APPLICATION_NAME, application); + } + String sessionId = stringTag(span, SESSION_ID_TAG); + if (sessionId != null) { + span.setTag(CONVERSATION_ID, sessionId); + } + + // Other kinds carry unrelated metrics that a gen_ai.usage.* key would misrepresent. + if (modelBacked) { + for (String[] metric : TOKEN_METRICS) { + Object value = span.getTag(metric[0]); + if (value instanceof Number) { + span.setMetric(metric[1], ((Number) value).doubleValue()); + } + } + } + } + + /** The value of {@code key} as a non-empty string, or null. */ + public static String stringTag(AgentSpan span, String key) { + Object value = span.getTag(key); + if (value == null) { + return null; + } + String string = String.valueOf(value); + return string.isEmpty() ? null : string; + } + + private static String firstNonEmpty(String preferred, String fallback) { + return preferred == null || preferred.isEmpty() ? fallback : preferred; + } + + private GenAiApmTags() {} +} diff --git a/internal-api/src/test/java/datadog/trace/api/llmobs/GenAiApmTagsTest.java b/internal-api/src/test/java/datadog/trace/api/llmobs/GenAiApmTagsTest.java new file mode 100644 index 00000000000..98783a6badd --- /dev/null +++ b/internal-api/src/test/java/datadog/trace/api/llmobs/GenAiApmTagsTest.java @@ -0,0 +1,207 @@ +package datadog.trace.api.llmobs; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyDouble; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.Tags; +import java.util.HashMap; +import java.util.Map; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class GenAiApmTagsTest { + private static final String LLMOBS_TAG_PREFIX = "_ml_obs_tag."; + private static final String LLMOBS_METRIC_PREFIX = "_ml_obs_metric."; + + private final Map tags = new HashMap<>(); + private AgentSpan span; + + @BeforeEach + void setUp() { + span = mock(AgentSpan.class); + when(span.getTag(anyString())).thenAnswer(call -> tags.get(call.getArgument(0))); + when(span.setTag(anyString(), anyString())) + .thenAnswer( + call -> { + tags.put(call.getArgument(0), call.getArgument(1)); + return span; + }); + when(span.setMetric(any(CharSequence.class), anyDouble())) + .thenAnswer( + call -> { + tags.put(call.getArgument(0).toString(), call.getArgument(1)); + return span; + }); + } + + @Test + void llmSpanEmitsEveryScalarAndTokenUsage() { + llmObsTag(Tags.SPAN_KIND, Tags.LLMOBS_LLM_SPAN_KIND); + llmObsTag(LLMObsTags.MODEL_NAME, "gpt-4"); + llmObsTag(LLMObsTags.MODEL_PROVIDER, "OpenAI"); + llmObsTag(LLMObsTags.ML_APP, "my-app"); + llmObsTag(LLMObsTags.SESSION_ID, "sess-1"); + llmObsMetric("input_tokens", 10); + llmObsMetric("output_tokens", 20); + llmObsMetric("total_tokens", 30); + llmObsMetric("cache_read_input_tokens", 4); + llmObsMetric("cache_write_input_tokens", 5); + llmObsMetric("reasoning_output_tokens", 6); + + GenAiApmTags.apply(span); + + assertEquals(Tags.LLMOBS_LLM_SPAN_KIND, tags.get(GenAiApmTags.OPERATION_NAME)); + assertEquals("gpt-4", tags.get(GenAiApmTags.REQUEST_MODEL)); + assertEquals("openai", tags.get(GenAiApmTags.PROVIDER_NAME)); + assertEquals("my-app", tags.get(GenAiApmTags.APPLICATION_NAME)); + assertEquals("sess-1", tags.get(GenAiApmTags.CONVERSATION_ID)); + assertEquals(10.0, tags.get(GenAiApmTags.USAGE_INPUT_TOKENS)); + assertEquals(20.0, tags.get(GenAiApmTags.USAGE_OUTPUT_TOKENS)); + assertEquals(30.0, tags.get(GenAiApmTags.USAGE_TOTAL_TOKENS)); + assertEquals(4.0, tags.get(GenAiApmTags.USAGE_CACHE_READ_INPUT_TOKENS)); + assertEquals(5.0, tags.get(GenAiApmTags.USAGE_CACHE_WRITE_INPUT_TOKENS)); + assertEquals(6.0, tags.get(GenAiApmTags.USAGE_REASONING_OUTPUT_TOKENS)); + } + + @Test + void modelBackedSpanWithoutModelFallsBackToCustom() { + llmObsTag(Tags.SPAN_KIND, Tags.LLMOBS_EMBEDDING_SPAN_KIND); + + GenAiApmTags.apply(span); + + assertEquals("custom", tags.get(GenAiApmTags.REQUEST_MODEL)); + assertEquals("custom", tags.get(GenAiApmTags.PROVIDER_NAME)); + } + + @Test + void emptyModelValuesAreTreatedAsAbsent() { + llmObsTag(Tags.SPAN_KIND, Tags.LLMOBS_LLM_SPAN_KIND); + llmObsTag(LLMObsTags.MODEL_NAME, ""); + llmObsTag(LLMObsTags.MODEL_PROVIDER, ""); + llmObsTag(LLMObsTags.ML_APP, ""); + llmObsTag(LLMObsTags.SESSION_ID, ""); + + GenAiApmTags.apply(span); + + assertEquals("custom", tags.get(GenAiApmTags.REQUEST_MODEL)); + assertEquals("custom", tags.get(GenAiApmTags.PROVIDER_NAME)); + assertNull(tags.get(GenAiApmTags.APPLICATION_NAME)); + assertNull(tags.get(GenAiApmTags.CONVERSATION_ID)); + } + + @Test + void nonModelBackedSpanKeepsExplicitModelFieldsWithoutFallbacks() { + llmObsTag(Tags.SPAN_KIND, Tags.LLMOBS_AGENT_SPAN_KIND); + llmObsTag(LLMObsTags.MODEL_NAME, "gpt-4"); + llmObsTag(LLMObsTags.MODEL_PROVIDER, "OpenAI"); + + GenAiApmTags.apply(span); + + assertEquals(Tags.LLMOBS_AGENT_SPAN_KIND, tags.get(GenAiApmTags.OPERATION_NAME)); + assertEquals("gpt-4", tags.get(GenAiApmTags.REQUEST_MODEL)); + assertEquals("openai", tags.get(GenAiApmTags.PROVIDER_NAME)); + } + + @Test + void nonModelBackedSpanWithoutModelEmitsNoModelFields() { + llmObsTag(Tags.SPAN_KIND, Tags.LLMOBS_WORKFLOW_SPAN_KIND); + + GenAiApmTags.apply(span); + + assertFalse(tags.containsKey(GenAiApmTags.REQUEST_MODEL)); + assertFalse(tags.containsKey(GenAiApmTags.PROVIDER_NAME)); + } + + @Test + void nonModelBackedSpanDropsTokenUsage() { + llmObsTag(Tags.SPAN_KIND, Tags.LLMOBS_TASK_SPAN_KIND); + llmObsMetric("input_tokens", 10); + + GenAiApmTags.apply(span); + + assertFalse(tags.containsKey(GenAiApmTags.USAGE_INPUT_TOKENS)); + } + + @Test + void nonNumericTokenMetricIsSkipped() { + llmObsTag(Tags.SPAN_KIND, Tags.LLMOBS_LLM_SPAN_KIND); + tags.put(LLMOBS_METRIC_PREFIX + "input_tokens", "not-a-number"); + + GenAiApmTags.apply(span); + + assertFalse(tags.containsKey(GenAiApmTags.USAGE_INPUT_TOKENS)); + } + + @Test + void spanWithoutLlmObsKindEmitsNothing() { + llmObsTag(LLMObsTags.MODEL_NAME, "gpt-4"); + + GenAiApmTags.apply(span); + + assertTrue(tags.keySet().stream().noneMatch(key -> key.startsWith("gen_ai."))); + } + + @Test + void nullSpanIsANoOp() { + GenAiApmTags.apply(null); + GenAiApmTags.apply(null, Tags.LLMOBS_LLM_SPAN_KIND, "gpt-4", "app"); + } + + @Test + void argumentsStandInForMissingLlmObsTags() { + llmObsTag(LLMObsTags.MODEL_PROVIDER, "OpenAI"); + + GenAiApmTags.apply(span, Tags.LLMOBS_LLM_SPAN_KIND, "gpt-4", "my-app"); + + assertEquals(Tags.LLMOBS_LLM_SPAN_KIND, tags.get(GenAiApmTags.OPERATION_NAME)); + assertEquals("gpt-4", tags.get(GenAiApmTags.REQUEST_MODEL)); + assertEquals("openai", tags.get(GenAiApmTags.PROVIDER_NAME)); + assertEquals("my-app", tags.get(GenAiApmTags.APPLICATION_NAME)); + assertFalse(tags.containsKey(GenAiApmTags.CONVERSATION_ID)); + assertFalse(tags.containsKey(GenAiApmTags.USAGE_INPUT_TOKENS)); + } + + @Test + void argumentsTakePrecedenceOverLlmObsTags() { + llmObsTag(Tags.SPAN_KIND, Tags.LLMOBS_WORKFLOW_SPAN_KIND); + llmObsTag(LLMObsTags.MODEL_NAME, "gpt-3.5"); + llmObsTag(LLMObsTags.ML_APP, "tagged-app"); + + GenAiApmTags.apply(span, Tags.LLMOBS_LLM_SPAN_KIND, "gpt-4", "my-app"); + + assertEquals(Tags.LLMOBS_LLM_SPAN_KIND, tags.get(GenAiApmTags.OPERATION_NAME)); + assertEquals("gpt-4", tags.get(GenAiApmTags.REQUEST_MODEL)); + assertEquals("my-app", tags.get(GenAiApmTags.APPLICATION_NAME)); + } + + @Test + void modelBackedArgumentsFallBackToCustom() { + GenAiApmTags.apply(span, Tags.LLMOBS_EMBEDDING_SPAN_KIND, null, "app"); + + assertEquals("custom", tags.get(GenAiApmTags.REQUEST_MODEL)); + assertEquals("custom", tags.get(GenAiApmTags.PROVIDER_NAME)); + } + + @Test + void spanWithNoResolvableOperationEmitsNothing() { + GenAiApmTags.apply(span, null, "gpt-4", "my-app"); + + assertTrue(tags.keySet().stream().noneMatch(key -> key.startsWith("gen_ai."))); + } + + private void llmObsTag(String key, String value) { + tags.put(LLMOBS_TAG_PREFIX + key, value); + } + + private void llmObsMetric(String key, Number value) { + tags.put(LLMOBS_METRIC_PREFIX + key, value); + } +}