set unified service tags on llm obs span start

Author: gary-huang

dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java

@@ -2,6 +2,7 @@
 
 import datadog.communication.ddagent.SharedCommunicationObjects;
 import datadog.trace.api.Config;
+import datadog.trace.api.WellKnownTags;
 import datadog.trace.api.llmobs.LLMObs;
 import datadog.trace.api.llmobs.LLMObsSpan;
 import datadog.trace.api.llmobs.LLMObsTags;
@@ -31,12 +32,9 @@ public static void start(Instrumentation inst, SharedCommunicationObjects sco) {
 
     sco.createRemaining(config);
 
-    LLMObsInternal.setLLMObsSpanFactory(
-        new LLMObsManualSpanFactory(config.getLlmObsMlApp(), config.getServiceName()));
-
     String mlApp = config.getLlmObsMlApp();
-    LLMObsInternal.setLLMObsSpanFactory(
-        new LLMObsManualSpanFactory(mlApp, config.getServiceName()));
+    WellKnownTags wellKnownTags = config.getWellKnownTags();
+    LLMObsInternal.setLLMObsSpanFactory(new LLMObsManualSpanFactory(mlApp, wellKnownTags));
 
     LLMObsInternal.setLLMObsEvalProcessor(new LLMObsCustomEvalProcessor(mlApp, sco, config));
   }
@@ -129,12 +127,14 @@ public void SubmitEvaluation(
 
   private static class LLMObsManualSpanFactory implements LLMObs.LLMObsSpanFactory {
 
-    private final String serviceName;
     private final String defaultMLApp;
+    private final String serviceName;
+    private final WellKnownTags wellKnownTags;
 
-    public LLMObsManualSpanFactory(String defaultMLApp, String serviceName) {
+    public LLMObsManualSpanFactory(String defaultMLApp, WellKnownTags wellKnownTags) {
       this.defaultMLApp = defaultMLApp;
-      this.serviceName = serviceName;
+      this.serviceName = wellKnownTags.getService().toString();
+      this.wellKnownTags = wellKnownTags;
     }
 
     @Override
@@ -147,7 +147,7 @@ public LLMObsSpan startLLMSpan(
 
       DDLLMObsSpan span =
           new DDLLMObsSpan(
-              Tags.LLMOBS_LLM_SPAN_KIND, spanName, getMLApp(mlApp), sessionId, serviceName);
+              Tags.LLMOBS_LLM_SPAN_KIND, spanName, getMLApp(mlApp), sessionId, serviceName, wellKnownTags);
 
       if (modelName == null || modelName.isEmpty()) {
         modelName = CUSTOM_MODEL_VAL;
@@ -165,28 +165,28 @@ public LLMObsSpan startLLMSpan(
     public LLMObsSpan startAgentSpan(
         String spanName, @Nullable String mlApp, @Nullable String sessionId) {
       return new DDLLMObsSpan(
-          Tags.LLMOBS_AGENT_SPAN_KIND, spanName, getMLApp(mlApp), sessionId, serviceName);
+          Tags.LLMOBS_AGENT_SPAN_KIND, spanName, getMLApp(mlApp), sessionId, serviceName, wellKnownTags);
     }
 
     @Override
     public LLMObsSpan startToolSpan(
         String spanName, @Nullable String mlApp, @Nullable String sessionId) {
       return new DDLLMObsSpan(
-          Tags.LLMOBS_TOOL_SPAN_KIND, spanName, getMLApp(mlApp), sessionId, serviceName);
+          Tags.LLMOBS_TOOL_SPAN_KIND, spanName, getMLApp(mlApp), sessionId, serviceName, wellKnownTags);
     }
 
     @Override
     public LLMObsSpan startTaskSpan(
         String spanName, @Nullable String mlApp, @Nullable String sessionId) {
       return new DDLLMObsSpan(
-          Tags.LLMOBS_TASK_SPAN_KIND, spanName, getMLApp(mlApp), sessionId, serviceName);
+          Tags.LLMOBS_TASK_SPAN_KIND, spanName, getMLApp(mlApp), sessionId, serviceName, wellKnownTags);
     }
 
     @Override
     public LLMObsSpan startWorkflowSpan(
         String spanName, @Nullable String mlApp, @Nullable String sessionId) {
       return new DDLLMObsSpan(
-          Tags.LLMOBS_WORKFLOW_SPAN_KIND, spanName, getMLApp(mlApp), sessionId, serviceName);
+          Tags.LLMOBS_WORKFLOW_SPAN_KIND, spanName, getMLApp(mlApp), sessionId, serviceName, wellKnownTags);
     }
 
     @Override
@@ -201,7 +201,7 @@ public LLMObsSpan startEmbeddingSpan(
       }
       DDLLMObsSpan embeddingSpan =
           new DDLLMObsSpan(
-              Tags.LLMOBS_EMBEDDING_SPAN_KIND, spanName, getMLApp(mlApp), sessionId, serviceName);
+              Tags.LLMOBS_EMBEDDING_SPAN_KIND, spanName, getMLApp(mlApp), sessionId, serviceName, wellKnownTags);
       embeddingSpan.setTag(LLMObsTags.MODEL_PROVIDER, modelProvider);
       embeddingSpan.setTag(LLMObsTags.MODEL_NAME, modelName);
       return embeddingSpan;
