Skip to content

Commit

Permalink
Merge pull request #1010 from jmartisk/issue-1009
Browse files Browse the repository at this point in the history
Don't include tools from removed beans in the ToolsMetadata
  • Loading branch information
geoand authored Oct 24, 2024
2 parents fb8bf24 + 3b236a8 commit c03cb43
Show file tree
Hide file tree
Showing 2 changed files with 60 additions and 10 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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();

Expand Down Expand Up @@ -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) {
Expand Down
Original file line number Diff line number Diff line change
@@ -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;
}

}

0 comments on commit c03cb43

Please sign in to comment.