Instrument openai async client

Author: anuraaga

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

@@ -0,0 +1,40 @@
+/*
+ * 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.OpenAIClientAsync;
+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;
+
+public class OpenAiClientAsyncInstrumentation implements TypeInstrumentation {
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("com.openai.client.okhttp.OpenAIOkHttpClientAsync$Builder");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        named("build").and(returns(named("com.openai.client.OpenAIClientAsync"))),
+        OpenAiClientAsyncInstrumentation.class.getName() + "$BuildAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class BuildAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    @Advice.AssignReturned.ToReturned
+    public static OpenAIClientAsync onExit(@Advice.Return OpenAIClientAsync client) {
+      return TELEMETRY.wrap(client);
+    }
+  }
+}

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

@@ -6,7 +6,7 @@
 package io.opentelemetry.javaagent.instrumentation.openai.v1_1;
 
 import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
-import static java.util.Collections.singletonList;
+import static java.util.Arrays.asList;
 
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
@@ -27,6 +27,6 @@ public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
 
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
-    return singletonList(new OpenAiClientInstrumentation());
+    return asList(new OpenAiClientInstrumentation(), new OpenAiClientAsyncInstrumentation());
   }
 }

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

@@ -6,6 +6,7 @@
 package io.opentelemetry.javaagent.instrumentation.openai.v1_1;
 
 import com.openai.client.OpenAIClient;
+import com.openai.client.OpenAIClientAsync;
 import io.opentelemetry.instrumentation.openai.v1_1.AbstractChatTest;
 import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
 import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
@@ -31,6 +32,11 @@ protected OpenAIClient wrap(OpenAIClient client) {
     return client;
   }
 
+  @Override
+  protected OpenAIClientAsync wrap(OpenAIClientAsync client) {
+    return client;
+  }
+
   @Override
   protected final List<Consumer<SpanDataAssert>> maybeWithTransportSpan(
       Consumer<SpanDataAssert> span) {

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

@@ -29,14 +29,16 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.stream.Collectors;
-import javax.annotation.Nullable;
 
 final class ChatCompletionEventsHelper {
 
   private static final AttributeKey<String> EVENT_NAME = stringKey("event.name");
 
   public static void emitPromptLogEvents(
-      Logger eventLogger, ChatCompletionCreateParams request, boolean captureMessageContent) {
+      Context ctx,
+      Logger eventLogger,
+      ChatCompletionCreateParams request,
+      boolean captureMessageContent) {
     for (ChatCompletionMessageParam msg : request.messages()) {
       String eventType;
       Map<String, Value<?>> body = new HashMap<>();
@@ -84,7 +86,7 @@ public static void emitPromptLogEvents(
       } else {
         continue;
       }
-      newEvent(eventLogger, eventType).setBody(Value.of(body)).emit();
+      newEvent(eventLogger, eventType).setContext(ctx).setBody(Value.of(body)).emit();
     }
   }
 
@@ -160,7 +162,7 @@ private static String joinContentParts(List<ChatCompletionContentPartText> conte
   }
 
   public static void emitCompletionLogEvents(
-      Logger eventLogger, ChatCompletion completion, boolean captureMessageContent) {
+      Context ctx, Logger eventLogger, ChatCompletion completion, boolean captureMessageContent) {
     for (ChatCompletion.Choice choice : completion.choices()) {
       ChatCompletionMessage choiceMsg = choice.message();
       Map<String, Value<?>> message = new HashMap<>();
@@ -179,25 +181,21 @@ public static void emitCompletionLogEvents(
                             .collect(Collectors.toList())));
               });
       emitCompletionLogEvent(
-          eventLogger, choice.index(), choice.finishReason().toString(), Value.of(message), null);
+          ctx, eventLogger, choice.index(), choice.finishReason().toString(), Value.of(message));
     }
   }
 
   public static void emitCompletionLogEvent(
+      Context ctx,
       Logger eventLogger,
       long index,
       String finishReason,
-      Value<?> eventMessageObject,
-      @Nullable Context contextOverride) {
+      Value<?> eventMessageObject) {
     Map<String, Value<?>> body = new HashMap<>();
     body.put("finish_reason", Value.of(finishReason));
     body.put("index", Value.of(index));
     body.put("message", eventMessageObject);
-    LogRecordBuilder builder = newEvent(eventLogger, "gen_ai.choice").setBody(Value.of(body));
-    if (contextOverride != null) {
-      builder.setContext(contextOverride);
-    }
-    builder.emit();
+    newEvent(eventLogger, "gen_ai.choice").setContext(ctx).setBody(Value.of(body)).emit();
   }
 
   private static LogRecordBuilder newEvent(Logger eventLogger, String name) {

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

@@ -0,0 +1,30 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.openai.v1_1;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import java.util.concurrent.CompletableFuture;
+
+final class CompletableFutureWrapper {
+  private CompletableFutureWrapper() {}
+
+  public static <T> CompletableFuture<T> wrap(CompletableFuture<T> future, Context context) {
+    CompletableFuture<T> result = new CompletableFuture<>();
+    future.whenComplete(
+        (T value, Throwable throwable) -> {
+          try (Scope ignored = context.makeCurrent()) {
+            if (throwable != null) {
+              result.completeExceptionally(throwable);
+            } else {
+              result.complete(value);
+            }
+          }
+        });
+
+    return result;
+  }
+}

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

@@ -75,13 +75,13 @@ private ChatCompletion create(
       ChatCompletionCreateParams chatCompletionCreateParams, RequestOptions requestOptions) {
     Context parentCtx = Context.current();
     if (!instrumenter.shouldStart(parentCtx, chatCompletionCreateParams)) {
-      return createWithLogs(chatCompletionCreateParams, requestOptions);
+      return createWithLogs(parentCtx, chatCompletionCreateParams, requestOptions);
     }
 
     Context ctx = instrumenter.start(parentCtx, chatCompletionCreateParams);
     ChatCompletion completion;
     try (Scope ignored = ctx.makeCurrent()) {
-      completion = createWithLogs(chatCompletionCreateParams, requestOptions);
+      completion = createWithLogs(ctx, chatCompletionCreateParams, requestOptions);
     } catch (Throwable t) {
       instrumenter.end(ctx, chatCompletionCreateParams, null, t);
       throw t;
@@ -92,46 +92,50 @@ private ChatCompletion create(
   }
 
   private ChatCompletion createWithLogs(
-      ChatCompletionCreateParams chatCompletionCreateParams, RequestOptions requestOptions) {
+      Context ctx,
+      ChatCompletionCreateParams chatCompletionCreateParams,
+      RequestOptions requestOptions) {
     ChatCompletionEventsHelper.emitPromptLogEvents(
-        eventLogger, chatCompletionCreateParams, captureMessageContent);
+        ctx, eventLogger, chatCompletionCreateParams, captureMessageContent);
     ChatCompletion result = delegate.create(chatCompletionCreateParams, requestOptions);
-    ChatCompletionEventsHelper.emitCompletionLogEvents(eventLogger, result, captureMessageContent);
+    ChatCompletionEventsHelper.emitCompletionLogEvents(
+        ctx, eventLogger, result, captureMessageContent);
     return result;
   }
 
   private StreamResponse<ChatCompletionChunk> createStreaming(
       ChatCompletionCreateParams chatCompletionCreateParams, RequestOptions requestOptions) {
     Context parentCtx = Context.current();
     if (!instrumenter.shouldStart(parentCtx, chatCompletionCreateParams)) {
-      return createStreamingWithLogs(chatCompletionCreateParams, requestOptions, parentCtx, false);
+      return createStreamingWithLogs(parentCtx, chatCompletionCreateParams, requestOptions, false);
     }
 
     Context ctx = instrumenter.start(parentCtx, chatCompletionCreateParams);
     try (Scope ignored = ctx.makeCurrent()) {
-      return createStreamingWithLogs(chatCompletionCreateParams, requestOptions, ctx, true);
+      return createStreamingWithLogs(ctx, chatCompletionCreateParams, requestOptions, true);
     } catch (Throwable t) {
       instrumenter.end(ctx, chatCompletionCreateParams, null, t);
       throw t;
     }
   }
 
   private StreamResponse<ChatCompletionChunk> createStreamingWithLogs(
+      Context ctx,
       ChatCompletionCreateParams chatCompletionCreateParams,
       RequestOptions requestOptions,
-      Context parentCtx,
       boolean newSpan) {
     ChatCompletionEventsHelper.emitPromptLogEvents(
-        eventLogger, chatCompletionCreateParams, captureMessageContent);
+        ctx, eventLogger, chatCompletionCreateParams, captureMessageContent);
     StreamResponse<ChatCompletionChunk> result =
         delegate.createStreaming(chatCompletionCreateParams, requestOptions);
     return new TracingStreamedResponse(
         result,
-        parentCtx,
-        chatCompletionCreateParams,
-        instrumenter,
-        eventLogger,
-        captureMessageContent,
-        newSpan);
+        new StreamListener(
+            ctx,
+            chatCompletionCreateParams,
+            instrumenter,
+            eventLogger,
+            captureMessageContent,
+            newSpan));
   }
 }