centralize type

This commit is contained in:
zzstoatzz
2025-02-03 14:27:07 -06:00
parent 00a44692e7
commit ca060014bb
2 changed files with 10 additions and 11 deletions

View File

@@ -12,7 +12,6 @@ import uvicorn
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from pydantic.networks import AnyUrl from pydantic.networks import AnyUrl
from pydantic_settings import BaseSettings, SettingsConfigDict from pydantic_settings import BaseSettings, SettingsConfigDict
from typing_extensions import TypeAlias
from mcp.server.fastmcp.exceptions import ResourceError from mcp.server.fastmcp.exceptions import ResourceError
from mcp.server.fastmcp.prompts import Prompt, PromptManager from mcp.server.fastmcp.prompts import Prompt, PromptManager
@@ -26,6 +25,7 @@ from mcp.server.sse import SseServerTransport
from mcp.server.stdio import stdio_server from mcp.server.stdio import stdio_server
from mcp.shared.context import RequestContext from mcp.shared.context import RequestContext
from mcp.types import ( from mcp.types import (
AnyFunction,
EmbeddedResource, EmbeddedResource,
GetPromptResult, GetPromptResult,
ImageContent, ImageContent,
@@ -49,8 +49,6 @@ from mcp.types import (
logger = get_logger(__name__) logger = get_logger(__name__)
_Function: TypeAlias = Callable[..., Any]
class Settings(BaseSettings): class Settings(BaseSettings):
"""FastMCP server settings. """FastMCP server settings.
@@ -217,7 +215,7 @@ class FastMCP:
def add_tool( def add_tool(
self, self,
fn: _Function, fn: AnyFunction,
name: str | None = None, name: str | None = None,
description: str | None = None, description: str | None = None,
) -> None: ) -> None:
@@ -235,7 +233,7 @@ class FastMCP:
def tool( def tool(
self, name: str | None = None, description: str | None = None self, name: str | None = None, description: str | None = None
) -> Callable[[_Function], _Function]: ) -> Callable[[AnyFunction], AnyFunction]:
"""Decorator to register a tool. """Decorator to register a tool.
Tools can optionally request a Context object by adding a parameter with the Tools can optionally request a Context object by adding a parameter with the
@@ -268,7 +266,7 @@ class FastMCP:
"Did you forget to call it? Use @tool() instead of @tool" "Did you forget to call it? Use @tool() instead of @tool"
) )
def decorator(fn: _Function) -> _Function: def decorator(fn: AnyFunction) -> AnyFunction:
self.add_tool(fn, name=name, description=description) self.add_tool(fn, name=name, description=description)
return fn return fn
@@ -289,7 +287,7 @@ class FastMCP:
name: str | None = None, name: str | None = None,
description: str | None = None, description: str | None = None,
mime_type: str | None = None, mime_type: str | None = None,
) -> Callable[[_Function], _Function]: ) -> Callable[[AnyFunction], AnyFunction]:
"""Decorator to register a function as a resource. """Decorator to register a function as a resource.
The function will be called when the resource is read to generate its content. The function will be called when the resource is read to generate its content.
@@ -333,7 +331,7 @@ class FastMCP:
"Did you forget to call it? Use @resource('uri') instead of @resource" "Did you forget to call it? Use @resource('uri') instead of @resource"
) )
def decorator(fn: _Function) -> _Function: def decorator(fn: AnyFunction) -> AnyFunction:
# Check if this should be a template # Check if this should be a template
has_uri_params = "{" in uri and "}" in uri has_uri_params = "{" in uri and "}" in uri
has_func_params = bool(inspect.signature(fn).parameters) has_func_params = bool(inspect.signature(fn).parameters)
@@ -381,7 +379,7 @@ class FastMCP:
def prompt( def prompt(
self, name: str | None = None, description: str | None = None self, name: str | None = None, description: str | None = None
) -> Callable[[_Function], _Function]: ) -> Callable[[AnyFunction], AnyFunction]:
"""Decorator to register a prompt. """Decorator to register a prompt.
Args: Args:
@@ -422,7 +420,7 @@ class FastMCP:
"Did you forget to call it? Use @prompt() instead of @prompt" "Did you forget to call it? Use @prompt() instead of @prompt"
) )
def decorator(func: _Function) -> _Function: def decorator(func: AnyFunction) -> AnyFunction:
prompt = Prompt.from_function(func, name=name, description=description) prompt = Prompt.from_function(func, name=name, description=description)
self.add_prompt(prompt) self.add_prompt(prompt)
return func return func

View File

@@ -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 import BaseModel, ConfigDict, Field, FileUrl, RootModel
from pydantic.networks import AnyUrl from pydantic.networks import AnyUrl
@@ -27,6 +27,7 @@ ProgressToken = str | int
Cursor = str Cursor = str
Role = Literal["user", "assistant"] Role = Literal["user", "assistant"]
RequestId = str | int RequestId = str | int
AnyFunction: TypeAlias = Callable[..., Any]
class RequestParams(BaseModel): class RequestParams(BaseModel):