mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-23 22:14:18 +02:00
Return objects from decorators
This commit is contained in:
parent
31fd9abee0
commit
d26a5dffc3
3 changed files with 61 additions and 67 deletions
|
|
@ -14,7 +14,7 @@ from contextlib import (
|
|||
)
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Generic, Literal
|
||||
from typing import TYPE_CHECKING, Any, Generic, Literal, overload
|
||||
|
||||
import anyio
|
||||
import httpx
|
||||
|
|
@ -45,6 +45,7 @@ import fastmcp.server
|
|||
import fastmcp.settings
|
||||
from fastmcp.exceptions import NotFoundError
|
||||
from fastmcp.prompts import Prompt, PromptManager
|
||||
from fastmcp.prompts.prompt import FunctionPrompt
|
||||
from fastmcp.resources import Resource, ResourceManager
|
||||
from fastmcp.resources.template import ResourceTemplate
|
||||
from fastmcp.server.auth.auth import OAuthProvider
|
||||
|
|
@ -55,7 +56,7 @@ from fastmcp.server.http import (
|
|||
create_streamable_http_app,
|
||||
)
|
||||
from fastmcp.tools import ToolManager
|
||||
from fastmcp.tools.tool import Tool
|
||||
from fastmcp.tools.tool import FunctionTool, Tool
|
||||
from fastmcp.utilities.cache import TimedCache
|
||||
from fastmcp.utilities.logging import get_logger
|
||||
from fastmcp.utilities.mcp_config import MCPConfig
|
||||
|
|
@ -510,6 +511,30 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
self._tool_manager.remove_tool(name)
|
||||
self._cache.clear()
|
||||
|
||||
@overload
|
||||
def tool(
|
||||
self,
|
||||
name_or_fn: AnyFunction,
|
||||
*,
|
||||
name: str | None = None,
|
||||
description: str | None = None,
|
||||
tags: set[str] | None = None,
|
||||
annotations: ToolAnnotations | dict[str, Any] | None = None,
|
||||
exclude_args: list[str] | None = None,
|
||||
) -> FunctionTool: ...
|
||||
|
||||
@overload
|
||||
def tool(
|
||||
self,
|
||||
name_or_fn: str | None = None,
|
||||
*,
|
||||
name: str | None = None,
|
||||
description: str | None = None,
|
||||
tags: set[str] | None = None,
|
||||
annotations: ToolAnnotations | dict[str, Any] | None = None,
|
||||
exclude_args: list[str] | None = None,
|
||||
) -> Callable[[AnyFunction], FunctionTool]: ...
|
||||
|
||||
def tool(
|
||||
self,
|
||||
name_or_fn: str | AnyFunction | None = None,
|
||||
|
|
@ -519,7 +544,7 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
tags: set[str] | None = None,
|
||||
annotations: ToolAnnotations | dict[str, Any] | None = None,
|
||||
exclude_args: list[str] | None = None,
|
||||
) -> Callable[[AnyFunction], AnyFunction] | AnyFunction:
|
||||
) -> Callable[[AnyFunction], FunctionTool] | FunctionTool:
|
||||
"""Decorator to register a tool.
|
||||
|
||||
Tools can optionally request a Context object by adding a parameter with the
|
||||
|
|
@ -571,7 +596,7 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
fn = name_or_fn
|
||||
tool_name = name # Use keyword name if provided, otherwise None
|
||||
|
||||
# Register the tool immediately and return the function
|
||||
# Register the tool immediately and return the tool object
|
||||
tool = Tool.from_function(
|
||||
fn,
|
||||
name=tool_name,
|
||||
|
|
@ -582,7 +607,7 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
serializer=self._tool_serializer,
|
||||
)
|
||||
self.add_tool(tool)
|
||||
return fn
|
||||
return tool
|
||||
|
||||
elif isinstance(name_or_fn, str):
|
||||
# Case 3: @tool("custom_name") - name passed as first argument
|
||||
|
|
@ -674,7 +699,7 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
description: str | None = None,
|
||||
mime_type: str | None = None,
|
||||
tags: set[str] | None = None,
|
||||
) -> Callable[[AnyFunction], AnyFunction]:
|
||||
) -> Callable[[AnyFunction], Resource | ResourceTemplate]:
|
||||
"""Decorator to register a function as a resource.
|
||||
|
||||
The function will be called when the resource is read to generate its content.
|
||||
|
|
@ -728,7 +753,7 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
"Did you forget to call it? Use @resource('uri') instead of @resource"
|
||||
)
|
||||
|
||||
def decorator(fn: AnyFunction) -> AnyFunction:
|
||||
def decorator(fn: AnyFunction) -> Resource | ResourceTemplate:
|
||||
from fastmcp.server.context import Context
|
||||
|
||||
# Check if this should be a template
|
||||
|
|
@ -750,6 +775,7 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
tags=tags,
|
||||
)
|
||||
self.add_template(template)
|
||||
return template
|
||||
elif not has_uri_params and not has_func_params:
|
||||
resource = Resource.from_function(
|
||||
fn=fn,
|
||||
|
|
@ -760,14 +786,13 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
tags=tags,
|
||||
)
|
||||
self.add_resource(resource)
|
||||
return resource
|
||||
else:
|
||||
raise ValueError(
|
||||
"Invalid resource or template definition due to a "
|
||||
"mismatch between URI parameters and function parameters."
|
||||
)
|
||||
|
||||
return fn
|
||||
|
||||
return decorator
|
||||
|
||||
def add_prompt(self, prompt: Prompt) -> None:
|
||||
|
|
@ -779,6 +804,26 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
self._prompt_manager.add_prompt(prompt)
|
||||
self._cache.clear()
|
||||
|
||||
@overload
|
||||
def prompt(
|
||||
self,
|
||||
name_or_fn: AnyFunction,
|
||||
*,
|
||||
name: str | None = None,
|
||||
description: str | None = None,
|
||||
tags: set[str] | None = None,
|
||||
) -> FunctionPrompt: ...
|
||||
|
||||
@overload
|
||||
def prompt(
|
||||
self,
|
||||
name_or_fn: str | None = None,
|
||||
*,
|
||||
name: str | None = None,
|
||||
description: str | None = None,
|
||||
tags: set[str] | None = None,
|
||||
) -> Callable[[AnyFunction], FunctionPrompt]: ...
|
||||
|
||||
def prompt(
|
||||
self,
|
||||
name_or_fn: str | AnyFunction | None = None,
|
||||
|
|
@ -786,7 +831,7 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
name: str | None = None,
|
||||
description: str | None = None,
|
||||
tags: set[str] | None = None,
|
||||
) -> Callable[[AnyFunction], AnyFunction] | AnyFunction:
|
||||
) -> Callable[[AnyFunction], FunctionPrompt] | FunctionPrompt:
|
||||
"""Decorator to register a prompt.
|
||||
|
||||
Prompts can optionally request a Context object by adding a parameter with the
|
||||
|
|
@ -867,7 +912,7 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
)
|
||||
self.add_prompt(prompt)
|
||||
|
||||
return fn
|
||||
return prompt
|
||||
|
||||
elif isinstance(name_or_fn, str):
|
||||
# Case 3: @prompt("custom_name") - name passed as first argument
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from pydantic import Field
|
|||
|
||||
from fastmcp import Client, FastMCP
|
||||
from fastmcp.exceptions import NotFoundError
|
||||
from fastmcp.prompts.prompt import Prompt
|
||||
from fastmcp.prompts.prompt import FunctionPrompt, Prompt
|
||||
from fastmcp.resources import Resource, ResourceTemplate
|
||||
from fastmcp.server.server import (
|
||||
MountedServer,
|
||||
|
|
@ -179,7 +179,6 @@ class TestToolDecorator:
|
|||
def __init__(self, x: int):
|
||||
self.x = x
|
||||
|
||||
@mcp.tool
|
||||
def add(self, y: int) -> int:
|
||||
return self.x + y
|
||||
|
||||
|
|
@ -326,11 +325,11 @@ class TestToolDecorator:
|
|||
result_fn = mcp.tool(standalone_function, name="direct_call_tool")
|
||||
|
||||
# The function should be returned unchanged
|
||||
assert result_fn is standalone_function
|
||||
assert isinstance(result_fn, FunctionTool)
|
||||
|
||||
# Verify the tool was registered correctly
|
||||
tools = await mcp.get_tools()
|
||||
assert "direct_call_tool" in tools
|
||||
assert tools["direct_call_tool"] is result_fn
|
||||
|
||||
# Verify it can be called
|
||||
result = await mcp._mcp_call_tool("direct_call_tool", {"x": 5, "y": 3})
|
||||
|
|
@ -857,11 +856,11 @@ class TestPromptDecorator:
|
|||
result_fn = mcp.prompt(standalone_function, name="direct_call_prompt")
|
||||
|
||||
# The function should be returned unchanged
|
||||
assert result_fn is standalone_function
|
||||
assert isinstance(result_fn, FunctionPrompt)
|
||||
|
||||
# Verify the prompt was registered correctly
|
||||
prompts = await mcp.get_prompts()
|
||||
assert "direct_call_prompt" in prompts
|
||||
assert prompts["direct_call_prompt"] is result_fn
|
||||
|
||||
# Verify it can be called
|
||||
async with Client(mcp) as client:
|
||||
|
|
|
|||
|
|
@ -936,56 +936,6 @@ class TestResourceTemplates:
|
|||
result = await client.read_resource(AnyUrl("resource://test/data"))
|
||||
assert result[0].text == "Data for test" # type: ignore[attr-defined]
|
||||
|
||||
async def test_stacked_resource_template_decorators(self):
|
||||
"""Test that resource template decorators can be stacked."""
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.resource("users://email/{email}")
|
||||
@mcp.resource("users://name/{name}")
|
||||
def lookup_user(name: str | None = None, email: str | None = None) -> dict:
|
||||
if name:
|
||||
return {
|
||||
"id": "123",
|
||||
"name": name,
|
||||
"email": "dummy@example.com",
|
||||
"lookup": "name",
|
||||
}
|
||||
elif email:
|
||||
return {
|
||||
"id": "123",
|
||||
"name": "Test User",
|
||||
"email": email,
|
||||
"lookup": "email",
|
||||
}
|
||||
else:
|
||||
raise ValueError("Either name or email must be provided")
|
||||
|
||||
# Verify both templates are registered
|
||||
templates_dict = await mcp.get_resource_templates()
|
||||
templates = list(templates_dict.values())
|
||||
assert len(templates) == 2
|
||||
template_uris = {t.uri_template for t in templates}
|
||||
assert "users://email/{email}" in template_uris
|
||||
assert "users://name/{name}" in template_uris
|
||||
|
||||
# Test lookup by email
|
||||
async with Client(mcp) as client:
|
||||
email_result = await client.read_resource(
|
||||
AnyUrl("users://email/user@example.com")
|
||||
)
|
||||
assert email_result[0].text # type: ignore[attr-defined]
|
||||
email_data = json.loads(email_result[0].text) # type: ignore[attr-defined]
|
||||
assert email_data["lookup"] == "email"
|
||||
assert email_data["email"] == "user@example.com"
|
||||
|
||||
# Test lookup by name
|
||||
name_result = await client.read_resource(AnyUrl("users://name/John"))
|
||||
assert name_result[0].text # type: ignore[attr-defined]
|
||||
name_data = json.loads(name_result[0].text) # type: ignore[attr-defined]
|
||||
assert name_data["lookup"] == "name"
|
||||
assert name_data["name"] == "John"
|
||||
assert name_data["email"] == "dummy@example.com"
|
||||
|
||||
async def test_template_decorator_with_tags(self):
|
||||
mcp = FastMCP()
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue