types: allow items and defs for tools (#501)

Author: ParthSareen

ollama/_types.py

@@ -307,15 +307,19 @@ class Function(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 Property(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.'
 

tests/test_utils.py

@@ -187,7 +187,7 @@ def 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 only_description_with_args(x: int, y: int):
   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'],
   }