mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-23 05:54:19 +02:00
Add Audio helper
This commit is contained in:
parent
f269c074f4
commit
113778514f
2 changed files with 83 additions and 4 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue