Add telemetry for LLMObs manual api span finish calls (#10616)

llmobs: add telemetry for manual api span finish calls

Add a test for a non-root span

Add has_session_id tag to the llmobs "span.finished" metric

Co-authored-by: yury.gribkov <yury.gribkov@datadoghq.com>

Author: ygree

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

@@ -8,6 +8,7 @@
 import datadog.trace.api.llmobs.LLMObsContext;
 import datadog.trace.api.llmobs.LLMObsSpan;
 import datadog.trace.api.llmobs.LLMObsTags;
+import datadog.trace.api.telemetry.LLMObsMetricCollector;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
@@ -47,6 +48,7 @@ public class DDLLMObsSpan implements LLMObsSpan {
   private final AgentSpan span;
   private final String spanKind;
   private final ContextScope scope;
+  private final boolean hasSessionId;
 
   private boolean finished = false;
 
@@ -68,51 +70,52 @@ public DDLLMObsSpan(
             .withServiceName(serviceName)
             .withSpanType(DDSpanTypes.LLMOBS);
 
-    this.span = spanBuilder.start();
+    span = spanBuilder.start();
 
     // set UST (unified service tags, env, service, version)
-    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);
-    if (sessionId != null && !sessionId.isEmpty()) {
-      this.span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID, sessionId);
+    span.setTag(ENV, wellKnownTags.getEnv());
+    span.setTag(SERVICE, wellKnownTags.getService());
+    span.setTag(VERSION, wellKnownTags.getVersion());
+
+    span.setTag(SPAN_KIND, kind);
+    spanKind = kind;
+    span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.ML_APP, mlApp);
+    this.hasSessionId = sessionId != null && !sessionId.isEmpty();
+    if (this.hasSessionId) {
+      span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID, sessionId);
     }
 
     AgentSpanContext parent = LLMObsContext.current();
     String parentSpanID = LLMObsContext.ROOT_SPAN_ID;
     if (null != parent) {
-      if (parent.getTraceId() != this.span.getTraceId()) {
+      if (parent.getTraceId() != span.getTraceId()) {
         LOGGER.error(
             "trace ID mismatch, retrieved parent from context trace_id={}, span_id={}, started span trace_id={}, span_id={}",
             parent.getTraceId(),
             parent.getSpanId(),
-            this.span.getTraceId(),
-            this.span.getSpanId());
+            span.getTraceId(),
+            span.getSpanId());
       } else {
         parentSpanID = String.valueOf(parent.getSpanId());
       }
     }
-    this.span.setTag(LLMOBS_TAG_PREFIX + PARENT_ID_TAG_INTERNAL, parentSpanID);
-    this.scope = LLMObsContext.attach(this.span.context());
+    span.setTag(LLMOBS_TAG_PREFIX + PARENT_ID_TAG_INTERNAL, parentSpanID);
+    scope = LLMObsContext.attach(span.context());
   }
 
   @Override
   public String toString() {
     return super.toString()
         + ", trace_id="
-        + this.span.context().getTraceId()
+        + span.context().getTraceId()
         + ", span_id="
-        + this.span.context().getSpanId()
+        + span.context().getSpanId()
         + ", ml_app="
-        + this.span.getTag(LLMObsTags.ML_APP)
+        + span.getTag(LLMObsTags.ML_APP)
         + ", service="
-        + this.span.getServiceName()
+        + span.getServiceName()
         + ", span_kind="
-        + this.span.getTag(SPAN_KIND);
+        + span.getTag(SPAN_KIND);
   }
 
   @Override
@@ -121,10 +124,10 @@ public void annotateIO(List<LLMObs.LLMMessage> inputData, List<LLMObs.LLMMessage
       return;
     }
     if (inputData != null && !inputData.isEmpty()) {
-      this.span.setTag(INPUT, inputData);
+      span.setTag(INPUT, inputData);
     }
     if (outputData != null && !outputData.isEmpty()) {
-      this.span.setTag(OUTPUT, outputData);
+      span.setTag(OUTPUT, outputData);
     }
   }
 
