Python: Preserve MCP array items schema in Pydantic field generation (#2382)

* mcp tool parameter fix

* small fixes
This commit is contained in:
Giles Odigwe
2025-11-24 09:16:30 -08:00
committed by GitHub
Unverified
parent 9f43108ef1
commit bcbf1b33e8
2 changed files with 50 additions and 8 deletions
+20 -8
View File
@@ -278,18 +278,30 @@ def _get_input_model_from_mcp_tool(tool: types.Tool) -> type[BaseModel]:
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)
@@ -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(