Skip to content

Commit

Permalink
Merge pull request #181 from zzstoatzz/update-fastmcp-server-typing
Browse files Browse the repository at this point in the history
satisfy `pyright` untyped decorators in `mcp.server.fastmcp.server`
  • Loading branch information
dsp-ant authored Feb 4, 2025
2 parents 51e65c8 + ca06001 commit 08042c3
Show file tree
Hide file tree
Showing 3 changed files with 16 additions and 9 deletions.
3 changes: 3 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -162,3 +162,6 @@ cython_debug/
# and can be added to the global gitignore or merged into this file. For a more nuclear
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
#.idea/

# vscode
.vscode/
19 changes: 11 additions & 8 deletions src/mcp/server/fastmcp/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
from mcp.server.stdio import stdio_server
from mcp.shared.context import RequestContext
from mcp.types import (
AnyFunction,
EmbeddedResource,
GetPromptResult,
ImageContent,
Expand Down Expand Up @@ -165,7 +166,7 @@ def get_context(self) -> "Context":
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()
Expand Down Expand Up @@ -214,7 +215,7 @@ async def read_resource(self, uri: AnyUrl | str) -> ReadResourceContents:

def add_tool(
self,
fn: Callable,
fn: AnyFunction,
name: str | None = None,
description: str | None = None,
) -> None:
Expand All @@ -230,7 +231,9 @@ def add_tool(
"""
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[[AnyFunction], AnyFunction]:
"""Decorator to register a tool.
Tools can optionally request a Context object by adding a parameter with the
Expand Down Expand Up @@ -263,7 +266,7 @@ async def async_tool(x: int, context: Context) -> str:
"Did you forget to call it? Use @tool() instead of @tool"
)

def decorator(fn: Callable) -> Callable:
def decorator(fn: AnyFunction) -> AnyFunction:
self.add_tool(fn, name=name, description=description)
return fn

Expand All @@ -284,7 +287,7 @@ def resource(
name: str | None = None,
description: str | None = None,
mime_type: str | None = None,
) -> Callable:
) -> Callable[[AnyFunction], AnyFunction]:
"""Decorator to register a function as a resource.
The function will be called when the resource is read to generate its content.
Expand Down Expand Up @@ -328,7 +331,7 @@ async def get_weather(city: str) -> str:
"Did you forget to call it? Use @resource('uri') instead of @resource"
)

def decorator(fn: Callable) -> Callable:
def decorator(fn: AnyFunction) -> AnyFunction:
# Check if this should be a template
has_uri_params = "{" in uri and "}" in uri
has_func_params = bool(inspect.signature(fn).parameters)
Expand Down Expand Up @@ -376,7 +379,7 @@ def add_prompt(self, prompt: Prompt) -> None:

def prompt(
self, name: str | None = None, description: str | None = None
) -> Callable:
) -> Callable[[AnyFunction], AnyFunction]:
"""Decorator to register a prompt.
Args:
Expand Down Expand Up @@ -417,7 +420,7 @@ async def analyze_file(path: str) -> list[Message]:
"Did you forget to call it? Use @prompt() instead of @prompt"
)

def decorator(func: Callable) -> Callable:
def decorator(func: AnyFunction) -> AnyFunction:
prompt = Prompt.from_function(func, name=name, description=description)
self.add_prompt(prompt)
return func
Expand Down
3 changes: 2 additions & 1 deletion src/mcp/types.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import Annotated, Any, Generic, Literal, TypeVar
from typing import Annotated, Any, Callable, Generic, Literal, TypeAlias, TypeVar

from pydantic import BaseModel, ConfigDict, Field, FileUrl, RootModel
from pydantic.networks import AnyUrl
Expand Down Expand Up @@ -27,6 +27,7 @@
Cursor = str
Role = Literal["user", "assistant"]
RequestId = str | int
AnyFunction: TypeAlias = Callable[..., Any]


class RequestParams(BaseModel):
Expand Down

0 comments on commit 08042c3

Please sign in to comment.