Skip to content

Commit caaf1b5

Browse files
anuraagaanuraaga
committed
Instrument openai embeddings
1 parent 1115cda commit caaf1b5

File tree

8 files changed

+308
-2
lines changed

8 files changed

+308
-2
lines changed

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

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
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 io.opentelemetry.context.Context;
9+
import io.opentelemetry.context.Scope;
10+
import java.util.concurrent.CompletableFuture;
11+
12+
final class CompletableFutureWrapper {
13+
private CompletableFutureWrapper() {}
14+
15+
static <T> CompletableFuture<T> wrap(CompletableFuture<T> future, Context context) {
16+
CompletableFuture<T> result = new CompletableFuture<>();
17+
future.whenComplete(
18+
(T value, Throwable throwable) -> {
19+
try (Scope ignored = context.makeCurrent()) {
20+
if (throwable != null) {
21+
result.completeExceptionally(throwable);
22+
} else {
23+
result.complete(value);
24+
}
25+
}
26+
});
27+
28+
return result;
29+
}
30+
}

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

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

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

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

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

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

Lines changed: 5 additions & 0 deletions
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
}

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

Lines changed: 6 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.OpenTelemetry;
1214
import io.opentelemetry.api.logs.Logger;
1315
import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
@@ -28,24 +30,27 @@ public static OpenAITelemetryBuilder builder(OpenTelemetry openTelemetry) {
2830
}
2931

3032
private final Instrumenter<ChatCompletionCreateParams, ChatCompletion> chatInstrumenter;
33+
private final Instrumenter<EmbeddingCreateParams, CreateEmbeddingResponse> embeddingsInstrumenter;
3134

3235
private final Logger eventLogger;
3336

3437
private final boolean captureMessageContent;
3538

3639
OpenAITelemetry(
3740
Instrumenter<ChatCompletionCreateParams, ChatCompletion> chatInstrumenter,
41+
Instrumenter<EmbeddingCreateParams, CreateEmbeddingResponse> embeddingsInstrumenter,
3842
Logger eventLogger,
3943
boolean captureMessageContent) {
4044
this.chatInstrumenter = chatInstrumenter;
45+
this.embeddingsInstrumenter = embeddingsInstrumenter;
4146
this.eventLogger = eventLogger;
4247
this.captureMessageContent = captureMessageContent;
4348
}
4449

4550
/** Wraps the provided OpenAIClient, enabling telemetry for it. */
4651
public OpenAIClient wrap(OpenAIClient client) {
4752
return new InstrumentedOpenAiClient(
48-
client, chatInstrumenter, eventLogger, captureMessageContent)
53+
client, chatInstrumenter, embeddingsInstrumenter, eventLogger, captureMessageContent)
4954
.createProxy();
5055
}
5156
}

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

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
import com.google.errorprone.annotations.CanIgnoreReturnValue;
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.OpenTelemetry;
1214
import io.opentelemetry.api.logs.Logger;
1315
import io.opentelemetry.instrumentation.api.incubator.semconv.genai.GenAiAttributesExtractor;
@@ -53,7 +55,16 @@ public OpenAITelemetry build() {
5355
.addOperationMetrics(GenAiClientMetrics.get())
5456
.buildInstrumenter();
5557

58+
Instrumenter<EmbeddingCreateParams, CreateEmbeddingResponse> embeddingsInstrumenter =
59+
Instrumenter.<EmbeddingCreateParams, CreateEmbeddingResponse>builder(
60+
openTelemetry,
61+
INSTRUMENTATION_NAME,
62+
GenAiSpanNameExtractor.create(EmbeddingAttributesGetter.INSTANCE))
63+
.addAttributesExtractor(GenAiAttributesExtractor.create(EmbeddingAttributesGetter.INSTANCE))
64+
.addOperationMetrics(GenAiClientMetrics.get())
65+
.buildInstrumenter();
66+
5667
Logger eventLogger = openTelemetry.getLogsBridge().get(INSTRUMENTATION_NAME);
57-
return new OpenAITelemetry(chatInstrumenter, eventLogger, captureMessageContent);
68+
return new OpenAITelemetry(chatInstrumenter, embeddingsInstrumenter, eventLogger, captureMessageContent);
5869
}
5970
}

0 commit comments

Comments
 (0)