Merge pull request #9 from jlowin/image

Move image to separate module
This commit is contained in:
Jeremiah Lowin 2024-11-30 15:19:35 -05:00 committed by GitHub
commit fe6be28bbe
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 62 additions and 53 deletions

View file

@ -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"]

View file

@ -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__)

View file

@ -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."""

View file

@ -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)

View file

@ -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