diff --git a/mcp/common/src/main/java/org/springframework/ai/mcp/McpToolUtils.java b/mcp/common/src/main/java/org/springframework/ai/mcp/McpToolUtils.java index 8e9a682ae..c3c58accb 100644 --- a/mcp/common/src/main/java/org/springframework/ai/mcp/McpToolUtils.java +++ b/mcp/common/src/main/java/org/springframework/ai/mcp/McpToolUtils.java @@ -18,13 +18,15 @@ package org.springframework.ai.mcp; import java.util.List; import java.util.Map; +import com.fasterxml.jackson.annotation.JsonAlias; +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; import io.micrometer.common.util.StringUtils; import io.modelcontextprotocol.client.McpAsyncClient; import io.modelcontextprotocol.client.McpSyncClient; import io.modelcontextprotocol.server.McpServerFeatures; -import io.modelcontextprotocol.server.McpSyncServerExchange; import io.modelcontextprotocol.server.McpServerFeatures.AsyncToolRegistration; import io.modelcontextprotocol.server.McpServerFeatures.AsyncToolSpecification; +import io.modelcontextprotocol.server.McpSyncServerExchange; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpSchema.Role; import reactor.core.publisher.Mono; @@ -33,6 +35,8 @@ import reactor.core.scheduler.Schedulers; import org.springframework.ai.chat.model.ToolContext; import org.springframework.ai.model.ModelOptionsUtils; import org.springframework.ai.tool.ToolCallback; +import org.springframework.ai.util.json.JsonParser; +import org.springframework.lang.Nullable; import org.springframework.util.CollectionUtils; import org.springframework.util.MimeType; @@ -234,9 +238,22 @@ public final class McpToolUtils { return new McpServerFeatures.SyncToolRegistration(tool, request -> { try { String callResult = toolCallback.call(ModelOptionsUtils.toJsonString(request)); - if (mimeType != null && mimeType.toString().startsWith("image")) { - return new McpSchema.CallToolResult(List - .of(new McpSchema.ImageContent(List.of(Role.ASSISTANT), null, callResult, mimeType.toString())), + String imgData = callResult; + if (mimeType != null && "image".equals(mimeType.getType())) { + String imgType = mimeType.toString(); + if (callResult.startsWith("{") && callResult.endsWith("}")) { + // This is most likely a JSON structure: + // let's try to parse it as a base64 wrapper. + var b64Struct = JsonParser.fromJson(callResult, Base64Wrapper.class); + if (b64Struct.mimeType() != null && b64Struct.data() != null + && b64Struct.mimeType.getType().equals("image")) { + // Get the base64 encoded image as is. + imgType = b64Struct.mimeType().toString(); + imgData = b64Struct.data(); + } + } + return new McpSchema.CallToolResult( + List.of(new McpSchema.ImageContent(List.of(Role.ASSISTANT), null, imgData, imgType)), false); } return new McpSchema.CallToolResult(List.of(new McpSchema.TextContent(callResult)), false); @@ -547,4 +564,9 @@ public final class McpToolUtils { return List.of((new AsyncMcpToolCallbackProvider(asyncMcpClients).getToolCallbacks())); } + @JsonIgnoreProperties(ignoreUnknown = true) + private record Base64Wrapper(@JsonAlias("mimetype") @Nullable MimeType mimeType, @JsonAlias( { + "base64", "b64", "imageData" }) @Nullable String data){ + } + } diff --git a/spring-ai-core/src/test/java/org/springframework/ai/tool/execution/DefaultToolCallResultConverterTests.java b/spring-ai-core/src/test/java/org/springframework/ai/tool/execution/DefaultToolCallResultConverterTests.java new file mode 100644 index 000000000..69c081da5 --- /dev/null +++ b/spring-ai-core/src/test/java/org/springframework/ai/tool/execution/DefaultToolCallResultConverterTests.java @@ -0,0 +1,132 @@ +package org.springframework.ai.tool.execution; + +import org.junit.jupiter.api.Test; +import org.springframework.ai.util.json.JsonParser; +import org.springframework.util.MimeType; +import org.springframework.util.MimeTypeUtils; + +import javax.imageio.ImageIO; +import java.awt.Color; +import java.awt.image.BufferedImage; +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.util.Base64; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Unit tests for {@link DefaultToolCallResultConverter}. + * + * @author Thomas Vitale + */ +class DefaultToolCallResultConverterTests { + + private final DefaultToolCallResultConverter converter = new DefaultToolCallResultConverter(); + + @Test + void convertWithNullReturnTypeShouldReturn() { + String result = converter.convert(null, null); + assertThat(result).isEqualTo("null"); + } + + @Test + void convertVoidReturnTypeShouldReturnDone() { + String result = converter.convert(null, void.class); + assertThat(result).isEqualTo("Done"); + } + + @Test + void convertStringReturnTypeShouldReturnJson() { + String result = converter.convert("test", String.class); + assertThat(result).isEqualTo("\"test\""); + } + + @Test + void convertNullReturnValueShouldReturnNullJson() { + String result = converter.convert(null, String.class); + assertThat(result).isEqualTo("null"); + } + + @Test + void convertObjectReturnTypeShouldReturnJson() { + TestObject testObject = new TestObject("test", 42); + String result = converter.convert(testObject, TestObject.class); + assertThat(result).containsIgnoringWhitespaces(""" + "name": "test" + """).containsIgnoringWhitespaces(""" + "value": 42 + """); + } + + @Test + void convertCollectionReturnTypeShouldReturnJson() { + List testList = List.of("one", "two", "three"); + String result = converter.convert(testList, List.class); + assertThat(result).isEqualTo(""" + ["one","two","three"] + """.trim()); + } + + @Test + void convertMapReturnTypeShouldReturnJson() { + Map testMap = Map.of("one", 1, "two", 2); + String result = converter.convert(testMap, Map.class); + assertThat(result).containsIgnoringWhitespaces(""" + "one": 1 + """).containsIgnoringWhitespaces(""" + "two": 2 + """); + } + + @Test + void convertImageShouldReturnBase64Image() throws IOException { + // We don't want any AWT windows. + System.setProperty("java.awt.headless", "true"); + + var img = new BufferedImage(64, 64, BufferedImage.TYPE_4BYTE_ABGR); + var g = img.createGraphics(); + g.setColor(Color.WHITE); + g.fillRect(0, 0, 64, 64); + g.dispose(); + String result = converter.convert(img, BufferedImage.class); + + var b64Struct = JsonParser.fromJson(result, Base64Wrapper.class); + assertThat(b64Struct.mimeType).isEqualTo(MimeTypeUtils.IMAGE_PNG); + assertThat(b64Struct.data).isNotNull(); + + var imgData = Base64.getDecoder().decode(b64Struct.data); + assertThat(imgData.length).isNotZero(); + + var imgRes = ImageIO.read(new ByteArrayInputStream(imgData)); + assertThat(imgRes.getWidth()).isEqualTo(64); + assertThat(imgRes.getHeight()).isEqualTo(64); + assertThat(imgRes.getRGB(0, 0)).isEqualTo(img.getRGB(0, 0)); + } + + record Base64Wrapper(MimeType mimeType, String data) { + } + + static class TestObject { + + private final String name; + + private final int value; + + TestObject(String name, int value) { + this.name = name; + this.value = value; + } + + public String getName() { + return name; + } + + public int getValue() { + return value; + } + + } + +} diff --git a/spring-ai-model/src/main/java/org/springframework/ai/tool/execution/DefaultToolCallResultConverter.java b/spring-ai-model/src/main/java/org/springframework/ai/tool/execution/DefaultToolCallResultConverter.java index 1c975291a..a70316d9e 100644 --- a/spring-ai-model/src/main/java/org/springframework/ai/tool/execution/DefaultToolCallResultConverter.java +++ b/spring-ai-model/src/main/java/org/springframework/ai/tool/execution/DefaultToolCallResultConverter.java @@ -16,14 +16,19 @@ package org.springframework.ai.tool.execution; -import java.lang.reflect.Type; - import org.slf4j.Logger; import org.slf4j.LoggerFactory; - import org.springframework.ai.util.json.JsonParser; import org.springframework.lang.Nullable; +import javax.imageio.ImageIO; +import java.awt.image.RenderedImage; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.lang.reflect.Type; +import java.util.Base64; +import java.util.Map; + /** * A default implementation of {@link ToolCallResultConverter}. * @@ -40,6 +45,17 @@ public final class DefaultToolCallResultConverter implements ToolCallResultConve logger.debug("The tool has no return type. Converting to conventional response."); return "Done"; } + if (result instanceof RenderedImage) { + final var buf = new ByteArrayOutputStream(1024 * 4); + try { + ImageIO.write((RenderedImage) result, "PNG", buf); + } + catch (IOException e) { + return "Failed to convert tool result to a base64 image: " + e.getMessage(); + } + final var imgB64 = Base64.getEncoder().encodeToString(buf.toByteArray()); + return JsonParser.toJson(Map.of("mimeType", "image/png", "data", imgB64)); + } else { logger.debug("Converting tool result to JSON."); return JsonParser.toJson(result);