-
Notifications
You must be signed in to change notification settings - Fork 5
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #44 from ehsavoie/opentelemetry
Using telemetry instrumentation in injection module to track LLM Usage.
- Loading branch information
Showing
19 changed files
with
367 additions
and
17 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
121 changes: 121 additions & 0 deletions
121
...g/wildfly/extension/ai/injection/observability/OpenTelemetryMetricsChatModelListener.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,121 @@ | ||
/* | ||
* Copyright The WildFly Authors | ||
* SPDX-License-Identifier: Apache-2.0 | ||
*/ | ||
package org.wildfly.extension.ai.injection.observability; | ||
|
||
import dev.langchain4j.data.message.AiMessage; | ||
import dev.langchain4j.model.chat.listener.ChatModelErrorContext; | ||
import dev.langchain4j.model.chat.listener.ChatModelListener; | ||
import dev.langchain4j.model.chat.listener.ChatModelRequest; | ||
import dev.langchain4j.model.chat.listener.ChatModelRequestContext; | ||
import dev.langchain4j.model.chat.listener.ChatModelResponse; | ||
import dev.langchain4j.model.chat.listener.ChatModelResponseContext; | ||
import io.opentelemetry.api.common.AttributeKey; | ||
import io.opentelemetry.api.common.Attributes; | ||
import io.opentelemetry.api.metrics.DoubleHistogram; | ||
import io.opentelemetry.api.metrics.LongHistogram; | ||
import io.opentelemetry.api.metrics.Meter; | ||
import jakarta.annotation.PostConstruct; | ||
import jakarta.enterprise.context.Dependent; | ||
import jakarta.inject.Inject; | ||
import java.util.List; | ||
import java.util.concurrent.TimeUnit; | ||
|
||
/** | ||
* Creates metrics following the <a href="https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-metrics/">Semantic | ||
* Conventions for GenAI Metrics</a>. | ||
*/ | ||
@Dependent | ||
public class OpenTelemetryMetricsChatModelListener implements ChatModelListener { | ||
|
||
private static final String MP_AI_METRIC_START_TIME_NAME = "MP_AI_METRIC_START_TIME"; | ||
|
||
private static final String METRIC_CLIENT_TOKEN_USAGE_NAME = "gen_ai.client.token.usage"; | ||
private static final String METRIC_CLIENT_OPERATION_DURATION_NAME = "gen_ai.client.operation.duration"; | ||
|
||
private LongHistogram clientTokenUsage; | ||
private DoubleHistogram clientOperationDuration; | ||
|
||
@Inject | ||
private Meter meter; | ||
|
||
@PostConstruct | ||
private void init() { | ||
clientTokenUsage = meter.histogramBuilder(METRIC_CLIENT_TOKEN_USAGE_NAME) | ||
.ofLongs() | ||
.setDescription("Measures number of input and output tokens used") | ||
.setExplicitBucketBoundariesAdvice(List.of(1L, 4L, 16L, 64L, 256L, 1024L, 4096L, 16384L, 65536L, 262144L, | ||
1048576L, 4194304L, 16777216L, 67108864L)) | ||
.build(); | ||
|
||
clientOperationDuration = meter.histogramBuilder(METRIC_CLIENT_OPERATION_DURATION_NAME) | ||
.setDescription("GenAI operation duration") | ||
.setExplicitBucketBoundariesAdvice( | ||
List.of(0.01, 0.02, 0.04, 0.08, 0.16, 0.32, 0.64, 1.28, 2.56, 5.12, 10.24, 20.48, 40.96, 81.92)) | ||
.setUnit("s") | ||
.build(); | ||
} | ||
|
||
@Override | ||
public void onRequest(ChatModelRequestContext requestContext) { | ||
requestContext.attributes().put(MP_AI_METRIC_START_TIME_NAME, System.nanoTime()); | ||
} | ||
|
||
@Override | ||
public void onResponse(ChatModelResponseContext responseContext) { | ||
final long endTime = System.nanoTime(); | ||
final long startTime = (Long) responseContext.attributes().get(MP_AI_METRIC_START_TIME_NAME); | ||
|
||
final ChatModelRequest request = responseContext.request(); | ||
final ChatModelResponse response = responseContext.response(); | ||
|
||
Attributes inputTokenCountAttributes = Attributes.of(AttributeKey.stringKey("gen_ai.operation.name"), "chat", | ||
AttributeKey.stringKey("gen_ai.request.model"), request.model(), | ||
AttributeKey.stringKey("gen_ai.response.model"), response.model(), | ||
AttributeKey.stringKey("gen_ai.token.type"), "input"); | ||
//Record | ||
clientTokenUsage.record(response.tokenUsage().inputTokenCount(), inputTokenCountAttributes); | ||
|
||
Attributes outputTokenCountAttributes = Attributes.of(AttributeKey.stringKey("gen_ai.operation.name"), "chat", | ||
AttributeKey.stringKey("gen_ai.request.model"), request.model(), | ||
AttributeKey.stringKey("gen_ai.response.model"), response.model(), | ||
AttributeKey.stringKey("gen_ai.token.type"), "output"); | ||
|
||
//Record | ||
clientTokenUsage.record(response.tokenUsage().outputTokenCount(), outputTokenCountAttributes); | ||
|
||
//Record duration | ||
Attributes durationAttributes = Attributes.of(AttributeKey.stringKey("gen_ai.operation.name"), "chat", | ||
AttributeKey.stringKey("gen_ai.request.model"), request.model(), | ||
AttributeKey.stringKey("gen_ai.response.model"), response.model()); | ||
recordClientOperationDuration(startTime, endTime, durationAttributes); | ||
} | ||
|
||
@Override | ||
public void onError(ChatModelErrorContext errorContext) { | ||
final long endTime = System.nanoTime(); | ||
final long startTime = (Long) errorContext.attributes().get(MP_AI_METRIC_START_TIME_NAME); | ||
final ChatModelRequest request = errorContext.request(); | ||
final ChatModelResponse response = errorContext.partialResponse(); | ||
|
||
StringBuilder sb = new StringBuilder() | ||
.append(errorContext.error().getClass().getName()); | ||
|
||
AiMessage aiMessage = errorContext.partialResponse().aiMessage(); | ||
if (aiMessage != null) { | ||
sb.append(";").append(aiMessage.text()); | ||
} | ||
|
||
//Record duration | ||
Attributes durationAttributes = Attributes.of(AttributeKey.stringKey("gen_ai.operation.name"), "chat", | ||
AttributeKey.stringKey("gen_ai.request.model"), request.model(), | ||
AttributeKey.stringKey("gen_ai.response.model"), response.model(), | ||
AttributeKey.stringKey("error.type"), sb.toString()); | ||
recordClientOperationDuration(startTime, endTime, durationAttributes); | ||
} | ||
|
||
private void recordClientOperationDuration(final long startTime, long endTime, final Attributes attributes) { | ||
clientOperationDuration.record(TimeUnit.SECONDS.convert(endTime - startTime, TimeUnit.NANOSECONDS), attributes); | ||
} | ||
} |
Oops, something went wrong.