mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: added more complete parsing for mcp tool arguments (#2756)
* added more complete parsing for mcp tool arguments * fixed mypy * added nonlocal model counter, and some fixes * fixes in naming logic * extracted json parsing function, added parametrized test and checked coverage
This commit is contained in:
Eduard van Valkenburg
committed by
GitHub
Unverified
GitHub
parent
8bb9927f3c
commit
c376868ec9
@@ -9,7 +9,7 @@ import pytest
|
||||
from mcp import types
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.shared.exceptions import McpError
|
||||
from pydantic import AnyUrl, ValidationError
|
||||
from pydantic import AnyUrl, BaseModel, ValidationError
|
||||
|
||||
from agent_framework import (
|
||||
ChatMessage,
|
||||
@@ -357,122 +357,360 @@ def test_chat_message_to_mcp_types():
|
||||
assert isinstance(mcp_contents[1], types.ImageContent)
|
||||
|
||||
|
||||
def test_get_input_model_from_mcp_tool():
|
||||
"""Test creation of input model from MCP tool."""
|
||||
tool = types.Tool(
|
||||
name="test_tool",
|
||||
description="A test tool",
|
||||
inputSchema={
|
||||
"type": "object",
|
||||
"properties": {"param1": {"type": "string"}, "param2": {"type": "number"}},
|
||||
"required": ["param1"],
|
||||
},
|
||||
)
|
||||
model = _get_input_model_from_mcp_tool(tool)
|
||||
|
||||
# Create an instance to verify the model works
|
||||
instance = model(param1="test", param2=42)
|
||||
assert instance.param1 == "test"
|
||||
assert instance.param2 == 42
|
||||
|
||||
# Test validation
|
||||
with pytest.raises(ValidationError): # Missing required param1
|
||||
model(param2=42)
|
||||
|
||||
|
||||
def test_get_input_model_from_mcp_tool_with_nested_object():
|
||||
"""Test creation of input model from MCP tool with nested object property."""
|
||||
tool = types.Tool(
|
||||
name="get_customer_detail",
|
||||
description="Get customer details",
|
||||
inputSchema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"params": {
|
||||
"type": "object",
|
||||
"properties": {"customer_id": {"type": "integer"}},
|
||||
"required": ["customer_id"],
|
||||
@pytest.mark.parametrize(
|
||||
"test_id,input_schema,valid_data,expected_values,invalid_data,validation_check",
|
||||
[
|
||||
# Basic types with required/optional fields
|
||||
(
|
||||
"basic_types",
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {"param1": {"type": "string"}, "param2": {"type": "number"}},
|
||||
"required": ["param1"],
|
||||
},
|
||||
{"param1": "test", "param2": 42},
|
||||
{"param1": "test", "param2": 42},
|
||||
{"param2": 42}, # Missing required param1
|
||||
None,
|
||||
),
|
||||
# Nested object
|
||||
(
|
||||
"nested_object",
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"params": {
|
||||
"type": "object",
|
||||
"properties": {"customer_id": {"type": "integer"}},
|
||||
"required": ["customer_id"],
|
||||
}
|
||||
},
|
||||
"required": ["params"],
|
||||
},
|
||||
{"params": {"customer_id": 251}},
|
||||
{"params.customer_id": 251},
|
||||
{"params": {}}, # Missing required customer_id
|
||||
lambda instance: isinstance(instance.params, BaseModel),
|
||||
),
|
||||
# $ref resolution
|
||||
(
|
||||
"ref_schema",
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {"params": {"$ref": "#/$defs/CustomerIdParam"}},
|
||||
"required": ["params"],
|
||||
"$defs": {
|
||||
"CustomerIdParam": {
|
||||
"type": "object",
|
||||
"properties": {"customer_id": {"type": "integer"}},
|
||||
"required": ["customer_id"],
|
||||
}
|
||||
},
|
||||
},
|
||||
{"params": {"customer_id": 251}},
|
||||
{"params.customer_id": 251},
|
||||
{"params": {}}, # Missing required customer_id
|
||||
lambda instance: isinstance(instance.params, BaseModel),
|
||||
),
|
||||
# Array of strings (typed)
|
||||
(
|
||||
"array_of_strings",
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"tags": {
|
||||
"type": "array",
|
||||
"description": "List of tags",
|
||||
"items": {"type": "string"},
|
||||
}
|
||||
},
|
||||
"required": ["tags"],
|
||||
},
|
||||
{"tags": ["tag1", "tag2", "tag3"]},
|
||||
{"tags": ["tag1", "tag2", "tag3"]},
|
||||
None, # No validation error test for this case
|
||||
None,
|
||||
),
|
||||
# Array of integers (typed)
|
||||
(
|
||||
"array_of_integers",
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"numbers": {
|
||||
"type": "array",
|
||||
"description": "List of integers",
|
||||
"items": {"type": "integer"},
|
||||
}
|
||||
},
|
||||
"required": ["numbers"],
|
||||
},
|
||||
{"numbers": [1, 2, 3]},
|
||||
{"numbers": [1, 2, 3]},
|
||||
None,
|
||||
None,
|
||||
),
|
||||
# Array of objects (complex nested)
|
||||
(
|
||||
"array_of_objects",
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"users": {
|
||||
"type": "array",
|
||||
"description": "List of users",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"id": {"type": "integer", "description": "User ID"},
|
||||
"name": {"type": "string", "description": "User name"},
|
||||
},
|
||||
"required": ["id", "name"],
|
||||
},
|
||||
}
|
||||
},
|
||||
"required": ["users"],
|
||||
},
|
||||
{"users": [{"id": 1, "name": "Alice"}, {"id": 2, "name": "Bob"}]},
|
||||
{"users[0].id": 1, "users[0].name": "Alice", "users[1].id": 2, "users[1].name": "Bob"},
|
||||
{"users": [{"id": 1}]}, # Missing required 'name'
|
||||
lambda instance: all(isinstance(user, BaseModel) for user in instance.users),
|
||||
),
|
||||
# Deeply nested objects (3+ levels)
|
||||
(
|
||||
"deeply_nested",
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"filters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"date_range": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"start": {"type": "string"},
|
||||
"end": {"type": "string"},
|
||||
},
|
||||
"required": ["start", "end"],
|
||||
},
|
||||
"categories": {"type": "array", "items": {"type": "string"}},
|
||||
},
|
||||
"required": ["date_range"],
|
||||
}
|
||||
},
|
||||
"required": ["filters"],
|
||||
}
|
||||
},
|
||||
"required": ["query"],
|
||||
},
|
||||
{
|
||||
"query": {
|
||||
"filters": {
|
||||
"date_range": {"start": "2024-01-01", "end": "2024-12-31"},
|
||||
"categories": ["tech", "science"],
|
||||
}
|
||||
}
|
||||
},
|
||||
"required": ["params"],
|
||||
},
|
||||
)
|
||||
model = _get_input_model_from_mcp_tool(tool)
|
||||
{
|
||||
"query.filters.date_range.start": "2024-01-01",
|
||||
"query.filters.date_range.end": "2024-12-31",
|
||||
"query.filters.categories": ["tech", "science"],
|
||||
},
|
||||
{"query": {"filters": {"date_range": {}}}}, # Missing required start and end
|
||||
None,
|
||||
),
|
||||
# Complex $ref with nested structure
|
||||
(
|
||||
"ref_nested_structure",
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {"order": {"$ref": "#/$defs/OrderParams"}},
|
||||
"required": ["order"],
|
||||
"$defs": {
|
||||
"OrderParams": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"customer": {"$ref": "#/$defs/Customer"},
|
||||
"items": {"type": "array", "items": {"$ref": "#/$defs/OrderItem"}},
|
||||
},
|
||||
"required": ["customer", "items"],
|
||||
},
|
||||
"Customer": {
|
||||
"type": "object",
|
||||
"properties": {"id": {"type": "integer"}, "email": {"type": "string"}},
|
||||
"required": ["id", "email"],
|
||||
},
|
||||
"OrderItem": {
|
||||
"type": "object",
|
||||
"properties": {"product_id": {"type": "string"}, "quantity": {"type": "integer"}},
|
||||
"required": ["product_id", "quantity"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"order": {
|
||||
"customer": {"id": 123, "email": "test@example.com"},
|
||||
"items": [{"product_id": "prod1", "quantity": 2}],
|
||||
}
|
||||
},
|
||||
{
|
||||
"order.customer.id": 123,
|
||||
"order.customer.email": "test@example.com",
|
||||
"order.items[0].product_id": "prod1",
|
||||
"order.items[0].quantity": 2,
|
||||
},
|
||||
{"order": {"customer": {"id": 123}, "items": []}}, # Missing email
|
||||
lambda instance: isinstance(instance.order.customer, BaseModel),
|
||||
),
|
||||
# Mixed types (primitives, arrays, nested objects)
|
||||
(
|
||||
"mixed_types",
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"simple_string": {"type": "string"},
|
||||
"simple_number": {"type": "integer"},
|
||||
"string_array": {"type": "array", "items": {"type": "string"}},
|
||||
"nested_config": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"enabled": {"type": "boolean"},
|
||||
"options": {"type": "array", "items": {"type": "string"}},
|
||||
},
|
||||
"required": ["enabled"],
|
||||
},
|
||||
},
|
||||
"required": ["simple_string", "nested_config"],
|
||||
},
|
||||
{
|
||||
"simple_string": "test",
|
||||
"simple_number": 42,
|
||||
"string_array": ["a", "b"],
|
||||
"nested_config": {"enabled": True, "options": ["opt1", "opt2"]},
|
||||
},
|
||||
{
|
||||
"simple_string": "test",
|
||||
"simple_number": 42,
|
||||
"string_array": ["a", "b"],
|
||||
"nested_config.enabled": True,
|
||||
"nested_config.options": ["opt1", "opt2"],
|
||||
},
|
||||
None,
|
||||
None,
|
||||
),
|
||||
# Empty schema (no properties)
|
||||
(
|
||||
"empty_schema",
|
||||
{"type": "object", "properties": {}},
|
||||
{},
|
||||
{},
|
||||
None,
|
||||
None,
|
||||
),
|
||||
# All primitive types
|
||||
(
|
||||
"all_primitives",
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"string_field": {"type": "string"},
|
||||
"integer_field": {"type": "integer"},
|
||||
"number_field": {"type": "number"},
|
||||
"boolean_field": {"type": "boolean"},
|
||||
},
|
||||
},
|
||||
{"string_field": "test", "integer_field": 42, "number_field": 3.14, "boolean_field": True},
|
||||
{"string_field": "test", "integer_field": 42, "number_field": 3.14, "boolean_field": True},
|
||||
None,
|
||||
None,
|
||||
),
|
||||
# Edge case: unresolvable $ref (fallback to dict)
|
||||
(
|
||||
"unresolvable_ref",
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {"data": {"$ref": "#/$defs/NonExistent"}},
|
||||
"$defs": {},
|
||||
},
|
||||
{"data": {"key": "value"}},
|
||||
{"data": {"key": "value"}},
|
||||
None,
|
||||
None,
|
||||
),
|
||||
# Edge case: array without items schema (fallback to bare list)
|
||||
(
|
||||
"array_no_items",
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {"items": {"type": "array"}},
|
||||
},
|
||||
{"items": [1, "two", 3.0]},
|
||||
{"items": [1, "two", 3.0]},
|
||||
None,
|
||||
None,
|
||||
),
|
||||
# Edge case: object without properties (fallback to dict)
|
||||
(
|
||||
"object_no_properties",
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {"config": {"type": "object"}},
|
||||
},
|
||||
{"config": {"arbitrary": "data", "nested": {"key": "value"}}},
|
||||
{"config": {"arbitrary": "data", "nested": {"key": "value"}}},
|
||||
None,
|
||||
None,
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_get_input_model_from_mcp_tool_parametrized(
|
||||
test_id, input_schema, valid_data, expected_values, invalid_data, validation_check
|
||||
):
|
||||
"""Parametrized test for JSON schema to Pydantic model conversion.
|
||||
|
||||
# Create an instance to verify the model works with nested objects
|
||||
instance = model(params={"customer_id": 251})
|
||||
assert instance.params == {"customer_id": 251}
|
||||
assert isinstance(instance.params, dict)
|
||||
This test covers various edge cases including:
|
||||
- Basic types with required/optional fields
|
||||
- Nested objects
|
||||
- $ref resolution
|
||||
- Typed arrays (strings, integers, objects)
|
||||
- Deeply nested structures
|
||||
- Complex $ref with nested structures
|
||||
- Mixed types
|
||||
|
||||
# Verify model_dump produces the correct nested structure
|
||||
dumped = instance.model_dump()
|
||||
assert dumped == {"params": {"customer_id": 251}}
|
||||
|
||||
|
||||
def test_get_input_model_from_mcp_tool_with_ref_schema():
|
||||
"""Test creation of input model from MCP tool with $ref schema.
|
||||
|
||||
This simulates a FastMCP tool that uses Pydantic models with $ref in the schema.
|
||||
The schema should be resolved and nested objects should be preserved.
|
||||
To add a new test case, add a tuple to the parametrize decorator with:
|
||||
- test_id: A descriptive name for the test case
|
||||
- input_schema: The JSON schema (inputSchema dict)
|
||||
- valid_data: Valid data to instantiate the model
|
||||
- expected_values: Dict of expected values (supports dot notation for nested access)
|
||||
- invalid_data: Invalid data to test validation errors (None to skip)
|
||||
- validation_check: Optional callable to perform additional validation checks
|
||||
"""
|
||||
# This is similar to what FastMCP generates when you have:
|
||||
# async def get_customer_detail(params: CustomerIdParam) -> CustomerDetail
|
||||
tool = types.Tool(
|
||||
name="get_customer_detail",
|
||||
description="Get customer details",
|
||||
inputSchema={
|
||||
"type": "object",
|
||||
"properties": {"params": {"$ref": "#/$defs/CustomerIdParam"}},
|
||||
"required": ["params"],
|
||||
"$defs": {
|
||||
"CustomerIdParam": {
|
||||
"type": "object",
|
||||
"properties": {"customer_id": {"type": "integer"}},
|
||||
"required": ["customer_id"],
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
tool = types.Tool(name="test_tool", description="A test tool", inputSchema=input_schema)
|
||||
model = _get_input_model_from_mcp_tool(tool)
|
||||
|
||||
# Create an instance to verify the model works with $ref schemas
|
||||
instance = model(params={"customer_id": 251})
|
||||
assert instance.params == {"customer_id": 251}
|
||||
assert isinstance(instance.params, dict)
|
||||
# Test valid data
|
||||
instance = model(**valid_data)
|
||||
|
||||
# Verify model_dump produces the correct nested structure
|
||||
dumped = instance.model_dump()
|
||||
assert dumped == {"params": {"customer_id": 251}}
|
||||
# Check expected values
|
||||
for field_path, expected_value in expected_values.items():
|
||||
# Support dot notation and array indexing for nested access
|
||||
current = instance
|
||||
parts = field_path.replace("]", "").replace("[", ".").split(".")
|
||||
for part in parts:
|
||||
current = current[int(part)] if part.isdigit() else getattr(current, part)
|
||||
assert current == expected_value, f"Field {field_path} = {current}, expected {expected_value}"
|
||||
|
||||
# Run additional validation checks if provided
|
||||
if validation_check:
|
||||
assert validation_check(instance), f"Validation check failed for {test_id}"
|
||||
|
||||
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"
|
||||
# Test invalid data if provided
|
||||
if invalid_data is not None:
|
||||
with pytest.raises(ValidationError):
|
||||
model(**invalid_data)
|
||||
|
||||
|
||||
def test_get_input_model_from_mcp_prompt():
|
||||
|
||||
Reference in New Issue
Block a user