Enhanced function capabilities through @tool annotation improvements

loong-coder · ilayaperumalg · commit 22fc730661a2 · 2025-06-20T11:27:08.000+01:00
Signed-off-by: loong-coder &lt;garen.mao@gmail.com&gt;
diff --git a/spring-ai-model/src/main/java/org/springframework/ai/tool/method/MethodToolCallbackProvider.java b/spring-ai-model/src/main/java/org/springframework/ai/tool/method/MethodToolCallbackProvider.java
@@ -19,6 +19,7 @@
 import java.lang.reflect.Method;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Objects;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
@@ -35,6 +36,7 @@
 import org.springframework.ai.tool.support.ToolDefinitions;
 import org.springframework.ai.tool.support.ToolUtils;
 import org.springframework.aop.support.AopUtils;
+import org.springframework.core.annotation.AnnotationUtils;
 import org.springframework.util.Assert;
 import org.springframework.util.ClassUtils;
 import org.springframework.util.ReflectionUtils;
@@ -67,7 +69,7 @@ private void assertToolAnnotatedMethodsPresent(List<Object> toolObjects) {
 			List<Method> toolMethods = Stream
 				.of(ReflectionUtils.getDeclaredMethods(
 						AopUtils.isAopProxy(toolObject) ? AopUtils.getTargetClass(toolObject) : toolObject.getClass()))
-				.filter(toolMethod -> toolMethod.isAnnotationPresent(Tool.class))
+				.filter(this::isToolAnnotatedMethod)
 				.filter(toolMethod -> !isFunctionalType(toolMethod))
 				.toList();
 
@@ -84,7 +86,7 @@ public ToolCallback[] getToolCallbacks() {
 			.map(toolObject -> Stream
 				.of(ReflectionUtils.getDeclaredMethods(
 						AopUtils.isAopProxy(toolObject) ? AopUtils.getTargetClass(toolObject) : toolObject.getClass()))
-				.filter(toolMethod -> toolMethod.isAnnotationPresent(Tool.class))
+				.filter(this::isToolAnnotatedMethod)
 				.filter(toolMethod -> !isFunctionalType(toolMethod))
 				.map(toolMethod -> MethodToolCallback.builder()
 					.toolDefinition(ToolDefinitions.from(toolMethod))
@@ -115,6 +117,11 @@ private boolean isFunctionalType(Method toolMethod) {
 		return isFunction;
 	}
 
+	private boolean isToolAnnotatedMethod(Method method) {
+		Tool annotation = AnnotationUtils.findAnnotation(method, Tool.class);
+		return Objects.nonNull(annotation);
+	}
+
 	private void validateToolCallbacks(ToolCallback[] toolCallbacks) {
 		List<String> duplicateToolNames = ToolUtils.getDuplicateToolNames(toolCallbacks);
 		if (!duplicateToolNames.isEmpty()) {
diff --git a/spring-ai-model/src/main/java/org/springframework/ai/tool/support/ToolUtils.java b/spring-ai-model/src/main/java/org/springframework/ai/tool/support/ToolUtils.java
@@ -27,6 +27,7 @@
 import org.springframework.ai.tool.execution.DefaultToolCallResultConverter;
 import org.springframework.ai.tool.execution.ToolCallResultConverter;
 import org.springframework.ai.util.ParsingUtils;
+import org.springframework.core.annotation.AnnotatedElementUtils;
 import org.springframework.util.Assert;
 import org.springframework.util.StringUtils;
 
@@ -42,7 +43,7 @@ private ToolUtils() {
 
 	public static String getToolName(Method method) {
 		Assert.notNull(method, "method cannot be null");
-		var tool = method.getAnnotation(Tool.class);
+		var tool = AnnotatedElementUtils.findMergedAnnotation(method, Tool.class);
 		if (tool == null) {
 			return method.getName();
 		}
@@ -56,7 +57,7 @@ public static String getToolDescriptionFromName(String toolName) {
 
 	public static String getToolDescription(Method method) {
 		Assert.notNull(method, "method cannot be null");
-		var tool = method.getAnnotation(Tool.class);
+		var tool = AnnotatedElementUtils.findMergedAnnotation(method, Tool.class);
 		if (tool == null) {
 			return ParsingUtils.reConcatenateCamelCase(method.getName(), " ");
 		}
@@ -65,13 +66,13 @@ public static String getToolDescription(Method method) {
 
 	public static boolean getToolReturnDirect(Method method) {
 		Assert.notNull(method, "method cannot be null");
-		var tool = method.getAnnotation(Tool.class);
+		var tool = AnnotatedElementUtils.findMergedAnnotation(method, Tool.class);
 		return tool != null && tool.returnDirect();
 	}
 
 	public static ToolCallResultConverter getToolCallResultConverter(Method method) {
 		Assert.notNull(method, "method cannot be null");
-		var tool = method.getAnnotation(Tool.class);
+		var tool = AnnotatedElementUtils.findMergedAnnotation(method, Tool.class);
 		if (tool == null) {
 			return new DefaultToolCallResultConverter();
 		}
diff --git a/spring-ai-model/src/test/java/org/springframework/ai/aot/ToolBeanRegistrationAotProcessorTests.java b/spring-ai-model/src/test/java/org/springframework/ai/aot/ToolBeanRegistrationAotProcessorTests.java
@@ -19,14 +19,19 @@
 import org.junit.jupiter.api.Test;
 
 import org.springframework.ai.tool.annotation.Tool;
+import org.springframework.ai.tool.execution.DefaultToolCallResultConverter;
+import org.springframework.ai.tool.execution.ToolCallResultConverter;
 import org.springframework.aot.generate.GenerationContext;
 import org.springframework.aot.hint.RuntimeHints;
 import org.springframework.beans.factory.aot.BeanRegistrationAotContribution;
 import org.springframework.beans.factory.support.DefaultListableBeanFactory;
 import org.springframework.beans.factory.support.RegisteredBean;
 import org.springframework.beans.factory.support.RootBeanDefinition;
+import org.springframework.core.annotation.AliasFor;
 import org.springframework.lang.Nullable;
 
+import java.lang.annotation.*;
+
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -55,6 +60,12 @@ void shouldProcessAnnotatedClass() {
 		assertThat(reflection().onType(TestTools.class)).accepts(this.runtimeHints);
 	}
 
+	@Test
+	void shouldProcessEnhanceAnnotatedClass() {
+		process(TestEnhanceToolTools.class);
+		assertThat(reflection().onType(TestEnhanceToolTools.class)).accepts(this.runtimeHints);
+	}
+
 	private void process(Class<?> beanClass) {
 		when(this.generationContext.getRuntimeHints()).thenReturn(this.runtimeHints);
 		BeanRegistrationAotContribution contribution = createContribution(beanClass);
@@ -87,4 +98,36 @@ String nonTool() {
 
 	}
 
+	@Target({ ElementType.METHOD, ElementType.ANNOTATION_TYPE })
+	@Retention(RetentionPolicy.RUNTIME)
+	@Documented
+	@Tool
+	@Inherited
+	@interface EnhanceTool {
+
+		@AliasFor(annotation = Tool.class)
+		String name() default "";
+
+		@AliasFor(annotation = Tool.class)
+		String description() default "";
+
+		@AliasFor(annotation = Tool.class)
+		boolean returnDirect() default false;
+
+		@AliasFor(annotation = Tool.class)
+		Class<? extends ToolCallResultConverter> resultConverter() default DefaultToolCallResultConverter.class;
+
+		String enhanceValue() default "";
+
+	}
+
+	static class TestEnhanceToolTools {
+
+		@EnhanceTool
+		String testTool() {
+			return "Testing EnhanceTool";
+		}
+
+	}
+
 }
diff --git a/spring-ai-model/src/test/java/org/springframework/ai/tool/method/MethodToolCallbackProviderTests.java b/spring-ai-model/src/test/java/org/springframework/ai/tool/method/MethodToolCallbackProviderTests.java
@@ -16,13 +16,17 @@
 
 package org.springframework.ai.tool.method;
 
+import java.lang.annotation.*;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
 
 import org.junit.jupiter.api.Test;
 
 import org.springframework.ai.tool.annotation.Tool;
+import org.springframework.ai.tool.execution.DefaultToolCallResultConverter;
+import org.springframework.ai.tool.execution.ToolCallResultConverter;
+import org.springframework.core.annotation.AliasFor;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -87,6 +91,27 @@ void whenToolObjectHasObjectTypeMethodThenSuccess() {
 		assertThat(provider.getToolCallbacks()[0].getToolDefinition().name()).isEqualTo("objectTool");
 	}
 
+	@Test
+	void whenToolObjectHasEnhanceToolAnnotatedMethodThenSucceed() {
+		MethodToolCallbackProvider provider = MethodToolCallbackProvider.builder()
+			.toolObjects(new ToolUseEnhanceToolObject())
+			.build();
+
+		assertThat(provider.getToolCallbacks()).hasSize(1);
+		assertThat(provider.getToolCallbacks()[0].getToolDefinition().name()).isEqualTo("enhanceTool");
+		assertThat(provider.getToolCallbacks()[0].getToolDefinition().description()).isEqualTo("enhance tool");
+	}
+
+	@Test
+	void whenEnhanceToolObjectHasMixOfValidAndFunctionalTypeToolMethodsThenSucceed() {
+		MethodToolCallbackProvider provider = MethodToolCallbackProvider.builder()
+			.toolObjects(new UseEnhanceToolMixedToolMethodsObject())
+			.build();
+
+		assertThat(provider.getToolCallbacks()).hasSize(1);
+		assertThat(provider.getToolCallbacks()[0].getToolDefinition().name()).isEqualTo("validTool");
+	}
+
 	static class ValidToolObject {
 
 		@Tool
@@ -155,4 +180,50 @@ public Object objectTool() {
 
 	}
 
+	@Target({ ElementType.METHOD, ElementType.ANNOTATION_TYPE })
+	@Retention(RetentionPolicy.RUNTIME)
+	@Documented
+	@Tool
+	@Inherited
+	@interface EnhanceTool {
+
+		@AliasFor(annotation = Tool.class)
+		String name() default "";
+
+		@AliasFor(annotation = Tool.class)
+		String description() default "";
+
+		@AliasFor(annotation = Tool.class)
+		boolean returnDirect() default false;
+
+		@AliasFor(annotation = Tool.class)
+		Class<? extends ToolCallResultConverter> resultConverter() default DefaultToolCallResultConverter.class;
+
+		String enhanceValue() default "";
+
+	}
+
+	static class ToolUseEnhanceToolObject {
+
+		@EnhanceTool(description = "enhance tool")
+		public String enhanceTool() {
+			return "enhance tool result";
+		}
+
+	}
+
+	static class UseEnhanceToolMixedToolMethodsObject {
+
+		@EnhanceTool
+		public String validTool() {
+			return "Valid tool result";
+		}
+
+		@EnhanceTool
+		public Function<String, String> functionTool() {
+			return input -> "Function result: " + input;
+		}
+
+	}
+
 }
