Spaces:
Running
Running
Jeremiah Lowin commited on
Commit ·
c52b60b
1
Parent(s): 0d751b2
Improve handling of image results from tools
Browse files- src/fastmcp/server.py +17 -0
- tests/test_server.py +33 -0
src/fastmcp/server.py
CHANGED
|
@@ -160,6 +160,23 @@ class FastMCP:
|
|
| 160 |
if isinstance(value, (list, tuple)):
|
| 161 |
if all(isinstance(x, (TextContent, ImageContent)) for x in value):
|
| 162 |
return value
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 163 |
|
| 164 |
# Single content type
|
| 165 |
if isinstance(value, (TextContent, ImageContent)):
|
|
|
|
| 160 |
if isinstance(value, (list, tuple)):
|
| 161 |
if all(isinstance(x, (TextContent, ImageContent)) for x in value):
|
| 162 |
return value
|
| 163 |
+
# Handle mixed content including Image objects
|
| 164 |
+
result = []
|
| 165 |
+
for item in value:
|
| 166 |
+
if isinstance(item, (TextContent, ImageContent)):
|
| 167 |
+
result.append(item)
|
| 168 |
+
elif isinstance(item, Image):
|
| 169 |
+
result.append(item.to_image_content())
|
| 170 |
+
else:
|
| 171 |
+
result.append(
|
| 172 |
+
TextContent(
|
| 173 |
+
type="text",
|
| 174 |
+
text=json.dumps(
|
| 175 |
+
item, indent=2, default=pydantic.json.pydantic_encoder
|
| 176 |
+
),
|
| 177 |
+
)
|
| 178 |
+
)
|
| 179 |
+
return result
|
| 180 |
|
| 181 |
# Single content type
|
| 182 |
if isinstance(value, (TextContent, ImageContent)):
|
tests/test_server.py
CHANGED
|
@@ -161,6 +161,39 @@ class TestServerTools:
|
|
| 161 |
assert result.content[1].mimeType == "image/png"
|
| 162 |
assert result.content[1].data == "abc"
|
| 163 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 164 |
|
| 165 |
class TestServerResources:
|
| 166 |
async def test_text_resource(self):
|
|
|
|
| 161 |
assert result.content[1].mimeType == "image/png"
|
| 162 |
assert result.content[1].data == "abc"
|
| 163 |
|
| 164 |
+
async def test_tool_mixed_list_with_image(self, tmp_path: Path):
|
| 165 |
+
"""Test that lists containing Image objects and other types are handled correctly"""
|
| 166 |
+
# Create a test image
|
| 167 |
+
image_path = tmp_path / "test.png"
|
| 168 |
+
image_path.write_bytes(b"test image data")
|
| 169 |
+
|
| 170 |
+
def mixed_list_fn() -> list:
|
| 171 |
+
return [
|
| 172 |
+
"text message",
|
| 173 |
+
Image(image_path),
|
| 174 |
+
{"key": "value"},
|
| 175 |
+
TextContent(type="text", text="direct content"),
|
| 176 |
+
]
|
| 177 |
+
|
| 178 |
+
mcp = FastMCP()
|
| 179 |
+
mcp.add_tool(mixed_list_fn)
|
| 180 |
+
async with client_session(mcp._mcp_server) as client:
|
| 181 |
+
result = await client.call_tool("mixed_list_fn", {})
|
| 182 |
+
assert len(result.content) == 4
|
| 183 |
+
# Check text conversion
|
| 184 |
+
assert result.content[0].type == "text"
|
| 185 |
+
assert '"text message"' in result.content[0].text
|
| 186 |
+
# Check image conversion
|
| 187 |
+
assert result.content[1].type == "image"
|
| 188 |
+
assert result.content[1].mimeType == "image/png"
|
| 189 |
+
assert base64.b64decode(result.content[1].data) == b"test image data"
|
| 190 |
+
# Check dict conversion
|
| 191 |
+
assert result.content[2].type == "text"
|
| 192 |
+
assert '"key": "value"' in result.content[2].text
|
| 193 |
+
# Check direct TextContent
|
| 194 |
+
assert result.content[3].type == "text"
|
| 195 |
+
assert result.content[3].text == "direct content"
|
| 196 |
+
|
| 197 |
|
| 198 |
class TestServerResources:
|
| 199 |
async def test_text_resource(self):
|