ensure toolObjects are populated after all beans are initialized

Signed-off-by: jitokim <[email protected]>

Author: jitokim

auto-configurations/common/spring-ai-autoconfigure-tool/src/main/java/org/springframework/ai/tool/autoconfigure/ToolAnnotatedBeanProcessor.java

@@ -7,18 +7,17 @@
 import org.springframework.beans.factory.SmartInitializingSingleton;
 import org.springframework.beans.factory.config.BeanPostProcessor;
 import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
+import org.springframework.beans.factory.support.BeanDefinitionRegistry;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.ApplicationContextAware;
 import org.springframework.context.ConfigurableApplicationContext;
 
-import java.util.Arrays;
-import java.util.HashSet;
-import java.util.Set;
+import java.util.*;
 
 /**
  * {@code ToolAnnotatedBeanProcessor} scans beans after initialization and collects those
  * that have methods annotated with {@link Tool} within the specified base packages.
- *
+ * <p>
  * The collected beans are then registered with a {@link ToolCallbackProvider}.
  */
 public class ToolAnnotatedBeanProcessor
@@ -28,7 +27,7 @@ public class ToolAnnotatedBeanProcessor
 
 	private ApplicationContext applicationContext;
 
-	private Set<Object> methodToolBeans = new HashSet<>();
+	private List<Object> methodToolBeans = new ArrayList<>();
 
 	public ToolAnnotatedBeanProcessor(Set<String> basePackages) {
 		this.basePackages = basePackages;
@@ -96,13 +95,19 @@ private boolean hasToolAnnotatedMethod(Class<?> clazz) {
 	public void afterSingletonsInstantiated() {
 
 		if (!methodToolBeans.isEmpty()) {
-			MethodToolCallbackProvider.Builder builder = MethodToolCallbackProvider.builder();
-			builder.toolObjects(methodToolBeans.toArray());
-			MethodToolCallbackProvider provider = builder.build();
+
 			ConfigurableListableBeanFactory factory = ((ConfigurableApplicationContext) applicationContext)
 				.getBeanFactory();
 
-			if (!factory.containsBean("methodToolCallbackProvider")) {
+			if (factory.containsBean("methodToolCallbackProvider")) {
+				BeanDefinitionRegistry registry = (BeanDefinitionRegistry) applicationContext;
+				MethodToolCallbackProvider bean = factory.getBean(MethodToolCallbackProvider.class);
+				bean.setToolObjects(methodToolBeans);
+			}
+			else {
+				MethodToolCallbackProvider.Builder builder = MethodToolCallbackProvider.builder();
+				builder.toolObjects(methodToolBeans.toArray());
+				MethodToolCallbackProvider provider = builder.build();
 				factory.registerSingleton("methodToolCallbackProvider", provider);
 			}
 		}

auto-configurations/common/spring-ai-autoconfigure-tool/src/main/java/org/springframework/ai/tool/autoconfigure/ToolCallbackAutoRegistrar.java

@@ -3,11 +3,13 @@
 import org.springframework.ai.tool.ToolCallbackProvider;
 import org.springframework.ai.tool.annotation.Tool;
 import org.springframework.ai.tool.autoconfigure.annotation.EnableToolCallbackAutoRegistration;
+import org.springframework.ai.tool.method.MethodToolCallbackProvider;
 import org.springframework.beans.factory.config.BeanDefinition;
 import org.springframework.beans.factory.config.ConstructorArgumentValues;
 import org.springframework.beans.factory.support.BeanDefinitionRegistry;
 import org.springframework.beans.factory.support.BeanNameGenerator;
 import org.springframework.beans.factory.support.GenericBeanDefinition;
+import org.springframework.beans.factory.support.RootBeanDefinition;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.context.annotation.ImportBeanDefinitionRegistrar;
 import org.springframework.core.type.AnnotationMetadata;
@@ -36,6 +38,21 @@ public class ToolCallbackAutoRegistrar implements ImportBeanDefinitionRegistrar
 	public void registerBeanDefinitions(AnnotationMetadata importingClassMetadata, BeanDefinitionRegistry registry,
 			BeanNameGenerator importBeanNameGenerator) {
 
+		registerToolAnnotatedBeanProcessor(importingClassMetadata, registry);
+
+		registerMethodToolCallbackProvider(registry);
+
+	}
+
+	private void registerMethodToolCallbackProvider(BeanDefinitionRegistry registry) {
+		RootBeanDefinition callbackProviderDef = new RootBeanDefinition(MethodToolCallbackProvider.class);
+		callbackProviderDef
+			.setInstanceSupplier(() -> MethodToolCallbackProvider.builder().toolObjects(new Object[0]).build());
+		registry.registerBeanDefinition("methodToolCallbackProvider", callbackProviderDef);
+	}
+
+	private void registerToolAnnotatedBeanProcessor(AnnotationMetadata importingClassMetadata,
+			BeanDefinitionRegistry registry) {
 		Set<String> basePackages = getBasePackages(importingClassMetadata);
 
 		GenericBeanDefinition beanDefinition = new GenericBeanDefinition();

spring-ai-model/src/main/java/org/springframework/ai/tool/method/MethodToolCallbackProvider.java

@@ -51,7 +51,7 @@ public final class MethodToolCallbackProvider implements ToolCallbackProvider {
 
 	private static final Logger logger = LoggerFactory.getLogger(MethodToolCallbackProvider.class);
 
-	private final List<Object> toolObjects;
+	private List<Object> toolObjects;
 
 	private MethodToolCallbackProvider(List<Object> toolObjects) {
 		Assert.notNull(toolObjects, "toolObjects cannot be null");
@@ -78,6 +78,10 @@ private void assertToolAnnotatedMethodsPresent(List<Object> toolObjects) {
 		}
 	}
 
+	public void setToolObjects(List<Object> toolObjects) {
+		this.toolObjects = toolObjects;
+	}
+
 	@Override
 	public ToolCallback[] getToolCallbacks() {
 		var toolCallbacks = this.toolObjects.stream()