mirror of
https://github.com/apache/superset.git
synced 2026-04-19 08:04:53 +00:00
feat(mcp): Add flexible input parsing to handle double-serialized requests (#36249)
This commit is contained in:
Amin Ghadersohi
444
superset/mcp_service/utils/schema_utils.py
Normal file
444
superset/mcp_service/utils/schema_utils.py
Normal file
@@ -0,0 +1,444 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
"""
|
||||
Generic utilities for flexible schema input handling in MCP tools.
|
||||
|
||||
This module provides utilities to accept both JSON string and object formats
|
||||
for input parameters, making MCP tools more flexible for different clients.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from functools import wraps
|
||||
from typing import Any, Callable, List, Type, TypeVar
|
||||
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class JSONParseError(ValueError):
|
||||
"""Raised when JSON parsing fails with helpful context."""
|
||||
|
||||
def __init__(self, value: Any, error: Exception, param_name: str = "parameter"):
|
||||
self.value = value
|
||||
self.original_error = error
|
||||
self.param_name = param_name
|
||||
super().__init__(
|
||||
f"Failed to parse {param_name} from JSON string: {error}. "
|
||||
f"Received value: {value!r}"
|
||||
)
|
||||
|
||||
|
||||
def parse_json_or_passthrough(
|
||||
value: Any, param_name: str = "parameter", strict: bool = False
|
||||
) -> Any:
|
||||
"""
|
||||
Parse a value that can be either a JSON string or a native Python object.
|
||||
|
||||
This function handles the common pattern where API parameters can be provided
|
||||
as either:
|
||||
- A JSON string (e.g., from CLI tools or tests): '{"key": "value"}'
|
||||
- A native Python object (e.g., from LLM clients): {"key": "value"}
|
||||
|
||||
Args:
|
||||
value: The input value to parse. Can be a string, list, dict, or any JSON-
|
||||
serializable type.
|
||||
param_name: Name of the parameter for error messages (default: "parameter")
|
||||
strict: If True, raises JSONParseError on parse failures. If False, logs
|
||||
warning and returns original value (default: False)
|
||||
|
||||
Returns:
|
||||
Parsed Python object if value was a JSON string, otherwise returns value
|
||||
unchanged.
|
||||
|
||||
Raises:
|
||||
JSONParseError: If strict=True and JSON parsing fails.
|
||||
|
||||
Examples:
|
||||
>>> parse_json_or_passthrough('[1, 2, 3]', 'numbers')
|
||||
[1, 2, 3]
|
||||
|
||||
>>> parse_json_or_passthrough([1, 2, 3], 'numbers')
|
||||
[1, 2, 3]
|
||||
|
||||
>>> parse_json_or_passthrough('{"key": "value"}', 'config')
|
||||
{'key': 'value'}
|
||||
|
||||
>>> parse_json_or_passthrough({'key': 'value'}, 'config')
|
||||
{'key': 'value'}
|
||||
"""
|
||||
# If not a string, return as-is (already in object form)
|
||||
if not isinstance(value, str):
|
||||
return value
|
||||
|
||||
# Try to parse as JSON
|
||||
try:
|
||||
from superset.utils import json
|
||||
|
||||
parsed = json.loads(value)
|
||||
logger.debug("Successfully parsed %s from JSON string", param_name)
|
||||
return parsed
|
||||
except (ValueError, TypeError) as e:
|
||||
error_msg = (
|
||||
f"Failed to parse {param_name} from JSON string: {e}. Received: {value!r}"
|
||||
)
|
||||
|
||||
if strict:
|
||||
raise JSONParseError(value, e, param_name) from None
|
||||
|
||||
logger.warning("%s. Returning original value.", error_msg)
|
||||
return value
|
||||
|
||||
|
||||
def parse_json_or_list(
|
||||
value: Any, param_name: str = "parameter", item_separator: str = ","
|
||||
) -> List[Any]:
|
||||
"""
|
||||
Parse a value into a list, accepting JSON string, list, or comma-separated string.
|
||||
|
||||
This function provides maximum flexibility for list parameters by accepting:
|
||||
- JSON array string: '["item1", "item2"]'
|
||||
- Python list: ["item1", "item2"]
|
||||
- Comma-separated string: "item1, item2, item3"
|
||||
- Empty/None: returns empty list
|
||||
|
||||
Args:
|
||||
value: Input value to parse into a list
|
||||
param_name: Name of the parameter for error messages
|
||||
item_separator: Separator for comma-separated strings (default: ",")
|
||||
|
||||
Returns:
|
||||
List of items. Returns empty list if value is None or empty.
|
||||
|
||||
Examples:
|
||||
>>> parse_json_or_list('["a", "b"]', 'items')
|
||||
['a', 'b']
|
||||
|
||||
>>> parse_json_or_list(['a', 'b'], 'items')
|
||||
['a', 'b']
|
||||
|
||||
>>> parse_json_or_list('a, b, c', 'items')
|
||||
['a', 'b', 'c']
|
||||
|
||||
>>> parse_json_or_list(None, 'items')
|
||||
[]
|
||||
"""
|
||||
# Handle None and empty values
|
||||
if value is None or value == "":
|
||||
return []
|
||||
|
||||
# Already a list, return as-is
|
||||
if isinstance(value, list):
|
||||
return value
|
||||
|
||||
# Try to parse as JSON
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
from superset.utils import json
|
||||
|
||||
parsed = json.loads(value)
|
||||
# If successfully parsed and it's a list, return it
|
||||
if isinstance(parsed, list):
|
||||
logger.debug("Successfully parsed %s from JSON string", param_name)
|
||||
return parsed
|
||||
|
||||
# If parsed to non-list (e.g., single value), wrap in list
|
||||
logger.debug(
|
||||
"Parsed %s from JSON to non-list, wrapping in list", param_name
|
||||
)
|
||||
return [parsed]
|
||||
except (ValueError, TypeError):
|
||||
# Not valid JSON, try comma-separated parsing
|
||||
logger.debug(
|
||||
"Could not parse %s as JSON, trying comma-separated", param_name
|
||||
)
|
||||
items = [
|
||||
item.strip() for item in value.split(item_separator) if item.strip()
|
||||
]
|
||||
return items
|
||||
|
||||
# For any other type, wrap in a list
|
||||
logger.debug("Wrapping %s value in list", param_name)
|
||||
return [value]
|
||||
|
||||
|
||||
def parse_json_or_model(
|
||||
value: Any, model_class: Type[BaseModel], param_name: str = "parameter"
|
||||
) -> BaseModel:
|
||||
"""
|
||||
Parse a value into a Pydantic model, accepting JSON string or dict.
|
||||
|
||||
Args:
|
||||
value: Input value to parse (JSON string, dict, or model instance)
|
||||
model_class: Pydantic model class to validate against
|
||||
param_name: Name of the parameter for error messages
|
||||
|
||||
Returns:
|
||||
Validated Pydantic model instance
|
||||
|
||||
Raises:
|
||||
ValidationError: If the value cannot be parsed or validated
|
||||
|
||||
Examples:
|
||||
>>> class MyModel(BaseModel):
|
||||
... name: str
|
||||
... value: int
|
||||
|
||||
>>> parse_json_or_model('{"name": "test", "value": 42}', MyModel)
|
||||
MyModel(name='test', value=42)
|
||||
|
||||
>>> parse_json_or_model({"name": "test", "value": 42}, MyModel)
|
||||
MyModel(name='test', value=42)
|
||||
"""
|
||||
# If already an instance of the model, return as-is
|
||||
if isinstance(value, model_class):
|
||||
return value
|
||||
|
||||
# Parse JSON string if needed
|
||||
parsed_value = parse_json_or_passthrough(value, param_name, strict=True)
|
||||
|
||||
# Validate and construct the model
|
||||
try:
|
||||
return model_class.model_validate(parsed_value)
|
||||
except ValidationError:
|
||||
logger.error(
|
||||
"Failed to validate %s against %s", param_name, model_class.__name__
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
def parse_json_or_model_list(
|
||||
value: Any,
|
||||
model_class: Type[BaseModel],
|
||||
param_name: str = "parameter",
|
||||
) -> List[BaseModel]:
|
||||
"""
|
||||
Parse a value into a list of Pydantic models, accepting JSON string or list.
|
||||
|
||||
Args:
|
||||
value: Input value to parse (JSON string, list of dicts, or list of models)
|
||||
model_class: Pydantic model class for list items
|
||||
param_name: Name of the parameter for error messages
|
||||
|
||||
Returns:
|
||||
List of validated Pydantic model instances
|
||||
|
||||
Raises:
|
||||
ValidationError: If any item cannot be parsed or validated
|
||||
|
||||
Examples:
|
||||
>>> class Item(BaseModel):
|
||||
... name: str
|
||||
|
||||
>>> parse_json_or_model_list('[{"name": "a"}, {"name": "b"}]', Item)
|
||||
[Item(name='a'), Item(name='b')]
|
||||
|
||||
>>> parse_json_or_model_list([{"name": "a"}, {"name": "b"}], Item)
|
||||
[Item(name='a'), Item(name='b')]
|
||||
"""
|
||||
# Handle None and empty
|
||||
if value is None or value == "":
|
||||
return []
|
||||
|
||||
# Parse to list first
|
||||
items = parse_json_or_list(value, param_name)
|
||||
|
||||
# Validate each item against the model
|
||||
validated_items = []
|
||||
for i, item in enumerate(items):
|
||||
try:
|
||||
if isinstance(item, model_class):
|
||||
validated_items.append(item)
|
||||
else:
|
||||
validated_items.append(model_class.model_validate(item))
|
||||
except ValidationError:
|
||||
logger.error(
|
||||
"Failed to validate %s[%s] against %s",
|
||||
param_name,
|
||||
i,
|
||||
model_class.__name__,
|
||||
)
|
||||
# Re-raise original validation error
|
||||
raise
|
||||
|
||||
return validated_items
|
||||
|
||||
|
||||
# Pydantic validator decorators for common use cases
|
||||
def json_or_passthrough_validator(
|
||||
param_name: str | None = None, strict: bool = False
|
||||
) -> Callable[[Type[BaseModel], Any, Any], Any]:
|
||||
"""
|
||||
Decorator factory for Pydantic field validators that accept JSON or objects.
|
||||
|
||||
This creates a validator that can be used with Pydantic's @field_validator
|
||||
decorator to automatically parse JSON strings.
|
||||
|
||||
Args:
|
||||
param_name: Parameter name for error messages (uses field name if None)
|
||||
strict: Whether to raise errors on parse failures
|
||||
|
||||
Returns:
|
||||
Validator function compatible with @field_validator
|
||||
|
||||
Example:
|
||||
>>> class MySchema(BaseModel):
|
||||
... config: dict
|
||||
...
|
||||
... @field_validator('config', mode='before')
|
||||
... @classmethod
|
||||
... def parse_config(cls, v):
|
||||
... return parse_json_or_passthrough(v, 'config')
|
||||
"""
|
||||
|
||||
def validator(cls: Type[BaseModel], v: Any, info: Any = None) -> Any:
|
||||
# Use field name from validation info if param_name not provided
|
||||
field_name = param_name or (info.field_name if info else "field")
|
||||
return parse_json_or_passthrough(v, field_name, strict)
|
||||
|
||||
return validator
|
||||
|
||||
|
||||
def json_or_list_validator(
|
||||
param_name: str | None = None, item_separator: str = ","
|
||||
) -> Callable[[Type[BaseModel], Any, Any], List[Any]]:
|
||||
"""
|
||||
Decorator factory for Pydantic validators that parse values into lists.
|
||||
|
||||
Args:
|
||||
param_name: Parameter name for error messages
|
||||
item_separator: Separator for comma-separated strings
|
||||
|
||||
Returns:
|
||||
Validator function compatible with @field_validator
|
||||
|
||||
Example:
|
||||
>>> class MySchema(BaseModel):
|
||||
... items: List[str]
|
||||
...
|
||||
... @field_validator('items', mode='before')
|
||||
... @classmethod
|
||||
... def parse_items(cls, v):
|
||||
... return parse_json_or_list(v, 'items')
|
||||
"""
|
||||
|
||||
def validator(cls: Type[BaseModel], v: Any, info: Any = None) -> List[Any]:
|
||||
field_name = param_name or (info.field_name if info else "field")
|
||||
return parse_json_or_list(v, field_name, item_separator)
|
||||
|
||||
return validator
|
||||
|
||||
|
||||
def json_or_model_list_validator(
|
||||
model_class: Type[BaseModel], param_name: str | None = None
|
||||
) -> Callable[[Type[BaseModel], Any, Any], List[BaseModel]]:
|
||||
"""
|
||||
Decorator factory for Pydantic validators that parse lists of models.
|
||||
|
||||
Args:
|
||||
model_class: Pydantic model class for list items
|
||||
param_name: Parameter name for error messages
|
||||
|
||||
Returns:
|
||||
Validator function compatible with @field_validator
|
||||
|
||||
Example:
|
||||
>>> class FilterModel(BaseModel):
|
||||
... col: str
|
||||
... value: str
|
||||
...
|
||||
>>> class MySchema(BaseModel):
|
||||
... filters: List[FilterModel]
|
||||
...
|
||||
... @field_validator('filters', mode='before')
|
||||
... @classmethod
|
||||
... def parse_filters(cls, v):
|
||||
... return parse_json_or_model_list(v, FilterModel, 'filters')
|
||||
"""
|
||||
|
||||
def validator(cls: Type[BaseModel], v: Any, info: Any = None) -> List[BaseModel]:
|
||||
field_name = param_name or (info.field_name if info else "field")
|
||||
return parse_json_or_model_list(v, model_class, field_name)
|
||||
|
||||
return validator
|
||||
|
||||
|
||||
def parse_request(
|
||||
request_class: Type[BaseModel],
|
||||
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
|
||||
"""
|
||||
Decorator to handle Claude Code bug where requests are double-serialized as strings.
|
||||
|
||||
Automatically parses string requests to Pydantic models before calling
|
||||
the tool function.
|
||||
This eliminates the need for manual parsing code in every tool function.
|
||||
|
||||
See: https://github.com/anthropics/claude-code/issues/5504
|
||||
|
||||
Args:
|
||||
request_class: The Pydantic model class for the request
|
||||
|
||||
Returns:
|
||||
Decorator function that wraps the tool function
|
||||
|
||||
Usage:
|
||||
@mcp.tool
|
||||
@mcp_auth_hook
|
||||
@parse_request(ListChartsRequest)
|
||||
async def list_charts(
|
||||
request: ListChartsRequest, ctx: Context
|
||||
) -> ChartList:
|
||||
# Decorator handles string conversion automatically
|
||||
await ctx.info(f"Listing charts: page={request.page}")
|
||||
...
|
||||
|
||||
Note:
|
||||
- Works with both async and sync functions
|
||||
- Request must be the first positional argument
|
||||
- If request is already a model instance, it passes through unchanged
|
||||
- Handles JSON string parsing with helpful error messages
|
||||
"""
|
||||
|
||||
def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
|
||||
if asyncio.iscoroutinefunction(func):
|
||||
|
||||
@wraps(func)
|
||||
async def async_wrapper(request: Any, *args: Any, **kwargs: Any) -> Any:
|
||||
# Parse if string, otherwise pass through
|
||||
# (parse_json_or_model handles both)
|
||||
parsed_request = parse_json_or_model(request, request_class, "request")
|
||||
return await func(parsed_request, *args, **kwargs)
|
||||
|
||||
return async_wrapper
|
||||
else:
|
||||
|
||||
@wraps(func)
|
||||
def sync_wrapper(request: Any, *args: Any, **kwargs: Any) -> Any:
|
||||
# Parse if string, otherwise pass through
|
||||
# (parse_json_or_model handles both)
|
||||
parsed_request = parse_json_or_model(request, request_class, "request")
|
||||
return func(parsed_request, *args, **kwargs)
|
||||
|
||||
return sync_wrapper
|
||||
|
||||
return decorator
|
||||
Reference in New Issue
Block a user