Add initial instrumentation of OpenAI client

Author: anuraaga

buildscripts/checkstyle.xml

@@ -254,6 +254,7 @@
       <property name="arrayInitIndent" value="2"/>
     </module>
     -->
+    <!-- handled by error prone
     <module name="AbbreviationAsWordInName">
       <property name="ignoreFinal" value="false"/>
       <property name="allowedAbbreviationLength" value="0"/>
@@ -262,6 +263,7 @@
                     PARAMETER_DEF, VARIABLE_DEF, METHOD_DEF, PATTERN_VARIABLE_DEF, RECORD_DEF,
                     RECORD_COMPONENT_DEF"/>
     </module>
+    -->
     <module name="OverloadMethodsDeclarationOrder"/>
     <!-- there are only a few violations of this, and they all appear to be for good reasons
     <module name="VariableDeclarationUsageDistance"/>

conventions/src/main/kotlin/otel.scala-conventions.gradle.kts

@@ -5,6 +5,7 @@ plugins {
 }
 
 dependencies {
+  compileOnly("org.scala-lang:scala-library:2.11.12")
   testCompileOnly("org.scala-lang:scala-library:2.11.12")
 }
 

instrumentation/openai/openai-java-1.1/javaagent/build.gradle.kts

@@ -0,0 +1,33 @@
+plugins {
+  id("otel.javaagent-instrumentation")
+}
+
+muzzle {
+  pass {
+    group.set("com.openai")
+    module.set("openai-java")
+    versions.set("[1.1.0,)")
+    assertInverse.set(true)
+  }
+}
+
+dependencies {
+  implementation(project(":instrumentation:openai:openai-java-1.1:library"))
+
+  testInstrumentation(project(":instrumentation:okhttp:okhttp-3.0:javaagent"))
+
+  library("com.openai:openai-java:1.1.0")
+
+  testImplementation(project(":instrumentation:openai:openai-java-1.1:testing"))
+
+  // needed for latest dep tests
+  testCompileOnly("com.google.errorprone:error_prone_annotations")
+}
+
+tasks {
+  withType<Test>().configureEach {
+    systemProperty("testLatestDeps", findProperty("testLatestDeps") as Boolean)
+    // TODO run tests both with and without genai message capture
+    systemProperty("otel.instrumentation.genai.capture-message-content", "true")
+  }
+}

instrumentation/openai/openai-java-1.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/openai/v1_1/OpenAIClientInstrumentation.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.openai.v1_1;
+
+import static io.opentelemetry.javaagent.instrumentation.openai.v1_1.OpenAISingletons.TELEMETRY;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+
+import com.openai.client.OpenAIClient;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@SuppressWarnings("IdentifierName") // Want to match library's convention
+public class OpenAIClientInstrumentation implements TypeInstrumentation {
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("com.openai.client.okhttp.OpenAIOkHttpClient$Builder");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        named("build").and(returns(named("com.openai.client.OpenAIClient"))),
+        OpenAIClientInstrumentation.class.getName() + "$BuildAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class BuildAdvice {
+    @Advice.OnMethodExit
+    @Advice.AssignReturned.ToReturned
+    public static OpenAIClient onExit(@Advice.Return OpenAIClient client) {
+      return TELEMETRY.wrap(client);
+    }
+  }
+}

instrumentation/openai/openai-java-1.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/openai/v1_1/OpenAIInstrumentationModule.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.openai.v1_1;
+
+import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
+import static java.util.Collections.singletonList;
+
+import com.google.auto.service.AutoService;
+import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import java.util.List;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(InstrumentationModule.class)
+@SuppressWarnings("IdentifierName") // Want to match library's convention
+public class OpenAIInstrumentationModule extends InstrumentationModule {
+  public OpenAIInstrumentationModule() {
+    super("openai", "openai-java", "openai-java-1.1");
+  }
+
+  @Override
+  public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
+    return hasClassesNamed("com.openai.client.OpenAIClient");
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return singletonList(new OpenAIClientInstrumentation());
+  }
+}

instrumentation/openai/openai-java-1.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/openai/v1_1/OpenAISingletons.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.openai.v1_1;
+
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.instrumentation.openai.v1_1.OpenAITelemetry;
+import io.opentelemetry.javaagent.bootstrap.internal.AgentInstrumentationConfig;
+
+@SuppressWarnings("IdentifierName") // Want to match library's convention
+public final class OpenAISingletons {
+  public static final OpenAITelemetry TELEMETRY =
+      OpenAITelemetry.builder(GlobalOpenTelemetry.get())
+          .setCaptureMessageContent(
+              AgentInstrumentationConfig.get()
+                  .getBoolean("otel.instrumentation.genai.capture-message-content", false))
+          .build();
+
+  private OpenAISingletons() {}
+}

instrumentation/openai/openai-java-1.1/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/openai/v1_1/ChatTest.java

@@ -0,0 +1,43 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.openai.v1_1;
+
+import com.openai.client.OpenAIClient;
+import io.opentelemetry.instrumentation.openai.v1_1.AbstractChatTest;
+import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
+import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
+import io.opentelemetry.sdk.testing.assertj.SpanDataAssert;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.Consumer;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+class ChatTest extends AbstractChatTest {
+
+  @RegisterExtension
+  private static final AgentInstrumentationExtension testing =
+      AgentInstrumentationExtension.create();
+
+  @Override
+  protected InstrumentationExtension getTesting() {
+    return testing;
+  }
+
+  @Override
+  protected OpenAIClient wrap(OpenAIClient client) {
+    return client;
+  }
+
+  @Override
+  protected final List<Consumer<SpanDataAssert>> maybeWithTransportSpan(
+      Consumer<SpanDataAssert> span) {
+    List<Consumer<SpanDataAssert>> result = new ArrayList<>();
+    result.add(span);
+    // Do a very simple assertion since the telemetry is not part of this library.
+    result.add(s -> s.hasName("POST"));
+    return result;
+  }
+}

