Merge pull request #1848 from sgdesmet/fix/tool-tracing

Fix tool call instrumentation regression

Author: geoand

core/runtime/src/main/java/io/quarkiverse/langchain4j/runtime/tool/QuarkusToolExecutor.java

@@ -34,8 +34,8 @@ public record Context(Object tool, String toolInvokerName, String methodName, St
 
     public interface Wrapper {
 
-        String wrap(ToolExecutionRequest toolExecutionRequest, Object memoryId,
-                BiFunction<ToolExecutionRequest, Object, String> fun);
+        ToolExecutionResult wrap(ToolExecutionRequest toolExecutionRequest, InvocationContext invocationContext,
+                BiFunction<ToolExecutionRequest, InvocationContext, ToolExecutionResult> fun);
     }
 
     public QuarkusToolExecutor(Context context) {
@@ -46,6 +46,7 @@ public ReturnBehavior returnBehavior() {
         return context.returnBehavior;
     }
 
+    @Override
     public String execute(ToolExecutionRequest toolExecutionRequest, Object memoryId) {
         return executeWithContext(toolExecutionRequest, InvocationContext.builder().chatMemoryId(memoryId).build())
                 .resultText();

core/runtime/src/main/java/io/quarkiverse/langchain4j/runtime/tool/QuarkusToolExecutorFactory.java

@@ -7,6 +7,8 @@
 import jakarta.inject.Singleton;
 
 import dev.langchain4j.agent.tool.ToolExecutionRequest;
+import dev.langchain4j.invocation.InvocationContext;
+import dev.langchain4j.service.tool.ToolExecutionResult;
 import io.quarkus.arc.All;
 import io.quarkus.arc.Unremovable;
 
@@ -29,27 +31,21 @@ public QuarkusToolExecutor create(QuarkusToolExecutor.Context 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);
-                            }
-                        });
+            public ToolExecutionResult executeWithContext(ToolExecutionRequest toolExecutionRequest,
+                    InvocationContext invocationContext) {
+
+                AtomicReference<BiFunction<ToolExecutionRequest, InvocationContext, ToolExecutionResult>> funRef = new AtomicReference<>(
+                        originalTool::executeWithContext);
 
                 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);
-                        }
-                    };
+                    BiFunction<ToolExecutionRequest, InvocationContext, ToolExecutionResult> newFunction = (
+                            toolExecutionRequest1,
+                            invocationContext1) -> wrapper.wrap(toolExecutionRequest1, invocationContext1, currentFun);
                     funRef.set(newFunction);
                 }
 
-                return funRef.get().apply(toolExecutionRequest, memoryId);
+                return funRef.get().apply(toolExecutionRequest, invocationContext);
             }
         };
     }

core/runtime/src/main/java/io/quarkiverse/langchain4j/runtime/tool/ToolSpanWrapper.java

@@ -7,6 +7,8 @@
 import org.eclipse.microprofile.config.inject.ConfigProperty;
 
 import dev.langchain4j.agent.tool.ToolExecutionRequest;
+import dev.langchain4j.invocation.InvocationContext;
+import dev.langchain4j.service.tool.ToolExecutionResult;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.api.trace.Tracer;
@@ -29,8 +31,8 @@ public ToolSpanWrapper(Tracer tracer,
     }
 
     @Override
-    public String wrap(ToolExecutionRequest toolExecutionRequest, Object memoryId,
-            BiFunction<ToolExecutionRequest, Object, String> fun) {
+    public ToolExecutionResult wrap(ToolExecutionRequest toolExecutionRequest, InvocationContext invocationContext,
+            BiFunction<ToolExecutionRequest, InvocationContext, ToolExecutionResult> fun) {
 
         // from https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/gen-ai-spans.md#execute-tool-span
         Span span = tracer.spanBuilder("langchain4j.tools." + toolExecutionRequest.name())
@@ -43,11 +45,11 @@ public String wrap(ToolExecutionRequest toolExecutionRequest, Object memoryId,
         if (includeArguments) {
             span.setAttribute("gen_ai.tool.call.arguments", toolExecutionRequest.arguments());
         }
-        try (Scope scope = span.makeCurrent()) {
+        try (Scope ignored = span.makeCurrent()) {
             // TODO Handle async method here.
-            var result = fun.apply(toolExecutionRequest, memoryId);
-            if (includeResult) {
-                span.setAttribute("gen_ai.tool.call.result", result);
+            var result = fun.apply(toolExecutionRequest, invocationContext);
+            if (includeResult && result != null) {
+                span.setAttribute("gen_ai.tool.call.result", result.resultText());
             }
             return result;
         } catch (Throwable t) {