Skip to content

Commit 0775f23

Browse files
committed
feat(core): Add support for map parameter type
This PR introduces support for the `map` parameter type in the Python SDKs, aligning it with the recent MCP Toolbox server enhancements ([#928](googleapis/genai-toolbox#928)). This allows tools to define parameters that are key-value pairs, either as a generic map with mixed-type values or as a strictly typed map. ## Changes * The `ParameterSchema` in `protocol.py` has been updated to include an optional `valueType` field. * The internal `__get_type` method now correctly resolves the `map` type to the appropriate Python `dict` type hint: * `dict[str, Any]` for generic maps (when `valueType` is omitted). * `dict[str, <type>]` for typed maps (e.g., `dict[str, str]`, `dict[str, int]`). ## Testing * Added extensive unit tests to `test_protocol.py` to ensure the `map` type is handled correctly. * Test coverage includes: * Generic (untyped) maps. * Typed maps for `string`, `integer`, `float`, and `boolean`. * Optional maps. * Error handling for unsupported `valueType`. > [!NOTE] > No modifications were required in `toolbox-langchain` or `toolbox-llamaindex`. They inherit this new functionality directly from `toolbox-core`.
1 parent 0c0d94e commit 0775f23

File tree

2 files changed

+115
-2
lines changed

2 files changed

+115
-2
lines changed

packages/toolbox-core/src/toolbox_core/protocol.py

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
# limitations under the License.
1414

1515
from inspect import Parameter
16-
from typing import Optional, Type
16+
from typing import Any, Optional, Type
1717

1818
from pydantic import BaseModel
1919

@@ -29,6 +29,7 @@ class ParameterSchema(BaseModel):
2929
description: str
3030
authSources: Optional[list[str]] = None
3131
items: Optional["ParameterSchema"] = None
32+
valueType: Optional[str] = None
3233

3334
def __get_type(self) -> Type:
3435
base_type: Type
@@ -44,6 +45,20 @@ def __get_type(self) -> Type:
4445
if self.items is None:
4546
raise Exception("Unexpected value: type is 'list' but items is None")
4647
base_type = list[self.items.__get_type()] # type: ignore
48+
elif self.type == "map":
49+
if self.valueType:
50+
if self.valueType == "string":
51+
base_type = dict[str, str]
52+
elif self.valueType == "integer":
53+
base_type = dict[str, int]
54+
elif self.valueType == "float":
55+
base_type = dict[str, float]
56+
elif self.valueType == "boolean":
57+
base_type = dict[str, bool]
58+
else:
59+
raise ValueError(f"Unsupported map value type: {self.valueType}")
60+
else:
61+
base_type = dict[str, Any]
4762
else:
4863
raise ValueError(f"Unsupported schema type: {self.type}")
4964

packages/toolbox-core/tests/test_protocol.py

Lines changed: 99 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414

1515

1616
from inspect import Parameter
17-
from typing import Optional
17+
from typing import Any, Optional
1818

1919
import pytest
2020

@@ -170,3 +170,101 @@ def test_parameter_schema_array_optional():
170170
assert param.annotation == expected_type
171171
assert param.kind == Parameter.POSITIONAL_OR_KEYWORD
172172
assert param.default is None
173+
174+
def test_parameter_schema_map_generic():
175+
"""Tests ParameterSchema with a generic 'map' type."""
176+
schema = ParameterSchema(name="metadata", type="map", description="Some metadata")
177+
expected_type = dict[str, Any]
178+
assert schema._ParameterSchema__get_type() == expected_type
179+
180+
param = schema.to_param()
181+
assert isinstance(param, Parameter)
182+
assert param.name == "metadata"
183+
assert param.annotation == expected_type
184+
assert param.kind == Parameter.POSITIONAL_OR_KEYWORD
185+
186+
187+
def test_parameter_schema_map_typed_string():
188+
"""Tests ParameterSchema with a typed 'map' type (string values)."""
189+
schema = ParameterSchema(
190+
name="headers",
191+
type="map",
192+
description="HTTP headers",
193+
valueType="string",
194+
)
195+
expected_type = dict[str, str]
196+
assert schema._ParameterSchema__get_type() == expected_type
197+
198+
param = schema.to_param()
199+
assert param.annotation == expected_type
200+
201+
202+
def test_parameter_schema_map_typed_integer():
203+
"""Tests ParameterSchema with a typed 'map' type (integer values)."""
204+
schema = ParameterSchema(
205+
name="user_scores",
206+
type="map",
207+
description="User scores",
208+
valueType="integer",
209+
)
210+
expected_type = dict[str, int]
211+
assert schema._ParameterSchema__get_type() == expected_type
212+
param = schema.to_param()
213+
assert param.annotation == expected_type
214+
215+
216+
def test_parameter_schema_map_typed_float():
217+
"""Tests ParameterSchema with a typed 'map' type (float values)."""
218+
schema = ParameterSchema(
219+
name="item_prices",
220+
type="map",
221+
description="Item prices",
222+
valueType="float",
223+
)
224+
expected_type = dict[str, float]
225+
assert schema._ParameterSchema__get_type() == expected_type
226+
param = schema.to_param()
227+
assert param.annotation == expected_type
228+
229+
230+
def test_parameter_schema_map_typed_boolean():
231+
"""Tests ParameterSchema with a typed 'map' type (boolean values)."""
232+
schema = ParameterSchema(
233+
name="feature_flags",
234+
type="map",
235+
description="Feature flags",
236+
valueType="boolean",
237+
)
238+
expected_type = dict[str, bool]
239+
assert schema._ParameterSchema__get_type() == expected_type
240+
param = schema.to_param()
241+
assert param.annotation == expected_type
242+
243+
244+
def test_parameter_schema_map_optional():
245+
"""Tests an optional ParameterSchema with a 'map' type."""
246+
schema = ParameterSchema(
247+
name="optional_metadata",
248+
type="map",
249+
description="Optional metadata",
250+
required=False,
251+
)
252+
expected_type = Optional[dict[str, Any]]
253+
assert schema._ParameterSchema__get_type() == expected_type
254+
param = schema.to_param()
255+
assert param.annotation == expected_type
256+
assert param.default is None
257+
258+
259+
def test_parameter_schema_map_unsupported_value_type_error():
260+
"""Tests that an unsupported map valueType raises ValueError."""
261+
unsupported_type = "custom_object"
262+
schema = ParameterSchema(
263+
name="custom_data",
264+
type="map",
265+
description="Custom data map",
266+
valueType=unsupported_type,
267+
)
268+
expected_error_msg = f"Unsupported map value type: {unsupported_type}"
269+
with pytest.raises(ValueError, match=expected_error_msg):
270+
schema._ParameterSchema__get_type()

0 commit comments

Comments
 (0)