Merge pull request #165 from jon-fox/feature/mcp-input-schema-fix-164

Feature/mcp input schema fix 164

Author: KennyVaneetvelde

atomic-agents/atomic_agents/connectors/mcp/schema_transformer.py

@@ -1,7 +1,7 @@
 """Module for transforming JSON schemas to Pydantic models."""
 
 import logging
-from typing import Any, Dict, List, Optional, Type, Tuple, Literal, cast
+from typing import Any, Dict, List, Optional, Type, Tuple, Literal, Union, cast
 
 from pydantic import Field, create_model
 
@@ -24,33 +24,113 @@ class SchemaTransformer:
     """Class for transforming JSON schemas to Pydantic models."""
 
     @staticmethod
-    def json_to_pydantic_field(prop_schema: Dict[str, Any], required: bool) -> Tuple[Type, Field]:
+    def _resolve_ref(ref_path: str, root_schema: Dict[str, Any], model_cache: Dict[str, Type]) -> Type:
+        """Resolve a $ref to a Pydantic model."""
+        # Extract ref name from path like "#/$defs/MyObject" or "#/definitions/ANode"
+        ref_name = ref_path.split("/")[-1]
+
+        if ref_name in model_cache:
+            return model_cache[ref_name]
+
+        # Look for the referenced schema in $defs or definitions
+        defs = root_schema.get("$defs", root_schema.get("definitions", {}))
+        if ref_name in defs:
+            ref_schema = defs[ref_name]
+            # Create model for the referenced schema
+            model_name = ref_schema.get("title", ref_name)
+            # Avoid infinite recursion by adding placeholder first
+            model_cache[ref_name] = Any
+            model = SchemaTransformer._create_nested_model(ref_schema, model_name, root_schema, model_cache)
+            model_cache[ref_name] = model
+            return model
+
+        logger.warning(f"Could not resolve $ref: {ref_path}")
+        return Any
+
+    @staticmethod
+    def _create_nested_model(
+        schema: Dict[str, Any], model_name: str, root_schema: Dict[str, Any], model_cache: Dict[str, Type]
+    ) -> Type:
+        """Create a nested Pydantic model from a schema."""
+        fields = {}
+        required_fields = set(schema.get("required", []))
+        properties = schema.get("properties", {})
+
+        for prop_name, prop_schema in properties.items():
+            is_required = prop_name in required_fields
+            fields[prop_name] = SchemaTransformer.json_to_pydantic_field(prop_schema, is_required, root_schema, model_cache)
+
+        return create_model(model_name, **fields)
+
+    @staticmethod
+    def json_to_pydantic_field(
+        prop_schema: Dict[str, Any],
+        required: bool,
+        root_schema: Optional[Dict[str, Any]] = None,
+        model_cache: Optional[Dict[str, Type]] = None,
+    ) -> Tuple[Type, Field]:
         """
         Convert a JSON schema property to a Pydantic field.
 
         Args:
             prop_schema: JSON schema for the property
             required: Whether the field is required
+            root_schema: Full root schema for resolving $refs
+            model_cache: Cache for resolved models
 
         Returns:
             Tuple of (type, Field)
         """
-        json_type = prop_schema.get("type")
+        if root_schema is None:
+            root_schema = {}
+        if model_cache is None:
+            model_cache = {}
+
         description = prop_schema.get("description")
         default = prop_schema.get("default")
         python_type: Any = Any
 
-        if json_type in JSON_TYPE_MAP:
-            python_type = JSON_TYPE_MAP[json_type]
-            if json_type == "array":
-                items_schema = prop_schema.get("items", {})
-                item_type_str = items_schema.get("type")
-                if item_type_str in JSON_TYPE_MAP:
-                    python_type = List[JSON_TYPE_MAP[item_type_str]]
+        # Handle $ref
+        if "$ref" in prop_schema:
+            python_type = SchemaTransformer._resolve_ref(prop_schema["$ref"], root_schema, model_cache)
+        # Handle oneOf/anyOf (unions)
+        elif "oneOf" in prop_schema or "anyOf" in prop_schema:
+            union_schemas = prop_schema.get("oneOf", prop_schema.get("anyOf", []))
+            if union_schemas:
+                union_types = []
+                for union_schema in union_schemas:
+                    if "$ref" in union_schema:
+                        union_types.append(SchemaTransformer._resolve_ref(union_schema["$ref"], root_schema, model_cache))
+                    else:
+                        # Recursively resolve the union member
+                        member_type, _ = SchemaTransformer.json_to_pydantic_field(union_schema, True, root_schema, model_cache)
+                        union_types.append(member_type)
+
+                if len(union_types) == 1:
+                    python_type = union_types[0]
                 else:
