mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-21 04:54:17 +02:00
Improve client return types
This commit is contained in:
parent
2df89204f4
commit
ea13dc114f
5 changed files with 122 additions and 73 deletions
|
|
@ -1,7 +1,7 @@
|
|||
import datetime
|
||||
from contextlib import AbstractAsyncContextManager
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from typing import Any, Literal, cast, overload
|
||||
|
||||
import mcp.types
|
||||
from mcp import ClientSession
|
||||
|
|
@ -24,6 +24,10 @@ from .transports import ClientTransport, SessionKwargs, infer_transport
|
|||
__all__ = ["Client", "RootsHandler", "RootsList"]
|
||||
|
||||
|
||||
class ClientError(ValueError):
|
||||
"""Base class for errors raised by the client."""
|
||||
|
||||
|
||||
class Client:
|
||||
"""
|
||||
MCP client that delegates connection management to a Transport instance.
|
||||
|
|
@ -122,60 +126,101 @@ class Client:
|
|||
"""Send a logging/setLevel request."""
|
||||
await self.session.set_logging_level(level)
|
||||
|
||||
async def list_resources(self) -> mcp.types.ListResourcesResult:
|
||||
async def send_roots_list_changed(self) -> None:
|
||||
"""Send a roots/list_changed notification."""
|
||||
await self.session.send_roots_list_changed()
|
||||
|
||||
async def list_resources(self) -> list[mcp.types.Resource]:
|
||||
"""Send a resources/list request."""
|
||||
return await self.session.list_resources()
|
||||
result = await self.session.list_resources()
|
||||
return result.resources
|
||||
|
||||
async def list_resource_templates(self) -> mcp.types.ListResourceTemplatesResult:
|
||||
async def list_resource_templates(self) -> list[mcp.types.ResourceTemplate]:
|
||||
"""Send a resources/listResourceTemplates request."""
|
||||
return await self.session.list_resource_templates()
|
||||
result = await self.session.list_resource_templates()
|
||||
return result.resourceTemplates
|
||||
|
||||
async def read_resource(self, uri: AnyUrl | str) -> mcp.types.ReadResourceResult:
|
||||
async def read_resource(
|
||||
self, uri: AnyUrl | str
|
||||
) -> list[mcp.types.TextResourceContents | mcp.types.BlobResourceContents]:
|
||||
"""Send a resources/read request."""
|
||||
if isinstance(uri, str):
|
||||
uri = AnyUrl(uri) # Ensure AnyUrl
|
||||
return await self.session.read_resource(uri)
|
||||
result = await self.session.read_resource(uri)
|
||||
return result.contents
|
||||
|
||||
async def subscribe_resource(self, uri: AnyUrl | str) -> None:
|
||||
"""Send a resources/subscribe request."""
|
||||
if isinstance(uri, str):
|
||||
uri = AnyUrl(uri)
|
||||
await self.session.subscribe_resource(uri)
|
||||
# async def subscribe_resource(self, uri: AnyUrl | str) -> None:
|
||||
# """Send a resources/subscribe request."""
|
||||
# if isinstance(uri, str):
|
||||
# uri = AnyUrl(uri)
|
||||
# await self.session.subscribe_resource(uri)
|
||||
|
||||
async def unsubscribe_resource(self, uri: AnyUrl | str) -> None:
|
||||
"""Send a resources/unsubscribe request."""
|
||||
if isinstance(uri, str):
|
||||
uri = AnyUrl(uri)
|
||||
await self.session.unsubscribe_resource(uri)
|
||||
# async def unsubscribe_resource(self, uri: AnyUrl | str) -> None:
|
||||
# """Send a resources/unsubscribe request."""
|
||||
# if isinstance(uri, str):
|
||||
# uri = AnyUrl(uri)
|
||||
# await self.session.unsubscribe_resource(uri)
|
||||
|
||||
async def list_prompts(self) -> mcp.types.ListPromptsResult:
|
||||
async def list_prompts(self) -> list[mcp.types.Prompt]:
|
||||
"""Send a prompts/list request."""
|
||||
return await self.session.list_prompts()
|
||||
result = await self.session.list_prompts()
|
||||
return result.prompts
|
||||
|
||||
async def get_prompt(
|
||||
self, name: str, arguments: dict[str, str] | None = None
|
||||
) -> mcp.types.GetPromptResult:
|
||||
"""Send a prompts/get request."""
|
||||
return await self.session.get_prompt(name, arguments)
|
||||
result = await self.session.get_prompt(name, arguments)
|
||||
return result
|
||||
|
||||
async def complete(
|
||||
self,
|
||||
ref: mcp.types.ResourceReference | mcp.types.PromptReference,
|
||||
argument: dict[str, str],
|
||||
) -> mcp.types.CompleteResult:
|
||||
"""Send a completion/complete request."""
|
||||
return await self.session.complete(ref, argument)
|
||||
) -> mcp.types.Completion:
|
||||
"""Send a completion request."""
|
||||
result = await self.session.complete(ref, argument)
|
||||
return result.completion
|
||||
|
||||
async def list_tools(self) -> mcp.types.ListToolsResult:
|
||||
async def list_tools(self) -> list[mcp.types.Tool]:
|
||||
"""Send a tools/list request."""
|
||||
return await self.session.list_tools()
|
||||
result = await self.session.list_tools()
|
||||
return result.tools
|
||||
|
||||
@overload
|
||||
async def call_tool(
|
||||
self,
|
||||
name: str,
|
||||
arguments: dict[str, Any] | None = None,
|
||||
_return_raw_result: Literal[False] = False,
|
||||
) -> list[
|
||||
mcp.types.TextContent | mcp.types.ImageContent | mcp.types.EmbeddedResource
|
||||
]: ...
|
||||
|
||||
@overload
|
||||
async def call_tool(
|
||||
self,
|
||||
name: str,
|
||||
arguments: dict[str, Any] | None = None,
|
||||
_return_raw_result: Literal[True] = True,
|
||||
) -> mcp.types.CallToolResult: ...
|
||||
|
||||
async def call_tool(
|
||||
self, name: str, arguments: dict[str, Any] | None = None
|
||||
) -> mcp.types.CallToolResult:
|
||||
self,
|
||||
name: str,
|
||||
arguments: dict[str, Any] | None = None,
|
||||
_return_raw_result: bool = False,
|
||||
) -> (
|
||||
list[
|
||||
mcp.types.TextContent | mcp.types.ImageContent | mcp.types.EmbeddedResource
|
||||
]
|
||||
| mcp.types.CallToolResult
|
||||
):
|
||||
"""Send a tools/call request."""
|
||||
return await self.session.call_tool(name, arguments)
|
||||
|
||||
async def send_roots_list_changed(self) -> None:
|
||||
"""Send a roots/list_changed notification."""
|
||||
await self.session.send_roots_list_changed()
|
||||
result = await self.session.call_tool(name, arguments)
|
||||
if _return_raw_result:
|
||||
return result
|
||||
elif result.isError:
|
||||
msg = cast(mcp.types.TextContent, result.content[0]).text
|
||||
raise ClientError(msg)
|
||||
return result.content
|
||||
|
|
|
|||
|
|
@ -40,11 +40,15 @@ class ProxyTool(Tool):
|
|||
async def run(
|
||||
self, arguments: dict[str, Any], context: Context | None = None
|
||||
) -> Any:
|
||||
# the client context manager will swallow any exceptions inside a TaskGroup
|
||||
# so we return the raw result and raise an exception ourselves
|
||||
async with self._client:
|
||||
result = await self._client.call_tool(self.name, arguments)
|
||||
result = await self._client.call_tool(
|
||||
self.name, arguments, _return_raw_result=True
|
||||
)
|
||||
if result.isError:
|
||||
raise ValueError(cast(mcp.types.TextContent, result.content[0]).text)
|
||||
return result.content[0]
|
||||
return result.content
|
||||
|
||||
|
||||
class ProxyResource(Resource):
|
||||
|
|
@ -73,12 +77,12 @@ class ProxyResource(Resource):
|
|||
|
||||
async with self._client:
|
||||
result = await self._client.read_resource(self.uri)
|
||||
if isinstance(result.contents[0], TextResourceContents):
|
||||
return result.contents[0].text
|
||||
elif isinstance(result.contents[0], BlobResourceContents):
|
||||
return result.contents[0].blob
|
||||
if isinstance(result[0], TextResourceContents):
|
||||
return result[0].text
|
||||
elif isinstance(result[0], BlobResourceContents):
|
||||
return result[0].blob
|
||||
else:
|
||||
raise ValueError(f"Unsupported content type: {type(result.contents[0])}")
|
||||
raise ValueError(f"Unsupported content type: {type(result[0])}")
|
||||
|
||||
|
||||
class ProxyTemplate(ResourceTemplate):
|
||||
|
|
@ -103,20 +107,20 @@ class ProxyTemplate(ResourceTemplate):
|
|||
async with self._client:
|
||||
result = await self._client.read_resource(uri)
|
||||
|
||||
if isinstance(result.contents[0], TextResourceContents):
|
||||
value = result.contents[0].text
|
||||
elif isinstance(result.contents[0], BlobResourceContents):
|
||||
value = result.contents[0].blob
|
||||
if isinstance(result[0], TextResourceContents):
|
||||
value = result[0].text
|
||||
elif isinstance(result[0], BlobResourceContents):
|
||||
value = result[0].blob
|
||||
else:
|
||||
raise ValueError(f"Unsupported content type: {type(result.contents[0])}")
|
||||
raise ValueError(f"Unsupported content type: {type(result[0])}")
|
||||
|
||||
return ProxyResource(
|
||||
client=self._client,
|
||||
uri=uri,
|
||||
name=self.name,
|
||||
description=self.description,
|
||||
mime_type=result.contents[0].mimeType,
|
||||
contents=result.contents,
|
||||
mime_type=result[0].mimeType,
|
||||
contents=result,
|
||||
_value=value,
|
||||
)
|
||||
|
||||
|
|
@ -177,15 +181,15 @@ class FastMCPProxy(FastMCP):
|
|||
|
||||
async with client:
|
||||
# Register proxies for client tools
|
||||
tools_result = await client.list_tools()
|
||||
for tool in tools_result.tools:
|
||||
tools = await client.list_tools()
|
||||
for tool in tools:
|
||||
tool_proxy = await ProxyTool.from_client(client, tool)
|
||||
server._tool_manager._tools[tool_proxy.name] = tool_proxy
|
||||
logger.debug(f"Created proxy for tool: {tool_proxy.name}")
|
||||
|
||||
# Register proxies for client resources
|
||||
resources_result = await client.list_resources()
|
||||
for resource in resources_result.resources:
|
||||
resources = await client.list_resources()
|
||||
for resource in resources:
|
||||
resource_proxy = await ProxyResource.from_client(client, resource)
|
||||
server._resource_manager._resources[str(resource_proxy.uri)] = (
|
||||
resource_proxy
|
||||
|
|
@ -193,8 +197,8 @@ class FastMCPProxy(FastMCP):
|
|||
logger.debug(f"Created proxy for resource: {resource_proxy.uri}")
|
||||
|
||||
# Register proxies for client resource templates
|
||||
templates_result = await client.list_resource_templates()
|
||||
for template in templates_result.resourceTemplates:
|
||||
templates = await client.list_resource_templates()
|
||||
for template in templates:
|
||||
template_proxy = await ProxyTemplate.from_client(client, template)
|
||||
server._resource_manager._templates[template_proxy.uri_template] = (
|
||||
template_proxy
|
||||
|
|
@ -204,8 +208,8 @@ class FastMCPProxy(FastMCP):
|
|||
)
|
||||
|
||||
# Register proxies for client prompts
|
||||
prompts_result = await client.list_prompts()
|
||||
for prompt in prompts_result.prompts:
|
||||
prompts = await client.list_prompts()
|
||||
for prompt in prompts:
|
||||
prompt_proxy = await ProxyPrompt.from_client(client, prompt)
|
||||
server._prompt_manager._prompts[prompt_proxy.name] = prompt_proxy
|
||||
logger.debug(f"Created proxy for prompt: {prompt_proxy.name}")
|
||||
|
|
|
|||
|
|
@ -51,8 +51,8 @@ async def test_list_tools(fastmcp_server):
|
|||
result = await client.list_tools()
|
||||
|
||||
# Check that our tools are available
|
||||
assert len(result.tools) == 2
|
||||
assert set(tool.name for tool in result.tools) == {"greet", "add"}
|
||||
assert len(result) == 2
|
||||
assert set(tool.name for tool in result) == {"greet", "add"}
|
||||
|
||||
|
||||
async def test_call_tool(fastmcp_server):
|
||||
|
|
@ -63,7 +63,7 @@ async def test_call_tool(fastmcp_server):
|
|||
result = await client.call_tool("greet", {"name": "World"})
|
||||
|
||||
# The result content should contain our greeting
|
||||
content_str = str(result.content[0])
|
||||
content_str = str(result[0])
|
||||
assert "Hello, World!" in content_str
|
||||
|
||||
|
||||
|
|
@ -75,8 +75,8 @@ async def test_list_resources(fastmcp_server):
|
|||
result = await client.list_resources()
|
||||
|
||||
# Check that our resource is available
|
||||
assert len(result.resources) == 1
|
||||
assert str(result.resources[0].uri) == "data://users"
|
||||
assert len(result) == 1
|
||||
assert str(result[0].uri) == "data://users"
|
||||
|
||||
|
||||
async def test_list_prompts(fastmcp_server):
|
||||
|
|
@ -87,8 +87,8 @@ async def test_list_prompts(fastmcp_server):
|
|||
result = await client.list_prompts()
|
||||
|
||||
# Check that our prompt is available
|
||||
assert len(result.prompts) == 1
|
||||
assert result.prompts[0].name == "welcome"
|
||||
assert len(result) == 1
|
||||
assert result[0].name == "welcome"
|
||||
|
||||
|
||||
async def test_get_prompt(fastmcp_server):
|
||||
|
|
@ -115,7 +115,7 @@ async def test_read_resource(fastmcp_server):
|
|||
result = await client.read_resource(uri)
|
||||
|
||||
# The contents should include our user list
|
||||
contents_str = str(result.contents[0])
|
||||
contents_str = str(result[0])
|
||||
assert "Alice" in contents_str
|
||||
assert "Bob" in contents_str
|
||||
assert "Charlie" in contents_str
|
||||
|
|
@ -145,15 +145,15 @@ async def test_resource_template(fastmcp_server):
|
|||
result = await client.list_resource_templates()
|
||||
|
||||
# Check that our template is available
|
||||
assert len(result.resourceTemplates) == 1
|
||||
assert "data://user/{user_id}" in result.resourceTemplates[0].uriTemplate
|
||||
assert len(result) == 1
|
||||
assert "data://user/{user_id}" in result[0].uriTemplate
|
||||
|
||||
# Now use the template with a specific user_id
|
||||
uri = cast(AnyUrl, "data://user/123")
|
||||
result = await client.read_resource(uri)
|
||||
|
||||
# Check the content matches what we expect for the provided user_id
|
||||
content_str = str(result.contents[0])
|
||||
content_str = str(result[0])
|
||||
assert '"id": "123"' in content_str
|
||||
assert '"name": "User 123"' in content_str
|
||||
assert '"active": true' in content_str
|
||||
|
|
@ -41,8 +41,8 @@ class TestClientRoots:
|
|||
async def test_valid_roots(self, fastmcp_server: FastMCP, roots: list[str]):
|
||||
async with Client(fastmcp_server, roots=roots) as client:
|
||||
result = await client.call_tool("list_roots", {})
|
||||
assert isinstance(result.content[0], TextContent)
|
||||
assert json.loads(result.content[0].text) == [
|
||||
assert isinstance(result[0], TextContent)
|
||||
assert json.loads(result[0].text) == [
|
||||
"file://x/y/z",
|
||||
"file://x/y/z",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -47,7 +47,7 @@ async def test_simple_sampling(fastmcp_server: FastMCP):
|
|||
|
||||
async with Client(fastmcp_server, sampling_handler=sampling_handler) as client:
|
||||
result = await client.call_tool("simple_sample", {"message": "Hello, world!"})
|
||||
reply = cast(TextContent, result.content[0])
|
||||
reply = cast(TextContent, result[0])
|
||||
assert reply.text == "This is the sample message!"
|
||||
|
||||
|
||||
|
|
@ -62,7 +62,7 @@ async def test_sampling_with_system_prompt(fastmcp_server: FastMCP):
|
|||
result = await client.call_tool(
|
||||
"sample_with_system_prompt", {"message": "Hello, world!"}
|
||||
)
|
||||
reply = cast(TextContent, result.content[0])
|
||||
reply = cast(TextContent, result[0])
|
||||
assert reply.text == "You love FastMCP"
|
||||
|
||||
|
||||
|
|
@ -81,5 +81,5 @@ async def test_sampling_with_messages(fastmcp_server: FastMCP):
|
|||
result = await client.call_tool(
|
||||
"sample_with_messages", {"message": "Hello, world!"}
|
||||
)
|
||||
reply = cast(TextContent, result.content[0])
|
||||
reply = cast(TextContent, result[0])
|
||||
assert reply.text == "I need to think."
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue