mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-21 21:14:17 +02:00
Added to_data_uri method for Image class. (#2227)
* Added to_data_uri method and path_to_data_uri classmethod for Image class. * Removed path_to_data_uri classmethod and modified _get_mime_type to use mimetypes.guess_type function instead of hardcoded dictionary. * Register image/webp with mimetypes before guess_type to support WEBP mimetype detection on Python 3.10. * Improved branch coverage for Image.to_data_uri. * Added Image._to_data_uri example in the docs.
This commit is contained in:
parent
08c49e62e8
commit
321f404046
3 changed files with 75 additions and 30 deletions
|
|
@ -115,6 +115,7 @@ For small icons or when you want to embed the icon directly, use data URIs:
|
|||
|
||||
```python
|
||||
from mcp.types import Icon
|
||||
from fastmcp.utilities.types import Image
|
||||
|
||||
# SVG icon as data URI
|
||||
svg_icon = Icon(
|
||||
|
|
@ -126,4 +127,13 @@ svg_icon = Icon(
|
|||
def my_tool() -> str:
|
||||
"""A tool with an embedded SVG icon."""
|
||||
return "result"
|
||||
|
||||
# Generating a data URI from a local image file.
|
||||
img = Image(path="./assets/brand/favicon.png")
|
||||
icon = Icon(src=img.to_data_uri())
|
||||
|
||||
@mcp.tool(icons=[icon])
|
||||
def file_icon_tool() -> str:
|
||||
"""A tool with an icon generated from a local file."""
|
||||
return "result"
|
||||
```
|
||||
|
|
|
|||
|
|
@ -190,34 +190,33 @@ class Image:
|
|||
if path is not None and data is not None:
|
||||
raise ValueError("Only one of path or data can be provided")
|
||||
|
||||
self.path = Path(os.path.expandvars(str(path))).expanduser() if path else None
|
||||
self.path = self._get_expanded_path(path)
|
||||
self.data = data
|
||||
self._format = format
|
||||
self._mime_type = self._get_mime_type()
|
||||
self.annotations = annotations
|
||||
|
||||
@staticmethod
|
||||
def _get_expanded_path(path: str | Path | None) -> Path | None:
|
||||
"""Expand environment variables and user home in path."""
|
||||
return Path(os.path.expandvars(str(path))).expanduser() if path else None
|
||||
|
||||
def _get_mime_type(self) -> str:
|
||||
"""Get MIME type from format or guess from file extension."""
|
||||
if self._format:
|
||||
return f"image/{self._format.lower()}"
|
||||
|
||||
if self.path:
|
||||
suffix = self.path.suffix.lower()
|
||||
return {
|
||||
".png": "image/png",
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".gif": "image/gif",
|
||||
".webp": "image/webp",
|
||||
}.get(suffix, "application/octet-stream")
|
||||
# Workaround for WEBP in Py3.10
|
||||
mimetypes.add_type("image/webp", ".webp")
|
||||
resp = mimetypes.guess_type(self.path, strict=False)
|
||||
if resp and resp[0] is not None:
|
||||
return resp[0]
|
||||
return "application/octet-stream"
|
||||
return "image/png" # default for raw binary data
|
||||
|
||||
def to_image_content(
|
||||
self,
|
||||
mime_type: str | None = None,
|
||||
annotations: Annotations | None = None,
|
||||
) -> mcp.types.ImageContent:
|
||||
"""Convert to MCP ImageContent."""
|
||||
def _get_data(self) -> str:
|
||||
"""Get raw image data as base64-encoded string."""
|
||||
if self.path:
|
||||
with open(self.path, "rb") as f:
|
||||
data = base64.b64encode(f.read()).decode()
|
||||
|
|
@ -225,6 +224,15 @@ class Image:
|
|||
data = base64.b64encode(self.data).decode()
|
||||
else:
|
||||
raise ValueError("No image data available")
|
||||
return data
|
||||
|
||||
def to_image_content(
|
||||
self,
|
||||
mime_type: str | None = None,
|
||||
annotations: Annotations | None = None,
|
||||
) -> mcp.types.ImageContent:
|
||||
"""Convert to MCP ImageContent."""
|
||||
data = self._get_data()
|
||||
|
||||
return mcp.types.ImageContent(
|
||||
type="image",
|
||||
|
|
@ -233,6 +241,11 @@ class Image:
|
|||
annotations=annotations or self.annotations,
|
||||
)
|
||||
|
||||
def to_data_uri(self, mime_type: str | None = None) -> str:
|
||||
"""Get image as a data URI."""
|
||||
data = self._get_data()
|
||||
return f"data:{mime_type or self._mime_type};base64,{data}"
|
||||
|
||||
|
||||
class Audio:
|
||||
"""Helper class for returning audio from tools."""
|
||||
|
|
|
|||
|
|
@ -177,22 +177,23 @@ class TestImage:
|
|||
):
|
||||
Image(path="test.png", data=b"test")
|
||||
|
||||
def test_get_mime_type_from_path(self, tmp_path):
|
||||
@pytest.mark.parametrize(
|
||||
"extension,mime_type",
|
||||
[
|
||||
(".png", "image/png"),
|
||||
(".jpg", "image/jpeg"),
|
||||
(".jpeg", "image/jpeg"),
|
||||
(".gif", "image/gif"),
|
||||
(".webp", "image/webp"),
|
||||
(".unknown", "application/octet-stream"),
|
||||
],
|
||||
)
|
||||
def test_get_mime_type_from_path(self, tmp_path, extension, mime_type):
|
||||
"""Test MIME type detection from file extension."""
|
||||
extensions = {
|
||||
".png": "image/png",
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".gif": "image/gif",
|
||||
".webp": "image/webp",
|
||||
".unknown": "application/octet-stream",
|
||||
}
|
||||
|
||||
for ext, mime in extensions.items():
|
||||
path = tmp_path / f"test{ext}"
|
||||
path.write_bytes(b"fake image data")
|
||||
img = Image(path=path)
|
||||
assert img._mime_type == mime
|
||||
path = tmp_path / f"test{extension}"
|
||||
path.write_bytes(b"fake image data")
|
||||
img = Image(path=path)
|
||||
assert img._mime_type == mime_type
|
||||
|
||||
def test_to_image_content(self, tmp_path, monkeypatch):
|
||||
"""Test conversion to ImageContent."""
|
||||
|
|
@ -227,6 +228,27 @@ class TestImage:
|
|||
with pytest.raises(ValueError, match="No image data available"):
|
||||
img.to_image_content()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"mime_type,fname,expected_mime",
|
||||
[
|
||||
(None, "test.png", "image/png"),
|
||||
("image/jpeg", "test.unknown", "image/jpeg"),
|
||||
],
|
||||
)
|
||||
def test_to_data_uri(self, tmp_path, mime_type, fname, expected_mime):
|
||||
"""Test conversion to data URI."""
|
||||
img_path = tmp_path / fname
|
||||
test_data = b"fake image data"
|
||||
img_path.write_bytes(test_data)
|
||||
|
||||
img = Image(path=img_path)
|
||||
data_uri = img.to_data_uri(mime_type=mime_type)
|
||||
|
||||
expected_data_uri = (
|
||||
f"data:{expected_mime};base64,{base64.b64encode(test_data).decode()}"
|
||||
)
|
||||
assert data_uri == expected_data_uri
|
||||
|
||||
|
||||
class TestAudio:
|
||||
def test_audio_initialization_with_path(self):
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue