From 1b3406887d6e3cf7cd4f70e2edeb5f69c8a9a9c4 Mon Sep 17 00:00:00 2001
From: Parth Sareen <parth.sareen@ollama.com>
Date: Tue, 8 Apr 2025 15:05:46 -0700
Subject: [PATCH] types: relax enum type for tool (#498)

---
 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 0df0ddb..03ccec7 100644
--- a/ollama/_types.py
+++ b/ollama/_types.py
@@ -315,6 +315,7 @@ class Tool(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 cef7bf2..77fdc29 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -200,8 +200,8 @@ def test_function_with_only_description():
   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'],
   }