pydantic_to_mcp_schema.py

def pydantic_to_mcp_schema(model_class):
    """
    Convert a Pydantic model to a MCP Tool.inputSchema format,
    including support for nested models and complex types.

    Args:
        model_class: A Pydantic model class

    Returns:
        dict: A simplified JSON schema compatible with MCP Tool.inputSchema specification
    """
    # Get the full JSON schema from Pydantic
    full_schema = model_class.model_json_schema()

    # Process the schema recursively
    return _process_schema_node(full_schema, full_schema.get("$defs", {}))

def _process_schema_node(node, defs):
    # Handle $ref references
    if "$ref" in node:
        ref_path = node["$ref"]
        if ref_path.startswith("#/$defs/"):
            ref_name = ref_path.split("/")[-1]
            if ref_name in defs:
                # Process the referenced definition
                return _process_schema_node(defs[ref_name], defs)

    # Start with a new schema object
    result = {}

    # Copy the basic properties
    if "type" in node:
        result["type"] = node["type"]

    # Handle anyOf (often used for optional fields with None)
    if "anyOf" in node:
        non_null_types = [t for t in node["anyOf"] if t.get("type") != "null"]
        if non_null_types:
            # Process the first non-null type
            processed = _process_schema_node(non_null_types[0], defs)
            result.update(processed)

    # Handle description
    if "description" in node:
        result["description"] = node["description"]

    # Handle object properties recursively
    if node.get("type") == "object" and "properties" in node:
        result["type"] = "object"
        result["properties"] = {}

        # Process each property
        for prop_name, prop_schema in node["properties"].items():
            result["properties"][prop_name] = _process_schema_node(prop_schema, defs)

        # Add required fields if present
        if "required" in node:
            result["required"] = node["required"]

    # Handle arrays
    if node.get("type") == "array" and "items" in node:
        result["type"] = "array"
        result["items"] = _process_schema_node(node["items"], defs)

    return result