Skip to content

Commit

Permalink
Improve handling of image results from tools
Browse files Browse the repository at this point in the history
  • Loading branch information
jlowin committed Nov 30, 2024
1 parent 0eaa841 commit be6e0c1
Show file tree
Hide file tree
Showing 2 changed files with 50 additions and 0 deletions.
17 changes: 17 additions & 0 deletions src/fastmcp/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -160,6 +160,23 @@ def _convert_to_content(
if isinstance(value, (list, tuple)):
if all(isinstance(x, (TextContent, ImageContent)) for x in value):
return value
# Handle mixed content including Image objects
result = []
for item in value:
if isinstance(item, (TextContent, ImageContent)):
result.append(item)
elif isinstance(item, Image):
result.append(item.to_image_content())
else:
result.append(
TextContent(
type="text",
text=json.dumps(
item, indent=2, default=pydantic.json.pydantic_encoder
),
)
)
return result

# Single content type
if isinstance(value, (TextContent, ImageContent)):
Expand Down
33 changes: 33 additions & 0 deletions tests/test_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,39 @@ async def test_tool_mixed_content(self):
assert result.content[1].mimeType == "image/png"
assert result.content[1].data == "abc"

async def test_tool_mixed_list_with_image(self, tmp_path: Path):
"""Test that lists containing Image objects and other types are handled correctly"""
# Create a test image
image_path = tmp_path / "test.png"
image_path.write_bytes(b"test image data")

def mixed_list_fn() -> list:
return [
"text message",
Image(image_path),
{"key": "value"},
TextContent(type="text", text="direct content"),
]

mcp = FastMCP()
mcp.add_tool(mixed_list_fn)
async with client_session(mcp._mcp_server) as client:
result = await client.call_tool("mixed_list_fn", {})
assert len(result.content) == 4
# Check text conversion
assert result.content[0].type == "text"
assert '"text message"' in result.content[0].text
# Check image conversion
assert result.content[1].type == "image"
assert result.content[1].mimeType == "image/png"
assert base64.b64decode(result.content[1].data) == b"test image data"
# Check dict conversion
assert result.content[2].type == "text"
assert '"key": "value"' in result.content[2].text
# Check direct TextContent
assert result.content[3].type == "text"
assert result.content[3].text == "direct content"


class TestServerResources:
async def test_text_resource(self):
Expand Down

0 comments on commit be6e0c1

Please sign in to comment.