diff --git a/ollama/_types.py b/ollama/_types.py index b9bb53cb..8176a48d 100644 --- a/ollama/_types.py +++ b/ollama/_types.py @@ -376,6 +376,7 @@ class Property(SubscriptableBaseModel): items: Optional[Any] = None description: Optional[str] = None enum: Optional[Sequence[Any]] = None + properties: Optional[Mapping[str, 'Tool.Function.Parameters.Property']] = None properties: Optional[Mapping[str, Property]] = None diff --git a/tests/test_type_serialization.py b/tests/test_type_serialization.py index 02a69e95..fa502e8b 100644 --- a/tests/test_type_serialization.py +++ b/tests/test_type_serialization.py @@ -1,10 +1,52 @@ +import json import tempfile from base64 import b64encode from pathlib import Path import pytest -from ollama._types import CreateRequest, Image +from ollama._types import ChatRequest, CreateRequest, Image, Tool + + +def test_tool_nested_properties_serialization(): + parameters = { + 'type': 'object', + 'properties': { + 'address': { + 'type': 'object', + 'description': 'Delivery address', + 'properties': { + 'location': { + 'type': 'object', + 'properties': {'city': {'type': 'string', 'enum': ['London', 'Paris']}}, + }, + 'lines': {'type': 'array', 'items': {'type': 'string'}}, + }, + }, + }, + } + tool = Tool(function=Tool.Function(name='deliver', parameters=parameters)) + assert tool.function.parameters.model_dump(exclude_none=True) == parameters + assert isinstance(tool.function.parameters.properties['address'].properties['location'].properties['city'], Tool.Function.Parameters.Property) + + request = ChatRequest(model='test-model', tools=[tool]) + assert request.model_dump(exclude_none=True)['tools'][0]['function']['parameters'] == parameters + assert json.loads(request.model_dump_json(exclude_none=True))['tools'][0]['function']['parameters'] == parameters + + +def test_tool_flat_properties_serialization(): + parameters = { + 'type': 'object', + '$defs': {'Label': {'type': 'string'}}, + 'required': ['name'], + 'properties': { + 'name': {'type': ['string', 'null'], 'description': 'Name', 'enum': ['a', 'b', None]}, + 'labels': {'type': 'array', 'items': {'$ref': '#/$defs/Label'}}, + 'empty': {'type': 'object', 'properties': {}}, + }, + } + request = ChatRequest(model='test-model', tools=[{'function': {'name': 'label', 'parameters': parameters}}]) + assert request.model_dump(exclude_none=True)['tools'][0]['function']['parameters'] == parameters def test_image_serialization_bytes(): diff --git a/tests/test_utils.py b/tests/test_utils.py index cb9e0d4f..a8a45c08 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -193,8 +193,8 @@ def only_description_with_args(x: int, y: int): 'defs': None, 'items': None, 'properties': { - 'x': {'type': 'integer', 'description': '', 'enum': None, 'items': None}, - 'y': {'type': 'integer', 'description': '', 'enum': None, 'items': None}, + 'x': {'type': 'integer', 'description': '', 'enum': None, 'items': None, 'properties': None}, + 'y': {'type': 'integer', 'description': '', 'enum': None, 'items': None, 'properties': None}, }, 'required': ['x', 'y'], }