From f08c6c9b981fd9ba09aebc20a7861c9c436707ed Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sun, 4 May 2025 13:20:34 -0400 Subject: [PATCH 1/4] Dry out retrieving context kwarg --- src/fastmcp/prompts/prompt.py | 28 ++++++++++------------------ src/fastmcp/resources/template.py | 11 ++--------- src/fastmcp/tools/tool.py | 15 ++++----------- src/fastmcp/utilities/types.py | 26 +++++++++++++++++++++++++- 4 files changed, 41 insertions(+), 39 deletions(-) diff --git a/src/fastmcp/prompts/prompt.py b/src/fastmcp/prompts/prompt.py index 0ce8f17e0..39c346942 100644 --- a/src/fastmcp/prompts/prompt.py +++ b/src/fastmcp/prompts/prompt.py @@ -12,9 +12,11 @@ from mcp.types import Prompt as MCPPrompt from mcp.types import PromptArgument as MCPPromptArgument from pydantic import BaseModel, BeforeValidator, Field, TypeAdapter, validate_call +from fastmcp.utilities.json_schema import prune_params from fastmcp.utilities.types import ( _convert_set_defaults, - is_class_member_of_type, + find_kwarg_by_type, + get_cached_typeadapter, ) if TYPE_CHECKING: @@ -110,34 +112,24 @@ class Prompt(BaseModel): if func_name == "": raise ValueError("You must provide a name for lambda functions") + type_adapter = get_cached_typeadapter(fn) + parameters = type_adapter.json_schema() + # 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 is_class_member_of_type(param.annotation, Context): - context_kwarg = param_name - break - - # Get schema from TypeAdapter - will fail if function isn't properly typed - parameters = TypeAdapter(fn).json_schema() + context_kwarg = find_kwarg_by_type(fn, kwarg_type=Context) + if context_kwarg: + parameters = prune_params(parameters, params=[context_kwarg]) # Convert parameters to PromptArguments 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( name=param_name, description=param.get("description"), - required=required, + required=param_name in parameters.get("required", []), ) ) diff --git a/src/fastmcp/resources/template.py b/src/fastmcp/resources/template.py index 6cb6ac3a3..f789076d9 100644 --- a/src/fastmcp/resources/template.py +++ b/src/fastmcp/resources/template.py @@ -22,7 +22,7 @@ from pydantic import ( from fastmcp.resources.types import FunctionResource, Resource from fastmcp.utilities.types import ( _convert_set_defaults, - is_class_member_of_type, + find_kwarg_by_type, ) if TYPE_CHECKING: @@ -111,14 +111,7 @@ class ResourceTemplate(BaseModel): # 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 is_class_member_of_type(param.annotation, Context): - context_kwarg = param_name - break + context_kwarg = find_kwarg_by_type(fn, kwarg_type=Context) # Validate that URI params match function params uri_params = set(re.findall(r"{(\w+)(?:\*)?}", uri_template)) diff --git a/src/fastmcp/tools/tool.py b/src/fastmcp/tools/tool.py index 4f178a6e8..1d1a09c8b 100644 --- a/src/fastmcp/tools/tool.py +++ b/src/fastmcp/tools/tool.py @@ -16,8 +16,8 @@ from fastmcp.utilities.logging import get_logger from fastmcp.utilities.types import ( Image, _convert_set_defaults, + find_kwarg_by_type, get_cached_typeadapter, - is_class_member_of_type, ) if TYPE_CHECKING: @@ -74,18 +74,11 @@ class Tool(BaseModel): func_doc = description or fn.__doc__ or "" - if inspect.ismethod(fn) and hasattr(fn, "__func__"): - sig = inspect.signature(fn.__func__) - else: - sig = inspect.signature(fn) - if context_kwarg is None: - for param_name, param in sig.parameters.items(): - if is_class_member_of_type(param.annotation, Context): - context_kwarg = param_name - break - type_adapter = get_cached_typeadapter(fn) schema = type_adapter.json_schema() + + if context_kwarg is None: + context_kwarg = find_kwarg_by_type(fn, kwarg_type=Context) if context_kwarg: schema = prune_params(schema, params=[context_kwarg]) diff --git a/src/fastmcp/utilities/types.py b/src/fastmcp/utilities/types.py index eb355a943..0ad371db3 100644 --- a/src/fastmcp/utilities/types.py +++ b/src/fastmcp/utilities/types.py @@ -1,6 +1,8 @@ """Common types used across FastMCP.""" import base64 +import inspect +from collections.abc import Callable from functools import lru_cache from pathlib import Path from types import UnionType @@ -34,7 +36,12 @@ def issubclass_safe(cls: type, base: type) -> bool: def is_class_member_of_type(cls: type, base: type) -> bool: - """Check if cls is a member of base, even if cls is a type variable.""" + """ + Check if cls is a member of base, even if cls is a type variable. + + Base can be a type, a UnionType, or an Annotated type. Generic types are not + considered members (e.g. T is not a member of list[T]). + """ origin = get_origin(cls) # Handle both types of unions: UnionType (from types module, used with | syntax) # and typing.Union (used with Union[] syntax) @@ -50,6 +57,23 @@ def is_class_member_of_type(cls: type, base: type) -> bool: return issubclass_safe(cls, base) +def find_kwarg_by_type(fn: Callable, kwarg_type: type) -> str | None: + """ + Find the name of the kwarg that is of type kwarg_type. + + Includes union types that contain the kwarg_type, as well as Annotated types. + """ + if inspect.ismethod(fn) and hasattr(fn, "__func__"): + sig = inspect.signature(fn.__func__) + else: + sig = inspect.signature(fn) + + for name, param in sig.parameters.items(): + if is_class_member_of_type(param.annotation, kwarg_type): + return name + return None + + def _convert_set_defaults(maybe_set: set[T] | list[T] | None) -> set[T]: """Convert a set or list to a set, defaulting to an empty set if None.""" if maybe_set is None: From d8c881bf210b6d7f01816e9e1efa846fa1ac4f2f Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sun, 4 May 2025 14:31:21 -0400 Subject: [PATCH 2/4] Add tests --- tests/utilities/test_types.py | 104 +++++++++++++++++++++++++++++++++- 1 file changed, 103 insertions(+), 1 deletion(-) diff --git a/tests/utilities/test_types.py b/tests/utilities/test_types.py index e317e50fa..43f0a9725 100644 --- a/tests/utilities/test_types.py +++ b/tests/utilities/test_types.py @@ -2,7 +2,12 @@ from typing import Annotated, Any import pytest -from fastmcp.utilities.types import Image, is_class_member_of_type, issubclass_safe +from fastmcp.utilities.types import ( + Image, + find_kwarg_by_type, + is_class_member_of_type, + issubclass_safe, +) class BaseClass: @@ -143,3 +148,100 @@ class TestImage: ValueError, match="Only one of path or data can be provided" ): Image(path="test.png", data=b"test") + + +class TestFindKwargByType: + def test_exact_type_match(self): + """Test finding parameter with exact type match.""" + + def func(a: int, b: str, c: BaseClass): + pass + + assert find_kwarg_by_type(func, BaseClass) == "c" + + def test_no_matching_parameter(self): + """Test finding parameter when no match exists.""" + + def func(a: int, b: str, c: OtherClass): + pass + + assert find_kwarg_by_type(func, BaseClass) is None + + def test_parameter_with_no_annotation(self): + """Test with a parameter that has no type annotation.""" + + def func(a: int, b, c: BaseClass): + pass + + assert find_kwarg_by_type(func, BaseClass) == "c" + + def test_union_type_match_pipe_syntax(self): + """Test finding parameter with union type using pipe syntax.""" + + def func(a: int, b: str | BaseClass, c: str): + pass + + assert find_kwarg_by_type(func, BaseClass) == "b" + + def test_union_type_match_typing_union(self): + """Test finding parameter with union type using Union.""" + + def func(a: int, b: str | BaseClass, c: str): + pass + + assert find_kwarg_by_type(func, BaseClass) == "b" + + def test_annotated_type_match(self): + """Test finding parameter with Annotated type.""" + + def func(a: int, b: Annotated[BaseClass, "metadata"], c: str): + pass + + assert find_kwarg_by_type(func, BaseClass) == "b" + + def test_method_parameter(self): + """Test finding parameter in a class method.""" + + class TestClass: + def method(self, a: int, b: BaseClass): + pass + + instance = TestClass() + assert find_kwarg_by_type(instance.method, BaseClass) == "b" + + def test_static_method_parameter(self): + """Test finding parameter in a static method.""" + + class TestClass: + @staticmethod + def static_method(a: int, b: BaseClass, c: str): + pass + + assert find_kwarg_by_type(TestClass.static_method, BaseClass) == "b" + + def test_class_method_parameter(self): + """Test finding parameter in a class method.""" + + class TestClass: + @classmethod + def class_method(cls, a: int, b: BaseClass, c: str): + pass + + assert find_kwarg_by_type(TestClass.class_method, BaseClass) == "b" + + def test_multiple_matching_parameters(self): + """Test finding first parameter when multiple matches exist.""" + + def func(a: BaseClass, b: str, c: BaseClass): + pass + + # Should return the first match + assert find_kwarg_by_type(func, BaseClass) == "a" + + def test_subclass_match(self): + """Test finding parameter with a subclass of the target type.""" + + def func(a: int, b: ChildClass, c: str): + pass + + assert find_kwarg_by_type(func, BaseClass) == "b" From 49d7bc92ffac8178f73a7a55984e2481bb24c0bb Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sun, 4 May 2025 14:32:52 -0400 Subject: [PATCH 3/4] Update test_types.py --- tests/utilities/test_types.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/utilities/test_types.py b/tests/utilities/test_types.py index 43f0a9725..ef4c8716d 100644 --- a/tests/utilities/test_types.py +++ b/tests/utilities/test_types.py @@ -245,3 +245,14 @@ class TestFindKwargByType: pass assert find_kwarg_by_type(func, BaseClass) == "b" + + def test_nonstandard_annotation(self): + """Test finding parameter with a nonstandard annotation like an + instance. This is irregular.""" + + SENTINEL = object() + + def func(a: int, b: SENTINEL, c: str): # type: ignore + pass + + assert find_kwarg_by_type(func, SENTINEL) is None # type: ignore From 55e62a0c635fbf982a7bf5592f9dbea416cb917c Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sun, 4 May 2025 14:34:37 -0400 Subject: [PATCH 4/4] Update test_types.py --- tests/utilities/test_types.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/utilities/test_types.py b/tests/utilities/test_types.py index ef4c8716d..ffe71f0c2 100644 --- a/tests/utilities/test_types.py +++ b/tests/utilities/test_types.py @@ -256,3 +256,11 @@ class TestFindKwargByType: pass assert find_kwarg_by_type(func, SENTINEL) is None # type: ignore + + def test_missing_type_annotation(self): + """Test finding parameter with a missing type annotation.""" + + def func(a: int, b, c: str): + pass + + assert find_kwarg_by_type(func, str) == "c"