From 3cc71af8a0dfe05ac34da33da85dfd6512030de8 Mon Sep 17 00:00:00 2001
From: ParthSareen <parth.sareen@ollama.com>
Date: Mon, 7 Apr 2025 13:05:05 -0700
Subject: [PATCH] types: relax enum type for tool

---
 ollama/_types.py    | 1 +
 tests/test_utils.py | 4 ++--
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/ollama/_types.py b/ollama/_types.py
index 0df0ddbf..03ccec7f 100644
--- a/ollama/_types.py
+++ b/ollama/_types.py
@@ -315,6 +315,7 @@ class Property(SubscriptableBaseModel):
 
         type: Optional[str] = None
         description: Optional[str] = None
+        enum: Optional[Sequence] = None
 
       properties: Optional[Mapping[str, Property]] = None
 
diff --git a/tests/test_utils.py b/tests/test_utils.py
index cef7bf29..77fdc29c 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -200,8 +200,8 @@ def only_description_with_args(x: int, y: int):
   assert tool['function']['parameters'] == {
     'type': 'object',
     'properties': {
-      'x': {'type': 'integer', 'description': ''},
-      'y': {'type': 'integer', 'description': ''},
+      'x': {'type': 'integer', 'description': '', 'enum': None},
+      'y': {'type': 'integer', 'description': '', 'enum': None},
     },
     'required': ['x', 'y'],
   }