diff --git a/ollama/_utils.py b/ollama/_utils.py index 15f1cc0c..36a5874c 100644 --- a/ollama/_utils.py +++ b/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', diff --git a/tests/test_utils.py b/tests/test_utils.py index cb9e0d4f..9d89f3e7 100644 --- a/tests/test_utils.py +++ b/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):