-
Notifications
You must be signed in to change notification settings - Fork 101
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add OpenTelemetry support to tool execution
- Loading branch information
Showing
8 changed files
with
189 additions
and
25 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
56 changes: 56 additions & 0 deletions
56
...ime/src/main/java/io/quarkiverse/langchain4j/runtime/tool/QuarkusToolExecutorFactory.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,56 @@ | ||
| package io.quarkiverse.langchain4j.runtime.tool; | ||
|
|
||
| import java.util.List; | ||
| import java.util.concurrent.atomic.AtomicReference; | ||
| import java.util.function.BiFunction; | ||
|
|
||
| import jakarta.inject.Singleton; | ||
|
|
||
| import dev.langchain4j.agent.tool.ToolExecutionRequest; | ||
| import io.quarkus.arc.All; | ||
| import io.quarkus.arc.Unremovable; | ||
|
|
||
| @Singleton | ||
| @Unremovable | ||
| public class QuarkusToolExecutorFactory { | ||
|
|
||
| private final List<QuarkusToolExecutor.Wrapper> wrappers; | ||
|
|
||
| public QuarkusToolExecutorFactory(@All List<QuarkusToolExecutor.Wrapper> wrappers) { | ||
| this.wrappers = wrappers; | ||
| } | ||
|
|
||
| public QuarkusToolExecutor create(QuarkusToolExecutor.Context context) { | ||
| if (wrappers.isEmpty()) { | ||
| return new QuarkusToolExecutor(context); | ||
| } | ||
|
|
||
| return new QuarkusToolExecutor(context) { | ||
| final QuarkusToolExecutor originalTool = new QuarkusToolExecutor(context); | ||
|
|
||
| @Override | ||
| public String execute(ToolExecutionRequest toolExecutionRequest, Object memoryId) { | ||
| AtomicReference<BiFunction<ToolExecutionRequest, Object, String>> funRef = new AtomicReference<>( | ||
| new BiFunction<>() { | ||
| @Override | ||
| public String apply(ToolExecutionRequest toolExecutionRequest, Object o) { | ||
| return originalTool.execute(toolExecutionRequest, memoryId); | ||
| } | ||
| }); | ||
|
|
||
| for (QuarkusToolExecutor.Wrapper wrapper : wrappers) { | ||
| var currentFun = funRef.get(); | ||
| BiFunction<ToolExecutionRequest, Object, String> newFunction = new BiFunction<>() { | ||
| @Override | ||
| public String apply(ToolExecutionRequest toolExecutionRequest, Object memoryId) { | ||
| return wrapper.wrap(toolExecutionRequest, memoryId, currentFun); | ||
| } | ||
| }; | ||
| funRef.set(newFunction); | ||
| } | ||
|
|
||
| return funRef.get().apply(toolExecutionRequest, memoryId); | ||
| } | ||
| }; | ||
| } | ||
| } |
81 changes: 81 additions & 0 deletions
81
core/runtime/src/main/java/io/quarkiverse/langchain4j/runtime/tool/ToolSpanWrapper.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,81 @@ | ||
| package io.quarkiverse.langchain4j.runtime.tool; | ||
|
|
||
| import java.util.function.BiFunction; | ||
|
|
||
| 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.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; | ||
|
|
||
| @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; | ||
| } | ||
| }); | ||
| } | ||
|
|
||
| @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; | ||
| } catch (Throwable t) { | ||
| if (shouldStart) { | ||
| instrumenter.end(spanContext, toolExecutionRequest, null, 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(); | ||
| } | ||
| } | ||
| } |