Refactor LLMObs getting current LLMObs parentSpanID

Author: ygree

dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/CompletionServiceInstrumentation.java

@@ -12,6 +12,7 @@
 import com.openai.core.http.StreamResponse;
 import com.openai.models.completions.Completion;
 import com.openai.models.completions.CompletionCreateParams;
+import datadog.context.ContextScope;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -45,15 +46,20 @@ public static class CreateAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static AgentScope enter(
         @Advice.Argument(0) final CompletionCreateParams params,
-        @Advice.FieldValue("clientOptions") ClientOptions clientOptions) {
+        @Advice.FieldValue("clientOptions") ClientOptions clientOptions,
+        @Advice.Local("llmScope") ContextScope llmScope) {
       AgentSpan span = DECORATE.startSpan(clientOptions);
+      // llmScope = LLMObsContext.attach(span.context());
+      // TODO why would we ever need to activate llmScope in this instrumentation if we never expect
+      // inner llmobs spans
       CompletionDecorator.DECORATE.withCompletionCreateParams(span, params);
       return activateSpan(span);
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void exit(
         @Advice.Enter final AgentScope scope,
+        @Advice.Local("llmScope") ContextScope llmScope,
         @Advice.Return(readOnly = false) HttpResponseFor<Completion> response,
         @Advice.Thrown final Throwable err) {
       AgentSpan span = scope.span();
@@ -64,6 +70,7 @@ public static void exit(
             HttpResponseWrapper.wrap(response, span, CompletionDecorator.DECORATE::withCompletion);
       }
       scope.close();
+      // llmScope.close();
     }
   }
 

dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/OpenAiDecorator.java

@@ -2,10 +2,8 @@
 
 import com.openai.core.ClientOptions;
 import com.openai.core.http.Headers;
-import datadog.trace.api.DDSpanId;
 import datadog.trace.api.llmobs.LLMObsContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
-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.UTF8BytesString;
@@ -27,7 +25,9 @@ public class OpenAiDecorator extends ClientDecorator {
   public AgentSpan startSpan(ClientOptions clientOptions) {
     AgentSpan span = AgentTracer.startSpan(INSTRUMENTATION_NAME, SPAN_NAME);
     afterStart(span);
-    withClientOptions(span, clientOptions);
+    span.setTag("openai.api_base", clientOptions.baseUrl());
+    // TODO api_version (either last part of the URL, or api-version param if Azure)
+    // clientOptions.queryParams().values("api-version")
     return span;
   }
 
@@ -61,22 +61,11 @@ protected CharSequence component() {
 
   @Override
   public AgentSpan afterStart(AgentSpan span) {
-    span.setTag("_ml_obs_tag.parent_id", currentLlmParentSpanId());
+    // TODO only if llmobs enabled
+    span.setTag("_ml_obs_tag.parent_id", LLMObsContext.parentSpanId());
     return super.afterStart(span);
   }
 
-  private String currentLlmParentSpanId() {
-    AgentSpanContext parentLlmContext = LLMObsContext.current();
-    if (parentLlmContext == null) {
-      return LLMObsContext.ROOT_SPAN_ID;
-    }
-    long parentLlmSpanId = parentLlmContext.getSpanId();
-    if (parentLlmSpanId == DDSpanId.ZERO) {
-      return LLMObsContext.ROOT_SPAN_ID;
-    }
-    return Long.toString(parentLlmSpanId);
-  }
-
   public void withHttpResponse(AgentSpan span, Headers headers) {
     List<String> values = headers.values("openai-organization");
     if (!values.isEmpty()) {
@@ -116,11 +105,4 @@ private static void setMetricFromHeader(
       // ~
     }
   }
-
-  public void withClientOptions(AgentSpan span, ClientOptions clientOptions) {
-    span.setTag("openai.api_base", clientOptions.baseUrl());
-
-    // TODO api_version (either last part of the URL, or api-version param if Azure)
-    // clientOptions.queryParams().values("api-version")
-  }
 }

internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsContext.java

@@ -3,6 +3,7 @@
 import datadog.context.Context;
 import datadog.context.ContextKey;
 import datadog.context.ContextScope;
+import datadog.trace.api.DDSpanId;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 
 public final class LLMObsContext {
@@ -21,4 +22,16 @@ public static ContextScope attach(AgentSpanContext ctx) {
   public static AgentSpanContext current() {
     return Context.current().get(CONTEXT_KEY);
   }
+
+  public static String parentSpanId() {
+    AgentSpanContext parentLlmContext = current();
+    if (parentLlmContext == null) {
+      return ROOT_SPAN_ID;
+    }
+    long parentLlmSpanId = parentLlmContext.getSpanId();
+    if (parentLlmSpanId == DDSpanId.ZERO) {
+      return ROOT_SPAN_ID;
+    }
+    return Long.toString(parentLlmSpanId);
+  }
 }