Skip to content

Commit

Permalink
Merge pull request #77 from quarkiverse/opentelemetry
Browse files Browse the repository at this point in the history
Add initial OpenTelemetry support for declarative services
  • Loading branch information
geoand authored Dec 4, 2023
2 parents ab8c6c2 + 61ab6df commit 1ba5659
Show file tree
Hide file tree
Showing 14 changed files with 455 additions and 32 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -53,13 +53,16 @@
import io.quarkiverse.langchain4j.runtime.aiservice.AiServiceMethodImplementationSupport;
import io.quarkiverse.langchain4j.runtime.aiservice.DeclarativeAiServiceCreateInfo;
import io.quarkiverse.langchain4j.runtime.aiservice.MetricsWrapper;
import io.quarkiverse.langchain4j.runtime.aiservice.SpanWrapper;
import io.quarkus.arc.Arc;
import io.quarkus.arc.ArcContainer;
import io.quarkus.arc.InstanceHandle;
import io.quarkus.arc.deployment.AdditionalBeanBuildItem;
import io.quarkus.arc.deployment.SyntheticBeanBuildItem;
import io.quarkus.arc.deployment.UnremovableBeanBuildItem;
import io.quarkus.builder.item.MultiBuildItem;
import io.quarkus.deployment.Capabilities;
import io.quarkus.deployment.Capability;
import io.quarkus.deployment.GeneratedClassGizmoAdaptor;
import io.quarkus.deployment.annotations.BuildProducer;
import io.quarkus.deployment.annotations.BuildStep;
Expand Down Expand Up @@ -333,7 +336,8 @@ public void handleAiServices(AiServicesRecorder recorder,
BuildProducer<ReflectiveClassBuildItem> reflectiveClassProducer,
BuildProducer<AiServicesMethodBuildItem> aiServicesMethodProducer,
BuildProducer<AdditionalBeanBuildItem> additionalBeanProducer,
Optional<MetricsCapabilityBuildItem> metricsCapability) {
Optional<MetricsCapabilityBuildItem> metricsCapability,
Capabilities capabilities) {

IndexView index = indexBuildItem.getIndex();

Expand Down Expand Up @@ -400,11 +404,15 @@ public void handleAiServices(AiServicesRecorder recorder,

var addMicrometerMetrics = metricsCapability.isPresent()
&& metricsCapability.get().metricsSupported(MetricsFactory.MICROMETER);

if (addMicrometerMetrics) {
additionalBeanProducer.produce(AdditionalBeanBuildItem.builder().addBeanClass(MetricsWrapper.class).build());
}

var addOpenTelemetrySpan = capabilities.isPresent(Capability.OPENTELEMETRY_TRACER);
if (addOpenTelemetrySpan) {
additionalBeanProducer.produce(AdditionalBeanBuildItem.builder().addBeanClass(SpanWrapper.class).build());
}

Map<String, AiServiceClassCreateInfo> perClassMetadata = new HashMap<>();
if (!ifacesForCreate.isEmpty()) {
ClassOutput classOutput = new GeneratedClassGizmoAdaptor(generatedClassProducer, true);
Expand Down Expand Up @@ -439,7 +447,8 @@ public void handleAiServices(AiServicesRecorder recorder,
// MethodImplementationSupport#implement

String methodId = createMethodId(methodInfo);
perMethodMetadata.put(methodId, gatherMethodMetadata(methodInfo, addMicrometerMetrics));
perMethodMetadata.put(methodId,
gatherMethodMetadata(methodInfo, addMicrometerMetrics, addOpenTelemetrySpan));
MethodCreator constructor = classCreator.getMethodCreator(MethodDescriptor.INIT, "V",
AiServiceContext.class);
constructor.invokeSpecialMethod(OBJECT_CONSTRUCTOR, constructor.getThis());
Expand Down Expand Up @@ -521,7 +530,8 @@ private static void addCreatedAware(IndexView index, Set<String> detectedForCrea
}
}

private AiServiceMethodCreateInfo gatherMethodMetadata(MethodInfo method, boolean addMicrometerMetrics) {
private AiServiceMethodCreateInfo gatherMethodMetadata(MethodInfo method, boolean addMicrometerMetrics,
boolean addOpenTelemetrySpans) {
if (method.returnType().kind() == Type.Kind.VOID) {
throw illegalConfiguration("Return type of method '%s' cannot be void", method);
}
Expand All @@ -537,9 +547,10 @@ private AiServiceMethodCreateInfo gatherMethodMetadata(MethodInfo method, boolea
returnType);
Optional<Integer> memoryIdParamPosition = gatherMemoryIdParamName(method);
Optional<AiServiceMethodCreateInfo.MetricsInfo> metricsInfo = gatherMetricsInfo(method, addMicrometerMetrics);
Optional<AiServiceMethodCreateInfo.SpanInfo> spanInfo = gatherSpanInfo(method, addOpenTelemetrySpans);

return new AiServiceMethodCreateInfo(systemMessageInfo, userMessageInfo, memoryIdParamPosition, requiresModeration,
returnType, metricsInfo);
returnType, metricsInfo, spanInfo);
}

private List<TemplateParameterInfo> gatherTemplateParamInfo(List<MethodParameterInfo> params) {
Expand Down Expand Up @@ -712,10 +723,27 @@ private Optional<AiServiceMethodCreateInfo.MetricsInfo> gatherMetricsInfo(Method
return Optional.of(builder.build());
}

private Optional<AiServiceMethodCreateInfo.SpanInfo> gatherSpanInfo(MethodInfo method,
boolean addOpenTelemetrySpans) {
if (!addOpenTelemetrySpans) {
return Optional.empty();
}

String name = defaultAiServiceSpanName(method);

// TODO: add more

return Optional.of(new AiServiceMethodCreateInfo.SpanInfo(name));
}

private String defaultAiServiceMetricName(MethodInfo method) {
return "langchain4j.aiservices." + method.declaringClass().name().withoutPackagePrefix() + "." + method.name();
}

private String defaultAiServiceSpanName(MethodInfo method) {
return "langchain4j.aiservices." + method.declaringClass().name().withoutPackagePrefix() + "." + method.name();
}

private static class TemplateParameterInfo {
private final int position;
private final String name;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -43,8 +43,12 @@
import io.quarkiverse.langchain4j.runtime.tool.ToolInvoker;
import io.quarkiverse.langchain4j.runtime.tool.ToolMethodCreateInfo;
import io.quarkiverse.langchain4j.runtime.tool.ToolParametersObjectSubstitution;
import io.quarkiverse.langchain4j.runtime.tool.ToolSpanWrapper;
import io.quarkiverse.langchain4j.runtime.tool.ToolSpecificationObjectSubstitution;
import io.quarkus.arc.deployment.AdditionalBeanBuildItem;
import io.quarkus.arc.deployment.ValidationPhaseBuildItem;
import io.quarkus.deployment.Capabilities;
import io.quarkus.deployment.Capability;
import io.quarkus.deployment.GeneratedClassGizmoAdaptor;
import io.quarkus.deployment.annotations.BuildProducer;
import io.quarkus.deployment.annotations.BuildStep;
Expand Down Expand Up @@ -75,6 +79,14 @@ public class ToolProcessor {
public static final MethodDescriptor MAP_PUT = MethodDescriptor.ofMethod(Map.class, "put", Object.class, Object.class,
Object.class);

@BuildStep
public void telemetry(Capabilities capabilities, BuildProducer<AdditionalBeanBuildItem> additionalBeanProducer) {
var addOpenTelemetrySpan = capabilities.isPresent(Capability.OPENTELEMETRY_TRACER);
if (addOpenTelemetrySpan) {
additionalBeanProducer.produce(AdditionalBeanBuildItem.builder().addBeanClass(ToolSpanWrapper.class).build());
}
}

@BuildStep
@Record(ExecutionTime.STATIC_INIT)
public void handleTools(CombinedIndexBuildItem indexBuildItem,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -271,8 +271,9 @@ private ToolExecutor getToolExecutor(String methodName) {
ToolSpecification toolSpecification = methodCreateInfo.getToolSpecification();
if (methodName.equals(
toolSpecification.name())) { // this only works because TestTool does not contain overloaded methods
toolExecutor = new QuarkusToolExecutor(testTool, invokerClassName, methodCreateInfo.getMethodName(),
methodCreateInfo.getArgumentMapperClassName());
toolExecutor = new QuarkusToolExecutor(
new QuarkusToolExecutor.Context(testTool, invokerClassName, methodCreateInfo.getMethodName(),
methodCreateInfo.getArgumentMapperClassName()));
break;
}
}
Expand Down
10 changes: 10 additions & 0 deletions core/runtime/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,16 @@
<artifactId>micrometer-core</artifactId>
<optional>true</optional>
</dependency>
<dependency>
<groupId>io.micrometer</groupId>
<artifactId>micrometer-core</artifactId>
<optional>true</optional>
</dependency>
<dependency>
<groupId>io.opentelemetry.instrumentation</groupId>
<artifactId>opentelemetry-instrumentation-api</artifactId>
<optional>true</optional>
</dependency>

<dependency>
<groupId>dev.langchain4j</groupId>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,9 @@
import io.quarkiverse.langchain4j.runtime.aiservice.AiServiceClassCreateInfo;
import io.quarkiverse.langchain4j.runtime.aiservice.AiServiceMethodCreateInfo;
import io.quarkiverse.langchain4j.runtime.tool.QuarkusToolExecutor;
import io.quarkiverse.langchain4j.runtime.tool.QuarkusToolExecutorFactory;
import io.quarkiverse.langchain4j.runtime.tool.ToolMethodCreateInfo;
import io.quarkus.arc.Arc;
import io.quarkus.arc.ClientProxy;

public class QuarkusAiServicesFactory implements AiServicesFactory {
Expand All @@ -33,8 +35,12 @@ public static class InstanceHolder {
}

public static class QuarkusAiServices<T> extends AiServices<T> {

private final QuarkusToolExecutorFactory toolExecutorFactory;

public QuarkusAiServices(AiServiceContext context) {
super(context);
toolExecutorFactory = Arc.container().instance(QuarkusToolExecutorFactory.class).get();
}

@Override
Expand All @@ -54,9 +60,10 @@ public AiServices<T> tools(List<Object> objectsWithTools) {
String invokerClassName = methodCreateInfo.getInvokerClassName();
ToolSpecification toolSpecification = methodCreateInfo.getToolSpecification();
context.toolSpecifications.add(toolSpecification);
context.toolExecutors.put(toolSpecification.name(),
new QuarkusToolExecutor(objectWithTool, invokerClassName, methodCreateInfo.getMethodName(),
methodCreateInfo.getArgumentMapperClassName()));
QuarkusToolExecutor.Context executorContext = new QuarkusToolExecutor.Context(objectWithTool,
invokerClassName, methodCreateInfo.getMethodName(),
methodCreateInfo.getArgumentMapperClassName());
context.toolExecutors.put(toolSpecification.name(), toolExecutorFactory.create(executorContext));
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,16 +17,21 @@ public class AiServiceMethodCreateInfo {

private final Optional<MetricsInfo> metricsInfo;

private final Optional<SpanInfo> spanInfo;

@RecordableConstructor
public AiServiceMethodCreateInfo(Optional<TemplateInfo> systemMessageInfo, UserMessageInfo userMessageInfo,
Optional<Integer> memoryIdParamPosition,
boolean requiresModeration, Class<?> returnType, Optional<MetricsInfo> metricsInfo) {
boolean requiresModeration, Class<?> returnType,
Optional<MetricsInfo> metricsInfo,
Optional<SpanInfo> spanInfo) {
this.systemMessageInfo = systemMessageInfo;
this.userMessageInfo = userMessageInfo;
this.memoryIdParamPosition = memoryIdParamPosition;
this.requiresModeration = requiresModeration;
this.returnType = returnType;
this.metricsInfo = metricsInfo;
this.spanInfo = spanInfo;
}

public Optional<TemplateInfo> getSystemMessageInfo() {
Expand All @@ -53,6 +58,10 @@ public Optional<MetricsInfo> getMetricsInfo() {
return metricsInfo;
}

public Optional<SpanInfo> getSpanInfo() {
return spanInfo;
}

public static class UserMessageInfo {
private final Optional<TemplateInfo> template;
private final Optional<Integer> paramPosition;
Expand Down Expand Up @@ -205,4 +214,17 @@ public AiServiceMethodCreateInfo.MetricsInfo build() {
}
}
}

public static class SpanInfo {
private final String name;

@RecordableConstructor
public SpanInfo(String name) {
this.name = name;
}

public String getName() {
return name;
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -33,9 +33,8 @@ public Object apply(Input input) {
});

for (Wrapper wrapper : wrappers) {
Function<Input, Object> currentFun = funRef.get();
wrapper.wrap(input, currentFun);
Function<Input, Object> newFunction = new Function<Input, Object>() {
var currentFun = funRef.get();
var newFunction = new Function<Input, Object>() {
@Override
public Object apply(Input input) {
return wrapper.wrap(input, currentFun);
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,86 @@
package io.quarkiverse.langchain4j.runtime.aiservice;

import java.util.Optional;
import java.util.function.Function;

import jakarta.inject.Inject;

import io.opentelemetry.api.OpenTelemetry;
import io.opentelemetry.api.trace.SpanKind;
import io.opentelemetry.context.Context;
import io.opentelemetry.context.Scope;
import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
import io.opentelemetry.instrumentation.api.instrumenter.InstrumenterBuilder;
import io.opentelemetry.instrumentation.api.instrumenter.SpanKindExtractor;
import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;

public class SpanWrapper implements AiServiceMethodImplementationSupport.Wrapper {

private static final String INSTRUMENTATION_NAME = "io.quarkus.opentelemetry";

private final Instrumenter<AiServiceMethodImplementationSupport.Input, Void> instrumenter;

@Inject
public SpanWrapper(OpenTelemetry openTelemetry) {
InstrumenterBuilder<AiServiceMethodImplementationSupport.Input, Void> builder = Instrumenter.builder(
openTelemetry,
INSTRUMENTATION_NAME,
InputSpanNameExtractor.INSTANCE);

// TODO: there is probably more information here we need to set
this.instrumenter = builder
.buildInstrumenter(new SpanKindExtractor<>() {
@Override
public SpanKind extract(AiServiceMethodImplementationSupport.Input input) {
return SpanKind.INTERNAL;
}
});
}

@Override
public Object wrap(AiServiceMethodImplementationSupport.Input input,
Function<AiServiceMethodImplementationSupport.Input, Object> fun) {

Context parentContext = Context.current();
Context spanContext = null;
Scope scope = null;
boolean shouldStart = instrumenter.shouldStart(parentContext, input);
if (shouldStart) {
spanContext = instrumenter.start(parentContext, input);
scope = spanContext.makeCurrent();
}

try {
Object result = fun.apply(input);

if (shouldStart) {
instrumenter.end(spanContext, input, null, null);
}

return result;
} catch (Throwable t) {
if (shouldStart) {
instrumenter.end(spanContext, input, null, t);
}
throw t;
} finally {
if (scope != null) {
scope.close();
}
}
}

private static class InputSpanNameExtractor implements SpanNameExtractor<AiServiceMethodImplementationSupport.Input> {

private static final InputSpanNameExtractor INSTANCE = new InputSpanNameExtractor();

@Override
public String extract(AiServiceMethodImplementationSupport.Input input) {
Optional<AiServiceMethodCreateInfo.SpanInfo> spanInfoOpt = input.createInfo.getSpanInfo();
if (spanInfoOpt.isPresent()) {
return spanInfoOpt.get().getName();
}
return null;
}
}
}
Loading

0 comments on commit 1ba5659

Please sign in to comment.