Passing Functions as Tools

ollama/_client.py

@@ -10,6 +10,7 @@
 
 from typing import (
   Any,
+  Callable,
   Literal,
   Mapping,
   Optional,
@@ -22,6 +23,9 @@
 
 import sys
 
+
+from ollama._utils import convert_function_to_tool
+
 if sys.version_info < (3, 9):
   from typing import Iterator, AsyncIterator
 else:
@@ -284,7 +288,7 @@ def chat(
     model: str = '',
     messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None,
     *,
-    tools: Optional[Sequence[Union[Mapping[str, Any], Tool]]] = None,
+    tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None,
     stream: bool = False,
     format: Optional[Literal['', 'json']] = None,
     options: Optional[Union[Mapping[str, Any], Options]] = None,
@@ -750,7 +754,7 @@ async def chat(
     model: str = '',
     messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None,
     *,
-    tools: Optional[Sequence[Union[Mapping[str, Any], Tool]]] = None,
+    tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None,
     stream: Literal[True] = True,
     format: Optional[Literal['', 'json']] = None,
     options: Optional[Union[Mapping[str, Any], Options]] = None,
@@ -1075,9 +1079,15 @@ def _copy_messages(messages: Optional[Sequence[Union[Mapping[str, Any], Message]
     )
 
 
-def _copy_tools(tools: Optional[Sequence[Union[Mapping[str, Any], Tool]]]) -> Iterator[Tool]:
-  for tool in tools or []:
-    yield Tool.model_validate(tool)
+def _copy_tools(tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None) -> Iterator[Tool]:
+  if not tools:
+    return []
+
+  for unprocessed_tool in tools:
+    if callable(unprocessed_tool):
+      yield convert_function_to_tool(unprocessed_tool)
+    else:
+      yield Tool.model_validate(unprocessed_tool)
 
 
 def _as_path(s: Optional[Union[str, PathLike]]) -> Union[Path, None]:

ollama/_types.py

@@ -2,25 +2,19 @@
 from base64 import b64encode
 from pathlib import Path
 from datetime import datetime
-from typing import (
-  Any,
-  Literal,
-  Mapping,
-  Optional,
-  Sequence,
-  Union,
-)
-from typing_extensions import Annotated
+from typing import Any, Mapping, Optional, Union, Sequence
+
+from typing_extensions import Annotated, Literal
 
 from pydantic import (
   BaseModel,
   ByteSize,
+  ConfigDict,
   Field,
   FilePath,
   Base64Str,
   model_serializer,
 )
-from pydantic.json_schema import JsonSchemaValue
 
 
 class SubscriptableBaseModel(BaseModel):
@@ -229,9 +223,16 @@ class Function(SubscriptableBaseModel):
     description: str
 
     class Parameters(SubscriptableBaseModel):
-      type: str
+      type: Literal['object'] = 'object'
       required: Optional[Sequence[str]] = None
-      properties: Optional[JsonSchemaValue] = None
+
+      class Property(SubscriptableBaseModel):
+        model_config = ConfigDict(arbitrary_types_allowed=True)
+
+        type: str
+        description: str
+
+      properties: Optional[Mapping[str, Property]] = None
 
     parameters: Parameters
 
@@ -335,6 +336,7 @@ class ModelDetails(SubscriptableBaseModel):
 
 class ListResponse(SubscriptableBaseModel):
   class Model(SubscriptableBaseModel):
+    name: Optional[str] = None
     modified_at: Optional[datetime] = None
     digest: Optional[str] = None
     size: Optional[ByteSize] = None

ollama/_utils.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+import inspect
+from typing import Callable, Union
+
+import pydantic
+from ollama._types import Tool
+
+
+def _parse_docstring(doc_string: Union[str, None]) -> dict[str, str]:
+  parsed_docstring = {'description': ''}
+  if not doc_string:
+    return parsed_docstring
+
+  lowered_doc_string = doc_string.lower()
+
+  if 'args:' not in lowered_doc_string:
+    parsed_docstring['description'] = lowered_doc_string.strip()
+    return parsed_docstring
+
+  else:
+    parsed_docstring['description'] = lowered_doc_string.split('args:')[0].strip()
+    args_section = lowered_doc_string.split('args:')[1]
+
+  if 'returns:' in lowered_doc_string:
+    # Return section can be captured and used
+    args_section = args_section.split('returns:')[0]
+
+  if 'yields:' in lowered_doc_string:
+    args_section = args_section.split('yields:')[0]
+
+  cur_var = None
+  for line in args_section.split('\n'):
+    line = line.strip()
+    if not line:
+      continue
+    if ':' not in line:
+      # Continuation of the previous parameter's description
+      if cur_var:
+        parsed_docstring[cur_var] += f' {line}'
+      continue
+
+    # For the case with: `param_name (type)`: ...
+    if '(' in line:
+      param_name = line.split('(')[0]
+      param_desc = line.split('):')[1]
+
+    # For the case with: `param_name: ...`
+    else:
+      param_name, param_desc = line.split(':', 1)
+
+    parsed_docstring[param_name.strip()] = param_desc.strip()
+    cur_var = param_name.strip()
+
+  return parsed_docstring
+
+
+def convert_function_to_tool(func: Callable) -> Tool:
+  schema = type(
+    func.__name__,
+    (pydantic.BaseModel,),
+    {
+      '__annotations__': {k: v.annotation for k, v in inspect.signature(func).parameters.items()},
+      '__signature__': inspect.signature(func),
+      '__doc__': inspect.getdoc(func),
+    },
+  ).model_json_schema()
+
+  properties = {}
+  required = []
+  parsed_docstring = _parse_docstring(schema.get('description'))
+  for k, v in schema.get('properties', {}).items():
+    prop = {
+      'description': parsed_docstring.get(k, ''),
+      'type': v.get('type'),
+    }
+
+    if 'anyOf' in v:
+      is_optional = any(t.get('type') == 'null' for t in v['anyOf'])
+      types = [t.get('type', 'string') for t in v['anyOf'] if t.get('type') != 'null']
+      prop['type'] = types[0] if len(types) == 1 else str(types)
+      if not is_optional:
+        required.append(k)
+    else:
+      if prop['type'] != 'null':
+        required.append(k)
+
+    properties[k] = prop
+
+  schema['properties'] = properties
+
+  tool = Tool(
+    function=Tool.Function(
+      name=func.__name__,
+      description=parsed_docstring.get('description'),
+      parameters=Tool.Function.Parameters(
+        type='object',
+        properties=schema.get('properties', {}),
+        required=required,
+      ),
+    )
+  )
+
+  return Tool.model_validate(tool)

tests/test_client.py

@@ -1,14 +1,15 @@
 import os
 import io
 import json
+from pydantic import ValidationError
 import pytest
 import tempfile
 from pathlib import Path
 from pytest_httpserver import HTTPServer, URIPattern
 from werkzeug.wrappers import Request, Response
 from PIL import Image
 
-from ollama._client import Client, AsyncClient
+from ollama._client import Client, AsyncClient, _copy_tools
 
 
 class PrefixPattern(URIPattern):
@@ -982,3 +983,67 @@ def test_headers():
   )
   assert client._client.headers['x-custom'] == 'value'
   assert client._client.headers['content-type'] == 'application/json'
+
+
+def test_copy_tools():
+  def func1(x: int) -> str:
+    """Simple function 1.
+    Args:
+        x (integer): A number
+    """
+    pass
+
+  def func2(y: str) -> int:
+    """Simple function 2.
+    Args:
+        y (string): A string
+    """
+    pass
+
+  # Test with list of functions
+  tools = list(_copy_tools([func1, func2]))
+  assert len(tools) == 2
+  assert tools[0].function.name == 'func1'
+  assert tools[1].function.name == 'func2'
+
+  # Test with empty input
+  assert list(_copy_tools()) == []
+  assert list(_copy_tools(None)) == []
+  assert list(_copy_tools([])) == []
+
+  # Test with mix of functions and tool dicts
+  tool_dict = {
+    'type': 'function',
+    'function': {
+      'name': 'test',
+      'description': 'Test function',
+      'parameters': {
+        'type': 'object',
+        'properties': {'x': {'type': 'string', 'description': 'A string'}},
+        'required': ['x'],
+      },
+    },
+  }
+
+  tool_json = json.loads(json.dumps(tool_dict))
+  tools = list(_copy_tools([func1, tool_dict, tool_json]))
+  assert len(tools) == 3
+  assert tools[0].function.name == 'func1'
+  assert tools[1].function.name == 'test'
+  assert tools[2].function.name == 'test'
+
+
+def test_tool_validation():
+  # Test that malformed tool dictionaries are rejected
+  # 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]))
+
+  # Test missing required fields
+  incomplete_tool = {
+    'type': 'function',
+    'function': {'name': 'test'},  # missing description and parameters
+  }
+  with pytest.raises(ValidationError):
+    list(_copy_tools([incomplete_tool]))

tests/test_type_serialization.py

@@ -1,6 +1,9 @@
 from base64 import b64decode, b64encode
 
-from ollama._types import Image
+import pytest
+
+
+from ollama._types import Image, Tool
 
 
 def test_image_serialization():
@@ -13,3 +16,81 @@ def test_image_serialization():
   b64_str = 'dGVzdCBiYXNlNjQgc3RyaW5n'
   img = Image(value=b64_str)
   assert img.model_dump() == b64decode(b64_str).decode()
+
+
+def test_tool_serialization():
+  # Test valid tool serialization
+  tool = Tool(
+    function=Tool.Function(
+      name='add_two_numbers',
+      description='Add two numbers together.',
+      parameters=Tool.Function.Parameters(
+        type='object',
+        properties={
+          'a': Tool.Function.Parameters.Property(
+            type='integer',
+            description='The first number',
+          ),
+          'b': Tool.Function.Parameters.Property(
+            type='integer',
+            description='The second number',
+          ),
+        },
+        required=['a', 'b'],
+      ),
+    )
+  )
+  assert tool.model_dump() == {
+    'type': 'function',
+    'function': {
+      'name': 'add_two_numbers',
+      'description': 'Add two numbers together.',
+      'parameters': {
+        'type': 'object',
+        'properties': {
+          'a': {
+            'type': 'integer',
+            'description': 'The first number',
+          },
+          'b': {
+            'type': 'integer',
+            'description': 'The second number',
+          },
+        },
+        'required': ['a', 'b'],
+      },
+    },
+  }
+
+  # Test invalid type
+  with pytest.raises(ValueError):
+    property = Tool.Function.Parameters.Property(
+      type=lambda x: x,  # Invalid type
+      description='Invalid type',
+    )
+    Tool.model_validate(
+      Tool(
+        function=Tool.Function(
+          parameters=Tool.Function.Parameters(
+            properties={
+              'x': property,
+            }
+          )
+        )
+      )
+    )
+
+  # Test invalid parameters type
+  with pytest.raises(ValueError):
+    Tool.model_validate(
+      Tool(
+        function=Tool.Function(
+          name='test',
+          description='Test',
+          parameters=Tool.Function.Parameters(
+            type='invalid_type',  # Must be 'object'
+            properties={},
+          ),
+        )
+      )
+    )