added chatcompletion llmobs spans

Author: Nayeem Kamal

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

@@ -12,7 +12,7 @@
 import java.lang.instrument.Instrumentation;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
-import org.jetbrains.annotations.Nullable;
+import javax.annotation.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -22,7 +22,7 @@ public class LLMObsSystem {
 
   private static final String CUSTOM_MODEL_VAL = "custom";
 
-  public static void start(Instrumentation inst, SharedCommunicationObjects sco) {
+  public static void start(@Nullable Instrumentation inst, SharedCommunicationObjects sco) {
     Config config = Config.get();
     if (!config.isLlmObsEnabled()) {
       LOGGER.debug("LLM Observability is disabled");

dd-java-agent/instrumentation/openai/build.gradle

@@ -8,8 +8,6 @@ muzzle {
 }
 
 // Instrumentation is currently in preview. Additional testing will be implemented before enabling this by default.
-//minimumBranchCoverage = 0.0
-//minimumInstructionCoverage = 0.0
 
 apply from: "$rootDir/gradle/java.gradle"
 

dd-java-agent/instrumentation/openai/src/main/java/datadog/trace/instrumentation/openai/ChatCompletionServiceInstrumentation.java

@@ -5,12 +5,16 @@
 import static net.bytebuddy.matcher.ElementMatchers.*;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 
+import com.openai.models.chat.completions.ChatCompletion;
 import com.openai.models.chat.completions.ChatCompletionCreateParams;
 import com.openai.services.blocking.CompletionService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.InstrumenterConfig;
+import datadog.trace.api.llmobs.LLMObsSpan;
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
@@ -23,6 +27,11 @@ public ChatCompletionServiceInstrumentation() {
     super("openai-client");
   }
 
+  @Override
+  protected boolean defaultEnabled() {
+    return InstrumenterConfig.get().isIntegrationEnabled(Collections.singleton("openai"), false);
+  }
+
   @Override
   public String[] helperClassNames() {
     return new String[] {
@@ -34,7 +43,7 @@ public String[] helperClassNames() {
   public Map<String, String> contextStore() {
     Map<String, String> contextStores = new HashMap<>(1);
     contextStores.put(
-        "com.openai.services.blocking.CompletionService", OpenAIClientInfo.class.getName());
+        "com.openai.services.blocking.ChatCompletionService", OpenAIClientInfo.class.getName());
     return contextStores;
   }
 
@@ -47,7 +56,7 @@ public void methodAdvice(MethodTransformer transformer) {
             .and(
                 takesArgument(
                     0, named("com.openai.models.completions.ChatCompletionCreateParams"))),
-        getClass().getName() + "$CompletionServiceAdvice");
+        getClass().getName() + "$ChatCompletionServiceAdvice");
   }
 
   @Override
@@ -60,24 +69,37 @@ public ElementMatcher<TypeDescription> hierarchyMatcher() {
     return implementsInterface(named(hierarchyMarkerType()));
   }
 
-  public static class CompletionServiceAdvice {
+  public static class ChatCompletionServiceAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static AgentScope methodEnter(
+    public static Map<String, Object> methodEnter(
         @Advice.Argument(0) final ChatCompletionCreateParams params) {
-
-      return OpenAIClientDecorator.DECORATE.startChatCompletionSpan(params);
+      Map<String, Object> spans = new HashMap<>();
+      spans.put(
+          "datadog.trace.api.llmobs.LLMObsSpan",
+          OpenAIClientDecorator.DECORATE.startLLMObsChatCompletionSpan(params));
+      spans.put(
+          "datadog.trace.bootstrap.instrumentation.api.AgentScope",
+          OpenAIClientDecorator.DECORATE.startChatCompletionSpan(params));
+      return spans;
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void methodExit(
-        @Advice.Enter final AgentScope span,
+        @Advice.Enter final Map<String, Object> spans,
         @Advice.This final CompletionService completionService,
         @Advice.Return final Object result,
         @Advice.Thrown final Throwable throwable) {
       OpenAIClientInfo info =
           InstrumentationContext.get(CompletionService.class, OpenAIClientInfo.class)
               .get(completionService);
-      OpenAIClientDecorator.DECORATE.finishSpan(span, result, throwable);
+      OpenAIClientDecorator.DECORATE.finishLLMObsChatCompletionSpan(
+          (LLMObsSpan) spans.get("datadog.trace.api.llmobs.LLMObsSpan"),
+          (ChatCompletion) result,
+          throwable);
+      OpenAIClientDecorator.DECORATE.finishSpan(
+          (AgentScope) spans.get("datadog.trace.bootstrap.instrumentation.api.AgentScope"),
+          result,
+          throwable);
     }
   }
 }

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

@@ -9,8 +9,10 @@
 import com.openai.services.blocking.CompletionService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
@@ -23,6 +25,11 @@ public CompletionServiceInstrumentation() {
     super("openai-client");
   }
 
+  @Override
+  protected boolean defaultEnabled() {
+    return InstrumenterConfig.get().isIntegrationEnabled(Collections.singleton("openai"), false);
+  }
+
   @Override
   public String[] helperClassNames() {
     return new String[] {

dd-java-agent/instrumentation/openai/src/main/java/datadog/trace/instrumentation/openai/EmbeddingServiceInstrumentation.java

@@ -10,7 +10,9 @@
 import com.openai.models.embeddings.EmbeddingCreateParams;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import java.util.Collections;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -23,6 +25,11 @@ public EmbeddingServiceInstrumentation() {
     super("openai", "openai-java");
   }
 
+  @Override
+  protected boolean defaultEnabled() {
+    return InstrumenterConfig.get().isIntegrationEnabled(Collections.singleton("openai"), false);
+  }
+
   @Override
   public String[] helperClassNames() {
     return new String[] {

dd-java-agent/instrumentation/openai/src/main/java/datadog/trace/instrumentation/openai/OpenAIClientDecorator.java

@@ -10,6 +10,9 @@
 import com.openai.models.embeddings.CreateEmbeddingResponse;
 import com.openai.models.embeddings.Embedding;
 import com.openai.models.embeddings.EmbeddingCreateParams;
+import datadog.trace.api.Config;
+import datadog.trace.api.llmobs.LLMObs;
+import datadog.trace.api.llmobs.LLMObsSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes;
@@ -20,6 +23,8 @@
 import java.util.Optional;
 
 public class OpenAIClientDecorator extends ClientDecorator {
+  private static final String mlApp = Config.get().getLlmObsMlApp();
+  private static final String mlProvider = "openai";
   private static final String COMPONENT_NAME = "openai";
   private static final UTF8BytesString OPENAI_REQUEST = UTF8BytesString.create("openai.request");
 
@@ -42,14 +47,14 @@ protected CharSequence spanType() {
 
   @Override
   protected String service() {
-    return null; // Use default service name
+    return null;
   }
 
   public AgentScope startChatCompletionSpan(ChatCompletionCreateParams params) {
     AgentSpan span = startSpan(OPENAI_REQUEST);
     span.setTag("openai.request.endpoint", "/chat/completions");
     span.setResourceName("chat.completions.create");
-    span.setTag("openai.provider", "openai");
+    span.setTag("openai.request.provider", "openai");
     extractChatCompletionRequestData(span, params);
     afterStart(span);
     return activateSpan(span);
@@ -59,7 +64,7 @@ public AgentScope startLLMChatCompletionSpan(ChatCompletionCreateParams params)
     AgentSpan span = startSpan(OPENAI_REQUEST);
     span.setTag("openai.request.endpoint", "/chat/completions");
     span.setResourceName("chat.completions.create");
-    span.setTag("openai.provider", "openai");
+    span.setTag("openai.request.provider", "openai");
     extractChatCompletionRequestData(span, params);
     afterStart(span);
     return activateSpan(span);
@@ -88,6 +93,115 @@ public AgentScope startEmbeddingSpan(EmbeddingCreateParams params) {
     return activateSpan(span);
   }
 
+  public LLMObsSpan startLLMObsChatCompletionSpan(ChatCompletionCreateParams params) {
+    String modelName = params.model().toString();
+
+    LLMObsSpan llmObsSpan =
+        LLMObs.startLLMSpan("chat_completion", modelName, mlProvider, mlApp, null);
+
+    // Extract and set input data from chat completion params
+    // May need to be reformatted
+    StringBuilder inputData = new StringBuilder();
+    List<ChatCompletionMessageParam> messages = params.messages();
+    for (int i = 0; i < messages.size(); i++) {
+      ChatCompletionMessageParam messageParam = messages.get(i);
+      if (i > 0) {
+        inputData.append(" | ");
+      }
+
+      if (messageParam.isUser()) {
+        inputData.append("User: ").append(messageParam.asUser().content());
+      } else if (messageParam.isAssistant()) {
+        inputData.append("Assistant: ").append(messageParam.asAssistant().content());
+      } else if (messageParam.isDeveloper()) {
+        inputData.append("Developer: ").append(messageParam.asDeveloper().content());
+      } else if (messageParam.isSystem()) {
+        inputData.append("System: ").append(messageParam.asSystem().content());
+      } else if (messageParam.isTool()) {
+        inputData.append("Tool: ").append(messageParam.asTool().content());
+      }
+    }
+
+    if (inputData.length() > 0) {
+      llmObsSpan.annotateIO(inputData.toString(), null); // No output yet, will be set in response
+    }
+
+    java.util.Map<String, Object> metadata = new java.util.HashMap<>();
+    metadata.put("endpoint", "/chat/completions");
+    metadata.put("provider", "openai");
+    metadata.put("model", modelName);
+
+    params.maxTokens().ifPresent(tokens -> metadata.put("max_tokens", tokens));
+    params.temperature().ifPresent(temp -> metadata.put("temperature", temp));
+
+    llmObsSpan.setMetadata(metadata);
+
+    return llmObsSpan;
+  }
+
+  public void finishLLMObsChatCompletionSpan(
+      LLMObsSpan llmObsSpan, ChatCompletion response, Throwable throwable) {
+    try {
+      if (throwable != null) {
+        // Set error information
+        java.util.Map<String, Object> errorMetadata = new java.util.HashMap<>();
+        errorMetadata.put("error.type", throwable.getClass().getSimpleName());
+        errorMetadata.put("error.message", throwable.getMessage());
+        llmObsSpan.setMetadata(errorMetadata);
+      } else if (response != null) {
+        StringBuilder outputData = new StringBuilder();
+        List<ChatCompletion.Choice> choices = response.choices();
+
+        for (int i = 0; i < choices.size(); i++) {
+          ChatCompletion.Choice choice = choices.get(i);
+          ChatCompletionMessage message = choice.message();
+
+          if (i > 0) {
+            outputData.append(" | ");
+          }
+
+          // Extract content
+          Optional<String> content = message.content();
+          content.ifPresent(s -> outputData.append("Assistant: ").append(s));
+
+          // Extract tool calls if present
+          Optional<List<ChatCompletionMessageToolCall>> toolCalls = message.toolCalls();
+          if (toolCalls.isPresent() && !toolCalls.get().isEmpty()) {
+            content.ifPresent(s -> outputData.append(" | "));
+            outputData.append("Tool calls: ");
+            for (int j = 0; j < toolCalls.get().size(); j++) {
+              ChatCompletionMessageToolCall call = toolCalls.get().get(j);
+              if (j > 0) {
+                outputData.append(", ");
+              }
+              outputData
+                  .append(call.function().name())
+                  .append("(")
+                  .append(call.function().arguments())
+                  .append(")");
+            }
+          }
+        }
+
+        if (outputData.length() > 0) {
+          llmObsSpan.annotateIO(null, outputData.toString());
+        }
+        java.util.Map<String, Object> responseMetadata = new java.util.HashMap<>();
+        responseMetadata.put("response.choices_count", choices.size());
+
+        llmObsSpan.setMetadata(responseMetadata);
+      }
+    } catch (Exception e) {
+      java.util.Map<String, Object> errorMetadata = new java.util.HashMap<>();
+      errorMetadata.put("error.type", "ResponseProcessingError");
+      errorMetadata.put("error.message", "Failed to process response: " + e.getMessage());
+      llmObsSpan.setMetadata(errorMetadata);
+    } finally {
+      // Always finish the span
+      llmObsSpan.finish();
+    }
+  }
+
   public void finishSpan(AgentScope scope, Object result, Throwable throwable) {
 
     AgentSpan span = scope.span();
@@ -107,8 +221,6 @@ public void finishSpan(AgentScope scope, Object result, Throwable throwable) {
 
   private void extractChatCompletionRequestData(AgentSpan span, ChatCompletionCreateParams params) {
 
-    // Extract model
-
     span.setTag("openai.model.name", params.model().toString());
 
     // Extract messages
@@ -124,11 +236,7 @@ private void extractChatCompletionRequestData(AgentSpan span, ChatCompletionCrea
 
   private void extractCompletionRequestData(
       AgentSpan span, CompletionCreateParams completionParams) {
-    // Extract model
-    CompletionCreateParams.Model model = completionParams.model();
-    if (model != null) {
-      span.setTag("openai.model.name", model.toString());
-    }
+    span.setTag("openai.model.name", completionParams.model().toString());
 
     // Extract prompt
     Optional<CompletionCreateParams.Prompt> prompt = completionParams.prompt();
@@ -156,18 +264,13 @@ private void extractCompletionRequestData(
 
   private void extractEmbeddingRequestData(AgentSpan span, EmbeddingCreateParams embeddingParams) {
 
-    // Extract model
-    Object model = embeddingParams.model();
-    if (model != null) {
-      span.setTag("openai.model.name", model.toString());
-    }
-
+    span.setTag("openai.model.name", embeddingParams.model().toString());
     // Extract input
     EmbeddingCreateParams.Input input = embeddingParams.input();
     int inputIndex = 0;
     List<String> inputStrings = input.asArrayOfStrings();
     for (String inputItem : inputStrings) {
-      span.setTag("openai.request.input." + inputIndex, input.asString());
+      span.setTag("openai.request.input." + inputIndex, inputItem);
       inputIndex++;
     }
   }
@@ -205,25 +308,19 @@ private void extractMessageData(
 
   private void extractChatCompletionParameters(AgentSpan span, ChatCompletionCreateParams params) {
     // Extract max_tokens
-    Optional<Long> maxTokens = params.maxTokens();
-    maxTokens.ifPresent(tokens -> span.setTag("openai.request.max_tokens", tokens));
+    params
+        .maxCompletionTokens()
+        .ifPresent(tokens -> span.setTag("openai.request.max_tokens", tokens));
 
     // Extract temperature
-    Optional<Double> temperature = params.temperature();
-    temperature.ifPresent(temp -> span.setTag("openai.request.temperature", temp));
+    params.temperature().ifPresent(temp -> span.setTag("openai.request.temperature", temp));
   }
 
   private void extractCompletionParameters(AgentSpan span, CompletionCreateParams params) {
     // Extract max_tokens
-    if (params.maxTokens().isPresent()) {
-      span.setTag("openai.request.max_tokens", params.maxTokens().get());
-    }
-
+    params.maxTokens().ifPresent(tokens -> span.setTag("openai.request.max_tokens", tokens));
     // Extract temperature
-    Optional<Double> temperature = params.temperature();
-    if (temperature.isPresent()) {
-      span.setTag("openai.request.temperature", temperature.get());
-    }
+    params.temperature().ifPresent(temp -> span.setTag("openai.request.temperature", temp));
   }
 
   private void extractResponseData(AgentSpan span, Object result) {