Skip to content

Commit e2e6ac4

Browse files
anuraagaanuraaga
authored
Instrument embeddings in openai client (#14353)
1 parent 3ac6d09 commit e2e6ac4

File tree

16 files changed

+2451
-70
lines changed

16 files changed

+2451
-70
lines changed

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

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.javaagent.instrumentation.openai.v1_1;
7+
8+
import com.openai.client.OpenAIClient;
9+
import com.openai.client.OpenAIClientAsync;
10+
import io.opentelemetry.instrumentation.openai.v1_1.AbstractEmbeddingsTest;
11+
import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
12+
import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
13+
import io.opentelemetry.sdk.testing.assertj.SpanDataAssert;
14+
import java.util.ArrayList;
15+
import java.util.List;
16+
import java.util.function.Consumer;
17+
import org.junit.jupiter.api.extension.RegisterExtension;
18+
19+
class EmbeddingsTest extends AbstractEmbeddingsTest {
20+
21+
@RegisterExtension
22+
private static final AgentInstrumentationExtension testing =
23+
AgentInstrumentationExtension.create();
24+
25+
@Override
26+
protected InstrumentationExtension getTesting() {
27+
return testing;
28+
}
29+
30+
@Override
31+
protected OpenAIClient wrap(OpenAIClient client) {
32+
return client;
33+
}
34+
35+
@Override
36+
protected OpenAIClientAsync wrap(OpenAIClientAsync client) {
37+
return client;
38+
}
39+
40+
@Override
41+
protected final List<Consumer<SpanDataAssert>> maybeWithTransportSpan(
42+
Consumer<SpanDataAssert> span) {
43+
List<Consumer<SpanDataAssert>> result = new ArrayList<>();
44+
result.add(span);
45+
// Do a very simple assertion since the telemetry is not part of this library.
46+
result.add(s -> s.hasName("POST"));
47+
return result;
48+
}
49+
}

‎instrumentation/openai/openai-java-1.1/library/README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# Library Instrumentation for OpenAI Java SDK version 1.1.0 and higher
22

33
Provides OpenTelemetry instrumentation for [openai-java](https://github.com/openai/openai-java/).
4+
Versions 1.1 through 2.x are supported.
45

56
## Quickstart
67

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

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.instrumentation.openai.v1_1;
7+
8+
import static java.util.Collections.singletonList;
9+
10+
import com.openai.models.embeddings.CreateEmbeddingResponse;
11+
import com.openai.models.embeddings.EmbeddingCreateParams;
12+
import io.opentelemetry.instrumentation.api.incubator.semconv.genai.GenAiAttributesGetter;
13+
import java.util.Collections;
14+
import java.util.List;
15+
import javax.annotation.Nullable;
16+
17+
enum EmbeddingAttributesGetter
18+
implements GenAiAttributesGetter<EmbeddingCreateParams, CreateEmbeddingResponse> {
19+
INSTANCE;
20+
21+
@Override
22+
public String getOperationName(EmbeddingCreateParams request) {
23+
return GenAiAttributes.GenAiOperationNameIncubatingValues.EMBEDDINGS;
24+
}
25+
26+
@Override
27+
public String getSystem(EmbeddingCreateParams request) {
28+
return GenAiAttributes.GenAiSystemIncubatingValues.OPENAI;
29+
}
30+
31+
@Override
32+
public String getRequestModel(EmbeddingCreateParams request) {
33+
return request.model().asString();
34+
}
35+
36+
@Nullable
37+
@Override
38+
public Long getRequestSeed(EmbeddingCreateParams request) {
39+
return null;
40+
}
41+
42+
@Nullable
43+
@Override
44+
public List<String> getRequestEncodingFormats(EmbeddingCreateParams request) {
45+
return request.encodingFormat().map(f -> singletonList(f.asString())).orElse(null);
46+
}
47+
48+
@Nullable
49+
@Override
50+
public Double getRequestFrequencyPenalty(EmbeddingCreateParams request) {
51+
return null;
52+
}
53+
54+
@Nullable
55+
@Override
56+
public Long getRequestMaxTokens(EmbeddingCreateParams request) {
57+
return null;
58+
}
59+
60+
@Nullable
61+
@Override
62+
public Double getRequestPresencePenalty(EmbeddingCreateParams request) {
63+
return null;
64+
}
65+
66+
@Nullable
67+
@Override
68+
public List<String> getRequestStopSequences(EmbeddingCreateParams request) {
69+
return null;
70+
}
71+
72+
@Nullable
73+
@Override
74+
public Double getRequestTemperature(EmbeddingCreateParams request) {
75+
return null;
76+
}
77+
78+
@Nullable
79+
@Override
80+
public Double getRequestTopK(EmbeddingCreateParams request) {
81+
return null;
82+
}
83+
84+
@Nullable
85+
@Override
86+
public Double getRequestTopP(EmbeddingCreateParams request) {
87+
return null;
88+
}
89+
90+
@Override
91+
public List<String> getResponseFinishReasons(
92+
EmbeddingCreateParams request, @Nullable CreateEmbeddingResponse response) {
93+
return Collections.emptyList();
94+
}
95+
96+
@Nullable
97+
@Override
98+
public String getResponseId(
99+
EmbeddingCreateParams request, @Nullable CreateEmbeddingResponse response) {
100+
return null;
101+
}
102+
103+
@Nullable
104+
@Override
105+
public String getResponseModel(
106+
EmbeddingCreateParams request, @Nullable CreateEmbeddingResponse response) {
107+
if (response == null) {
108+
return null;
109+
}
110+
return response.model();
111+
}
112+
113+
@Nullable
114+
@Override
115+
public Long getUsageInputTokens(
116+
EmbeddingCreateParams request, @Nullable CreateEmbeddingResponse response) {
117+
if (response == null) {
118+
return null;
119+
}
120+
return response.usage().promptTokens();
121+
}
122+
123+
@Nullable
124+
@Override
125+
public Long getUsageOutputTokens(
126+
EmbeddingCreateParams request, @Nullable CreateEmbeddingResponse response) {
127+
return null;
128+
}
129+
}

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ final class GenAiAttributes {
1515

1616
static final class GenAiOperationNameIncubatingValues {
1717
static final String CHAT = "chat";
18+
static final String EMBEDDINGS = "embeddings";
1819

1920
private GenAiOperationNameIncubatingValues() {}
2021
}

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

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.instrumentation.openai.v1_1;
7+
8+
import com.openai.core.RequestOptions;
9+
import com.openai.models.embeddings.CreateEmbeddingResponse;
10+
import com.openai.models.embeddings.EmbeddingCreateParams;
11+
import com.openai.services.blocking.EmbeddingService;
12+
import io.opentelemetry.context.Context;
13+
import io.opentelemetry.context.Scope;
14+
import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
15+
import java.lang.reflect.Method;
16+
17+
final class InstrumentedEmbeddingService
18+
extends DelegatingInvocationHandler<EmbeddingService, InstrumentedEmbeddingService> {
19+
20+
private final Instrumenter<EmbeddingCreateParams, CreateEmbeddingResponse> instrumenter;
21+
22+
public InstrumentedEmbeddingService(
23+
EmbeddingService delegate,
24+
Instrumenter<EmbeddingCreateParams, CreateEmbeddingResponse> instrumenter) {
25+
super(delegate);
26+
this.instrumenter = instrumenter;
27+
}
28+
29+
@Override
30+
protected Class<EmbeddingService> getProxyType() {
31+
return EmbeddingService.class;
32+
}
33+
34+
@Override
35+
public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
36+
String methodName = method.getName();
37+
Class<?>[] parameterTypes = method.getParameterTypes();
38+
39+
if (methodName.equals("create")
40+
&& parameterTypes.length >= 1
41+
&& parameterTypes[0] == EmbeddingCreateParams.class) {
42+
if (parameterTypes.length == 1) {
43+
return create((EmbeddingCreateParams) args[0], RequestOptions.none());
44+
} else if (parameterTypes.length == 2 && parameterTypes[1] == RequestOptions.class) {
45+
return create((EmbeddingCreateParams) args[0], (RequestOptions) args[1]);
46+
}
47+
}
48+
49+
return super.invoke(proxy, method, args);
50+
}
51+
52+
private CreateEmbeddingResponse create(
53+
EmbeddingCreateParams request, RequestOptions requestOptions) {
54+
Context parentContext = Context.current();
55+
if (!instrumenter.shouldStart(parentContext, request)) {
56+
return delegate.create(request, requestOptions);
57+
}
58+
59+
Context context = instrumenter.start(parentContext, request);
60+
CreateEmbeddingResponse response;
61+
try (Scope ignored = context.makeCurrent()) {
62+
response = delegate.create(request, requestOptions);
63+
} catch (Throwable t) {
64+
instrumenter.end(context, request, null, t);
65+
throw t;
66+
}
67+
68+
instrumenter.end(context, request, response, null);
69+
return response;
70+
}
71+
}

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

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.instrumentation.openai.v1_1;
7+
8+
import com.openai.core.RequestOptions;
9+
import com.openai.models.embeddings.CreateEmbeddingResponse;
10+
import com.openai.models.embeddings.EmbeddingCreateParams;
11+
import com.openai.services.async.EmbeddingServiceAsync;
12+
import io.opentelemetry.context.Context;
13+
import io.opentelemetry.context.Scope;
14+
import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
15+
import java.lang.reflect.Method;
16+
import java.util.concurrent.CompletableFuture;
17+
18+
final class InstrumentedEmbeddingServiceAsync
19+
extends DelegatingInvocationHandler<EmbeddingServiceAsync, InstrumentedEmbeddingServiceAsync> {
20+
21+
private final Instrumenter<EmbeddingCreateParams, CreateEmbeddingResponse> instrumenter;
22+
23+
public InstrumentedEmbeddingServiceAsync(
24+
EmbeddingServiceAsync delegate,
25+
Instrumenter<EmbeddingCreateParams, CreateEmbeddingResponse> instrumenter) {
26+
super(delegate);
27+
this.instrumenter = instrumenter;
28+
}
29+
30+
@Override
31+
protected Class<EmbeddingServiceAsync> getProxyType() {
32+
return EmbeddingServiceAsync.class;
33+
}
34+
35+
@Override
36+
public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
37+
String methodName = method.getName();
38+
Class<?>[] parameterTypes = method.getParameterTypes();
39+
40+
if (methodName.equals("create")
41+
&& parameterTypes.length >= 1
42+
&& parameterTypes[0] == EmbeddingCreateParams.class) {
43+
if (parameterTypes.length == 1) {
44+
return create((EmbeddingCreateParams) args[0], RequestOptions.none());
45+
} else if (parameterTypes.length == 2 && parameterTypes[1] == RequestOptions.class) {
46+
return create((EmbeddingCreateParams) args[0], (RequestOptions) args[1]);
47+
}
48+
}
49+
50+
return super.invoke(proxy, method, args);
51+
}
52+
53+
private CompletableFuture<CreateEmbeddingResponse> create(
54+
EmbeddingCreateParams request, RequestOptions requestOptions) {
55+
Context parentContext = Context.current();
56+
if (!instrumenter.shouldStart(parentContext, request)) {
57+
return delegate.create(request, requestOptions);
58+
}
59+
60+
Context context = instrumenter.start(parentContext, request);
61+
CompletableFuture<CreateEmbeddingResponse> future;
62+
try (Scope ignored = context.makeCurrent()) {
63+
future = delegate.create(request, requestOptions);
64+
} catch (Throwable t) {
65+
instrumenter.end(context, request, null, t);
66+
throw t;
67+
}
68+
69+
future = future.whenComplete((res, t) -> instrumenter.end(context, request, res, t));
70+
return CompletableFutureWrapper.wrap(future, parentContext);
71+
}
72+
}

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

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
import com.openai.client.OpenAIClient;
99
import com.openai.models.chat.completions.ChatCompletion;
1010
import com.openai.models.chat.completions.ChatCompletionCreateParams;
11+
import com.openai.models.embeddings.CreateEmbeddingResponse;
12+
import com.openai.models.embeddings.EmbeddingCreateParams;
1113
import io.opentelemetry.api.logs.Logger;
1214
import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
1315
import java.lang.reflect.Method;
@@ -16,16 +18,19 @@ final class InstrumentedOpenAiClient
1618
extends DelegatingInvocationHandler<OpenAIClient, InstrumentedOpenAiClient> {
1719

1820
private final Instrumenter<ChatCompletionCreateParams, ChatCompletion> chatInstrumenter;
21+
private final Instrumenter<EmbeddingCreateParams, CreateEmbeddingResponse> embeddingInstrumenter;
1922
private final Logger eventLogger;
2023
private final boolean captureMessageContent;
2124

2225
InstrumentedOpenAiClient(
2326
OpenAIClient delegate,
2427
Instrumenter<ChatCompletionCreateParams, ChatCompletion> chatInstrumenter,
28+
Instrumenter<EmbeddingCreateParams, CreateEmbeddingResponse> embeddingInstrumenter,
2529
Logger eventLogger,
2630
boolean captureMessageContent) {
2731
super(delegate);
2832
this.chatInstrumenter = chatInstrumenter;
33+
this.embeddingInstrumenter = embeddingInstrumenter;
2934
this.eventLogger = eventLogger;
3035
this.captureMessageContent = captureMessageContent;
3136
}
@@ -44,9 +49,17 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl
4449
delegate.chat(), chatInstrumenter, eventLogger, captureMessageContent)
4550
.createProxy();
4651
}
52+
if (methodName.equals("embeddings") && parameterTypes.length == 0) {
53+
return new InstrumentedEmbeddingService(delegate.embeddings(), embeddingInstrumenter)
54+
.createProxy();
55+
}
4756
if (methodName.equals("async") && parameterTypes.length == 0) {
4857
return new InstrumentedOpenAiClientAsync(
49-
delegate.async(), chatInstrumenter, eventLogger, captureMessageContent)
58+
delegate.async(),
59+
chatInstrumenter,
60+
embeddingInstrumenter,
61+
eventLogger,
62+
captureMessageContent)
5063
.createProxy();
5164
}
5265
return super.invoke(proxy, method, args);

0 commit comments

Comments
 (0)