-                    python_type = List[Any]
-            elif json_type == "object":
-                python_type = Dict[str, Any]
+                    python_type = Union[tuple(union_types)]
+        # Handle regular types
+        else:
+            json_type = prop_schema.get("type")
+            if json_type in JSON_TYPE_MAP:
+                python_type = JSON_TYPE_MAP[json_type]
+
+                if json_type == "array":
+                    items_schema = prop_schema.get("items", {})
+                    if "$ref" in items_schema:
+                        item_type = SchemaTransformer._resolve_ref(items_schema["$ref"], root_schema, model_cache)
+                    elif "oneOf" in items_schema or "anyOf" in items_schema:
+                        # Handle arrays of unions
+                        item_type, _ = SchemaTransformer.json_to_pydantic_field(items_schema, True, root_schema, model_cache)
+                    elif items_schema.get("type") in JSON_TYPE_MAP:
+                        item_type = JSON_TYPE_MAP[items_schema["type"]]
+                    else:
+                        item_type = Any
+                    python_type = List[item_type]
+
+                elif json_type == "object":
+                    python_type = Dict[str, Any]
 
         field_kwargs = {"description": description}
         if required:
@@ -85,11 +165,12 @@ def create_model_from_schema(
         fields = {}
         required_fields = set(schema.get("required", []))
         properties = schema.get("properties")
+        model_cache: Dict[str, Type] = {}
 
         if properties:
             for prop_name, prop_schema in properties.items():
                 is_required = prop_name in required_fields
-                fields[prop_name] = SchemaTransformer.json_to_pydantic_field(prop_schema, is_required)
+                fields[prop_name] = SchemaTransformer.json_to_pydantic_field(prop_schema, is_required, schema, model_cache)
         elif schema.get("type") == "object" and not properties:
             pass
         elif schema:

atomic-agents/tests/connectors/mcp/test_schema_transformer.py

@@ -1,5 +1,5 @@
 import pytest
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from atomic_agents import BaseIOSchema
 from atomic_agents.connectors.mcp import SchemaTransformer
@@ -134,3 +134,76 @@ def test_tool_name_field(self):
         # Test that an invalid tool_name raises an error
         with pytest.raises(ValueError):
             model(tool_name="wrong_tool")
+
+    def test_union_type_oneof(self):
+        """Test oneOf creates Union types."""
+        prop_schema = {"oneOf": [{"type": "string"}, {"type": "integer"}], "description": "A union field"}
+        result = SchemaTransformer.json_to_pydantic_field(prop_schema, True)
+        # Should create Union[str, int]
+        assert result[0] == Union[str, int]
+        assert result[1].description == "A union field"
+
+    def test_union_type_anyof(self):
+        """Test anyOf creates Union types."""
+        prop_schema = {"anyOf": [{"type": "boolean"}, {"type": "number"}], "description": "Another union field"}
+        result = SchemaTransformer.json_to_pydantic_field(prop_schema, True)
+        # Should create Union[bool, float]
+        assert result[0] == Union[bool, float]
+
+    def test_array_with_ref_items(self):
+        """Test arrays with $ref items are resolved."""
+        root_schema = {
+            "$defs": {"MyObject": {"type": "object", "properties": {"name": {"type": "string"}}, "title": "MyObject"}}
+        }
+        prop_schema = {"type": "array", "items": {"$ref": "#/$defs/MyObject"}, "description": "Array of MyObject"}
+        result = SchemaTransformer.json_to_pydantic_field(prop_schema, True, root_schema)
+        # Should be List[MyObject] not List[Any]
+        assert hasattr(result[0], "__origin__") and result[0].__origin__ is list
+        # The inner type should be the created model, not Any
+        inner_type = result[0].__args__[0]
+        assert inner_type != Any
+        assert hasattr(inner_type, "model_fields")
+
+    def test_array_with_union_items(self):
+        """Test arrays with oneOf items."""
+        prop_schema = {
+            "type": "array",
+            "items": {"oneOf": [{"type": "string"}, {"type": "integer"}]},
+            "description": "Array of union items",
+        }
+        result = SchemaTransformer.json_to_pydantic_field(prop_schema, True)
+        # Should be List[Union[str, int]]
+        assert hasattr(result[0], "__origin__") and result[0].__origin__ is list
+        inner_type = result[0].__args__[0]
+        assert inner_type == Union[str, int]
+
+    def test_model_with_complex_types(self):
+        """Test create_model_from_schema with complex types."""
+        schema = {
+            "type": "object",
+            "properties": {
+                "expr": {"oneOf": [{"$ref": "#/$defs/ANode"}, {"$ref": "#/$defs/BNode"}], "description": "Expression node"},
+                "objects": {"type": "array", "items": {"$ref": "#/$defs/MyObject"}, "description": "List of objects"},
+            },
+            "required": ["expr", "objects"],
+            "$defs": {
+                "ANode": {"type": "object", "properties": {"a_value": {"type": "string"}}, "title": "ANode"},
+                "BNode": {"type": "object", "properties": {"b_value": {"type": "integer"}}, "title": "BNode"},
+                "MyObject": {"type": "object", "properties": {"name": {"type": "string"}}, "title": "MyObject"},
+            },
+        }
+
+        model = SchemaTransformer.create_model_from_schema(schema, "ComplexModel", "complex_tool")
+
+        # Check that expr is a Union, not Any
+        expr_field = model.model_fields["expr"]
+        assert expr_field.annotation != Any
+        # Should be Union[ANode, BNode]
+        assert hasattr(expr_field.annotation, "__origin__") and expr_field.annotation.__origin__ is Union
+
+        # Check that objects is List[MyObject], not List[Any]
+        objects_field = model.model_fields["objects"]
+        assert objects_field.annotation != List[Any]
+        assert hasattr(objects_field.annotation, "__origin__") and objects_field.annotation.__origin__ is list
+        inner_type = objects_field.annotation.__args__[0]
+        assert inner_type != Any

atomic-examples/mcp-agent/example-client/example_client/main_fastapi.py

@@ -20,7 +20,7 @@ class MCPConfig:
     """Configuration for the MCP Agent system using HTTP Stream transport."""
 
     mcp_server_url: str = "http://localhost:6969"
-    openai_model: str = "gpt-4o"
+    openai_model: str = "gpt-5-mini"
     openai_api_key: str = os.getenv("OPENAI_API_KEY") or ""
 
     def __post_init__(self):

atomic-examples/mcp-agent/example-client/example_client/main_http.py

@@ -23,7 +23,7 @@ class MCPConfig:
     """Configuration for the MCP Agent system using HTTP Stream transport."""
 
     mcp_server_url: str = "http://localhost:6969"
-    openai_model: str = "gpt-4o"
+    openai_model: str = "gpt-5-mini"
     openai_api_key: str = os.getenv("OPENAI_API_KEY")
 
     def __post_init__(self):

atomic-examples/mcp-agent/example-client/example_client/main_sse.py

@@ -21,11 +21,11 @@ class MCPConfig:
 
     mcp_server_url: str = "http://localhost:6969"
 
-    # NOTE: In contrast to other examples, we use gpt-4o and not gpt-4o-mini here.
-    # In my tests, gpt-4o-mini was not smart enough to deal with multiple tools like that
+    # NOTE: In contrast to other examples, we use gpt-5-mini and not gpt-4o-mini here.
+    # In my tests, gpt-5-mini was not smart enough to deal with multiple tools like that
     # and at the moment MCP does not yet allow for adding sufficient metadata to
     # clarify tools even more and introduce more constraints.
-    openai_model: str = "gpt-4o"
+    openai_model: str = "gpt-5-mini"
     openai_api_key: str = os.getenv("OPENAI_API_KEY")
 
     def __post_init__(self):

atomic-examples/mcp-agent/example-client/example_client/main_stdio.py

@@ -22,11 +22,11 @@
 class MCPConfig:
     """Configuration for the MCP Agent system using STDIO transport."""
 
-    # NOTE: In contrast to other examples, we use gpt-4o and not gpt-4o-mini here.
-    # In my tests, gpt-4o-mini was not smart enough to deal with multiple tools like that
+    # NOTE: In contrast to other examples, we use gpt-5-mini and not gpt-5-mini here.
+    # In my tests, gpt-5-mini was not smart enough to deal with multiple tools like that
     # and at the moment MCP does not yet allow for adding sufficient metadata to
     # clarify tools even more and introduce more constraints.
-    openai_model: str = "gpt-4o"
+    openai_model: str = "gpt-5-mini"
     openai_api_key: str = os.getenv("OPENAI_API_KEY")
 
     # Command to run the STDIO server.

atomic-examples/mcp-agent/example-client/example_client/main_stdio_async.py

@@ -22,11 +22,11 @@
 class MCPConfig:
     """Configuration for the MCP Agent system using STDIO transport."""
 
-    # NOTE: In contrast to other examples, we use gpt-4o and not gpt-4o-mini here.
-    # In my tests, gpt-4o-mini was not smart enough to deal with multiple tools like that
+    # NOTE: In contrast to other examples, we use gpt-5-mini and not gpt-4o-mini here.
+    # In my tests, gpt-5-mini was not smart enough to deal with multiple tools like that
     # and at the moment MCP does not yet allow for adding sufficient metadata to
     # clarify tools even more and introduce more constraints.
-    openai_model: str = "gpt-4o"
+    openai_model: str = "gpt-5-mini"
     openai_api_key: str = os.getenv("OPENAI_API_KEY")
 
     # Command to run the STDIO server.

atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/server_http.py

@@ -17,6 +17,7 @@
     SubtractNumbersTool,
     MultiplyNumbersTool,
     DivideNumbersTool,
+    BatchCalculatorTool,
 )
 
 
