Stringify return type arrays to not break server

ollama · Nov 13, 2024 · 23548ce · 23548ce
1 parent 2efa54a
commit 23548ce
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 7 deletions.
diff --git a/ollama/_types.py b/ollama/_types.py
@@ -230,7 +230,7 @@ class Function(SubscriptableBaseModel):
     description: str
 
     class Parameters(SubscriptableBaseModel):
-      type: Union[str, List[str]]
+      type: str
       required: Optional[Sequence[str]] = None
       properties: Optional[JsonSchemaValue] = None
 

diff --git a/ollama/_utils.py b/ollama/_utils.py
@@ -66,8 +66,8 @@ def _get_json_type(python_type: Any) -> str | List[str]:
     if non_none_args := [arg for arg in args if arg not in (None, type(None))]:
       if len(non_none_args) == 1:
         return _get_json_type(non_none_args[0])
-      # For multiple types (e.g., int | str | None), return array of types
-      return [_get_json_type(arg) for arg in non_none_args]
+      # For multiple return types (e.g., int | str | None), return stringified array of types -> "['integer', 'string', 'null']"
+      return str([_get_json_type(arg) for arg in non_none_args]).replace(' ', '')
     return 'null'
 
   # Handle generic types (List[int], Dict[str, int], etc.)

diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -90,8 +90,8 @@ def all_types(
   assert tool.function.parameters.properties['y']['type'] == 'string'
   assert tool.function.parameters.properties['z']['type'] == 'array'
   assert tool.function.parameters.properties['w']['type'] == 'object'
-  assert set(tool.function.parameters.properties['v']['type']) == {'string', 'integer'}
-  assert set(tool.function.return_type) == {'string', 'integer', 'array', 'object'}
+  assert set(x.strip().strip("'") for x in tool.function.parameters.properties['v']['type'].removeprefix('[').removesuffix(']').split(',')) == {'string', 'integer'}
+  assert set(x.strip().strip("'") for x in tool.function.return_type.removeprefix('[').removesuffix(']').split(',')) == {'string', 'integer', 'array', 'object'}
 
 
 def test_process_tools():
@@ -136,7 +136,7 @@ def test_advanced_json_type_conversion():
   assert _get_json_type(Dict[str, List[int]]) == 'object'
 
   # Test multiple unions
-  assert set(_get_json_type(Union[int, str, float])) == {'integer', 'string', 'number'}
+  assert set(x.strip().strip("'") for x in _get_json_type(Union[int, str, float]).removeprefix('[').removesuffix(']').split(',')) == {'integer', 'string', 'number'}
 
   # Test collections.abc types
   assert _get_json_type(Sequence[int]) == 'array'

diff --git a/tests/test_utils_legacy.py b/tests/test_utils_legacy.py
@@ -69,4 +69,4 @@ def all_types(
   assert tool.function.parameters.properties['t']['type'] == 'array'
   assert tool.function.parameters.properties['l']['type'] == 'array'
   assert tool.function.parameters.properties['o']['type'] == 'integer'
-  assert set(tool.function.return_type) == {'string', 'object'}
+  assert set(x.strip().strip("'") for x in tool.function.return_type.removeprefix('[').removesuffix(']').split(',')) == {'string', 'object'}