diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 586ebd8df1..6e4c71fe49 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -278,18 +278,30 @@ def resolve_type(prop_details: dict[str, Any]) -> type: python_type = resolve_type(prop_details) description = prop_details.get("description", "") + # Build field kwargs (description, array items schema, etc.) + field_kwargs: dict[str, Any] = {} + if description: + field_kwargs["description"] = description + + # Preserve array items schema if present + if prop_details.get("type") == "array" and "items" in prop_details: + items_schema = prop_details["items"] + if items_schema and items_schema != {}: + field_kwargs["json_schema_extra"] = {"items": items_schema} + # Create field definition for create_model if prop_name in required: - field_definitions[prop_name] = ( - (python_type, Field(description=description)) if description else (python_type, ...) - ) + if field_kwargs: + field_definitions[prop_name] = (python_type, Field(**field_kwargs)) + else: + field_definitions[prop_name] = (python_type, ...) else: default_value = prop_details.get("default", None) - field_definitions[prop_name] = ( - (python_type, Field(default=default_value, description=description)) - if description - else (python_type, default_value) - ) + field_kwargs["default"] = default_value + if field_kwargs and any(k != "default" for k in field_kwargs): + field_definitions[prop_name] = (python_type, Field(**field_kwargs)) + else: + field_definitions[prop_name] = (python_type, default_value) return create_model(f"{tool.name}_input", **field_definitions) diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 34f4857dc2..b73136fd6b 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -483,6 +483,36 @@ def test_get_input_model_from_mcp_tool_with_ref_schema(): assert dumped == {"params": {"customer_id": 251}} +def test_get_input_model_from_mcp_tool_with_simple_array(): + """Test array with simple items schema (items schema should be preserved in json_schema_extra).""" + tool = types.Tool( + name="simple_array_tool", + description="Tool with simple array", + inputSchema={ + "type": "object", + "properties": { + "tags": { + "type": "array", + "description": "List of tags", + "items": {"type": "string"}, # Simple string array + } + }, + "required": ["tags"], + }, + ) + model = _get_input_model_from_mcp_tool(tool) + + # Create an instance + instance = model(tags=["tag1", "tag2", "tag3"]) + assert instance.tags == ["tag1", "tag2", "tag3"] + + # Verify JSON schema still preserves items for simple types + json_schema = model.model_json_schema() + tags_property = json_schema["properties"]["tags"] + assert "items" in tags_property + assert tags_property["items"]["type"] == "string" + + def test_get_input_model_from_mcp_prompt(): """Test creation of input model from MCP prompt.""" prompt = types.Prompt(