@@ -135,24 +138,24 @@ public void annotateIO(String inputData, String outputData) {
     }
     boolean wrongSpanKind = false;
     if (inputData != null && !inputData.isEmpty()) {
-      if (Tags.LLMOBS_LLM_SPAN_KIND.equals(this.spanKind)) {
+      if (Tags.LLMOBS_LLM_SPAN_KIND.equals(spanKind)) {
         wrongSpanKind = true;
         annotateIO(
             Collections.singletonList(LLMObs.LLMMessage.from(LLM_MESSAGE_UNKNOWN_ROLE, inputData)),
             null);
       } else {
-        this.span.setTag(INPUT, inputData);
+        span.setTag(INPUT, inputData);
       }
     }
     if (outputData != null && !outputData.isEmpty()) {
-      if (Tags.LLMOBS_LLM_SPAN_KIND.equals(this.spanKind)) {
+      if (Tags.LLMOBS_LLM_SPAN_KIND.equals(spanKind)) {
         wrongSpanKind = true;
         annotateIO(
             null,
             Collections.singletonList(
                 LLMObs.LLMMessage.from(LLM_MESSAGE_UNKNOWN_ROLE, outputData)));
       } else {
-        this.span.setTag(OUTPUT, outputData);
+        span.setTag(OUTPUT, outputData);
       }
     }
     if (wrongSpanKind) {
@@ -168,7 +171,7 @@ public void setMetadata(Map<String, Object> metadata) {
     }
     Object value = span.getTag(METADATA);
     if (value == null) {
-      this.span.setTag(METADATA, new HashMap<>(metadata));
+      span.setTag(METADATA, new HashMap<>(metadata));
       return;
     }
 
@@ -178,7 +181,7 @@ public void setMetadata(Map<String, Object> metadata) {
       LOGGER.debug(
           "unexpected instance type for metadata {}, overwriting for now",
           value.getClass().getName());
-      this.span.setTag(METADATA, new HashMap<>(metadata));
+      span.setTag(METADATA, new HashMap<>(metadata));
     }
   }
 
@@ -188,7 +191,7 @@ public void setMetrics(Map<String, Number> metrics) {
       return;
     }
     for (Map.Entry<String, Number> entry : metrics.entrySet()) {
-      this.span.setMetric(LLMOBS_METRIC_PREFIX + entry.getKey(), entry.getValue().doubleValue());
+      span.setMetric(LLMOBS_METRIC_PREFIX + entry.getKey(), entry.getValue().doubleValue());
     }
   }
 
@@ -197,23 +200,23 @@ public void setMetric(CharSequence key, int value) {
     if (finished) {
       return;
     }
-    this.span.setMetric(LLMOBS_METRIC_PREFIX + key, value);
+    span.setMetric(LLMOBS_METRIC_PREFIX + key, value);
   }
 
   @Override
   public void setMetric(CharSequence key, long value) {
     if (finished) {
       return;
     }
-    this.span.setMetric(LLMOBS_METRIC_PREFIX + key, value);
+    span.setMetric(LLMOBS_METRIC_PREFIX + key, value);
   }
 
   @Override
   public void setMetric(CharSequence key, double value) {
     if (finished) {
       return;
     }
-    this.span.setMetric(LLMOBS_METRIC_PREFIX + key, value);
+    span.setMetric(LLMOBS_METRIC_PREFIX + key, value);
   }
 
   @Override
@@ -223,7 +226,7 @@ public void setTags(Map<String, Object> tags) {
     }
     if (tags != null && !tags.isEmpty()) {
       for (Map.Entry<String, Object> entry : tags.entrySet()) {
-        this.span.setTag(LLMOBS_TAG_PREFIX + entry.getKey(), entry.getValue());
+        span.setTag(LLMOBS_TAG_PREFIX + entry.getKey(), entry.getValue());
       }
     }
   }
@@ -233,47 +236,47 @@ public void setTag(String key, String value) {
     if (finished) {
       return;
     }
-    this.span.setTag(LLMOBS_TAG_PREFIX + key, value);
+    span.setTag(LLMOBS_TAG_PREFIX + key, value);
   }
 
   @Override
   public void setTag(String key, boolean value) {
     if (finished) {
       return;
     }
-    this.span.setTag(LLMOBS_TAG_PREFIX + key, value);
+    span.setTag(LLMOBS_TAG_PREFIX + key, value);
   }
 
   @Override
   public void setTag(String key, int value) {
     if (finished) {
       return;
     }
-    this.span.setTag(LLMOBS_TAG_PREFIX + key, value);
+    span.setTag(LLMOBS_TAG_PREFIX + key, value);
   }
 
   @Override
   public void setTag(String key, long value) {
     if (finished) {
       return;
     }
-    this.span.setTag(LLMOBS_TAG_PREFIX + key, value);
+    span.setTag(LLMOBS_TAG_PREFIX + key, value);
   }
 
   @Override
   public void setTag(String key, double value) {
     if (finished) {
       return;
     }
-    this.span.setTag(LLMOBS_TAG_PREFIX + key, value);
+    span.setTag(LLMOBS_TAG_PREFIX + key, value);
   }
 
   @Override
   public void setError(boolean error) {
     if (finished) {
       return;
     }
-    this.span.setError(error);
+    span.setError(error);
   }
 
   @Override
@@ -284,8 +287,8 @@ public void setErrorMessage(String errorMessage) {
     if (errorMessage == null || errorMessage.isEmpty()) {
       return;
     }
-    this.span.setError(true);
-    this.span.setErrorMessage(errorMessage);
+    span.setError(true);
+    span.setErrorMessage(errorMessage);
   }
 
   @Override
@@ -296,27 +299,36 @@ public void addThrowable(Throwable throwable) {
     if (throwable == null) {
       return;
     }
-    this.span.setError(true);
-    this.span.addThrowable(throwable);
+    span.setError(true);
+    span.addThrowable(throwable);
   }
 
   @Override
   public void finish() {
     if (finished) {
       return;
     }
-    this.span.finish();
-    this.scope.close();
-    this.finished = true;
+    span.finish();
+    scope.close();
+    finished = true;
+    boolean isRootSpan = span.getLocalRootSpan() == span;
+    LLMObsMetricCollector.get()
+        .recordSpanFinished(
+            LLM_OBS_INSTRUMENTATION_NAME,
+            spanKind,
+            isRootSpan,
+            false,
+            span.isError(),
+            hasSessionId);
   }
 
   @Override
   public DDTraceId getTraceId() {
-    return this.span.getTraceId();
+    return span.getTraceId();
   }
 
   @Override
   public long getSpanId() {
-    return this.span.getSpanId();
+    return span.getSpanId();
   }
 }