Initial APM-only openai-java instrumentation with a unit test.

Author: ygree

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

@@ -0,0 +1,23 @@
+apply from: "$rootDir/gradle/java.gradle"
+apply plugin: 'idea'
+
+muzzle {
+  pass {
+    group = "com.openai"
+    module = "openai-java"
+    versions = "[1.0.0,)"
+    //TODO assertInverse = true
+  }
+}
+
+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')
+}
+

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

@@ -0,0 +1,55 @@
+package datadog.trace.instrumentation.openai_java;
+
+import com.openai.models.completions.Completion;
+import com.openai.models.completions.CompletionCreateParams;
+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;
+
+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.takesArgument;
+
+public class CompletionServiceInstrumentation implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+  @Override
+  public String instrumentedType() {
+    return "com.openai.services.blocking.CompletionServiceImpl";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("create"))
+            .and(takesArgument(0, named("com.openai.models.completions.CompletionCreateParams"))),
+        getClass().getName() + "$CreateAdvice");
+  }
+
+  public static class CreateAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentScope enter(@Advice.Argument(0) final CompletionCreateParams params) {
+      AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
+      DECORATE.afterStart(span);
+      DECORATE.decorate(span, params);
+      return activateSpan(span);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void exit(@Advice.Enter final AgentScope scope, @Advice.Return Completion result, @Advice.Thrown final Throwable err) {
+      final AgentSpan span = scope.span();
+      if (err != null) {
+        DECORATE.onError(span, err);
+      }
+      if (result != null) {
+        DECORATE.decorate(span, result);
+      }
+      DECORATE.beforeFinish(span);
+      scope.close();
+      span.finish();
+    }
+  }
+}

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

@@ -0,0 +1,58 @@
+package datadog.trace.instrumentation.openai_java;
+
+import com.openai.models.completions.Completion;
+import com.openai.models.completions.CompletionChoice;
+import com.openai.models.completions.CompletionCreateParams;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes;
+import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
+import datadog.trace.bootstrap.instrumentation.decorator.ClientDecorator;
+
+public class OpenAiDecorator extends ClientDecorator {
+  public static final OpenAiDecorator DECORATE = new OpenAiDecorator();
+
+  public static final String INSTRUMENTATION_NAME = "openai-java";
+  public static final CharSequence SPAN_NAME = UTF8BytesString.create("openai.request");
+
+  private static final CharSequence COMPLETIONS_CREATE = UTF8BytesString.create("completions.create");
+
+  @Override
+  protected String service() {
+    return null;
+  }
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[] {
+      INSTRUMENTATION_NAME
+    };
+  }
+
+  @Override
+  protected CharSequence spanType() {
+    return InternalSpanTypes.HTTP_CLIENT;
+  }
+
+  @Override
+  protected CharSequence component() {
+    return null;
+  }
+
+  public void decorate(AgentSpan span, CompletionCreateParams params) {
+    span.setResourceName(COMPLETIONS_CREATE);
+
+    if (params == null) {
+      return;
+    }
+
+    span.setTag("openai.request.model", params.model().toString());
+    //TODO set other tags: ai.provider, openai.request.endpoint, request.prompt?
+    //TODO max_tokens, temperature or these are LLMObs – non-APM?
+  }
+
+  public void decorate(AgentSpan span, Completion completion) {
+    for (CompletionChoice choice : completion.choices()) {
+      span.setTag("openai.response.choices." + choice.index() + ".text", choice.text()); // TODO
+    }
+  }
+}

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

@@ -0,0 +1,28 @@
+package datadog.trace.instrumentation.openai_java;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import java.util.Arrays;
+import java.util.List;
+
+@AutoService(InstrumenterModule.class)
+public class OpenAiModule extends InstrumenterModule.Tracing {
+  public OpenAiModule() {
+    super("openai-java");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+        packageName + ".OpenAiDecorator",
+    };
+  }
+
+  @Override
+  public List<Instrumenter> typeInstrumentations() {
+    return Arrays.asList(
+        new CompletionServiceInstrumentation()
+    );
+  }
+}

