add support for union types in mcp client

cheese-head · cheese-head · commit 256698f17822 · 2025-06-13T16:09:15.000Z
Signed-off-by: Patrick Riel &lt;priel@nvidia.com&gt;
diff --git a/src/aiq/tool/mcp/mcp_client.py b/src/aiq/tool/mcp/mcp_client.py
@@ -68,6 +68,16 @@ def _generate_field(field_name: str, field_properties: dict[str, Any]) -> tuple:
             else:
                 item_type = _type_map.get(item_properties.get("type", "string"), Any)
             field_type = list[item_type]
+        elif isinstance(json_type, list):
+            nullable = "null" in json_type
+            main_type = next((t for t in json_type if t != "null"), "string")
+            field_type = _type_map.get(main_type, Any)
+            if nullable:
+                field_type = field_type | None
+            return field_type, Field(
+                default=field_properties.get("default", None if nullable else ...),
+                description=field_properties.get("description", "")
+            )
         else:
             field_type = _type_map.get(json_type, Any)
 
