fix: preserve function tool schema constraints

ollama/_utils.py

@@ -73,10 +73,13 @@ def convert_function_to_tool(func: Callable) -> Tool:
       schema['required'].remove(k)
       types.discard('null')
 
-    schema['properties'][k] = {
-      'description': parsed_docstring[k],
-      'type': ', '.join(types),
-    }
+    property_schema = dict(v)
+    property_schema.pop('anyOf', None)
+    property_schema.update(
+      description=parsed_docstring[k],
+      type=', '.join(types),
+    )
+    schema['properties'][k] = property_schema
 
   tool = Tool(
     type='function',

tests/test_utils.py

@@ -1,6 +1,6 @@
 import json
 import sys
-from typing import Dict, List, Mapping, Sequence, Set, Tuple, Union
+from typing import Dict, List, Literal, Mapping, Sequence, Set, Tuple, Union
 
 from ollama._utils import convert_function_to_tool
 
@@ -44,6 +44,17 @@ def simple_func():
   assert tool['function']['parameters']['properties'] == {}
 
 
+def test_function_preserves_parameter_schema_constraints():
+  def process(ids: list[int], mode: Literal['fast', 'safe']): ...
+
+  properties = convert_function_to_tool(process).model_dump()['function']['parameters']['properties']
+
+  assert properties['ids']['type'] == 'array'
+  assert properties['ids']['items'] == {'type': 'integer'}
+  assert properties['mode']['type'] == 'string'
+  assert properties['mode']['enum'] == ['fast', 'safe']
+
+
 def test_function_with_all_types():
   if sys.version_info >= (3, 10):