instrumentation/openai/openai-java-1.1/library/build.gradle.kts

@@ -0,0 +1,19 @@
+plugins {
+  id("otel.library-instrumentation")
+  id("otel.nullaway-conventions")
+}
+
+dependencies {
+  library("com.openai:openai-java:1.1.0")
+
+  testImplementation(project(":instrumentation:openai:openai-java-1.1:testing"))
+
+  // needed for latest dep tests
+  testCompileOnly("com.google.errorprone:error_prone_annotations")
+}
+
+tasks {
+  withType<Test>().configureEach {
+    systemProperty("testLatestDeps", findProperty("testLatestDeps") as Boolean)
+  }
+}

instrumentation/openai/openai-java-1.1/library/src/main/java/io/opentelemetry/instrumentation/openai/v1_1/ChatAttributesGetter.java

@@ -0,0 +1,154 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.openai.v1_1;
+
+import static java.util.Collections.emptyList;
+import static java.util.Collections.singletonList;
+
+import com.openai.models.chat.completions.ChatCompletion;
+import com.openai.models.chat.completions.ChatCompletionCreateParams;
+import com.openai.models.completions.CompletionUsage;
+import io.opentelemetry.instrumentation.api.incubator.semconv.genai.GenAiAttributesGetter;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.jetbrains.annotations.Nullable;
+
+enum ChatAttributesGetter
+    implements GenAiAttributesGetter<ChatCompletionCreateParams, ChatCompletion> {
+  INSTANCE;
+
+  @Override
+  public String getOperationName(ChatCompletionCreateParams request) {
+    return GenAiAttributes.GenAiOperationNameIncubatingValues.CHAT;
+  }
+
+  @Override
+  public String getSystem(ChatCompletionCreateParams request) {
+    return GenAiAttributes.GenAiSystemIncubatingValues.OPENAI;
+  }
+
+  @Override
+  public String getRequestModel(ChatCompletionCreateParams request) {
+    return request.model().asString();
+  }
+
+  @Nullable
+  @Override
+  public Long getRequestSeed(ChatCompletionCreateParams request) {
+    return request.seed().orElse(null);
+  }
+
+  @Nullable
+  @Override
+  public List<String> getRequestEncodingFormats(ChatCompletionCreateParams request) {
+    return null;
+  }
+
+  @Nullable
+  @Override
+  public Double getRequestFrequencyPenalty(ChatCompletionCreateParams request) {
+    return request.frequencyPenalty().orElse(null);
+  }
+
+  @Nullable
+  @Override
+  public Long getRequestMaxTokens(ChatCompletionCreateParams request) {
+    return request.maxCompletionTokens().orElse(null);
+  }
+
+  @Nullable
+  @Override
+  public Double getRequestPresencePenalty(ChatCompletionCreateParams request) {
+    return request.presencePenalty().orElse(null);
+  }
+
+  @Nullable
+  @Override
+  public List<String> getRequestStopSequences(ChatCompletionCreateParams request) {
+    return request
+        .stop()
+        .map(
+            s -> {
+              if (s.isString()) {
+                return singletonList(s.asString());
+              }
+              if (s.isStrings()) {
+                return s.asStrings();
+              }
+              return null;
+            })
+        .orElse(null);
+  }
+
+  @Nullable
+  @Override
+  public Double getRequestTemperature(ChatCompletionCreateParams request) {
+    return request.temperature().orElse(null);
+  }
+
+  @Nullable
+  @Override
+  public Double getRequestTopK(ChatCompletionCreateParams request) {
+    return null;
+  }
+
+  @Nullable
+  @Override
+  public Double getRequestTopP(ChatCompletionCreateParams request) {
+    return request.topP().orElse(null);
+  }
+
+  @Override
+  public List<String> getResponseFinishReasons(
+      ChatCompletionCreateParams request, @Nullable ChatCompletion response) {
+    if (response == null) {
+      return emptyList();
+    }
+    return response.choices().stream()
+        .map(choice -> choice.finishReason().asString())
+        .collect(Collectors.toList());
+  }
+
+  @Override
+  @Nullable
+  public String getResponseId(
+      ChatCompletionCreateParams request, @Nullable ChatCompletion response) {
+    if (response == null) {
+      return null;
+    }
+    return response.id();
+  }
+
+  @Override
+  @Nullable
+  public String getResponseModel(
+      ChatCompletionCreateParams request, @Nullable ChatCompletion response) {
+    if (response == null) {
+      return null;
+    }
+    return response.model();
+  }
+
+  @Override
+  @Nullable
+  public Long getUsageInputTokens(
+      ChatCompletionCreateParams request, @Nullable ChatCompletion response) {
+    if (response == null) {
+      return null;
+    }
+    return response.usage().map(CompletionUsage::promptTokens).orElse(null);
+  }
+
+  @Override
+  @Nullable
+  public Long getUsageOutputTokens(
+      ChatCompletionCreateParams request, @Nullable ChatCompletion response) {
+    if (response == null) {
+      return null;
+    }
+    return response.usage().map(CompletionUsage::completionTokens).orElse(null);
+  }
+}