Add context support to all objects

This commit is contained in:
Jeremiah Lowin 2025-04-25 20:33:24 -04:00
commit 1480c8771b
9 changed files with 231 additions and 45 deletions

View file

@ -1,9 +1,11 @@
"""Base classes for FastMCP prompts."""
from __future__ import annotations as _annotations
import inspect
import json
from collections.abc import Awaitable, Callable, Sequence
from typing import Annotated, Any, Literal
from typing import TYPE_CHECKING, Annotated, Any, Literal
import pydantic_core
from mcp.types import EmbeddedResource, ImageContent, TextContent
@ -13,6 +15,12 @@ from pydantic import BaseModel, BeforeValidator, Field, TypeAdapter, validate_ca
from fastmcp.utilities.types import _convert_set_defaults
if TYPE_CHECKING:
from mcp.server.session import ServerSessionT
from mcp.shared.context import LifespanContextT
from fastmcp.server import Context
CONTENT_TYPES = TextContent | ImageContent | EmbeddedResource
@ -72,6 +80,9 @@ class Prompt(BaseModel):
None, description="Arguments that can be passed to the prompt"
)
fn: Callable[..., PromptResult | Awaitable[PromptResult]]
context_kwarg: str | None = Field(
None, description="Name of the kwarg that should receive context"
)
@classmethod
def from_function(
@ -80,7 +91,8 @@ class Prompt(BaseModel):
name: str | None = None,
description: str | None = None,
tags: set[str] | None = None,
) -> "Prompt":
context_kwarg: str | None = None,
) -> Prompt:
"""Create a Prompt from a function.
The function can return:
@ -89,11 +101,24 @@ class Prompt(BaseModel):
- A dict (converted to a message)
- A sequence of any of the above
"""
from fastmcp import Context
func_name = name or fn.__name__
if func_name == "<lambda>":
raise ValueError("You must provide a name for lambda functions")
# Auto-detect context parameter if not provided
if context_kwarg is None:
if inspect.ismethod(fn) and hasattr(fn, "__func__"):
sig = inspect.signature(fn.__func__)
else:
sig = inspect.signature(fn)
for param_name, param in sig.parameters.items():
if param.annotation is Context:
context_kwarg = param_name
break
# Get schema from TypeAdapter - will fail if function isn't properly typed
parameters = TypeAdapter(fn).json_schema()
@ -101,6 +126,10 @@ class Prompt(BaseModel):
arguments: list[PromptArgument] = []
if "properties" in parameters:
for param_name, param in parameters["properties"].items():
# Skip context parameter
if param_name == context_kwarg:
continue
required = param_name in parameters.get("required", [])
arguments.append(
PromptArgument(
@ -119,9 +148,14 @@ class Prompt(BaseModel):
arguments=arguments,
fn=fn,
tags=tags or set(),
context_kwarg=context_kwarg,
)
async def render(self, arguments: dict[str, Any] | None = None) -> list[Message]:
async def render(
self,
arguments: dict[str, Any] | None = None,
context: Context[ServerSessionT, LifespanContextT] | None = None,
) -> list[Message]:
"""Render the prompt with arguments."""
# Validate required arguments
if self.arguments:
@ -132,8 +166,13 @@ class Prompt(BaseModel):
raise ValueError(f"Missing required arguments: {missing}")
try:
# Prepare arguments with context
kwargs = arguments.copy() if arguments else {}
if self.context_kwarg is not None and context is not None:
kwargs[self.context_kwarg] = context
# Call function and check if result is a coroutine
result = self.fn(**(arguments or {}))
result = self.fn(**kwargs)
if inspect.iscoroutine(result):
result = await result

View file

@ -1,13 +1,21 @@
"""Prompt management functionality."""
from __future__ import annotations as _annotations
from collections.abc import Awaitable, Callable
from typing import Any
from typing import TYPE_CHECKING, Any
from fastmcp.exceptions import NotFoundError
from fastmcp.prompts.prompt import Message, Prompt, PromptResult
from fastmcp.settings import DuplicateBehavior
from fastmcp.utilities.logging import get_logger
if TYPE_CHECKING:
from mcp.server.session import ServerSessionT
from mcp.shared.context import LifespanContextT
from fastmcp.server import Context
logger = get_logger(__name__)
@ -69,14 +77,17 @@ class PromptManager:
return prompt
async def render_prompt(
self, name: str, arguments: dict[str, Any] | None = None
self,
name: str,
arguments: dict[str, Any] | None = None,
context: Context[ServerSessionT, LifespanContextT] | None = None,
) -> list[Message]:
"""Render a prompt by name with arguments."""
prompt = self.get_prompt(name)
if not prompt:
raise NotFoundError(f"Unknown prompt: {name}")
return await prompt.render(arguments)
return await prompt.render(arguments, context=context)
def has_prompt(self, key: str) -> bool:
"""Check if a prompt exists."""

View file

@ -1,7 +1,9 @@
"""Base classes and interfaces for FastMCP resources."""
from __future__ import annotations
import abc
from typing import Annotated, Any
from typing import TYPE_CHECKING, Annotated, Any
from mcp.types import Resource as MCPResource
from pydantic import (
@ -17,6 +19,12 @@ from pydantic import (
from fastmcp.utilities.types import _convert_set_defaults
if TYPE_CHECKING:
from mcp.server.session import ServerSessionT
from mcp.shared.context import LifespanContextT
from fastmcp.server import Context
class Resource(BaseModel, abc.ABC):
"""Base class for all resources."""
@ -58,7 +66,9 @@ class Resource(BaseModel, abc.ABC):
raise ValueError("Either name or uri must be provided")
@abc.abstractmethod
async def read(self) -> str | bytes:
async def read(
self, context: Context[ServerSessionT, LifespanContextT] | None = None
) -> str | bytes:
"""Read the resource content."""
pass

View file

@ -212,9 +212,13 @@ class ResourceManager:
return True
return False
async def get_resource(self, uri: AnyUrl | str) -> Resource:
async def get_resource(self, uri: AnyUrl | str, context=None) -> Resource:
"""Get resource by URI, checking concrete resources first, then templates.
Args:
uri: The URI of the resource to get
context: Optional context object to pass to template resources
Raises:
NotFoundError: If no resource or template matching the URI is found.
"""
@ -230,7 +234,9 @@ class ResourceManager:
# Try to match against the storage key (which might be a custom key)
if params := match_uri_template(uri_str, storage_key):
try:
return await template.create_resource(uri_str, params)
return await template.create_resource(
uri_str, params, context=context
)
except Exception as e:
raise ValueError(f"Error creating resource from template: {e}")

View file

@ -5,7 +5,7 @@ from __future__ import annotations
import inspect
import re
from collections.abc import Callable
from typing import Annotated, Any
from typing import TYPE_CHECKING, Annotated, Any
from urllib.parse import unquote
from mcp.types import ResourceTemplate as MCPResourceTemplate
@ -22,6 +22,12 @@ from pydantic import (
from fastmcp.resources.types import FunctionResource, Resource
from fastmcp.utilities.types import _convert_set_defaults
if TYPE_CHECKING:
from mcp.server.session import ServerSessionT
from mcp.shared.context import LifespanContextT
from fastmcp.server import Context
def build_regex(template: str) -> re.Pattern:
parts = re.split(r"(\{[^}]+\})", template)
@ -70,6 +76,9 @@ class ResourceTemplate(BaseModel):
parameters: dict[str, Any] = Field(
description="JSON schema for function parameters"
)
context_kwarg: str | None = Field(
None, description="Name of the kwarg that should receive context"
)
@field_validator("mime_type", mode="before")
@classmethod
@ -88,18 +97,34 @@ class ResourceTemplate(BaseModel):
description: str | None = None,
mime_type: str | None = None,
tags: set[str] | None = None,
context_kwarg: str | None = None,
) -> ResourceTemplate:
"""Create a template from a function."""
from fastmcp import Context
func_name = name or fn.__name__
if func_name == "<lambda>":
raise ValueError("You must provide a name for lambda functions")
# Auto-detect context parameter if not provided
if context_kwarg is None:
if inspect.ismethod(fn) and hasattr(fn, "__func__"):
sig = inspect.signature(fn.__func__)
else:
sig = inspect.signature(fn)
for param_name, param in sig.parameters.items():
if param.annotation is Context:
context_kwarg = param_name
break
# Validate that URI params match function params
uri_params = set(re.findall(r"{(\w+)(?:\*)?}", uri_template))
if not uri_params:
raise ValueError("URI template must contain at least one parameter")
func_params = set(inspect.signature(fn).parameters.keys())
if context_kwarg:
func_params.discard(context_kwarg)
# get the parameters that are required
required_params = {
@ -107,6 +132,8 @@ class ResourceTemplate(BaseModel):
for p in func_params
if inspect.signature(fn).parameters[p].default is inspect.Parameter.empty
}
if context_kwarg and context_kwarg in required_params:
required_params.discard(context_kwarg)
if not required_params.issubset(uri_params):
raise ValueError(
@ -132,17 +159,28 @@ class ResourceTemplate(BaseModel):
fn=fn,
parameters=parameters,
tags=tags or set(),
context_kwarg=context_kwarg,
)
def matches(self, uri: str) -> dict[str, Any] | None:
"""Check if URI matches template and extract parameters."""
return match_uri_template(uri, self.uri_template)
async def create_resource(self, uri: str, params: dict[str, Any]) -> Resource:
async def create_resource(
self,
uri: str,
params: dict[str, Any],
context: Context[ServerSessionT, LifespanContextT] | None = None,
) -> Resource:
"""Create a resource from the template with the given parameters."""
try:
# Add context to parameters if needed
kwargs = params.copy()
if self.context_kwarg is not None and context is not None:
kwargs[self.context_kwarg] = context
# Call function and check if result is a coroutine
result = self.fn(**params)
result = self.fn(**kwargs)
if inspect.iscoroutine(result):
result = await result
@ -153,6 +191,7 @@ class ResourceTemplate(BaseModel):
mime_type=self.mime_type,
fn=lambda: result, # Capture result in closure
tags=self.tags,
context_kwarg=self.context_kwarg,
)
except Exception as e:
raise ValueError(f"Error creating resource from template: {e}")

View file

@ -1,10 +1,12 @@
"""Concrete resource implementations."""
from __future__ import annotations
import inspect
import json
from collections.abc import Callable
from pathlib import Path
from typing import Any
from typing import TYPE_CHECKING, Any
import anyio
import anyio.to_thread
@ -15,13 +17,21 @@ from pydantic import Field, ValidationInfo
from fastmcp.resources.resource import Resource
if TYPE_CHECKING:
from mcp.server.session import ServerSessionT
from mcp.shared.context import LifespanContextT
from fastmcp.server import Context
class TextResource(Resource):
"""A resource that reads from a string."""
text: str = Field(description="Text content of the resource")
async def read(self) -> str:
async def read(
self, context: Context[ServerSessionT, LifespanContextT] | None = None
) -> str:
"""Read the text content."""
return self.text
@ -31,7 +41,9 @@ class BinaryResource(Resource):
data: bytes = Field(description="Binary content of the resource")
async def read(self) -> bytes:
async def read(
self, context: Context[ServerSessionT, LifespanContextT] | None = None
) -> bytes:
"""Read the binary content."""
return self.data
@ -50,13 +62,23 @@ class FunctionResource(Resource):
"""
fn: Callable[[], Any]
context_kwarg: str | None = Field(
default=None, description="Name of the kwarg that should receive context"
)
async def read(self) -> str | bytes:
async def read(
self, context: Context[ServerSessionT, LifespanContextT] | None = None
) -> str | bytes:
"""Read the resource by calling the wrapped function."""
try:
result = (
await self.fn() if inspect.iscoroutinefunction(self.fn) else self.fn()
)
kwargs = {}
if self.context_kwarg is not None:
kwargs[self.context_kwarg] = context
result = self.fn(**kwargs)
if inspect.iscoroutinefunction(self.fn):
result = await result
if isinstance(result, Resource):
return await result.read()
if isinstance(result, bytes):

View file

@ -1,11 +1,13 @@
"""FastMCP server implementation for OpenAPI integration."""
from __future__ import annotations
import enum
import json
import re
from dataclasses import dataclass
from re import Pattern
from typing import Any, Literal
from typing import TYPE_CHECKING, Any, Literal
import httpx
from mcp.types import TextContent
@ -22,6 +24,12 @@ from fastmcp.utilities.openapi import (
format_description_with_responses,
)
if TYPE_CHECKING:
from mcp.server.session import ServerSessionT
from mcp.shared.context import LifespanContextT
from fastmcp.server import Context
logger = get_logger(__name__)
HttpMethod = Literal["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"]
@ -347,11 +355,17 @@ class OpenAPIResourceTemplate(ResourceTemplate):
fn=lambda **kwargs: None,
parameters=parameters,
tags=tags,
context_kwarg=None,
)
self._client = client
self._route = route
async def create_resource(self, uri: str, params: dict[str, Any]) -> Resource:
async def create_resource(
self,
uri: str,
params: dict[str, Any],
context: Context[ServerSessionT, LifespanContextT] | None = None,
) -> Resource:
"""Create a resource with the given parameters."""
# Generate a URI for this resource instance
uri_parts = []

View file

@ -1,4 +1,6 @@
from typing import Any, cast
from __future__ import annotations
from typing import TYPE_CHECKING, Any, cast
from urllib.parse import quote
import mcp.types
@ -25,6 +27,12 @@ from fastmcp.tools.tool import Tool
from fastmcp.utilities.func_metadata import func_metadata
from fastmcp.utilities.logging import get_logger
if TYPE_CHECKING:
from mcp.server.session import ServerSessionT
from mcp.shared.context import LifespanContextT
from fastmcp.server import Context
logger = get_logger(__name__)
@ -33,12 +41,12 @@ def _proxy_passthrough():
class ProxyTool(Tool):
def __init__(self, client: "Client", **kwargs):
def __init__(self, client: Client, **kwargs):
super().__init__(**kwargs)
self._client = client
@classmethod
async def from_client(cls, client: "Client", tool: mcp.types.Tool) -> "ProxyTool":
async def from_client(cls, client: Client, tool: mcp.types.Tool) -> ProxyTool:
return cls(
client=client,
name=tool.name,
@ -50,7 +58,9 @@ class ProxyTool(Tool):
)
async def run(
self, arguments: dict[str, Any], context: Context | None = None
self,
arguments: dict[str, Any],
context: Context[ServerSessionT, LifespanContextT] | 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
@ -64,17 +74,15 @@ class ProxyTool(Tool):
class ProxyResource(Resource):
def __init__(
self, client: "Client", *, _value: str | bytes | None = None, **kwargs
):
def __init__(self, client: Client, *, _value: str | bytes | None = None, **kwargs):
super().__init__(**kwargs)
self._client = client
self._value = _value
@classmethod
async def from_client(
cls, client: "Client", resource: mcp.types.Resource
) -> "ProxyResource":
cls, client: Client, resource: mcp.types.Resource
) -> ProxyResource:
return cls(
client=client,
uri=resource.uri,
@ -83,7 +91,9 @@ class ProxyResource(Resource):
mime_type=resource.mimeType,
)
async def read(self) -> str | bytes:
async def read(
self, context: Context[ServerSessionT, LifespanContextT] | None = None
) -> str | bytes:
if self._value is not None:
return self._value
@ -98,14 +108,14 @@ class ProxyResource(Resource):
class ProxyTemplate(ResourceTemplate):
def __init__(self, client: "Client", **kwargs):
def __init__(self, client: Client, **kwargs):
super().__init__(**kwargs)
self._client = client
@classmethod
async def from_client(
cls, client: "Client", template: mcp.types.ResourceTemplate
) -> "ProxyTemplate":
cls, client: Client, template: mcp.types.ResourceTemplate
) -> ProxyTemplate:
return cls(
client=client,
uri_template=template.uriTemplate,
@ -115,7 +125,12 @@ class ProxyTemplate(ResourceTemplate):
parameters={},
)
async def create_resource(self, uri: str, params: dict[str, Any]) -> ProxyResource:
async def create_resource(
self,
uri: str,
params: dict[str, Any],
context: Context[ServerSessionT, LifespanContextT] | None = None,
) -> ProxyResource:
# dont use the provided uri, because it may not be the same as the
# uri_template on the remote server.
# quote params to ensure they are valid for the uri_template
@ -144,14 +159,12 @@ class ProxyTemplate(ResourceTemplate):
class ProxyPrompt(Prompt):
def __init__(self, client: "Client", **kwargs):
def __init__(self, client: Client, **kwargs):
super().__init__(**kwargs)
self._client = client
@classmethod
async def from_client(
cls, client: "Client", prompt: mcp.types.Prompt
) -> "ProxyPrompt":
async def from_client(cls, client: Client, prompt: mcp.types.Prompt) -> ProxyPrompt:
return cls(
client=client,
name=prompt.name,
@ -160,14 +173,18 @@ class ProxyPrompt(Prompt):
fn=_proxy_passthrough,
)
async def render(self, arguments: dict[str, Any]) -> list[Message]:
async def render(
self,
arguments: dict[str, Any],
context: Context[ServerSessionT, LifespanContextT] | None = None,
) -> list[Message]:
async with self._client:
result = await self._client.get_prompt(self.name, arguments)
return [Message(role=m.role, content=m.content) for m in result]
class FastMCPProxy(FastMCP):
def __init__(self, client: "Client", **kwargs):
def __init__(self, client: Client, **kwargs):
super().__init__(**kwargs)
self.client = client

View file

@ -398,7 +398,8 @@ class FastMCP(Generic[LifespanResultT]):
server.
"""
if self._resource_manager.has_resource(uri):
resource = await self._resource_manager.get_resource(uri)
context = self.get_context()
resource = await self._resource_manager.get_resource(uri, context=context)
try:
content = await resource.read()
return [
@ -424,7 +425,10 @@ class FastMCP(Generic[LifespanResultT]):
"""
if self._prompt_manager.has_prompt(name):
messages = await self._prompt_manager.render_prompt(name, arguments)
context = self.get_context()
messages = await self._prompt_manager.render_prompt(
name, arguments, context=context
)
return GetPromptResult(messages=pydantic_core.to_jsonable_python(messages))
else:
for server in self._mounted_servers.values():
@ -562,6 +566,10 @@ class FastMCP(Generic[LifespanResultT]):
- bytes for binary content
- other types will be converted to JSON
Resources 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 session information.
If the URI contains parameters (e.g. "resource://{param}") or the function
has parameters, it will be registered as a template resource.
@ -586,6 +594,11 @@ class FastMCP(Generic[LifespanResultT]):
def get_weather(city: str) -> str:
return f"Weather for {city}"
@server.resource("resource://{city}/weather")
def get_weather_with_context(city: str, ctx: Context) -> str:
ctx.info(f"Fetching weather for {city}")
return f"Weather for {city}"
@server.resource("resource://{city}/weather")
async def get_weather(city: str) -> str:
data = await fetch_weather(city)
@ -639,6 +652,10 @@ class FastMCP(Generic[LifespanResultT]):
) -> Callable[[AnyFunction], AnyFunction]:
"""Decorator to register a prompt.
Prompts 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 session information.
Args:
name: Optional name for the prompt (defaults to function name)
description: Optional description of what the prompt does
@ -655,6 +672,17 @@ class FastMCP(Generic[LifespanResultT]):
}
]
@server.prompt()
def analyze_with_context(table_name: str, ctx: Context) -> list[Message]:
ctx.info(f"Analyzing table {table_name}")
schema = read_table_schema(table_name)
return [
{
"role": "user",
"content": f"Analyze this schema:\n{schema}"
}
]
@server.prompt()
async def analyze_file(path: str) -> list[Message]:
content = await read_file(path)