@@ -27,6 +28,7 @@ def get_available_tools() -> List[Tool]:
         SubtractNumbersTool(),
         MultiplyNumbersTool(),
         DivideNumbersTool(),
+        BatchCalculatorTool(),
     ]
 
 

atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/services/tool_service.py

@@ -3,7 +3,6 @@
 from typing import Dict, List, Any
 from mcp.server.fastmcp import FastMCP
 from example_mcp_server.interfaces.tool import Tool, ToolResponse, ToolContent
-from pydantic import Field
 
 
 class ToolService:
@@ -43,8 +42,8 @@ async def execute_tool(self, tool_name: str, input_data: Dict[str, Any]) -> Tool
         """
         tool = self.get_tool(tool_name)
 
-        # Convert input dictionary to the tool's input model
-        input_model = tool.input_model(**input_data)
+        # Use model_validate to handle complex nested objects properly
+        input_model = tool.input_model.model_validate(input_data)
 
         # Execute the tool with validated input
         return await tool.execute(input_model)
@@ -90,41 +89,19 @@ def _serialize_response(self, response: ToolResponse) -> Any:
     def register_mcp_handlers(self, mcp: FastMCP) -> None:
         """Register all tools as MCP handlers."""
         for tool in self._tools.values():
-            # Get the tool's schema
-            schema = tool.input_model.model_json_schema()
-            properties = schema.get("properties", {})
-
-            # Create a function signature that matches the schema with parameter descriptions
-            params = []
-
-            for name, info in properties.items():
-                type_hint = "str"  # Default to str
-                if info.get("type") == "integer":
-                    type_hint = "int"
-                elif info.get("type") == "number":
-                    type_hint = "float"
-                elif info.get("type") == "boolean":
-                    type_hint = "bool"
-
-                default = info.get("default", "...")
-                description = info.get("description", "")
-
-                # Create parameter string for function definition with Field for descriptions
-                if default == "...":
-                    params.append(f"{name}: {type_hint} = Field(description='{description}')")
-                else:
-                    params.append(f"{name}: {type_hint} = Field(description='{description}', default={repr(default)})")
-
-            # Create the function definition
-            fn_def = f"async def {tool.name}({', '.join(params)}):\n"
-            fn_def += f'    """{tool.description}"""\n'
-            fn_def += "    result = await self.execute_tool(tool.name, locals())\n"
-            fn_def += "    return self._serialize_response(result)"
-
-            # Create the function
-            namespace = {"self": self, "tool": tool, "Field": Field}
-            exec(fn_def, namespace)
-            handler = namespace[tool.name]
-
-            # Register the handler
+            # Create a handler that uses the tool's input model directly for schema generation
+            def create_handler(tool_instance):
+                # Use the actual Pydantic model as the function parameter
+                # This ensures FastMCP gets the complete schema including nested objects
+                async def handler(input_data: tool_instance.input_model):
+                    f'"""{tool_instance.description}"""'
+                    result = await self.execute_tool(tool_instance.name, input_data.model_dump())
+                    return self._serialize_response(result)
+
+                return handler
+
+            # Create the handler
+            handler = create_handler(tool)
+
+            # Register with FastMCP - it should auto-detect the schema from the type annotation
             mcp.tool(name=tool.name, description=tool.description)(handler)