Skip to content

Commit 885e824

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. * 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]`). * 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 885e824

File tree

2 files changed

+116
-2
lines changed

2 files changed

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

0 commit comments

Comments
 (0)