mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-22 05:24:18 +02:00
Support enabled/disabled prompts
This commit is contained in:
parent
27259a7c49
commit
17a71278aa
5 changed files with 103 additions and 4 deletions
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -178,7 +178,9 @@ class TestResources:
|
|||
assert json.loads(result[0].text) == USERS # type: ignore[attr-defined]
|
||||
|
||||
async def test_read_resource_returns_none_if_not_found(self, proxy_server):
|
||||
with pytest.raises(McpError, match="Unknown resource: resource://nonexistent"):
|
||||
with pytest.raises(
|
||||
McpError, match="Unknown resource: 'resource://nonexistent'"
|
||||
):
|
||||
async with Client(proxy_server) as client:
|
||||
await client.read_resource("resource://nonexistent")
|
||||
|
||||
|
|
|
|||
|
|
@ -1220,6 +1220,93 @@ class TestPrompts:
|
|||
assert prompt.tags == {"example", "test-tag"}
|
||||
|
||||
|
||||
class TestPromptEnabled:
|
||||
async def test_toggle_enabled(self):
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.prompt
|
||||
def sample_prompt() -> str:
|
||||
return "Hello, world!"
|
||||
|
||||
assert sample_prompt.enabled
|
||||
|
||||
prompt = await mcp.get_prompt("sample_prompt")
|
||||
assert prompt.enabled
|
||||
|
||||
prompt.disable()
|
||||
|
||||
assert not prompt.enabled
|
||||
assert not sample_prompt.enabled
|
||||
|
||||
prompt.enable()
|
||||
assert prompt.enabled
|
||||
assert sample_prompt.enabled
|
||||
|
||||
async def test_prompt_disabled_in_decorator(self):
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.prompt(enabled=False)
|
||||
def sample_prompt() -> str:
|
||||
return "Hello, world!"
|
||||
|
||||
async with Client(mcp) as client:
|
||||
prompts = await client.list_prompts()
|
||||
assert len(prompts) == 0
|
||||
|
||||
async def test_prompt_toggle_enabled(self):
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.prompt(enabled=False)
|
||||
def sample_prompt() -> str:
|
||||
return "Hello, world!"
|
||||
|
||||
sample_prompt.enable()
|
||||
|
||||
async with Client(mcp) as client:
|
||||
prompts = await client.list_prompts()
|
||||
assert len(prompts) == 1
|
||||
|
||||
async def test_prompt_toggle_disabled(self):
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.prompt
|
||||
def sample_prompt() -> str:
|
||||
return "Hello, world!"
|
||||
|
||||
sample_prompt.disable()
|
||||
|
||||
async with Client(mcp) as client:
|
||||
prompts = await client.list_prompts()
|
||||
assert len(prompts) == 0
|
||||
|
||||
async def test_get_prompt_and_disable(self):
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.prompt
|
||||
def sample_prompt() -> str:
|
||||
return "Hello, world!"
|
||||
|
||||
prompt = await mcp.get_prompt("sample_prompt")
|
||||
assert prompt.enabled
|
||||
|
||||
sample_prompt.disable()
|
||||
|
||||
async with Client(mcp) as client:
|
||||
result = await client.list_prompts()
|
||||
assert len(result) == 0
|
||||
|
||||
async def test_cant_get_disabled_prompt(self):
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.prompt(enabled=False)
|
||||
def sample_prompt() -> str:
|
||||
return "Hello, world!"
|
||||
|
||||
with pytest.raises(McpError, match="Unknown prompt"):
|
||||
async with Client(mcp) as client:
|
||||
await client.get_prompt("sample_prompt")
|
||||
|
||||
|
||||
class TestPromptContext:
|
||||
async def test_prompt_context(self):
|
||||
mcp = FastMCP()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue