Merge pull request #123 from jlowin/client

Improve client return types
This commit is contained in:
Jeremiah Lowin 2025-04-12 12:18:18 -04:00 committed by GitHub
commit 4bc85e7cd1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 122 additions and 73 deletions

View file

@ -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

View file

@ -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}")

View file

@ -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

View file

@ -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",
]

View file

@ -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."