[WIP] openai-java instrumentation

dd-java-agent/agent-llmobs/llmobs-test-fixtures/build.gradle

@@ -0,0 +1,8 @@
+apply from: "$rootDir/gradle/java.gradle"
+apply from: "$rootDir/gradle/version.gradle"
+
+dependencies {
+  api project(':dd-java-agent:agent-llmobs')
+  api project(':dd-java-agent:instrumentation-testing')
+}
+

dd-java-agent/agent-llmobs/llmobs-test-fixtures/src/main/groovy/datadog/trace/llmobs/LlmObsSpecification.groovy

@@ -0,0 +1,26 @@
+package datadog.trace.llmobs
+
+import datadog.communication.ddagent.SharedCommunicationObjects
+import datadog.trace.agent.test.InstrumentationSpecification
+import datadog.trace.api.Config
+import datadog.trace.api.config.LlmObsConfig
+
+class LlmObsSpecification extends InstrumentationSpecification {
+
+  void setupSpec() {
+    def sco = new SharedCommunicationObjects()
+    def config = Config.get()
+    sco.createRemaining(config)
+    // assert sco.configurationPoller(config) == null
+    // assert sco.monitoring instanceof Monitoring.DisabledMonitoring
+
+    LLMObsSystem.start(null, sco)
+  }
+
+  @Override
+  void configurePreAgent() {
+    super.configurePreAgent()
+
+    injectSysConfig(LlmObsConfig.LLMOBS_ENABLED, "true") // TODO maybe extract to an override method similar to DSM/DBM (see the super impl)
+  }
+}

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java

@@ -107,7 +107,14 @@ public static ReferenceMatcher loadStaticMuzzleReferences(
     }
   }
 
-  /** @return Class names of helpers to inject into the user's classloader */
+  /**
+   * @return Class names of helpers to inject into the user's classloader.
+   *     <blockquote>
+   *     <p><b>NOTE:</b> The order matters. If the muzzle check fails with a NoClassDefFoundError
+   *     (as seen in build/reports/muzzle-*.txt), it may be because some helper classes depend on
+   *     each other. In this case, the order must be adjusted accordingly.
+   *     </blockquote>
+   */
   public String[] helperClassNames() {
     return NO_HELPERS;
   }

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

@@ -0,0 +1,25 @@
+apply from: "$rootDir/gradle/java.gradle"
+apply plugin: 'idea'
+
+muzzle {
+  pass {
+    group = "com.openai"
+    module = "openai-java"
+    versions = "[1.0.0,)"
+    // assertInverse = true // TODO exclude <1.0.0, now it passes the muzzle check for >=0.34.0 but the tests will fail
+  }
+}
+
+addTestSuiteForDir('latestDepTest', 'test')
+
+dependencies {
+  compileOnly group: 'com.openai', name: 'openai-java', version: '1.0.0'
+
+  testImplementation group: 'com.openai', name: 'openai-java', version: '1.0.0'
+  latestDepTestImplementation group: 'com.openai', name: 'openai-java', version: '+'
+
+  testImplementation project(':dd-java-agent:instrumentation:okhttp:okhttp-3.0')
+
+  testImplementation project(':dd-java-agent:agent-llmobs:llmobs-test-fixtures')
+}
+

dd-java-agent/instrumentation/openai-java/openai-java-1.0/src/main/java/datadog/trace/instrumentation/openai_java/ChatCompletionServiceAsyncInstrumentation.java

@@ -0,0 +1,122 @@
+package datadog.trace.instrumentation.openai_java;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.instrumentation.openai_java.OpenAiDecorator.DECORATE;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.openai.core.ClientOptions;
+import com.openai.core.http.HttpResponseFor;
+import com.openai.core.http.StreamResponse;
+import com.openai.models.chat.completions.ChatCompletion;
+import com.openai.models.chat.completions.ChatCompletionChunk;
+import com.openai.models.chat.completions.ChatCompletionCreateParams;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import java.util.concurrent.CompletableFuture;
+import net.bytebuddy.asm.Advice;
+
+public class ChatCompletionServiceAsyncInstrumentation
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+  @Override
+  public String instrumentedType() {
+    return "com.openai.services.async.chat.ChatCompletionServiceAsyncImpl$WithRawResponseImpl";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("create"))
+            .and(
+                takesArgument(
+                    0, named("com.openai.models.chat.completions.ChatCompletionCreateParams")))
+            .and(returns(named(CompletableFuture.class.getName()))),
+        getClass().getName() + "$CreateAdvice");
+
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("createStreaming"))
+            .and(
+                takesArgument(
+                    0, named("com.openai.models.chat.completions.ChatCompletionCreateParams")))
+            .and(returns(named(CompletableFuture.class.getName()))),
+        getClass().getName() + "$CreateStreamingAdvice");
+  }
+
+  public static class CreateAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentScope enter(
+        @Advice.Argument(0) final ChatCompletionCreateParams params,
+        @Advice.FieldValue("clientOptions") ClientOptions clientOptions) {
+      AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
+      DECORATE.afterStart(span);
+      DECORATE.decorateWithClientOptions(span, clientOptions);
+      DECORATE.decorateChatCompletion(span, params);
+      return activateSpan(span);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void exit(
+        @Advice.Enter final AgentScope scope,
+        @Advice.Return(readOnly = false) CompletableFuture<HttpResponseFor<ChatCompletion>> future,
+        @Advice.Thrown final Throwable err) {
+      final AgentSpan span = scope.span();
+      try {
+        if (err != null) {
+          DECORATE.onError(span, err);
+        }
+        if (future != null) {
+          future =
+              ResponseWrappers.wrapFutureResponse(
+                  future, span, DECORATE::decorateWithChatCompletion);
+        } else {
+          span.finish();
+        }
+      } finally {
+        scope.close();
+      }
+    }
+  }
+
+  public static class CreateStreamingAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentScope enter(
+        @Advice.Argument(0) final ChatCompletionCreateParams params,
+        @Advice.FieldValue("clientOptions") ClientOptions clientOptions) {
+      AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
+      DECORATE.afterStart(span);
+      DECORATE.decorateWithClientOptions(span, clientOptions);
+      DECORATE.decorateChatCompletion(span, params);
+      return activateSpan(span);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void exit(
+        @Advice.Enter final AgentScope scope,
+        @Advice.Return(readOnly = false)
+            CompletableFuture<HttpResponseFor<StreamResponse<ChatCompletionChunk>>> future,
+        @Advice.Thrown final Throwable err) {
+      final AgentSpan span = scope.span();
+      try {
+        if (err != null) {
+          DECORATE.onError(span, err);
+        }
+        if (future != null) {
+          future =
+              ResponseWrappers.wrapFutureStreamResponse(
+                  future, span, DECORATE::decorateWithChatCompletionChunks);
+        } else {
+          span.finish();
+        }
+        DECORATE.beforeFinish(span);
+      } finally {
+        scope.close();
+      }
+    }
+  }
+}

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