@@ -210,7 +210,7 @@ public LLMObsSpan startEmbeddingSpan(
     public LLMObsSpan startRetrievalSpan(
         String spanName, @Nullable String mlApp, @Nullable String sessionId) {
       return new DDLLMObsSpan(
-          Tags.LLMOBS_RETRIEVAL_SPAN_KIND, spanName, getMLApp(mlApp), sessionId, serviceName);
+          Tags.LLMOBS_RETRIEVAL_SPAN_KIND, spanName, getMLApp(mlApp), sessionId, serviceName, wellKnownTags);
     }
 
     private String getMLApp(String mlApp) {

dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/DDLLMObsSpan.java

@@ -3,6 +3,7 @@
 import datadog.context.ContextScope;
 import datadog.trace.api.DDSpanTypes;
 import datadog.trace.api.DDTraceId;
+import datadog.trace.api.WellKnownTags;
 import datadog.trace.api.llmobs.LLMObs;
 import datadog.trace.api.llmobs.LLMObsSpan;
 import datadog.trace.api.llmobs.LLMObsTags;
@@ -34,6 +35,10 @@ public class DDLLMObsSpan implements LLMObsSpan {
   private static final String METADATA = LLMOBS_TAG_PREFIX + LLMObsTags.METADATA;
   private static final String PARENT_ID_TAG_INTERNAL = "parent_id";
 
+  private static final String SERVICE = LLMOBS_TAG_PREFIX + "service";
+  private static final String VERSION = LLMOBS_TAG_PREFIX + "version";
+  private static final String ENV = LLMOBS_TAG_PREFIX + "env";
+
   private static final String LLM_OBS_INSTRUMENTATION_NAME = "llmobs";
 
   private static final Logger LOGGER = LoggerFactory.getLogger(DDLLMObsSpan.class);
@@ -49,7 +54,8 @@ public DDLLMObsSpan(
       String spanName,
       @Nonnull String mlApp,
       String sessionId,
-      @Nonnull String serviceName) {
+      @Nonnull String serviceName,
+      WellKnownTags wellKnownTags) {
 
     if (null == spanName || spanName.isEmpty()) {
       spanName = kind;
@@ -62,6 +68,12 @@ public DDLLMObsSpan(
             .withSpanType(DDSpanTypes.LLMOBS);
 
     this.span = spanBuilder.start();
+
+    // set UST
+    this.span.setTag(ENV, wellKnownTags.getEnv());
+    this.span.setTag(SERVICE, wellKnownTags.getService());
+    this.span.setTag(VERSION, wellKnownTags.getVersion());
+
     this.span.setTag(SPAN_KIND, kind);
     this.spanKind = kind;
     this.span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.ML_APP, mlApp);