diff --git a/.gitignore b/.gitignore index f27f895..bb25f9e 100644 --- a/.gitignore +++ b/.gitignore @@ -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/ \ No newline at end of file diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py index 06b16e0..aa7c79b 100644 --- a/src/mcp/server/fastmcp/server.py +++ b/src/mcp/server/fastmcp/server.py @@ -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, @@ -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() @@ -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: @@ -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 @@ -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 @@ -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. @@ -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) @@ -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: @@ -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 diff --git a/src/mcp/types.py b/src/mcp/types.py index a4fa254..d1157aa 100644 --- a/src/mcp/types.py +++ b/src/mcp/types.py @@ -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 @@ -27,6 +27,7 @@ Cursor = str Role = Literal["user", "assistant"] RequestId = str | int +AnyFunction: TypeAlias = Callable[..., Any] class RequestParams(BaseModel):