Add support for instrumenting the OpenAI Client

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-client/build.gradle

@@ -0,0 +1,27 @@
+muzzle {
+  pass {
+    group = "com.openai"
+    module = "openai-java"
+    versions = "[2.8.0,)"
+    assertInverse = true
+  }
+}
+
+// Instrumentation is currently in preview. Additional testing will be implemented before enabling this by default.
+
+apply from: "$rootDir/gradle/java.gradle"
+
+addTestSuiteForDir('latestDepTest', 'test')
+
+dependencies {
+  compileOnly(group: 'com.openai', name: 'openai-java', version: '2.8.0')
+
+  testImplementation(group: 'com.openai', name: 'openai-java') {
+    version {
+      strictly '[2.8.0,)'
+    }
+  }
+
+  latestDepTestImplementation group: 'com.openai', name: 'openai-java', version: '+'
+
+}

dd-java-agent/instrumentation/openai-client/gradle.lockfile

@@ -0,0 +1,4 @@
+# This is a Gradle generated file for dependency locking.
+# This file is expected to be part of source control.
+# Manual entries are allowed.
+empty= 

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

@@ -0,0 +1,105 @@
+package datadog.trace.instrumentation.openaiclient;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+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;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class ChatCompletionServiceInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
+  public ChatCompletionServiceInstrumentation() {
+    super("openai-client", "openai-java", "openai-2.8");
+  }
+
+  @Override
+  protected boolean defaultEnabled() {
+    return InstrumenterConfig.get().isIntegrationEnabled(Collections.singleton("openai"), false);
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".OpenAIClientInfo",
+    };
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    Map<String, String> contextStores = new HashMap<>(1);
+    contextStores.put(
+        "com.openai.services.blocking.ChatCompletionService", OpenAIClientInfo.class.getName());
+    return contextStores;
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod()
+            .and(isPublic())
+            .and(named("create"))
+            .and(
+                takesArgument(
+                    0, named("com.openai.models.completions.ChatCompletionCreateParams"))),
+        getClass().getName() + "$ChatCompletionServiceAdvice");
+  }
+
+  @Override
+  public String hierarchyMarkerType() {
+    return "com.openai.services.blocking.chat.ChatCompletionService";
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> hierarchyMatcher() {
+    return implementsInterface(named(hierarchyMarkerType()));
+  }
+
+  public static class ChatCompletionServiceAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static Map<String, Object> methodEnter(
+        @Advice.Argument(0) final ChatCompletionCreateParams 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 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.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-client/src/main/java/datadog/trace/instrumentation/openaiclient/CompletionServiceInstrumentation.java

@@ -0,0 +1,87 @@
+package datadog.trace.instrumentation.openaiclient;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.*;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+
+import com.openai.models.completions.CompletionCreateParams;
+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;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class CompletionServiceInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
+  public CompletionServiceInstrumentation() {
+    super("openai-client", "openai-java", "openai-2.8");
+  }
+
+  @Override
+  protected boolean defaultEnabled() {
+    return InstrumenterConfig.get().isIntegrationEnabled(Collections.singleton("openai"), false);
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".OpenAIClientInfo",
+    };
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    Map<String, String> contextStores = new HashMap<>(1);
+    contextStores.put(
+        "com.openai.services.blocking.CompletionService", OpenAIClientInfo.class.getName());
+    return contextStores;
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod()
+            .and(isPublic())
+            .and(named("create"))
+            .and(takesArgument(0, named("com.openai.models.completions.CompletionCreateParams"))),
+        getClass().getName() + "$CompletionServiceAdvice");
+  }
+
+  @Override
+  public String hierarchyMarkerType() {
+    return "com.openai.services.blocking.CompletionService";
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> hierarchyMatcher() {
+    return implementsInterface(named(hierarchyMarkerType()));
+  }
+
+  public static class CompletionServiceAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentScope methodEnter(@Advice.Argument(0) final CompletionCreateParams params) {
+
+      return OpenAIClientDecorator.DECORATE.startCompletionSpan(params);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void methodExit(
+        @Advice.Enter final AgentScope span,
+        @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);
+    }
+  }
+}

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

@@ -0,0 +1,78 @@
+package datadog.trace.instrumentation.openaiclient;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.google.auto.service.AutoService;
+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;
+
+@AutoService(InstrumenterModule.class)
+public class EmbeddingServiceInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
+
+  public EmbeddingServiceInstrumentation() {
+    super("openai", "openai-java", "openai-2.8");
+  }
+
+  @Override
+  protected boolean defaultEnabled() {
+    return InstrumenterConfig.get().isIntegrationEnabled(Collections.singleton("openai"), false);
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".OpenAIDecorator",
+    };
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    // Instrument embedding creation methods
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("create"))
+            .and(isPublic())
+            .and(takesArgument(0, named("com.openai.models.embeddings.EmbeddingCreateParams"))),
+        EmbeddingServiceInstrumentation.class.getName() + "$EmbeddingAdvice");
+  }
+
+  @Override
+  public String hierarchyMarkerType() {
+    return "com.openai.services.embedding.EmbeddingService";
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> hierarchyMatcher() {
+    return implementsInterface(named(hierarchyMarkerType()));
+  }
+
+  public static class EmbeddingAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentScope onEnter(
+        @Advice.Argument(0) final EmbeddingCreateParams embeddingParams) {
+
+      return OpenAIClientDecorator.DECORATE.startEmbeddingSpan(embeddingParams);
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+    public static void onExit(
+        @Advice.Enter final AgentScope span,
+        @Advice.Return final Object result,
+        @Advice.Thrown final Throwable throwable) {
+
+      OpenAIClientDecorator.DECORATE.finishSpan(span, result, throwable);
+    }
+  }
+}