Support enabled/disabled prompts

This commit is contained in:
Jeremiah Lowin 2025-06-10 10:24:40 -04:00
commit 17a71278aa
5 changed files with 103 additions and 4 deletions

View file

@ -96,6 +96,7 @@ class Prompt(FastMCPComponent, ABC):
name: str | None = None,
description: str | None = None,
tags: set[str] | None = None,
enabled: bool | None = None,
) -> FunctionPrompt:
"""Create a Prompt from a function.
@ -106,7 +107,7 @@ class Prompt(FastMCPComponent, ABC):
- A sequence of any of the above
"""
return FunctionPrompt.from_function(
fn=fn, name=name, description=description, tags=tags
fn=fn, name=name, description=description, tags=tags, enabled=enabled
)
@abstractmethod
@ -130,6 +131,7 @@ class FunctionPrompt(Prompt):
name: str | None = None,
description: str | None = None,
tags: set[str] | None = None,
enabled: bool | None = None,
) -> FunctionPrompt:
"""Create a Prompt from a function.
@ -195,6 +197,7 @@ class FunctionPrompt(Prompt):
description=description,
arguments=arguments,
tags=tags or set(),
enabled=enabled if enabled is not None else True,
fn=fn,
)

View file

@ -40,9 +40,11 @@ class PromptManager:
self.duplicate_behavior = duplicate_behavior
def get_prompt(self, key: str) -> Prompt | None:
def get_prompt(self, key: str) -> Prompt:
"""Get prompt by key."""
return self._prompts.get(key)
if key in self._prompts:
return self._prompts[key]
raise NotFoundError(f"Unknown prompt: {key}")
def get_prompts(self) -> dict[str, Prompt]:
"""Get all registered prompts, indexed by registered key."""

View file

@ -940,6 +940,7 @@ class FastMCP(Generic[LifespanResultT]):
name: str | None = None,
description: str | None = None,
tags: set[str] | None = None,
enabled: bool | None = None,
) -> FunctionPrompt: ...
@overload
@ -950,6 +951,7 @@ class FastMCP(Generic[LifespanResultT]):
name: str | None = None,
description: str | None = None,
tags: set[str] | None = None,
enabled: bool | None = None,
) -> Callable[[AnyFunction], FunctionPrompt]: ...
def prompt(
@ -959,6 +961,7 @@ class FastMCP(Generic[LifespanResultT]):
name: str | None = None,
description: str | None = None,
tags: set[str] | None = None,
enabled: bool | None = None,
) -> Callable[[AnyFunction], FunctionPrompt] | FunctionPrompt:
"""Decorator to register a prompt.
@ -1050,6 +1053,7 @@ class FastMCP(Generic[LifespanResultT]):
name=prompt_name,
description=description,
tags=tags,
enabled=enabled,
)
self.add_prompt(prompt)
@ -1077,6 +1081,7 @@ class FastMCP(Generic[LifespanResultT]):
name=prompt_name,
description=description,
tags=tags,
enabled=enabled,
)
async def run_stdio_async(self) -> None: