Skip to content

Commit 50d631b

Browse files
authored
fix mcp client schema creation in flat lists (NVIDIA#346)
For MCP tools that have an input that is a flat list, eg ``` { 'items': ['item', 'item'] } ``` The original code led to an invalid pydantic schema like this: ``` { 'items': {'items': {}, 'type': 'array'}, 'type': 'array' } ``` Now the schema is correctly created like this: ``` { 'items': {'type': 'string'}, 'type': 'array' } ``` Confirmed the fix using a JIRA MCP tool. Closes NVIDIA#342 ## By Submitting this PR I confirm: - I am familiar with the [Contributing Guidelines](https://github.com/NVIDIA/AIQToolkit/blob/develop/docs/source/resources/contributing.md). - We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license. - Any contribution which contains commits that are not Signed-Off will not be accepted. - When the PR is ready for review, new or existing tests cover these changes. - When the PR is ready for review, the documentation is up to date with these changes. Authors: - Sean Lopp (https://github.com/slopp) - Anuradha Karuppiah (https://github.com/AnuradhaKaruppiah) Approvers: - Anuradha Karuppiah (https://github.com/AnuradhaKaruppiah) URL: NVIDIA#346
1 parent febdb56 commit 50d631b

File tree

2 files changed

+30
-2
lines changed

2 files changed

+30
-2
lines changed

src/aiq/tool/mcp/mcp_client.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,9 +63,9 @@ def _generate_field(field_name: str, field_properties: dict[str, Any]) -> tuple:
6363
elif json_type == "array" and "items" in field_properties:
6464
item_properties = field_properties.get("items", {})
6565
if item_properties.get("type") == "object":
66-
item_type = model_from_mcp_schema(name=field_name, mcp_input_schema=field_properties)
66+
item_type = model_from_mcp_schema(name=field_name, mcp_input_schema=item_properties)
6767
else:
68-
item_type = _type_map.get(json_type, Any)
68+
item_type = _type_map.get(item_properties.get("type", "string"), Any)
6969
field_type = list[item_type]
7070
else:
7171
field_type = _type_map.get(json_type, Any)

tests/aiq/tools/test_mcp.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,30 @@ def _get_sample_schema():
6363
},
6464
'optional_bool_field': {
6565
'default': False, 'description': 'Optional Boolean Field.', 'title': 'Raw', 'type': 'boolean'
66+
},
67+
'optional_array_field': {
68+
'default': ['item'],
69+
'description': 'Optional Array Field.',
70+
'title': 'Array',
71+
'type': 'array',
72+
'items': {
73+
'type': 'string'
74+
}
75+
},
76+
'optional_array_object_field': {
77+
'default': [{
78+
'key': 'value'
79+
}],
80+
'description': 'Optional Array Field.',
81+
'title': 'Array',
82+
'type': 'array',
83+
'items': {
84+
'type': 'object', 'properties': {
85+
'key': {
86+
'type': 'string'
87+
}
88+
}
89+
}
6690
}
6791
},
6892
'required': [
@@ -100,6 +124,10 @@ def test_schema_generation(sample_schema):
100124
"required_string_field": "This is a string",
101125
"required_int_field": 4,
102126
"required_float_field": 5.5,
127+
"optional_array_field": ["item1"],
128+
"optional_array_object_field": [{
129+
'key': 'value1'
130+
}],
103131
}
104132

105133
m = _model.model_validate(test_input)

0 commit comments

Comments
 (0)