mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-23 14:04:18 +02:00
Merge pull request #12 from jlowin/context
Add support for request context, progress, logging, etc.
This commit is contained in:
Jeremiah Lowin
GitHub
commit
8626f40998
6 changed files with 341 additions and 18 deletions
|
|
@ -1,6 +1,6 @@
|
|||
"""FastMCP - A more ergonomic interface for MCP servers."""
|
||||
|
||||
from .server import FastMCP
|
||||
from .server import FastMCP, Context
|
||||
from .utilities.types import Image
|
||||
|
||||
__all__ = ["FastMCP", "Image"]
|
||||
__all__ = ["FastMCP", "Context", "Image"]
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
"""Concrete resource implementations."""
|
||||
|
||||
import pydantic_core
|
||||
import asyncio
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
|
@ -58,8 +59,8 @@ class FunctionResource(Resource):
|
|||
if isinstance(result, str):
|
||||
return result
|
||||
try:
|
||||
return json.dumps(result, default=pydantic.json.pydantic_encoder)
|
||||
except TypeError:
|
||||
return json.dumps(pydantic_core.to_jsonable_python(result))
|
||||
except (TypeError, pydantic_core.PydanticSerializationError):
|
||||
# If JSON serialization fails, try str()
|
||||
return str(result)
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -1,13 +1,20 @@
|
|||
"""FastMCP - A more ergonomic interface for MCP servers."""
|
||||
|
||||
import pydantic_core
|
||||
from typing import Any, Literal, Optional, Union
|
||||
|
||||
from mcp.server import RequestContext
|
||||
from pydantic import BaseModel
|
||||
from pydantic.networks import AnyUrl
|
||||
|
||||
from fastmcp.utilities.logging import get_logger
|
||||
import asyncio
|
||||
import functools
|
||||
import json
|
||||
from typing import Any, Callable, Optional, Sequence, Union, Literal
|
||||
from typing import Callable, Sequence
|
||||
import inspect
|
||||
import re
|
||||
|
||||
import pydantic.json
|
||||
from mcp.server import Server as MCPServer
|
||||
from mcp.server.stdio import stdio_server
|
||||
from mcp.server.sse import SseServerTransport
|
||||
|
|
@ -25,7 +32,7 @@ from fastmcp.exceptions import ResourceError
|
|||
from fastmcp.resources import Resource, ResourceManager
|
||||
from fastmcp.resources.types import FunctionResource
|
||||
from fastmcp.tools import ToolManager
|
||||
from fastmcp.utilities.logging import get_logger, configure_logging
|
||||
from fastmcp.utilities.logging import configure_logging
|
||||
from fastmcp.utilities.types import Image
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
|
@ -112,12 +119,22 @@ class FastMCP:
|
|||
for info in tools
|
||||
]
|
||||
|
||||
def get_context(self) -> Optional["Context"]:
|
||||
try:
|
||||
request_context = self._mcp_server.request_context
|
||||
return Context(request_context=request_context, fastmcp=self)
|
||||
except LookupError:
|
||||
return None
|
||||
|
||||
async def call_tool(
|
||||
self, name: str, arguments: dict
|
||||
) -> Sequence[Union[TextContent, ImageContent]]:
|
||||
"""Call a tool by name with arguments."""
|
||||
try:
|
||||
result = await self._tool_manager.call_tool(name, arguments)
|
||||
context = self.get_context()
|
||||
result = await self._tool_manager.call_tool(
|
||||
name, arguments, context=context
|
||||
)
|
||||
return _convert_to_content(result)
|
||||
except Exception as e:
|
||||
logger.error(f"Error calling tool {name}: {e}")
|
||||
|
|
@ -172,13 +189,45 @@ class FastMCP:
|
|||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
) -> None:
|
||||
"""Add a tool to the server."""
|
||||
"""Add a tool to the server.
|
||||
|
||||
The tool function can optionally request a Context object by adding a parameter
|
||||
with the Context type annotation. See the @tool decorator for examples.
|
||||
|
||||
Args:
|
||||
func: The function to register as a tool
|
||||
name: Optional name for the tool (defaults to function name)
|
||||
description: Optional description of what the tool does
|
||||
"""
|
||||
self._tool_manager.add_tool(func, name=name, description=description)
|
||||
|
||||
def tool(
|
||||
self, name: Optional[str] = None, description: Optional[str] = None
|
||||
) -> Callable:
|
||||
"""Decorator to register a tool."""
|
||||
"""Decorator to register a tool.
|
||||
|
||||
Tools can optionally request a Context object by adding a parameter with the Context type annotation.
|
||||
The context provides access to MCP capabilities like logging, progress reporting, and resource access.
|
||||
|
||||
Args:
|
||||
name: Optional name for the tool (defaults to function name)
|
||||
description: Optional description of what the tool does
|
||||
|
||||
Example:
|
||||
@server.tool()
|
||||
def my_tool(x: int) -> str:
|
||||
return str(x)
|
||||
|
||||
@server.tool()
|
||||
def tool_with_context(x: int, ctx: Context) -> str:
|
||||
ctx.info(f"Processing {x}")
|
||||
return str(x)
|
||||
|
||||
@server.tool()
|
||||
async def async_tool(x: int, context: Context) -> str:
|
||||
await context.report_progress(50, 100)
|
||||
return str(x)
|
||||
"""
|
||||
# Check if user passed function directly instead of calling decorator
|
||||
if callable(name):
|
||||
raise TypeError(
|
||||
|
|
@ -348,7 +397,7 @@ def _convert_to_content(value: Any) -> Sequence[Union[TextContent, ImageContent]
|
|||
result.append(
|
||||
TextContent(
|
||||
type="text",
|
||||
text=json.dumps(item, default=pydantic.json.pydantic_encoder),
|
||||
text=json.dumps(pydantic_core.to_jsonable_python(item)),
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
|
@ -365,6 +414,146 @@ def _convert_to_content(value: Any) -> Sequence[Union[TextContent, ImageContent]
|
|||
return [
|
||||
TextContent(
|
||||
type="text",
|
||||
text=json.dumps(value, indent=2, default=pydantic.json.pydantic_encoder),
|
||||
text=json.dumps(pydantic_core.to_jsonable_python(value)),
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
class Context(BaseModel):
|
||||
"""Context object providing access to MCP capabilities.
|
||||
|
||||
This provides a cleaner interface to MCP's RequestContext functionality.
|
||||
It gets injected into tool and resource functions that request it via type hints.
|
||||
|
||||
To use context in a tool function, add a parameter with the Context type annotation:
|
||||
|
||||
```python
|
||||
@server.tool()
|
||||
def my_tool(x: int, ctx: Context) -> str:
|
||||
# Log messages to the client
|
||||
ctx.info(f"Processing {x}")
|
||||
ctx.debug("Debug info")
|
||||
ctx.warning("Warning message")
|
||||
ctx.error("Error message")
|
||||
|
||||
# Report progress
|
||||
ctx.report_progress(50, 100)
|
||||
|
||||
# Access resources
|
||||
data = ctx.read_resource("resource://data")
|
||||
|
||||
# Get request info
|
||||
request_id = ctx.request_id
|
||||
client_id = ctx.client_id
|
||||
|
||||
return str(x)
|
||||
```
|
||||
|
||||
The context parameter name can be anything as long as it's annotated with Context.
|
||||
The context is optional - tools that don't need it can omit the parameter.
|
||||
"""
|
||||
|
||||
_request_context: RequestContext
|
||||
_fastmcp: FastMCP
|
||||
|
||||
def __init__(
|
||||
self, *, request_context: RequestContext, fastmcp: FastMCP, **kwargs: Any
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
self._request_context = request_context
|
||||
self._fastmcp = fastmcp
|
||||
|
||||
@property
|
||||
def fastmcp(self) -> FastMCP:
|
||||
"""Access to the FastMCP server."""
|
||||
return self._fastmcp
|
||||
|
||||
@property
|
||||
def request_context(self) -> RequestContext:
|
||||
"""Access to the underlying request context."""
|
||||
return self._request_context
|
||||
|
||||
async def report_progress(
|
||||
self, progress: float, total: Optional[float] = None
|
||||
) -> None:
|
||||
"""Report progress for the current operation.
|
||||
|
||||
Args:
|
||||
progress: Current progress value e.g. 24
|
||||
total: Optional total value e.g. 100
|
||||
"""
|
||||
|
||||
progress_token = (
|
||||
self.request_context.meta.progressToken
|
||||
if self.request_context.meta
|
||||
else None
|
||||
)
|
||||
|
||||
if not progress_token:
|
||||
return
|
||||
|
||||
await self.request_context.session.send_progress_notification(
|
||||
progress_token=progress_token, progress=progress, total=total
|
||||
)
|
||||
|
||||
async def read_resource(self, uri: Union[str, AnyUrl]) -> Union[str, bytes]:
|
||||
"""Read a resource by URI.
|
||||
|
||||
Args:
|
||||
uri: Resource URI to read
|
||||
|
||||
Returns:
|
||||
The resource content as either text or bytes
|
||||
"""
|
||||
return await self._fastmcp.read_resource(uri)
|
||||
|
||||
def log(
|
||||
self,
|
||||
level: Literal["debug", "info", "warning", "error"],
|
||||
message: str,
|
||||
*,
|
||||
logger_name: Optional[str] = None,
|
||||
) -> None:
|
||||
"""Send a log message to the client.
|
||||
|
||||
Args:
|
||||
level: Log level (debug, info, warning, error)
|
||||
message: Log message
|
||||
logger_name: Optional logger name
|
||||
**extra: Additional structured data to include
|
||||
"""
|
||||
self.request_context.session.send_log_message(
|
||||
level=level, data=message, logger=logger_name
|
||||
)
|
||||
|
||||
@property
|
||||
def client_id(self) -> Optional[str]:
|
||||
"""Get the client ID if available."""
|
||||
return self.request_context.meta.clientId if self.request_context.meta else None
|
||||
|
||||
@property
|
||||
def request_id(self) -> str:
|
||||
"""Get the unique ID for this request."""
|
||||
return self.request_context.request_id
|
||||
|
||||
@property
|
||||
def session(self):
|
||||
"""Access to the underlying session for advanced usage."""
|
||||
return self.request_context.session
|
||||
|
||||
# Convenience methods for common log levels
|
||||
def debug(self, message: str, **extra: Any) -> None:
|
||||
"""Send a debug log message."""
|
||||
self.log("debug", message, **extra)
|
||||
|
||||
def info(self, message: str, **extra: Any) -> None:
|
||||
"""Send an info log message."""
|
||||
self.log("info", message, **extra)
|
||||
|
||||
def warning(self, message: str, **extra: Any) -> None:
|
||||
"""Send a warning log message."""
|
||||
self.log("warning", message, **extra)
|
||||
|
||||
def error(self, message: str, **extra: Any) -> None:
|
||||
"""Send an error log message."""
|
||||
self.log("error", message, **extra)
|
||||
|
|
|
|||
|
|
@ -1,12 +1,16 @@
|
|||
"""Tool management for FastMCP."""
|
||||
|
||||
import inspect
|
||||
from typing import Any, Callable, Dict, Optional
|
||||
from typing import Any, Callable, Dict, Optional, TYPE_CHECKING
|
||||
|
||||
from pydantic import BaseModel, Field, TypeAdapter, validate_call
|
||||
|
||||
from .exceptions import ToolError
|
||||
from .utilities.logging import get_logger
|
||||
import fastmcp
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastmcp.server import Context
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
|
@ -19,6 +23,9 @@ class Tool(BaseModel):
|
|||
description: str = Field(description="Description of what the tool does")
|
||||
parameters: dict = Field(description="JSON schema for tool parameters")
|
||||
is_async: bool = Field(description="Whether the tool is async")
|
||||
context_kwarg: Optional[str] = Field(
|
||||
None, description="Name of the kwarg that should receive context"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_function(
|
||||
|
|
@ -26,6 +33,7 @@ class Tool(BaseModel):
|
|||
func: Callable,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
context_kwarg: Optional[str] = None,
|
||||
) -> "Tool":
|
||||
"""Create a Tool from a function."""
|
||||
func_name = name or func.__name__
|
||||
|
|
@ -39,6 +47,14 @@ class Tool(BaseModel):
|
|||
# Get schema from TypeAdapter - will fail if function isn't properly typed
|
||||
parameters = TypeAdapter(func).json_schema()
|
||||
|
||||
# Find context parameter if it exists
|
||||
if context_kwarg is None:
|
||||
sig = inspect.signature(func)
|
||||
for param_name, param in sig.parameters.items():
|
||||
if param.annotation is fastmcp.Context:
|
||||
context_kwarg = param_name
|
||||
break
|
||||
|
||||
# ensure the arguments are properly cast
|
||||
func = validate_call(func)
|
||||
|
||||
|
|
@ -48,11 +64,16 @@ class Tool(BaseModel):
|
|||
description=func_doc,
|
||||
parameters=parameters,
|
||||
is_async=is_async,
|
||||
context_kwarg=context_kwarg,
|
||||
)
|
||||
|
||||
async def run(self, arguments: dict) -> Any:
|
||||
async def run(self, arguments: dict, context: Optional["Context"] = None) -> Any:
|
||||
"""Run the tool with arguments."""
|
||||
try:
|
||||
# Inject context if needed
|
||||
if self.context_kwarg and context:
|
||||
arguments[self.context_kwarg] = context
|
||||
|
||||
# Call function with proper async handling
|
||||
if self.is_async:
|
||||
return await self.func(**arguments)
|
||||
|
|
@ -92,10 +113,12 @@ class ToolManager:
|
|||
self._tools[tool.name] = tool
|
||||
return tool
|
||||
|
||||
async def call_tool(self, name: str, arguments: dict) -> Any:
|
||||
async def call_tool(
|
||||
self, name: str, arguments: dict, context: Optional["Context"] = None
|
||||
) -> Any:
|
||||
"""Call a tool by name with arguments."""
|
||||
tool = self.get_tool(name)
|
||||
if not tool:
|
||||
raise ToolError(f"Unknown tool: {name}")
|
||||
|
||||
return await tool.run(arguments)
|
||||
return await tool.run(arguments, context=context)
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
from pydantic import BaseModel
|
||||
import pytest
|
||||
from fastmcp.resources import FunctionResource
|
||||
|
||||
|
|
@ -80,6 +81,20 @@ class TestFunctionResource:
|
|||
with pytest.raises(ValueError, match="Error reading resource function://test"):
|
||||
await resource.read()
|
||||
|
||||
async def test_basemodel_conversion(self):
|
||||
"""Test handling of BaseModel types."""
|
||||
|
||||
class MyModel(BaseModel):
|
||||
name: str
|
||||
|
||||
resource = FunctionResource(
|
||||
uri="function://test",
|
||||
name="test",
|
||||
func=lambda: MyModel(name="test"),
|
||||
)
|
||||
content = await resource.read()
|
||||
assert content == '{"name": "test"}'
|
||||
|
||||
async def test_custom_type_conversion(self):
|
||||
"""Test handling of custom types."""
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
from mcp.shared.memory import (
|
||||
create_connected_server_and_client_session as client_session,
|
||||
)
|
||||
from fastmcp import FastMCP
|
||||
from fastmcp import FastMCP, Context
|
||||
from fastmcp.resources import FileResource, FunctionResource
|
||||
from fastmcp.utilities.types import Image
|
||||
from mcp.types import TextContent, ImageContent
|
||||
|
|
@ -9,7 +9,10 @@ import pytest
|
|||
from pydantic import BaseModel
|
||||
from pathlib import Path
|
||||
import base64
|
||||
from typing import Union
|
||||
from typing import Union, TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastmcp import Context
|
||||
|
||||
|
||||
class TestServer:
|
||||
|
|
@ -368,3 +371,95 @@ class TestServerResourceTemplates:
|
|||
assert isinstance(resource, FunctionResource)
|
||||
result = await resource.read()
|
||||
assert result == "Data for test"
|
||||
|
||||
|
||||
class TestContextInjection:
|
||||
"""Test context injection in tools."""
|
||||
|
||||
async def test_context_detection(self):
|
||||
"""Test that context parameters are properly detected."""
|
||||
mcp = FastMCP()
|
||||
|
||||
def tool_with_context(x: int, ctx: Context) -> str:
|
||||
return f"Request {ctx.request_id}: {x}"
|
||||
|
||||
tool = mcp._tool_manager.add_tool(tool_with_context)
|
||||
assert tool.context_kwarg == "ctx"
|
||||
|
||||
async def test_context_injection(self):
|
||||
"""Test that context is properly injected into tool calls."""
|
||||
mcp = FastMCP()
|
||||
|
||||
def tool_with_context(x: int, ctx: Context) -> str:
|
||||
assert ctx.request_id is not None
|
||||
return f"Request {ctx.request_id}: {x}"
|
||||
|
||||
mcp.add_tool(tool_with_context)
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
result = await client.call_tool("tool_with_context", {"x": 42})
|
||||
assert len(result.content) == 1
|
||||
assert "Request" in result.content[0].text
|
||||
assert "42" in result.content[0].text
|
||||
|
||||
async def test_async_context(self):
|
||||
"""Test that context works in async functions."""
|
||||
mcp = FastMCP()
|
||||
|
||||
async def async_tool(x: int, ctx: Context) -> str:
|
||||
assert ctx.request_id is not None
|
||||
return f"Async request {ctx.request_id}: {x}"
|
||||
|
||||
mcp.add_tool(async_tool)
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
result = await client.call_tool("async_tool", {"x": 42})
|
||||
assert len(result.content) == 1
|
||||
assert "Async request" in result.content[0].text
|
||||
assert "42" in result.content[0].text
|
||||
|
||||
async def test_context_logging(self):
|
||||
"""Test that context logging methods work."""
|
||||
mcp = FastMCP()
|
||||
|
||||
def logging_tool(msg: str, ctx: Context) -> str:
|
||||
ctx.debug("Debug message")
|
||||
ctx.info("Info message")
|
||||
ctx.warning("Warning message")
|
||||
ctx.error("Error message")
|
||||
return f"Logged messages for {msg}"
|
||||
|
||||
mcp.add_tool(logging_tool)
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
result = await client.call_tool("logging_tool", {"msg": "test"})
|
||||
assert len(result.content) == 1
|
||||
assert "Logged messages for test" in result.content[0].text
|
||||
|
||||
async def test_optional_context(self):
|
||||
"""Test that context is optional."""
|
||||
mcp = FastMCP()
|
||||
|
||||
def no_context(x: int) -> int:
|
||||
return x * 2
|
||||
|
||||
mcp.add_tool(no_context)
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
result = await client.call_tool("no_context", {"x": 21})
|
||||
assert len(result.content) == 1
|
||||
assert result.content[0].text == "42"
|
||||
|
||||
async def test_context_resource_access(self):
|
||||
"""Test that context can access resources."""
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.resource("test://data")
|
||||
def test_resource() -> str:
|
||||
return "resource data"
|
||||
|
||||
@mcp.tool()
|
||||
async def tool_with_resource(ctx: Context) -> str:
|
||||
data = await ctx.read_resource("test://data")
|
||||
return f"Read resource: {data}"
|
||||
|
||||
async with client_session(mcp._mcp_server) as client:
|
||||
result = await client.call_tool("tool_with_resource", {})
|
||||
assert len(result.content) == 1
|
||||
assert "Read resource: resource data" in result.content[0].text
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue