Instrument embeddings in openai client

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

@@ -0,0 +1,49 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+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.AbstractEmbeddingsTest;
+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 EmbeddingsTest extends AbstractEmbeddingsTest {
+
+  @RegisterExtension
+  private static final AgentInstrumentationExtension testing =
+      AgentInstrumentationExtension.create();
+
+  @Override
+  protected InstrumentationExtension getTesting() {
+    return testing;
+  }
+
+  @Override
+  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) {
+    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/README.md

@@ -1,6 +1,7 @@
 # Library Instrumentation for OpenAI Java SDK version 1.1.0 and higher
 
 Provides OpenTelemetry instrumentation for [openai-java](https://github.com/openai/openai-java/).
+Versions 1.1 through 2.x are supported.
 
 ## Quickstart
 

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

@@ -0,0 +1,129 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.openai.v1_1;
+
+import static java.util.Collections.singletonList;
+
+import com.openai.models.embeddings.CreateEmbeddingResponse;
+import com.openai.models.embeddings.EmbeddingCreateParams;
+import io.opentelemetry.instrumentation.api.incubator.semconv.genai.GenAiAttributesGetter;
+import java.util.Collections;
+import java.util.List;
+import javax.annotation.Nullable;
+
+enum EmbeddingAttributesGetter
+    implements GenAiAttributesGetter<EmbeddingCreateParams, CreateEmbeddingResponse> {
+  INSTANCE;
+
+  @Override
+  public String getOperationName(EmbeddingCreateParams request) {
+    return GenAiAttributes.GenAiOperationNameIncubatingValues.EMBEDDINGS;
+  }
+
+  @Override
+  public String getSystem(EmbeddingCreateParams request) {
+    return GenAiAttributes.GenAiSystemIncubatingValues.OPENAI;
+  }
+
+  @Override
+  public String getRequestModel(EmbeddingCreateParams request) {
+    return request.model().asString();
+  }
+
+  @Nullable
+  @Override
+  public Long getRequestSeed(EmbeddingCreateParams request) {
+    return null;
+  }
+
+  @Nullable
+  @Override
+  public List<String> getRequestEncodingFormats(EmbeddingCreateParams request) {
+    return request.encodingFormat().map(f -> singletonList(f.asString())).orElse(null);
+  }
+
+  @Nullable
+  @Override
+  public Double getRequestFrequencyPenalty(EmbeddingCreateParams request) {
+    return null;
+  }
+
+  @Nullable
+  @Override
+  public Long getRequestMaxTokens(EmbeddingCreateParams request) {
+    return null;
+  }
+
+  @Nullable
+  @Override
+  public Double getRequestPresencePenalty(EmbeddingCreateParams request) {
+    return null;
+  }
+
+  @Nullable
+  @Override
+  public List<String> getRequestStopSequences(EmbeddingCreateParams request) {
+    return null;
+  }
+
+  @Nullable
+  @Override
+  public Double getRequestTemperature(EmbeddingCreateParams request) {
+    return null;
+  }
+
+  @Nullable
+  @Override
+  public Double getRequestTopK(EmbeddingCreateParams request) {
+    return null;
+  }
+
+  @Nullable
+  @Override
+  public Double getRequestTopP(EmbeddingCreateParams request) {
+    return null;
+  }
+
+  @Override
+  public List<String> getResponseFinishReasons(
+      EmbeddingCreateParams request, @Nullable CreateEmbeddingResponse response) {
+    return Collections.emptyList();
+  }
+
+  @Nullable
+  @Override
+  public String getResponseId(
+      EmbeddingCreateParams request, @Nullable CreateEmbeddingResponse response) {
+    return null;
+  }
+
+  @Nullable
+  @Override
+  public String getResponseModel(
+      EmbeddingCreateParams request, @Nullable CreateEmbeddingResponse response) {
+    if (response == null) {
+      return null;
+    }
+    return response.model();
+  }
+
+  @Nullable
+  @Override
+  public Long getUsageInputTokens(
+      EmbeddingCreateParams request, @Nullable CreateEmbeddingResponse response) {
+    if (response == null) {
+      return null;
+    }
+    return response.usage().promptTokens();
+  }
+
+  @Nullable
+  @Override
+  public Long getUsageOutputTokens(
+      EmbeddingCreateParams request, @Nullable CreateEmbeddingResponse response) {
+    return null;
+  }
+}

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

@@ -15,6 +15,7 @@ final class GenAiAttributes {
 
   static final class GenAiOperationNameIncubatingValues {
     static final String CHAT = "chat";
+    static final String EMBEDDINGS = "embeddings";
 
     private GenAiOperationNameIncubatingValues() {}
   }

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

@@ -0,0 +1,71 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.openai.v1_1;
+
+import com.openai.core.RequestOptions;
+import com.openai.models.embeddings.CreateEmbeddingResponse;
+import com.openai.models.embeddings.EmbeddingCreateParams;
+import com.openai.services.blocking.EmbeddingService;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
+import java.lang.reflect.Method;
+
+final class InstrumentedEmbeddingService
+    extends DelegatingInvocationHandler<EmbeddingService, InstrumentedEmbeddingService> {
+
+  private final Instrumenter<EmbeddingCreateParams, CreateEmbeddingResponse> instrumenter;
+
+  public InstrumentedEmbeddingService(
+      EmbeddingService delegate,
+      Instrumenter<EmbeddingCreateParams, CreateEmbeddingResponse> instrumenter) {
+    super(delegate);
+    this.instrumenter = instrumenter;
+  }
+
+  @Override
+  protected Class<EmbeddingService> getProxyType() {
+    return EmbeddingService.class;
+  }
+
+  @Override
+  public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
+    String methodName = method.getName();
+    Class<?>[] parameterTypes = method.getParameterTypes();
+
+    if (methodName.equals("create")
+        && parameterTypes.length >= 1
+        && parameterTypes[0] == EmbeddingCreateParams.class) {
+      if (parameterTypes.length == 1) {
+        return create((EmbeddingCreateParams) args[0], RequestOptions.none());
+      } else if (parameterTypes.length == 2 && parameterTypes[1] == RequestOptions.class) {
+        return create((EmbeddingCreateParams) args[0], (RequestOptions) args[1]);
+      }
+    }
+
+    return super.invoke(proxy, method, args);
+  }
+
+  private CreateEmbeddingResponse create(
+      EmbeddingCreateParams request, RequestOptions requestOptions) {
+    Context parentContext = Context.current();
+    if (!instrumenter.shouldStart(parentContext, request)) {
+      return delegate.create(request, requestOptions);
+    }
+
+    Context context = instrumenter.start(parentContext, request);
+    CreateEmbeddingResponse response;
+    try (Scope ignored = context.makeCurrent()) {
+      response = delegate.create(request, requestOptions);
+    } catch (Throwable t) {
+      instrumenter.end(context, request, null, t);
+      throw t;
+    }
+
+    instrumenter.end(context, request, response, null);
+    return response;
+  }
+}

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

@@ -0,0 +1,72 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.openai.v1_1;
+
+import com.openai.core.RequestOptions;
+import com.openai.models.embeddings.CreateEmbeddingResponse;
+import com.openai.models.embeddings.EmbeddingCreateParams;
+import com.openai.services.async.EmbeddingServiceAsync;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
+import java.lang.reflect.Method;
+import java.util.concurrent.CompletableFuture;
+
+final class InstrumentedEmbeddingServiceAsync
+    extends DelegatingInvocationHandler<EmbeddingServiceAsync, InstrumentedEmbeddingServiceAsync> {
+
+  private final Instrumenter<EmbeddingCreateParams, CreateEmbeddingResponse> instrumenter;
+
+  public InstrumentedEmbeddingServiceAsync(
+      EmbeddingServiceAsync delegate,
+      Instrumenter<EmbeddingCreateParams, CreateEmbeddingResponse> instrumenter) {
+    super(delegate);
+    this.instrumenter = instrumenter;
+  }
+
+  @Override
+  protected Class<EmbeddingServiceAsync> getProxyType() {
+    return EmbeddingServiceAsync.class;
+  }
+
+  @Override
+  public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
+    String methodName = method.getName();
+    Class<?>[] parameterTypes = method.getParameterTypes();
+
+    if (methodName.equals("create")
+        && parameterTypes.length >= 1
+        && parameterTypes[0] == EmbeddingCreateParams.class) {
+      if (parameterTypes.length == 1) {
+        return create((EmbeddingCreateParams) args[0], RequestOptions.none());
+      } else if (parameterTypes.length == 2 && parameterTypes[1] == RequestOptions.class) {
+        return create((EmbeddingCreateParams) args[0], (RequestOptions) args[1]);
+      }
+    }
+
+    return super.invoke(proxy, method, args);
+  }
+
+  private CompletableFuture<CreateEmbeddingResponse> create(
+      EmbeddingCreateParams request, RequestOptions requestOptions) {
+    Context parentContext = Context.current();
+    if (!instrumenter.shouldStart(parentContext, request)) {
+      return delegate.create(request, requestOptions);
+    }
+
+    Context context = instrumenter.start(parentContext, request);
+    CompletableFuture<CreateEmbeddingResponse> future;
+    try (Scope ignored = context.makeCurrent()) {
+      future = delegate.create(request, requestOptions);
+    } catch (Throwable t) {
+      instrumenter.end(context, request, null, t);
+      throw t;
+    }
+
+    future = future.whenComplete((res, t) -> instrumenter.end(context, request, res, t));
+    return CompletableFutureWrapper.wrap(future, parentContext);
+  }
+}

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

@@ -8,6 +8,8 @@
 import com.openai.client.OpenAIClient;
 import com.openai.models.chat.completions.ChatCompletion;
 import com.openai.models.chat.completions.ChatCompletionCreateParams;
+import com.openai.models.embeddings.CreateEmbeddingResponse;
+import com.openai.models.embeddings.EmbeddingCreateParams;
 import io.opentelemetry.api.logs.Logger;
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import java.lang.reflect.Method;
@@ -16,16 +18,19 @@ final class InstrumentedOpenAiClient
     extends DelegatingInvocationHandler<OpenAIClient, InstrumentedOpenAiClient> {
 
   private final Instrumenter<ChatCompletionCreateParams, ChatCompletion> chatInstrumenter;
+  private final Instrumenter<EmbeddingCreateParams, CreateEmbeddingResponse> embeddingInstrumenter;
   private final Logger eventLogger;
   private final boolean captureMessageContent;
 
   InstrumentedOpenAiClient(
       OpenAIClient delegate,
       Instrumenter<ChatCompletionCreateParams, ChatCompletion> chatInstrumenter,
+      Instrumenter<EmbeddingCreateParams, CreateEmbeddingResponse> embeddingInstrumenter,
       Logger eventLogger,
       boolean captureMessageContent) {
     super(delegate);
     this.chatInstrumenter = chatInstrumenter;
+    this.embeddingInstrumenter = embeddingInstrumenter;
     this.eventLogger = eventLogger;
     this.captureMessageContent = captureMessageContent;
   }
@@ -44,9 +49,17 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl
               delegate.chat(), chatInstrumenter, eventLogger, captureMessageContent)
           .createProxy();
     }
+    if (methodName.equals("embeddings") && parameterTypes.length == 0) {
+      return new InstrumentedEmbeddingService(delegate.embeddings(), embeddingInstrumenter)
+          .createProxy();
+    }
     if (methodName.equals("async") && parameterTypes.length == 0) {
       return new InstrumentedOpenAiClientAsync(
-              delegate.async(), chatInstrumenter, eventLogger, captureMessageContent)
+              delegate.async(),
+              chatInstrumenter,
+              embeddingInstrumenter,
+              eventLogger,
+              captureMessageContent)
           .createProxy();
     }
     return super.invoke(proxy, method, args);