Wrap Exception from toolFunction in ToolExecutionException and rethrow

- When `toolFunction` throws a `ToolExecutionException`, rethrow it directly.
- When `toolFunction` throws a `Exception`, wrap it in a `ToolExecutionException` and rethrow it.
- Add related tests

Signed-off-by: YunKui Lu <[email protected]>

Author: YunKuiLu

spring-ai-model/src/main/java/org/springframework/ai/tool/function/FunctionToolCallback.java

@@ -31,6 +31,7 @@
 import org.springframework.ai.tool.definition.ToolDefinition;
 import org.springframework.ai.tool.execution.DefaultToolCallResultConverter;
 import org.springframework.ai.tool.execution.ToolCallResultConverter;
+import org.springframework.ai.tool.execution.ToolExecutionException;
 import org.springframework.ai.tool.metadata.ToolMetadata;
 import org.springframework.ai.tool.support.ToolUtils;
 import org.springframework.ai.util.json.JsonParser;
@@ -99,13 +100,25 @@ public String call(String toolInput, @Nullable ToolContext toolContext) {
 		logger.debug("Starting execution of tool: {}", this.toolDefinition.name());
 
 		I request = JsonParser.fromJson(toolInput, this.toolInputType);
-		O response = this.toolFunction.apply(request, toolContext);
+		O response = callMethod(request, toolContext);
 
 		logger.debug("Successful execution of tool: {}", this.toolDefinition.name());
 
 		return this.toolCallResultConverter.convert(response, null);
 	}
 
