types: relax type for tools (#550)

Author: ParthSareen

ollama/_types.py

@@ -79,7 +79,7 @@ def __contains__(self, key: str) -> bool:
     if key in self.model_fields_set:
       return True
 
-    if value := self.model_fields.get(key):
+    if value := self.__class__.model_fields.get(key):
       return value.default is not None
 
     return False
@@ -313,7 +313,7 @@ class Function(SubscriptableBaseModel):
 
 
 class Tool(SubscriptableBaseModel):
-  type: Optional[Literal['function']] = 'function'
+  type: Optional[str] = 'function'
 
   class Function(SubscriptableBaseModel):
     name: Optional[str] = None

ollama/_utils.py

@@ -79,11 +79,12 @@ def convert_function_to_tool(func: Callable) -> Tool:
     }
 
   tool = Tool(
+    type='function',
     function=Tool.Function(
       name=func.__name__,
       description=schema.get('description', ''),
       parameters=Tool.Function.Parameters(**schema),
-    )
+    ),
   )
 
   return Tool.model_validate(tool)

tests/test_client.py

@@ -8,7 +8,7 @@
 
 import pytest
 from httpx import Response as httpxResponse
-from pydantic import BaseModel, ValidationError
+from pydantic import BaseModel
 from pytest_httpserver import HTTPServer, URIPattern
 from werkzeug.wrappers import Request, Response
 
@@ -1136,10 +1136,11 @@ def func2(y: str) -> int:
 
 
 def test_tool_validation():
-  # Raises ValidationError when used as it is a generator
-  with pytest.raises(ValidationError):
-    invalid_tool = {'type': 'invalid_type', 'function': {'name': 'test'}}
-    list(_copy_tools([invalid_tool]))
+  arbitrary_tool = {'type': 'custom_type', 'function': {'name': 'test'}}
+  tools = list(_copy_tools([arbitrary_tool]))
+  assert len(tools) == 1
+  assert tools[0].type == 'custom_type'
+  assert tools[0].function.name == 'test'
 
 
 def test_client_connection_error():