Merge pull request #1010 from jmartisk/issue-1009

Don't include tools from removed beans in the ToolsMetadata

Author: geoand

core/deployment/src/main/java/io/quarkiverse/langchain4j/deployment/ToolProcessor.java

@@ -89,20 +89,13 @@ public void telemetry(Capabilities capabilities, BuildProducer<AdditionalBeanBui
     }
 
     @BuildStep
-    @Record(ExecutionTime.STATIC_INIT)
     public void handleTools(CombinedIndexBuildItem indexBuildItem,
-            ToolsRecorder recorder,
             BuildProducer<AdditionalBeanBuildItem> additionalBeanProducer,
-            RecorderContext recorderContext,
             BuildProducer<BytecodeTransformerBuildItem> transformerProducer,
             BuildProducer<GeneratedClassBuildItem> generatedClassProducer,
             BuildProducer<ReflectiveClassBuildItem> reflectiveClassProducer,
             BuildProducer<ValidationPhaseBuildItem.ValidationErrorBuildItem> validation,
-            BuildProducer<ToolsMetadataBuildItem> toolsMetadataProducer) {
-        recorderContext.registerSubstitution(ToolSpecification.class, ToolSpecificationObjectSubstitution.Serialized.class,
-                ToolSpecificationObjectSubstitution.class);
-        recorderContext.registerSubstitution(ToolParameters.class, ToolParametersObjectSubstitution.Serialized.class,
-                ToolParametersObjectSubstitution.class);
+            BuildProducer<ToolsMetadataBeforeRemovalBuildItem> toolsMetadataProducer) {
 
         IndexView index = indexBuildItem.getIndex();
 
@@ -250,8 +243,38 @@ public void handleTools(CombinedIndexBuildItem indexBuildItem,
                     .build());
         }
 
-        toolsMetadataProducer.produce(new ToolsMetadataBuildItem(metadata));
-        recorder.setMetadata(metadata);
+        toolsMetadataProducer.produce(new ToolsMetadataBeforeRemovalBuildItem(metadata));
+    }
+
+    /**
+     * Transforms ToolsMetadataBeforeRemovalBuildItem into ToolsMetadataBuildItem by filtering
+     * out tools belonging to beans that have been removed by ArC.
+     */
+    @BuildStep
+    @Record(ExecutionTime.STATIC_INIT)
+    public ToolsMetadataBuildItem filterOutRemovedTools(ToolsMetadataBeforeRemovalBuildItem beforeRemoval,
+            ValidationPhaseBuildItem validationPhase,
+            RecorderContext recorderContext,
+            ToolsRecorder recorder) {
+        if (beforeRemoval != null) {
+            recorderContext.registerSubstitution(ToolSpecification.class, ToolSpecificationObjectSubstitution.Serialized.class,
+                    ToolSpecificationObjectSubstitution.class);
+            recorderContext.registerSubstitution(ToolParameters.class, ToolParametersObjectSubstitution.Serialized.class,
+                    ToolParametersObjectSubstitution.class);
+            Map<String, List<ToolMethodCreateInfo>> metadataWithoutRemovedBeans = beforeRemoval.getMetadata().entrySet()
+                    .stream()
+                    .filter(entry -> validationPhase.getContext().removedBeans().stream()
+                            .noneMatch(
+                                    retainedBean -> DotName.createSimple(entry.getKey()).equals(retainedBean.getBeanClass())))
+                    .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+            ToolsMetadataBuildItem toolsMetadata = new ToolsMetadataBuildItem(metadataWithoutRemovedBeans);
+            recorder.setMetadata(toolsMetadata.getMetadata());
+            log.debug("Tool classes before filtering out removed beans: " + beforeRemoval.getMetadata().keySet());
+            log.debug("Tool classes after filtering out removed beans: " + toolsMetadata.getMetadata().keySet());
+            return toolsMetadata;
+        } else {
+            return null;
+        }
     }
 
     private boolean ignoreToolMethod(MethodInfo toolMethod, IndexView indexView) {

core/deployment/src/main/java/io/quarkiverse/langchain4j/deployment/ToolsMetadataBeforeRemovalBuildItem.java

@@ -0,0 +1,27 @@
+package io.quarkiverse.langchain4j.deployment;
+
+import java.util.List;
+import java.util.Map;
+
+import io.quarkiverse.langchain4j.runtime.tool.ToolMethodCreateInfo;
+import io.quarkus.builder.item.SimpleBuildItem;
+
+/**
+ * Holds metadata about tools discovered at build time, but tools belonging
+ * to beans removed by ArC have not been filtered out yet
+ * (ToolsMetadataBuildItem is the build item that already has them filtered
+ * out).
+ */
+public final class ToolsMetadataBeforeRemovalBuildItem extends SimpleBuildItem {
+
+    Map<String, List<ToolMethodCreateInfo>> metadata;
+
+    public ToolsMetadataBeforeRemovalBuildItem(Map<String, List<ToolMethodCreateInfo>> metadata) {
+        this.metadata = metadata;
+    }
+
+    public Map<String, List<ToolMethodCreateInfo>> getMetadata() {
+        return metadata;
+    }
+
+}