Skip to content

Commit

Permalink
Move image to separate module
Browse files Browse the repository at this point in the history
  • Loading branch information
jlowin committed Nov 30, 2024
1 parent 8a1237b commit eadec9b
Show file tree
Hide file tree
Showing 3 changed files with 59 additions and 51 deletions.
6 changes: 5 additions & 1 deletion src/fastmcp/__init__.py
Original file line number Diff line number Diff line change
@@ -1,2 +1,6 @@
"""FastMCP - A more ergonomic interface for MCP servers."""

from .server import FastMCP
from .tools import Image
from .utilities.types import Image

__all__ = ["FastMCP", "Image"]
51 changes: 1 addition & 50 deletions src/fastmcp/tools.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,8 @@
"""Tool management for FastMCP."""

import base64
import inspect
from pathlib import Path
from typing import Any, Callable, Dict, Optional, Union
from typing import Any, Callable, Dict, Optional

from mcp.types import ImageContent
from pydantic import BaseModel, Field, TypeAdapter, validate_call

from .exceptions import ToolError
Expand All @@ -14,52 +11,6 @@
logger = get_logger(__name__)


class Image:
"""Helper class for returning images from tools."""

def __init__(
self,
path: Optional[Union[str, Path]] = None,
data: Optional[bytes] = None,
format: Optional[str] = None,
):
if path is None and data is None:
raise ValueError("Either path or data must be provided")
if path is not None and data is not None:
raise ValueError("Only one of path or data can be provided")

self.path = Path(path) if path else None
self.data = data
self._format = format
self._mime_type = self._get_mime_type()

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")
return "image/png" # default for raw binary data

def to_image_content(self) -> ImageContent:
"""Convert to MCP ImageContent."""
if self.path:
with open(self.path, "rb") as f:
data = base64.b64encode(f.read()).decode()
else:
data = base64.b64encode(self.data).decode()

return ImageContent(type="image", data=data, mimeType=self._mime_type)


class Tool(BaseModel):
"""Internal tool registration info."""

Expand Down
53 changes: 53 additions & 0 deletions src/fastmcp/utilities/types.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,53 @@
"""Common types used across FastMCP."""

import base64
from pathlib import Path
from typing import Optional, Union

from mcp.types import ImageContent


class Image:
"""Helper class for returning images from tools."""

def __init__(
self,
path: Optional[Union[str, Path]] = None,
data: Optional[bytes] = None,
format: Optional[str] = None,
):
if path is None and data is None:
raise ValueError("Either path or data must be provided")
if path is not None and data is not None:
raise ValueError("Only one of path or data can be provided")

self.path = Path(path) if path else None
self.data = data
self._format = format
self._mime_type = self._get_mime_type()

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")
return "image/png" # default for raw binary data

def to_image_content(self) -> ImageContent:
"""Convert to MCP ImageContent."""
if self.path:
with open(self.path, "rb") as f:
data = base64.b64encode(f.read()).decode()
else:
data = base64.b64encode(self.data).decode()

return ImageContent(type="image", data=data, mimeType=self._mime_type)

0 comments on commit eadec9b

Please sign in to comment.