+	private O callMethod(I request, @Nullable ToolContext toolContext) {
+		try {
+			return this.toolFunction.apply(request, toolContext);
+		}
+		catch (ToolExecutionException ex) {
+			throw ex;
+		}
+		catch (Exception ex) {
+			throw new ToolExecutionException(this.toolDefinition, ex);
+		}
+	}
+
 	@Override
 	public String toString() {
 		return "FunctionToolCallback{" + "toolDefinition=" + this.toolDefinition + ", toolMetadata=" + this.toolMetadata

spring-ai-model/src/test/java/org/springframework/ai/tool/function/FunctionToolCallbackTest.java

@@ -0,0 +1,171 @@
+package org.springframework.ai.tool.function;
+
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BiFunction;
+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.chat.model.ToolContext;
+import org.springframework.ai.tool.execution.ToolExecutionException;
+import org.springframework.ai.tool.metadata.ToolMetadata;
+
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.assertj.core.api.InstanceOfAssertFactories.type;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+/**
+ * @author YunKui Lu
+ */
+class FunctionToolCallbackTest {
+
+	@Test
+	void testConsumerToolCall() {
+		TestFunctionTool tool = new TestFunctionTool();
+		FunctionToolCallback<String, Void> callback = FunctionToolCallback.builder("testTool", tool.stringConsumer())
+			.toolMetadata(ToolMetadata.builder().returnDirect(true).build())
+			.description("test description")
+			.inputType(String.class)
+			.build();
+
+		callback.call("\"test string param\"");
+
+		assertEquals("test string param", tool.calledValue.get());
+	}
+
+	@Test
+	void testBiFunctionToolCall() {
+		TestFunctionTool tool = new TestFunctionTool();
+		FunctionToolCallback<String, String> callback = FunctionToolCallback
+			.builder("testTool", tool.stringBiFunction())
+			.toolMetadata(ToolMetadata.builder().returnDirect(true).build())
+			.description("test description")
+			.inputType(String.class)
+			.build();
+
+		ToolContext toolContext = new ToolContext(Map.of("foo", "bar"));
+
+		String callResult = callback.call("\"test string param\"", toolContext);
+
+		assertEquals("test string param", tool.calledValue.get());
+		assertEquals("\"return value = test string param\"", callResult);
+		assertEquals(toolContext, tool.calledToolContext.get());
+	}
+
+	@Test
+	void testFunctionToolCall() {
+		TestFunctionTool tool = new TestFunctionTool();
+		FunctionToolCallback<String, String> callback = FunctionToolCallback.builder("testTool", tool.stringFunction())
+			.toolMetadata(ToolMetadata.builder().returnDirect(true).build())
+			.description("test description")
+			.inputType(String.class)
+			.build();
+
+		ToolContext toolContext = new ToolContext(Map.of());
+
+		String callResult = callback.call("\"test string param\"", toolContext);
+
+		assertEquals("test string param", tool.calledValue.get());
+		assertEquals("\"return value = test string param\"", callResult);
+	}
+
+	@Test
+	void testSupplierToolCall() {
+		TestFunctionTool tool = new TestFunctionTool();
+
+		FunctionToolCallback<Void, String> callback = FunctionToolCallback.builder("testTool", tool.stringSupplier())
+			.toolMetadata(ToolMetadata.builder().returnDirect(true).build())
+			.description("test description")
+			.inputType(Void.class)
+			.build();
+
+		ToolContext toolContext = new ToolContext(Map.of());
+
+		String callResult = callback.call("\"test string param\"", toolContext);
+
+		assertEquals("not params", tool.calledValue.get());
+		assertEquals("\"return value = \"", callResult);
+	}
+
+	@Test
+	void testThrowRuntimeException() {
+		TestFunctionTool tool = new TestFunctionTool();
+		FunctionToolCallback<String, Void> callback = FunctionToolCallback
+			.builder("testTool", tool.throwRuntimeException())
+			.toolMetadata(ToolMetadata.builder().returnDirect(true).build())
+			.description("test description")
+			.inputType(String.class)
+			.build();
+
+		assertThatThrownBy(() -> callback.call("\"test string param\"")).hasMessage("test exception")
+			.hasCauseInstanceOf(RuntimeException.class)
+			.asInstanceOf(type(ToolExecutionException.class))
+			.extracting(ToolExecutionException::getToolDefinition)
+			.isEqualTo(callback.getToolDefinition());
+	}
+
+	@Test
+	void testThrowToolExecutionException() {
+		TestFunctionTool tool = new TestFunctionTool();
+		FunctionToolCallback<String, Void> callback = FunctionToolCallback
+			.builder("testTool", tool.throwToolExecutionException())
+			.toolMetadata(ToolMetadata.builder().returnDirect(true).build())
+			.description("test description")
+			.inputType(String.class)
+			.build();
+
+		assertThatThrownBy(() -> callback.call("\"test string param\"")).hasMessage("test exception")
+			.hasCauseInstanceOf(RuntimeException.class)
+			.isInstanceOf(ToolExecutionException.class);
+	}
+
+	static class TestFunctionTool {
+
+		AtomicReference<Object> calledValue = new AtomicReference<>();
+
+		AtomicReference<ToolContext> calledToolContext = new AtomicReference<>();
+
+		public Consumer<String> stringConsumer() {
+			return s -> {
+				calledValue.set(s);
+			};
+		}
+
+		public BiFunction<String, ToolContext, String> stringBiFunction() {
+			return (s, context) -> {
+				calledValue.set(s);
+				calledToolContext.set(context);
+				return "return value = " + s;
+			};
+		}
+
+		public Function<String, String> stringFunction() {
+			return s -> {
+				calledValue.set(s);
+				return "return value = " + s;
+			};
+		}
+
+		public Supplier<String> stringSupplier() {
+			calledValue.set("not params");
+			return () -> "return value = ";
+		}
+
+		public Consumer<String> throwRuntimeException() {
+			return s -> {
+				throw new RuntimeException("test exception");
+			};
+		}
+
+		public Consumer<String> throwToolExecutionException() {
+			return s -> {
+				throw new ToolExecutionException(null, new RuntimeException("test exception"));
+			};
+		}
+
+	}
+
+}