From 65f94b4fba2cc95dd718cbda683a326a8e5e670a Mon Sep 17 00:00:00 2001
From: Parth Sareen <parth.sareen@ollama.com>
Date: Wed, 16 Apr 2025 13:33:51 -0700
Subject: [PATCH] types: allow items and defs for tools (#501)

---
 ollama/_types.py    | 15 ++++++++++++++-
 tests/test_utils.py |  8 +++++---
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/ollama/_types.py b/ollama/_types.py
index 0d802e0..dfeafcf 100644
--- a/ollama/_types.py
+++ b/ollama/_types.py
@@ -307,15 +307,19 @@ class Tool(SubscriptableBaseModel):
     description: Optional[str] = None
 
     class Parameters(SubscriptableBaseModel):
+      model_config = ConfigDict(populate_by_name=True)
       type: Optional[Literal['object']] = 'object'
+      defs: Optional[Any] = Field(None, alias='$defs')
+      items: Optional[Any] = None
       required: Optional[Sequence[str]] = None
 
       class Property(SubscriptableBaseModel):
         model_config = ConfigDict(arbitrary_types_allowed=True)
 
         type: Optional[Union[str, Sequence[str]]] = None
+        items: Optional[Any] = None
         description: Optional[str] = None
-        enum: Optional[Sequence] = None
+        enum: Optional[Sequence[Any]] = None
 
       properties: Optional[Mapping[str, Property]] = None
 
@@ -325,6 +329,15 @@ class Tool(SubscriptableBaseModel):
 
 
 class ChatRequest(BaseGenerateRequest):
+  @model_serializer(mode='wrap')
+  def serialize_model(self, nxt):
+    output = nxt(self)
+    if 'tools' in output and output['tools']:
+      for tool in output['tools']:
+        if 'function' in tool and 'parameters' in tool['function'] and 'defs' in tool['function']['parameters']:
+          tool['function']['parameters']['$defs'] = tool['function']['parameters'].pop('defs')
+    return output
+
   messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None
   'Messages to chat with.'
 
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 77fdc29..e5e0743 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -187,7 +187,7 @@ def test_function_with_only_description():
 
   tool = convert_function_to_tool(only_description).model_dump()
   assert tool['function']['description'] == 'A function with only a description.'
-  assert tool['function']['parameters'] == {'type': 'object', 'properties': {}, 'required': None}
+  assert tool['function']['parameters'] == {'type': 'object', 'defs': None, 'items': None, 'required': None, 'properties': {}}
 
   def only_description_with_args(x: int, y: int):
     """
@@ -199,9 +199,11 @@ def test_function_with_only_description():
   assert tool['function']['description'] == 'A function with only a description.'
   assert tool['function']['parameters'] == {
     'type': 'object',
+    'defs': None,
+    'items': None,
     'properties': {
-      'x': {'type': 'integer', 'description': '', 'enum': None},
-      'y': {'type': 'integer', 'description': '', 'enum': None},
+      'x': {'type': 'integer', 'description': '', 'enum': None, 'items': None},
+      'y': {'type': 'integer', 'description': '', 'enum': None, 'items': None},
     },
     'required': ['x', 'y'],
   }