Skip to content

Commit

Permalink
Merge pull request #623 from quarkiverse/tracing
Browse files Browse the repository at this point in the history 
Simplify OpenTelemetry instrumentation code
  • Loading branch information
@geoand
geoand committed May 24, 2024
2 parents 00f0794 + 0a2b1db commit f55be2a
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 115 deletions.
71 changes: 13 additions & 58 deletions core/runtime/src/main/java/io/quarkiverse/langchain4j/runtime/aiservice/SpanWrapper.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,82 +5,37 @@

import jakarta.inject.Inject;

import io.opentelemetry.api.OpenTelemetry;
import io.opentelemetry.api.trace.SpanKind;
import io.opentelemetry.context.Context;
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.Tracer;
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;
private final Tracer tracer;

@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;
}
});
public SpanWrapper(Tracer tracer) {
this.tracer = tracer;
}

@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();
Optional<AiServiceMethodCreateInfo.SpanInfo> spanInfoOpt = input.createInfo.getSpanInfo();
if (spanInfoOpt.isEmpty()) {
return fun.apply(input);
}

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

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

return result;
Span span = tracer.spanBuilder(spanInfoOpt.get().getName()).startSpan();
try (Scope scope = span.makeCurrent()) {
return fun.apply(input);
} catch (Throwable t) {
if (shouldStart) {
instrumenter.end(spanContext, input, null, t);
}
span.recordException(t);
throw t;
} finally {
if (scope != null) {
scope.close();
}
span.end();
}
}

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;
}
}
}
67 changes: 10 additions & 57 deletions core/runtime/src/main/java/io/quarkiverse/langchain4j/runtime/tool/ToolSpanWrapper.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,77 +5,30 @@
import jakarta.inject.Inject;

import dev.langchain4j.agent.tool.ToolExecutionRequest;
import io.opentelemetry.api.OpenTelemetry;
import io.opentelemetry.api.trace.SpanKind;
import io.opentelemetry.context.Context;
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.Tracer;
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 ToolSpanWrapper implements QuarkusToolExecutor.Wrapper {

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

private final Instrumenter<ToolExecutionRequest, Void> instrumenter;
private final Tracer tracer;

@Inject
public ToolSpanWrapper(OpenTelemetry openTelemetry) {
InstrumenterBuilder<ToolExecutionRequest, 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(ToolExecutionRequest toolExecutionRequest) {
return SpanKind.INTERNAL;
}
});
public ToolSpanWrapper(Tracer tracer) {
this.tracer = tracer;
}

@Override
public String wrap(ToolExecutionRequest toolExecutionRequest, Object memoryId,
BiFunction<ToolExecutionRequest, Object, String> fun) {
Context parentContext = Context.current();
Context spanContext = null;
Scope scope = null;
boolean shouldStart = instrumenter.shouldStart(parentContext, toolExecutionRequest);
if (shouldStart) {
spanContext = instrumenter.start(parentContext, toolExecutionRequest);
scope = spanContext.makeCurrent();
}

try {
String result = fun.apply(toolExecutionRequest, memoryId);

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

return result;
Span span = tracer.spanBuilder("langchain4j.tools." + toolExecutionRequest.name()).startSpan();
try (Scope scope = span.makeCurrent()) {
return fun.apply(toolExecutionRequest, memoryId);
} catch (Throwable t) {
if (shouldStart) {
instrumenter.end(spanContext, toolExecutionRequest, null, t);
}
span.recordException(t);
throw t;
} finally {
if (scope != null) {
scope.close();
}
}
}

private static class InputSpanNameExtractor implements SpanNameExtractor<ToolExecutionRequest> {

private static final InputSpanNameExtractor INSTANCE = new InputSpanNameExtractor();

@Override
public String extract(ToolExecutionRequest toolExecutionRequest) {
return "langchain4j.tools." + toolExecutionRequest.name();
span.end();
}
}
}

0 comments on commit f55be2a

Please sign in to comment.