@@ -0,0 +1,122 @@
+package datadog.trace.instrumentation.openai_java;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.instrumentation.openai_java.OpenAiDecorator.DECORATE;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.openai.core.ClientOptions;
+import com.openai.core.http.HttpResponseFor;
+import com.openai.core.http.StreamResponse;
+import com.openai.models.chat.completions.ChatCompletion;
+import com.openai.models.chat.completions.ChatCompletionChunk;
+import com.openai.models.chat.completions.ChatCompletionCreateParams;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import net.bytebuddy.asm.Advice;
+
+public class ChatCompletionServiceInstrumentation
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+  @Override
+  public String instrumentedType() {
+    return "com.openai.services.blocking.chat.ChatCompletionServiceImpl$WithRawResponseImpl";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("create"))
+            .and(
+                takesArgument(
+                    0, named("com.openai.models.chat.completions.ChatCompletionCreateParams")))
+            .and(returns(named("com.openai.core.http.HttpResponseFor"))),
+        getClass().getName() + "$CreateAdvice");
+
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("createStreaming"))
+            .and(
+                takesArgument(
+                    0, named("com.openai.models.chat.completions.ChatCompletionCreateParams")))
+            .and(returns(named("com.openai.core.http.HttpResponseFor"))),
+        getClass().getName() + "$CreateStreamingAdvice");
+  }
+
+  public static class CreateAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentScope enter(
+        @Advice.Argument(0) final ChatCompletionCreateParams params,
+        @Advice.FieldValue("clientOptions") ClientOptions clientOptions) {
+      AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
+      DECORATE.afterStart(span);
+      DECORATE.decorateWithClientOptions(span, clientOptions);
+      DECORATE.decorateChatCompletion(span, params);
+      return activateSpan(span);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void exit(
+        @Advice.Enter final AgentScope scope,
+        @Advice.Return(readOnly = false) HttpResponseFor<ChatCompletion> response,
+        @Advice.Thrown final Throwable err) {
+      final AgentSpan span = scope.span();
+      try {
+        if (err != null) {
+          DECORATE.onError(span, err);
+        }
+        if (response != null) {
+          response =
+              ResponseWrappers.wrapResponse(
+                  response, span, OpenAiDecorator.DECORATE::decorateWithChatCompletion);
+        }
+        DECORATE.beforeFinish(span);
+      } finally {
+        scope.close();
+        span.finish();
+      }
+    }
+  }
+
+  public static class CreateStreamingAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentScope enter(
+        @Advice.Argument(0) final ChatCompletionCreateParams params,
+        @Advice.FieldValue("clientOptions") ClientOptions clientOptions) {
+      AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
+      DECORATE.afterStart(span);
+      DECORATE.decorateWithClientOptions(span, clientOptions);
+      DECORATE.decorateChatCompletion(span, params);
+      return activateSpan(span);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void exit(
+        @Advice.Enter final AgentScope scope,
+        @Advice.Return(readOnly = false)
+            HttpResponseFor<StreamResponse<ChatCompletionChunk>> response,
+        @Advice.Thrown final Throwable err) {
+      final AgentSpan span = scope.span();
+      try {
+        if (err != null) {
+          DECORATE.onError(span, err);
+        }
+        if (response != null) {
+          response =
+              ResponseWrappers.wrapStreamResponse(
+                  response, span, DECORATE::decorateWithChatCompletionChunks);
+        } else {
+          span.finish();
+        }
+        DECORATE.beforeFinish(span);
+      } finally {
+        scope.close();
+      }
+    }
+  }
+}