This commit is contained in:
zzstoatzz
2025-01-29 23:01:07 -06:00
parent 51e65c8310
commit 00a44692e7
2 changed files with 16 additions and 8 deletions

View File

@@ -12,6 +12,7 @@ import uvicorn
from pydantic import BaseModel, Field
from pydantic.networks import AnyUrl
from pydantic_settings import BaseSettings, SettingsConfigDict
from typing_extensions import TypeAlias
from mcp.server.fastmcp.exceptions import ResourceError
from mcp.server.fastmcp.prompts import Prompt, PromptManager
@@ -48,6 +49,8 @@ from mcp.types import (
logger = get_logger(__name__)
_Function: TypeAlias = Callable[..., Any]
class Settings(BaseSettings):
"""FastMCP server settings.
@@ -165,7 +168,7 @@ class FastMCP:
return Context(request_context=request_context, fastmcp=self)
async def call_tool(
self, name: str, arguments: dict
self, name: str, arguments: dict[str, Any]
) -> Sequence[TextContent | ImageContent | EmbeddedResource]:
"""Call a tool by name with arguments."""
context = self.get_context()
@@ -214,7 +217,7 @@ class FastMCP:
def add_tool(
self,
fn: Callable,
fn: _Function,
name: str | None = None,
description: str | None = None,
) -> None:
@@ -230,7 +233,9 @@ class FastMCP:
"""
self._tool_manager.add_tool(fn, name=name, description=description)
def tool(self, name: str | None = None, description: str | None = None) -> Callable:
def tool(
self, name: str | None = None, description: str | None = None
) -> Callable[[_Function], _Function]:
"""Decorator to register a tool.
Tools can optionally request a Context object by adding a parameter with the
@@ -263,7 +268,7 @@ class FastMCP:
"Did you forget to call it? Use @tool() instead of @tool"
)
def decorator(fn: Callable) -> Callable:
def decorator(fn: _Function) -> _Function:
self.add_tool(fn, name=name, description=description)
return fn
@@ -284,7 +289,7 @@ class FastMCP:
name: str | None = None,
description: str | None = None,
mime_type: str | None = None,
) -> Callable:
) -> Callable[[_Function], _Function]:
"""Decorator to register a function as a resource.
The function will be called when the resource is read to generate its content.
@@ -328,7 +333,7 @@ class FastMCP:
"Did you forget to call it? Use @resource('uri') instead of @resource"
)
def decorator(fn: Callable) -> Callable:
def decorator(fn: _Function) -> _Function:
# Check if this should be a template
has_uri_params = "{" in uri and "}" in uri
has_func_params = bool(inspect.signature(fn).parameters)
@@ -376,7 +381,7 @@ class FastMCP:
def prompt(
self, name: str | None = None, description: str | None = None
) -> Callable:
) -> Callable[[_Function], _Function]:
"""Decorator to register a prompt.
Args:
@@ -417,7 +422,7 @@ class FastMCP:
"Did you forget to call it? Use @prompt() instead of @prompt"
)
def decorator(func: Callable) -> Callable:
def decorator(func: _Function) -> _Function:
prompt = Prompt.from_function(func, name=name, description=description)
self.add_prompt(prompt)
return func