dd-java-agent/instrumentation/openai-java/openai-java-1.0/src/test/groovy/CompletionServiceTest.groovy

@@ -0,0 +1,43 @@
+import static datadog.trace.agent.test.utils.TraceUtils.runnableUnderTrace
+
+import com.openai.models.completions.Completion
+import com.openai.models.completions.CompletionCreateParams
+
+class CompletionServiceTest extends OpenAiTest {
+
+  def "test"() {
+    CompletionCreateParams createParams = CompletionCreateParams.builder()
+    .model(CompletionCreateParams.Model.GPT_3_5_TURBO_INSTRUCT)
+    .prompt("Tell me a story about building the best SDK!")
+    .build()
+
+    Completion completion = runnableUnderTrace("parent") {
+      openAiClient.completions().create(createParams)
+    }
+
+    expect:
+    assertTraces(1) {
+      trace(3) {
+        span(0) {
+          operationName "parent"
+          parent()
+          errored false
+        }
+        span(1) {
+          operationName "openai.request"
+          resourceName "completions.create"
+          childOf span(0)
+          errored false
+          spanType "http"
+        }
+        span(2) {
+          operationName "okhttp.request"
+          // resourceName "POST /v1/completions"
+          childOf span(1)
+          errored false
+          spanType "http"
+        }
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/openai-java/openai-java-1.0/src/test/groovy/OpenAiTest.groovy

@@ -0,0 +1,70 @@
+import com.google.common.base.Strings
+import com.openai.client.OpenAIClient
+import com.openai.client.okhttp.OpenAIOkHttpClient
+import com.openai.credential.BearerTokenCredential
+import datadog.trace.agent.test.InstrumentationSpecification
+import datadog.trace.agent.test.server.http.TestHttpServer
+import spock.lang.AutoCleanup
+import spock.lang.Shared
+
+
+class OpenAiTest extends InstrumentationSpecification {
+
+  // will use real openai backend when provided
+  static String openAiToken() {
+    return null
+  }
+
+  @AutoCleanup
+  @Shared
+  OpenAIClient openAiClient
+
+  @AutoCleanup
+  @Shared
+  def mockOpenAiBackend = TestHttpServer.httpServer {
+    handlers {
+      prefix("/completions") {
+        redirect("/v1/completions")
+      }
+      prefix("/v1/completions") {
+        response.status(200).send(
+            """
+                  {
+                    "id": "cmpl-CUJTd66qbuEe2vu9cSEUWoFKFKm6O",
+                    "object": "text_completion",
+                    "created": 1761340745,
+                    "model": "gpt-3.5-turbo-instruct:20230824-v2",
+                    "choices": [
+                      {
+                        "text": "\\n\\nOnce upon a time in a tech company called Innovix, a team of",
+                        "index": 0,
+                        "logprobs": null,
+                        "finish_reason": "length"
+                      }
+                    ],
+                    "usage": {
+                      "prompt_tokens": 10,
+                      "completion_tokens": 16,
+                      "total_tokens": 26
+                    }
+                  }
+                  """
+        )
+      }
+    }
+  }
+
+  def setup() {
+    OpenAIOkHttpClient.Builder b = OpenAIOkHttpClient.builder()
+    if (Strings.isNullOrEmpty(openAiToken())) {
+      // mock backend
+      b.baseUrl(mockOpenAiBackend.address.toURL().toString())
+      b.credential(BearerTokenCredential.create(""))
+    } else {
+      // real openai backend
+      b.credential(BearerTokenCredential.create(openAiToken()))
+    }
+    openAiClient = b.build()
+  }
+}
+

settings.gradle.kts

@@ -479,6 +479,7 @@ include(
   ":dd-java-agent:instrumentation:okhttp:okhttp-2.2",
   ":dd-java-agent:instrumentation:okhttp:okhttp-3.0",
   ":dd-java-agent:instrumentation:ognl-appsec",
+  ":dd-java-agent:instrumentation:openai-java:openai-java-1.0",
   ":dd-java-agent:instrumentation:opensearch",
   ":dd-java-agent:instrumentation:opensearch:rest",
   ":dd-java-agent:instrumentation:opensearch:transport",