Repository: jlowin/fastmcp
Pull request #9 (2024-11-30 20:06:44)

Move image to separate module


Patch:
diff --git a/src/fastmcp/__init__.py b/src/fastmcp/__init__.py
index 3946cf5..ebbb267 100644
--- a/src/fastmcp/__init__.py
+++ b/src/fastmcp/__init__.py
@@ -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"]
diff --git a/src/fastmcp/server.py b/src/fastmcp/server.py
index b12fb1c..8b7b33b 100644
--- a/src/fastmcp/server.py
+++ b/src/fastmcp/server.py
@@ -24,8 +24,9 @@ from pydantic.networks import _BaseUrl
 from fastmcp.exceptions import ResourceError
 from fastmcp.resources import Resource, ResourceManager
 from fastmcp.resources.types import FunctionResource
-from fastmcp.tools import ToolManager, Image
+from fastmcp.tools import ToolManager
 from fastmcp.utilities.logging import get_logger, configure_logging
+from fastmcp.utilities.types import Image
 
 logger = get_logger(__name__)
 
diff --git a/src/fastmcp/tools.py b/src/fastmcp/tools.py
index 7f0c55a..351d3eb 100644
--- a/src/fastmcp/tools.py
+++ b/src/fastmcp/tools.py
@@ -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
@@ -14,52 +11,6 @@ from .utilities.logging import get_logger
 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."""
 
diff --git a/src/fastmcp/utilities/types.py b/src/fastmcp/utilities/types.py
new file mode 100644
index 0000000..60bbf15
--- /dev/null
+++ b/src/fastmcp/utilities/types.py
@@ -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)


Test patch:
diff --git a/tests/test_server.py b/tests/test_server.py
index 04b1b11..2fa16b7 100644
--- a/tests/test_server.py
+++ b/tests/test_server.py
@@ -3,7 +3,7 @@ from mcp.shared.memory import (
 )
 from fastmcp import FastMCP
 from fastmcp.resources import FileResource, FunctionResource
-from fastmcp.tools import Image
+from fastmcp.utilities.types import Image
 from mcp.types import TextContent, ImageContent
 import pytest
 from pydantic import BaseModel