Add Audio helper

This commit is contained in:
Jeremiah Lowin 2025-06-14 09:12:16 -04:00
commit 113778514f
2 changed files with 83 additions and 4 deletions

View file

@ -17,6 +17,7 @@ from fastmcp.utilities.components import FastMCPComponent
from fastmcp.utilities.json_schema import compress_schema
from fastmcp.utilities.logging import get_logger
from fastmcp.utilities.types import (
Audio,
Image,
MCPContent,
find_kwarg_by_type,
@ -273,6 +274,9 @@ def _convert_to_content(
if isinstance(result, Image):
return [result.to_image_content()]
elif isinstance(result, Audio):
return [result.to_audio_content()]
if isinstance(result, list | tuple) and not _process_as_single_item:
# if the result is a list, then it could either be a list of MCP types,
# or a "regular" list that the tool is returning, or a mix of both.
@ -284,7 +288,7 @@ def _convert_to_content(
other_content = []
for item in result:
if isinstance(item, MCPContent | Image):
if isinstance(item, MCPContent | Image | Audio):
mcp_types.append(_convert_to_content(item)[0])
else:
other_content.append(item)

View file

@ -8,7 +8,13 @@ from pathlib import Path
from types import UnionType
from typing import Annotated, TypeAlias, TypeVar, Union, get_args, get_origin
from mcp.types import AudioContent, EmbeddedResource, ImageContent, TextContent
from mcp.types import (
Annotations,
AudioContent,
EmbeddedResource,
ImageContent,
TextContent,
)
from pydantic import BaseModel, ConfigDict, TypeAdapter
T = TypeVar("T")
@ -90,6 +96,7 @@ class Image:
path: str | Path | None = None,
data: bytes | None = None,
format: str | None = None,
annotations: Annotations | None = None,
):
if path is None and data is None:
raise ValueError("Either path or data must be provided")
@ -100,6 +107,7 @@ class Image:
self.data = data
self._format = format
self._mime_type = self._get_mime_type()
self.annotations = annotations
def _get_mime_type(self) -> str:
"""Get MIME type from format or guess from file extension."""
@ -117,7 +125,11 @@ class Image:
}.get(suffix, "application/octet-stream")
return "image/png" # default for raw binary data
def to_image_content(self) -> ImageContent:
def to_image_content(
self,
mime_type: str | None = None,
annotations: Annotations | None = None,
) -> ImageContent:
"""Convert to MCP ImageContent."""
if self.path:
with open(self.path, "rb") as f:
@ -127,4 +139,67 @@ class Image:
else:
raise ValueError("No image data available")
return ImageContent(type="image", data=data, mimeType=self._mime_type)
return ImageContent(
type="image",
data=data,
mimeType=mime_type or self._mime_type,
annotations=annotations or self.annotations,
)
class Audio:
"""Helper class for returning audio from tools."""
def __init__(
self,
path: str | Path | None = None,
data: bytes | None = None,
format: str | None = None,
annotations: Annotations | None = 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()
self.annotations = annotations
def _get_mime_type(self) -> str:
"""Get MIME type from format or guess from file extension."""
if self._format:
return f"audio/{self._format.lower()}"
if self.path:
suffix = self.path.suffix.lower()
return {
".wav": "audio/wav",
".mp3": "audio/mpeg",
".ogg": "audio/ogg",
".m4a": "audio/mp4",
".flac": "audio/flac",
}.get(suffix, "application/octet-stream")
return "audio/wav" # default for raw binary data
def to_audio_content(
self,
mime_type: str | None = None,
annotations: Annotations | None = None,
) -> AudioContent:
if self.path:
with open(self.path, "rb") as f:
data = base64.b64encode(f.read()).decode()
elif self.data is not None:
data = base64.b64encode(self.data).decode()
else:
raise ValueError("No audio data available")
return AudioContent(
type="audio",
data=data,
mimeType=mime_type or self._mime_type,
annotations=annotations or self.annotations,
)