From 004a3fb6f0163f700640a5e995d17352f758621f Mon Sep 17 00:00:00 2001 From: ramjibc Date: Sun, 11 May 2025 23:17:55 -0400 Subject: [PATCH 1/4] move Prompt instantiation from server to prompt manager --- src/mcp/server/fastmcp/prompts/manager.py | 21 ++++++++++++--- src/mcp/server/fastmcp/server.py | 31 +++++++++++++++-------- 2 files changed, 39 insertions(+), 13 deletions(-) diff --git a/src/mcp/server/fastmcp/prompts/manager.py b/src/mcp/server/fastmcp/prompts/manager.py index 7ccbdef36..a795a2ed1 100644 --- a/src/mcp/server/fastmcp/prompts/manager.py +++ b/src/mcp/server/fastmcp/prompts/manager.py @@ -1,6 +1,6 @@ """Prompt management functionality.""" -from typing import Any +from typing import Any, Callable from mcp.server.fastmcp.prompts.base import Message, Prompt from mcp.server.fastmcp.utilities.logging import get_logger @@ -25,9 +25,24 @@ def list_prompts(self) -> list[Prompt]: def add_prompt( self, - prompt: Prompt, + prompt_or_fn: Prompt | Callable[..., Any], + name: str | None = None, + description: str | None = None, ) -> Prompt: - """Add a prompt to the manager.""" + """Add a prompt to the manager. + + Args: + prompt_or_fn: Either a Prompt object or a function to create a prompt from + name: Optional name for the prompt (only used if prompt_or_fn is a function) + description: Optional description of the prompt (only used if prompt_or_fn is a function) + """ + # If a function was provided, create a Prompt object from it + if callable(prompt_or_fn) and not isinstance(prompt_or_fn, Prompt): + prompt = Prompt.from_function( + prompt_or_fn, name=name, description=description + ) + else: + prompt = prompt_or_fn # Check for duplicates existing = self._prompts.get(prompt.name) diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py index c31f29d4c..d1177aced 100644 --- a/src/mcp/server/fastmcp/server.py +++ b/src/mcp/server/fastmcp/server.py @@ -138,8 +138,9 @@ def __init__( self, name: str | None = None, instructions: str | None = None, - auth_server_provider: OAuthAuthorizationServerProvider[Any, Any, Any] - | None = None, + auth_server_provider: ( + OAuthAuthorizationServerProvider[Any, Any, Any] | None + ) = None, event_store: EventStore | None = None, **settings: Any, ): @@ -148,9 +149,11 @@ def __init__( self._mcp_server = MCPServer( name=name or "FastMCP", instructions=instructions, - lifespan=lifespan_wrapper(self, self.settings.lifespan) - if self.settings.lifespan - else default_lifespan, + lifespan=( + lifespan_wrapper(self, self.settings.lifespan) + if self.settings.lifespan + else default_lifespan + ), ) self._tool_manager = ToolManager( warn_on_duplicate_tools=self.settings.warn_on_duplicate_tools @@ -481,13 +484,22 @@ def decorator(fn: AnyFunction) -> AnyFunction: return decorator - def add_prompt(self, prompt: Prompt) -> None: + def add_prompt( + self, + prompt_or_fn: Prompt | Callable[..., Any], + name: str | None = None, + description: str | None = None, + ) -> None: """Add a prompt to the server. Args: - prompt: A Prompt instance to add + prompt_or_fn: Either a Prompt object or a function to create a prompt from + name: Optional name for the prompt (only used if prompt_or_fn is a function) + description: Optional description of the prompt (only used if prompt_or_fn is a function) """ - self._prompt_manager.add_prompt(prompt) + self._prompt_manager.add_prompt( + prompt_or_fn, name=name, description=description + ) def prompt( self, name: str | None = None, description: str | None = None @@ -533,8 +545,7 @@ async def analyze_file(path: str) -> list[Message]: ) def decorator(func: AnyFunction) -> AnyFunction: - prompt = Prompt.from_function(func, name=name, description=description) - self.add_prompt(prompt) + self.add_prompt(func, name=name, description=description) return func return decorator From 8e3a312b81dcc36bf60674093e7b68aabdeebc6d Mon Sep 17 00:00:00 2001 From: ramjibc Date: Sun, 11 May 2025 23:19:39 -0400 Subject: [PATCH 2/4] simplify add_prompt() args to always use fn --- src/mcp/server/fastmcp/prompts/manager.py | 16 +++++----------- src/mcp/server/fastmcp/server.py | 12 +++++------- 2 files changed, 10 insertions(+), 18 deletions(-) diff --git a/src/mcp/server/fastmcp/prompts/manager.py b/src/mcp/server/fastmcp/prompts/manager.py index a795a2ed1..0c60a2f46 100644 --- a/src/mcp/server/fastmcp/prompts/manager.py +++ b/src/mcp/server/fastmcp/prompts/manager.py @@ -25,24 +25,18 @@ def list_prompts(self) -> list[Prompt]: def add_prompt( self, - prompt_or_fn: Prompt | Callable[..., Any], + fn: Callable[..., Any], name: str | None = None, description: str | None = None, ) -> Prompt: """Add a prompt to the manager. Args: - prompt_or_fn: Either a Prompt object or a function to create a prompt from - name: Optional name for the prompt (only used if prompt_or_fn is a function) - description: Optional description of the prompt (only used if prompt_or_fn is a function) + fn: Function to create a prompt from + name: Optional name for the prompt + description: Optional description of the prompt """ - # If a function was provided, create a Prompt object from it - if callable(prompt_or_fn) and not isinstance(prompt_or_fn, Prompt): - prompt = Prompt.from_function( - prompt_or_fn, name=name, description=description - ) - else: - prompt = prompt_or_fn + prompt = Prompt.from_function(fn, name=name, description=description) # Check for duplicates existing = self._prompts.get(prompt.name) diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py index d1177aced..5444f6004 100644 --- a/src/mcp/server/fastmcp/server.py +++ b/src/mcp/server/fastmcp/server.py @@ -486,20 +486,18 @@ def decorator(fn: AnyFunction) -> AnyFunction: def add_prompt( self, - prompt_or_fn: Prompt | Callable[..., Any], + fn: Callable[..., Any], name: str | None = None, description: str | None = None, ) -> None: """Add a prompt to the server. Args: - prompt_or_fn: Either a Prompt object or a function to create a prompt from - name: Optional name for the prompt (only used if prompt_or_fn is a function) - description: Optional description of the prompt (only used if prompt_or_fn is a function) + fn: Function to create a prompt from + name: Optional name for the prompt + description: Optional description of the prompt """ - self._prompt_manager.add_prompt( - prompt_or_fn, name=name, description=description - ) + self._prompt_manager.add_prompt(fn, name=name, description=description) def prompt( self, name: str | None = None, description: str | None = None From f0f8beb5beb08e133581e72ed68dc87910b79264 Mon Sep 17 00:00:00 2001 From: ramjibc Date: Sun, 11 May 2025 23:26:27 -0400 Subject: [PATCH 3/4] update unit tests for prompt manager --- src/mcp/server/fastmcp/server.py | 2 +- tests/server/fastmcp/prompts/test_manager.py | 33 ++++++++------------ 2 files changed, 14 insertions(+), 21 deletions(-) diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py index 5444f6004..bcbebb19d 100644 --- a/src/mcp/server/fastmcp/server.py +++ b/src/mcp/server/fastmcp/server.py @@ -35,7 +35,7 @@ AuthSettings, ) from mcp.server.fastmcp.exceptions import ResourceError -from mcp.server.fastmcp.prompts import Prompt, PromptManager +from mcp.server.fastmcp.prompts import PromptManager from mcp.server.fastmcp.resources import FunctionResource, Resource, ResourceManager from mcp.server.fastmcp.tools import ToolManager from mcp.server.fastmcp.utilities.logging import configure_logging, get_logger diff --git a/tests/server/fastmcp/prompts/test_manager.py b/tests/server/fastmcp/prompts/test_manager.py index c64a4a564..59d0e5538 100644 --- a/tests/server/fastmcp/prompts/test_manager.py +++ b/tests/server/fastmcp/prompts/test_manager.py @@ -12,10 +12,10 @@ def fn() -> str: return "Hello, world!" manager = PromptManager() - prompt = Prompt.from_function(fn) - added = manager.add_prompt(prompt) - assert added == prompt - assert manager.get_prompt("fn") == prompt + added = manager.add_prompt(fn) + assert isinstance(added, Prompt) + assert added.name == "fn" + assert manager.get_prompt("fn") == added def test_add_duplicate_prompt(self, caplog): """Test adding the same prompt twice.""" @@ -24,9 +24,8 @@ def fn() -> str: return "Hello, world!" manager = PromptManager() - prompt = Prompt.from_function(fn) - first = manager.add_prompt(prompt) - second = manager.add_prompt(prompt) + first = manager.add_prompt(fn) + second = manager.add_prompt(fn) assert first == second assert "Prompt already exists" in caplog.text @@ -37,9 +36,8 @@ def fn() -> str: return "Hello, world!" manager = PromptManager(warn_on_duplicate_prompts=False) - prompt = Prompt.from_function(fn) - first = manager.add_prompt(prompt) - second = manager.add_prompt(prompt) + first = manager.add_prompt(fn) + second = manager.add_prompt(fn) assert first == second assert "Prompt already exists" not in caplog.text @@ -53,10 +51,8 @@ def fn2() -> str: return "Goodbye, world!" manager = PromptManager() - prompt1 = Prompt.from_function(fn1) - prompt2 = Prompt.from_function(fn2) - manager.add_prompt(prompt1) - manager.add_prompt(prompt2) + prompt1 = manager.add_prompt(fn1) + prompt2 = manager.add_prompt(fn2) prompts = manager.list_prompts() assert len(prompts) == 2 assert prompts == [prompt1, prompt2] @@ -69,8 +65,7 @@ def fn() -> str: return "Hello, world!" manager = PromptManager() - prompt = Prompt.from_function(fn) - manager.add_prompt(prompt) + manager.add_prompt(fn) messages = await manager.render_prompt("fn") assert messages == [ UserMessage(content=TextContent(type="text", text="Hello, world!")) @@ -84,8 +79,7 @@ def fn(name: str) -> str: return f"Hello, {name}!" manager = PromptManager() - prompt = Prompt.from_function(fn) - manager.add_prompt(prompt) + manager.add_prompt(fn) messages = await manager.render_prompt("fn", arguments={"name": "World"}) assert messages == [ UserMessage(content=TextContent(type="text", text="Hello, World!")) @@ -106,7 +100,6 @@ def fn(name: str) -> str: return f"Hello, {name}!" manager = PromptManager() - prompt = Prompt.from_function(fn) - manager.add_prompt(prompt) + manager.add_prompt(fn) with pytest.raises(ValueError, match="Missing required arguments"): await manager.render_prompt("fn") From 3b9d4fa1aec8294bfcb65c6fc0a402a75330bba0 Mon Sep 17 00:00:00 2001 From: ramjibc Date: Sun, 11 May 2025 23:40:47 -0400 Subject: [PATCH 4/4] update type hint --- src/mcp/server/fastmcp/prompts/manager.py | 5 +++-- src/mcp/server/fastmcp/server.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/mcp/server/fastmcp/prompts/manager.py b/src/mcp/server/fastmcp/prompts/manager.py index 0c60a2f46..71e55aa06 100644 --- a/src/mcp/server/fastmcp/prompts/manager.py +++ b/src/mcp/server/fastmcp/prompts/manager.py @@ -1,9 +1,10 @@ """Prompt management functionality.""" -from typing import Any, Callable +from typing import Any from mcp.server.fastmcp.prompts.base import Message, Prompt from mcp.server.fastmcp.utilities.logging import get_logger +from mcp.types import AnyFunction logger = get_logger(__name__) @@ -25,7 +26,7 @@ def list_prompts(self) -> list[Prompt]: def add_prompt( self, - fn: Callable[..., Any], + fn: AnyFunction, name: str | None = None, description: str | None = None, ) -> Prompt: diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py index bcbebb19d..68803e324 100644 --- a/src/mcp/server/fastmcp/server.py +++ b/src/mcp/server/fastmcp/server.py @@ -486,7 +486,7 @@ def decorator(fn: AnyFunction) -> AnyFunction: def add_prompt( self, - fn: Callable[..., Any], + fn: AnyFunction, name: str | None = None, description: str | None = None, ) -> None: