Merge pull request #31 from jurasofish/draft-handling-of-complex-inputs

Handle complex inputs

README.md

@@ -212,6 +212,27 @@ async def fetch_weather(city: str) -> str:
         return response.text
 ```
 
+Complex input handling example:
+```python
+from pydantic import BaseModel, Field
+from typing import Annotated
+
+class ShrimpTank(BaseModel):
+    class Shrimp(BaseModel):
+        name: Annotated[str, Field(max_length=10)]
+
+    shrimp: list[Shrimp]
+
+@mcp.tool()
+def name_shrimp(
+    tank: ShrimpTank,
+    # You can use pydantic Field in function signatures for validation.
+    extra_names: Annotated[list[str], Field(max_length=10)],
+) -> list[str]:
+    """List all shrimp names in the tank"""
+    return [shrimp.name for shrimp in tank.shrimp] + extra_names
+```
+
 ### Prompts
 
 Prompts are reusable templates that help LLMs interact with your server effectively. They're like "best practices" encoded into your server. A prompt can be as simple as a string:

examples/complex_inputs.py

@@ -0,0 +1,28 @@
+"""
+FastMCP Complex inputs Example
+
+Demonstrates validation via pydantic with complex models.
+"""
+
+from pydantic import BaseModel, Field
+from typing import Annotated
+from fastmcp.server import FastMCP
+
+mcp = FastMCP("Shrimp Tank")
+
+
+class ShrimpTank(BaseModel):
+    class Shrimp(BaseModel):
+        name: Annotated[str, Field(max_length=10)]
+
+    shrimp: list[Shrimp]
+
+
+@mcp.tool()
+def name_shrimp(
+    tank: ShrimpTank,
+    # You can use pydantic Field in function signatures for validation.
+    extra_names: Annotated[list[str], Field(max_length=10)],
+) -> list[str]:
+    """List all shrimp names in the tank"""
+    return [shrimp.name for shrimp in tank.shrimp] + extra_names

src/fastmcp/exceptions.py

@@ -15,3 +15,7 @@ class ResourceError(FastMCPError):
 
 class ToolError(FastMCPError):
     """Error in tool operations."""
+
+
+class InvalidSignature(Exception):
+    """Invalid signature for use with FastMCP."""

src/fastmcp/tools/base.py

@@ -1,8 +1,8 @@
 import fastmcp
 from fastmcp.exceptions import ToolError
 
-
-from pydantic import BaseModel, Field, TypeAdapter, validate_call
+from fastmcp.utilities.func_metadata import func_metadata, FuncMetadata
+from pydantic import BaseModel, Field
 
 
 import inspect
@@ -19,6 +19,9 @@ class Tool(BaseModel):
     name: str = Field(description="Name of the tool")
     description: str = Field(description="Description of what the tool does")
     parameters: dict = Field(description="JSON schema for tool parameters")
+    fn_metadata: FuncMetadata = Field(
+        description="Metadata about the function including a pydantic model for tool arguments"
+    )
     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"
@@ -41,9 +44,6 @@ def from_function(
         func_doc = description or fn.__doc__ or ""
         is_async = inspect.iscoroutinefunction(fn)
 
-        # Get schema from TypeAdapter - will fail if function isn't properly typed
-        parameters = TypeAdapter(fn).json_schema()
-
         # Find context parameter if it exists
         if context_kwarg is None:
             sig = inspect.signature(fn)
@@ -52,28 +52,32 @@ def from_function(
                     context_kwarg = param_name
                     break
 
-        # ensure the arguments are properly cast
-        fn = validate_call(fn)
+        func_arg_metadata = func_metadata(
+            fn,
+            skip_names=[context_kwarg] if context_kwarg is not None else [],
+        )
+        parameters = func_arg_metadata.arg_model.model_json_schema()
 
         return cls(
             fn=fn,
             name=func_name,
             description=func_doc,
             parameters=parameters,
+            fn_metadata=func_arg_metadata,
             is_async=is_async,
             context_kwarg=context_kwarg,
         )
 
     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:
-                arguments[self.context_kwarg] = context
-
-            # Call function with proper async handling
-            if self.is_async:
-                return await self.fn(**arguments)
-            return self.fn(**arguments)
+            return await self.fn_metadata.call_fn_with_arg_validation(
+                self.fn,
+                self.is_async,
+                arguments,
+                {self.context_kwarg: context}
+                if self.context_kwarg is not None
+                else None,
+            )
         except Exception as e:
             raise ToolError(f"Error executing tool {self.name}: {e}") from e

src/fastmcp/utilities/func_metadata.py

@@ -0,0 +1,200 @@
+import inspect
+from collections.abc import Callable, Sequence, Awaitable
+from typing import (
+    Annotated,
+    Any,
+    Dict,
+    ForwardRef,
+)
+from pydantic import Field
+from fastmcp.exceptions import InvalidSignature
+from pydantic._internal._typing_extra import try_eval_type
+import json
+from pydantic import BaseModel
+from pydantic.fields import FieldInfo
+from pydantic import ConfigDict, create_model
+from pydantic import WithJsonSchema
+from pydantic_core import PydanticUndefined
+from fastmcp.utilities.logging import get_logger
+
+
+logger = get_logger(__name__)
+
+
+class ArgModelBase(BaseModel):
+    """A model representing the arguments to a function."""
+
+    def model_dump_one_level(self) -> dict[str, Any]:
+        """Return a dict of the model's fields, one level deep.
+
+        That is, sub-models etc are not dumped - they are kept as pydantic models.
+        """
+        kwargs: dict[str, Any] = {}
+        for field_name in self.model_fields.keys():
+            kwargs[field_name] = getattr(self, field_name)
+        return kwargs
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+    )
+
+
+class FuncMetadata(BaseModel):
+    arg_model: Annotated[type[ArgModelBase], WithJsonSchema(None)]
+    # We can add things in the future like
+    #  - Maybe some args are excluded from attempting to parse from JSON
+    #  - Maybe some args are special (like context) for dependency injection
+
+    async def call_fn_with_arg_validation(
+        self,
+        fn: Callable | Awaitable,
+        fn_is_async: bool,
+        arguments_to_validate: dict[str, Any],
+        arguments_to_pass_directly: dict[str, Any] | None,
+    ) -> Any:
+        """Call the given function with arguments validated and injected.
+
+        Arguments are first attempted to be parsed from JSON, then validated against
+        the argument model, before being passed to the function.
+        """
+        arguments_pre_parsed = self.pre_parse_json(arguments_to_validate)
+        arguments_parsed_model = self.arg_model.model_validate(arguments_pre_parsed)
+        arguments_parsed_dict = arguments_parsed_model.model_dump_one_level()
+
+        arguments_parsed_dict |= arguments_to_pass_directly or {}
+
+        if fn_is_async:
+            return await fn(**arguments_parsed_dict)
+        return fn(**arguments_parsed_dict)
+
+    def pre_parse_json(self, data: dict[str, Any]) -> dict[str, Any]:
+        """Pre-parse data from JSON.
+
+        Return a dict with same keys as input but with values parsed from JSON
+        if appropriate.
+
+        This is to handle cases like `["a", "b", "c"]` being passed in as JSON inside
+        a string rather than an actual list. Claude desktop is prone to this - in fact
+        it seems incapable of NOT doing this. For sub-models, it tends to pass
+        dicts (JSON objects) as JSON strings, which can be pre-parsed here.
+        """
+        new_data = data.copy()  # Shallow copy
+        for field_name, field_info in self.arg_model.model_fields.items():
+            if field_name not in data.keys():
+                continue
+            if isinstance(data[field_name], str):
+                try:
+                    pre_parsed = json.loads(data[field_name])
+                except json.JSONDecodeError:
+                    continue  # Not JSON - skip
+                if isinstance(pre_parsed, str):
+                    # This is likely that the raw value is e.g. `"hello"` which we
+                    # Should really be parsed as '"hello"' in Python - but if we parse
+                    # it as JSON it'll turn into just 'hello'. So we skip it.
+                    continue
+                new_data[field_name] = pre_parsed
+        assert new_data.keys() == data.keys()
+        return new_data
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+    )
+
+
+def func_metadata(func: Callable, skip_names: Sequence[str] = ()) -> FuncMetadata:
+    """Given a function, return metadata including a pydantic model representing its signature.
+
+    The use case for this is
+    ```
+    meta = func_to_pyd(func)
+    validated_args = meta.arg_model.model_validate(some_raw_data_dict)
+    return func(**validated_args.model_dump_one_level())
+    ```
+
+    **critically** it also provides pre-parse helper to attempt to parse things from JSON.
+
+    Args:
+        func: The function to convert to a pydantic model
+        skip_names: A list of parameter names to skip. These will not be included in
+            the model.
+    Returns:
+        A pydantic model representing the function's signature.
+    """
+    sig = _get_typed_signature(func)
+    params = sig.parameters
+    dynamic_pydantic_model_params: dict[str, Any] = {}
+    for param in params.values():
+        if param.name.startswith("_"):
+            raise InvalidSignature(
+                f"Parameter {param.name} of {func.__name__} may not start with an underscore"
+            )
+        if param.name in skip_names:
+            continue
+        annotation = param.annotation
+
+        # `x: None` / `x: None = None`
+        if annotation is None:
+            annotation = Annotated[
+                None,
+                Field(
+                    default=param.default
+                    if param.default is not inspect.Parameter.empty
+                    else PydanticUndefined
+                ),
+            ]
+
+        # Untyped field
+        if annotation is inspect.Parameter.empty:
+            annotation = Annotated[
+                Any,
+                Field(),
+                # 🤷
+                WithJsonSchema({"title": param.name, "type": "string"}),
+            ]
+
+        field_info = FieldInfo.from_annotated_attribute(
+            annotation,
+            param.default
+            if param.default is not inspect.Parameter.empty
+            else PydanticUndefined,
+        )
+        dynamic_pydantic_model_params[param.name] = (field_info.annotation, field_info)
+        continue
+
+    arguments_model = create_model(
+        f"{func.__name__}Arguments",
+        **dynamic_pydantic_model_params,
+        __base__=ArgModelBase,
+    )
+    resp = FuncMetadata(arg_model=arguments_model)
+    return resp
+
+
+def _get_typed_annotation(annotation: Any, globalns: Dict[str, Any]) -> Any:
+    if isinstance(annotation, str):
+        annotation = ForwardRef(annotation)
+        annotation, status = try_eval_type(annotation, globalns, globalns)
+
+        # This check and raise could perhaps be skipped, and we (FastMCP) just call
+        # model_rebuild right before using it 🤷
+        if status is False:
+            raise InvalidSignature(f"Unable to evaluate type annotation {annotation}")
+
+    return annotation
+
+
+def _get_typed_signature(call: Callable[..., Any]) -> inspect.Signature:
+    """Get function signature while evaluating forward references"""
+    signature = inspect.signature(call)
+    globalns = getattr(call, "__globals__", {})
+    typed_params = [
+        inspect.Parameter(
+            name=param.name,
+            kind=param.kind,
+            default=param.default,
+            annotation=_get_typed_annotation(param.annotation, globalns),
+        )
+        for param in signature.parameters.values()
+    ]
+    typed_signature = inspect.Signature(typed_params)
+    return typed_signature