Passing Functions as Tools (#321)

* Functions can now be passed as tools

tests/test_client.py

@@ -1,6 +1,7 @@
 import os
 import io
 import json
+from pydantic import ValidationError
 import pytest
 import tempfile
 from pathlib import Path
@@ -8,7 +9,7 @@ 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,56 @@ 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'],
+      },
+    },
+  }
+
+  tools = list(_copy_tools([func1, tool_dict]))
+  assert len(tools) == 2
+  assert tools[0].function.name == 'func1'
+  assert tools[1].function.name == 'test'
+
+
+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]))

tests/test_type_serialization.py

@@ -1,15 +1,48 @@
-from base64 import b64decode, b64encode
+from base64 import b64encode
+from pathlib import Path
 
+import pytest
 from ollama._types import Image
+import tempfile
 
 
-def test_image_serialization():
-  # Test bytes serialization
+def test_image_serialization_bytes():
   image_bytes = b'test image bytes'
+  encoded_string = b64encode(image_bytes).decode()
   img = Image(value=image_bytes)
-  assert img.model_dump() == b64encode(image_bytes).decode()
+  assert img.model_dump() == encoded_string
 
-  # Test base64 string serialization
+
+def test_image_serialization_base64_string():
   b64_str = 'dGVzdCBiYXNlNjQgc3RyaW5n'
   img = Image(value=b64_str)
-  assert img.model_dump() == b64decode(b64_str).decode()
+  assert img.model_dump() == b64_str  # Should return as-is if valid base64
+
+
+def test_image_serialization_plain_string():
+  img = Image(value='not a path or base64')
+  assert img.model_dump() == 'not a path or base64'  # Should return as-is
+
+
+def test_image_serialization_path():
+  with tempfile.NamedTemporaryFile() as temp_file:
+    temp_file.write(b'test file content')
+    temp_file.flush()
+    img = Image(value=Path(temp_file.name))
+    assert img.model_dump() == b64encode(b'test file content').decode()
+
+
+def test_image_serialization_string_path():
+  with tempfile.NamedTemporaryFile() as temp_file:
+    temp_file.write(b'test file content')
+    temp_file.flush()
+    img = Image(value=temp_file.name)
+    assert img.model_dump() == b64encode(b'test file content').decode()
+
+  with pytest.raises(ValueError):
+    img = Image(value='some_path/that/does/not/exist.png')
+    img.model_dump()
+
+  with pytest.raises(ValueError):
+    img = Image(value='not an image')
+    img.model_dump()

tests/test_utils.py

@@ -0,0 +1,270 @@
+import json
+import sys
+from typing import Dict, List, Mapping, Sequence, Set, Tuple, Union
+
+
+from ollama._utils import convert_function_to_tool
+
+
+def test_function_to_tool_conversion():
+  def add_numbers(x: int, y: Union[int, None] = None) -> int:
+    """Add two numbers together.
+    args:
+        x (integer): The first number
+        y (integer, optional): The second number
+
+    Returns:
+        integer: The sum of x and y
+    """
+    return x + y
+
+  tool = convert_function_to_tool(add_numbers).model_dump()
+
+  assert tool['type'] == 'function'
+  assert tool['function']['name'] == 'add_numbers'
+  assert tool['function']['description'] == 'Add two numbers together.'
+  assert tool['function']['parameters']['type'] == 'object'
+  assert tool['function']['parameters']['properties']['x']['type'] == 'integer'
+  assert tool['function']['parameters']['properties']['x']['description'] == 'The first number'
+  assert tool['function']['parameters']['required'] == ['x']
+
+
+def test_function_with_no_args():
+  def simple_func():
+    """
+    A simple function with no arguments.
+    Args:
+        None
+    Returns:
+        None
+    """
+    pass
+
+  tool = convert_function_to_tool(simple_func).model_dump()
+  assert tool['function']['name'] == 'simple_func'
+  assert tool['function']['description'] == 'A simple function with no arguments.'
+  assert tool['function']['parameters']['properties'] == {}
+
+
+def test_function_with_all_types():
+  if sys.version_info >= (3, 10):
+
+    def all_types(
+      x: int,
+      y: str,
+      z: list[int],
+      w: dict[str, int],
+      v: int | str | None,
+    ) -> int | dict[str, int] | str | list[int] | None:
+      """
+      A function with all types.
+      Args:
+          x (integer): The first number
+          y (string): The second number
+          z (array): The third number
+          w (object): The fourth number
+          v (integer | string | None): The fifth number
+      """
+      pass
+  else:
+
+    def all_types(
+      x: int,
+      y: str,
+      z: Sequence,
+      w: Mapping[str, int],
+      d: Dict[str, int],
+      s: Set[int],
+      t: Tuple[int, str],
+      l: List[int],  # noqa: E741
+      o: Union[int, None],
+    ) -> Union[Mapping[str, int], str, None]:
+      """
+      A function with all types.
+      Args:
+          x (integer): The first number
+          y (string): The second number
+          z (array): The third number
+          w (object): The fourth number
+          d (object): The fifth number
+          s (array): The sixth number
+          t (array): The seventh number
+          l (array): The eighth number
+          o (integer | None): The ninth number
+      """
+      pass
+
+  tool_json = convert_function_to_tool(all_types).model_dump_json()
+  tool = json.loads(tool_json)
+  assert tool['function']['parameters']['properties']['x']['type'] == 'integer'
+  assert tool['function']['parameters']['properties']['y']['type'] == 'string'
+
+  if sys.version_info >= (3, 10):
+    assert tool['function']['parameters']['properties']['z']['type'] == 'array'
+    assert tool['function']['parameters']['properties']['w']['type'] == 'object'
+    assert set(x.strip().strip("'") for x in tool['function']['parameters']['properties']['v']['type'].removeprefix('[').removesuffix(']').split(',')) == {'string', 'integer'}
+    assert tool['function']['parameters']['properties']['v']['type'] != 'null'
+    assert tool['function']['parameters']['required'] == ['x', 'y', 'z', 'w']
+  else:
+    assert tool['function']['parameters']['properties']['z']['type'] == 'array'
+    assert tool['function']['parameters']['properties']['w']['type'] == 'object'
+    assert tool['function']['parameters']['properties']['d']['type'] == 'object'
+    assert tool['function']['parameters']['properties']['s']['type'] == 'array'
+    assert tool['function']['parameters']['properties']['t']['type'] == 'array'
+    assert tool['function']['parameters']['properties']['l']['type'] == 'array'
+    assert tool['function']['parameters']['properties']['o']['type'] == 'integer'
+    assert tool['function']['parameters']['properties']['o']['type'] != 'null'
+    assert tool['function']['parameters']['required'] == ['x', 'y', 'z', 'w', 'd', 's', 't', 'l']
+
+
+def test_function_docstring_parsing():
+  from typing import List, Dict, Any
+
+  def func_with_complex_docs(x: int, y: List[str]) -> Dict[str, Any]:
+    """
+    Test function with complex docstring.
+
+    Args:
+        x (integer): A number
+            with multiple lines
+        y (array of string): A list
+            with multiple lines
+
+    Returns:
+        object: A dictionary
+            with multiple lines
+    """
+    pass
+
+  tool = convert_function_to_tool(func_with_complex_docs).model_dump()
+  assert tool['function']['description'] == 'Test function with complex docstring.'
+  assert tool['function']['parameters']['properties']['x']['description'] == 'A number with multiple lines'
+  assert tool['function']['parameters']['properties']['y']['description'] == 'A list with multiple lines'
+
+
+def test_skewed_docstring_parsing():
+  def add_two_numbers(x: int, y: int) -> int:
+    """
+    Add two numbers together.
+    Args:
+        x (integer): : The first number
+
+
+
+
+        y (integer ): The second number
+    Returns:
+        integer: The sum of x and y
+    """
+    pass
+
+  tool = convert_function_to_tool(add_two_numbers).model_dump()
+  assert tool['function']['parameters']['properties']['x']['description'] == ': The first number'
+  assert tool['function']['parameters']['properties']['y']['description'] == 'The second number'
+
+
+def test_function_with_no_docstring():
+  def no_docstring():
+    pass
+
+  def no_docstring_with_args(x: int, y: int):
+    pass
+
+  tool = convert_function_to_tool(no_docstring).model_dump()
+  assert tool['function']['description'] == ''
+
+  tool = convert_function_to_tool(no_docstring_with_args).model_dump()
+  assert tool['function']['description'] == ''
+  assert tool['function']['parameters']['properties']['x']['description'] == ''
+  assert tool['function']['parameters']['properties']['y']['description'] == ''
+
+
+def test_function_with_only_description():
+  def only_description():
+    """
+    A function with only a description.
+    """
+    pass
+
+  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}
+
+  def only_description_with_args(x: int, y: int):
+    """
+    A function with only a description.
+    """
+    pass
+
+  tool = convert_function_to_tool(only_description_with_args).model_dump()
+  assert tool['function']['description'] == 'A function with only a description.'
+  assert tool['function']['parameters'] == {
+    'type': 'object',
+    'properties': {
+      'x': {'type': 'integer', 'description': ''},
+      'y': {'type': 'integer', 'description': ''},
+    },
+    'required': ['x', 'y'],
+  }
+
+
+def test_function_with_yields():
+  def function_with_yields(x: int, y: int):
+    """
+    A function with yields section.
+
+    Args:
+      x: the first number
+      y: the second number
+
+    Yields:
+      The sum of x and y
+    """
+    pass
+
+  tool = convert_function_to_tool(function_with_yields).model_dump()
+  assert tool['function']['description'] == 'A function with yields section.'
+  assert tool['function']['parameters']['properties']['x']['description'] == 'the first number'
+  assert tool['function']['parameters']['properties']['y']['description'] == 'the second number'
+
+
+def test_function_with_no_types():
+  def no_types(a, b):
+    """
+    A function with no types.
+    """
+    pass
+
+  tool = convert_function_to_tool(no_types).model_dump()
+  assert tool['function']['parameters']['properties']['a']['type'] == 'string'
+  assert tool['function']['parameters']['properties']['b']['type'] == 'string'
+
+
+def test_function_with_parentheses():
+  def func_with_parentheses(a: int, b: int) -> int:
+    """
+    A function with parentheses.
+    Args:
+        a: First (:thing) number to add
+        b: Second number to add
+    Returns:
+        int: The sum of a and b
+    """
+    pass
+
+  def func_with_parentheses_and_args(a: int, b: int):
+    """
+    A function with parentheses and args.
+    Args:
+        a(integer) : First (:thing) number to add
+        b(integer) :Second number to add
+    """
+    pass
+
+  tool = convert_function_to_tool(func_with_parentheses).model_dump()
+  assert tool['function']['parameters']['properties']['a']['description'] == 'First (:thing) number to add'
+  assert tool['function']['parameters']['properties']['b']['description'] == 'Second number to add'
+
+  tool = convert_function_to_tool(func_with_parentheses_and_args).model_dump()
+  assert tool['function']['parameters']['properties']['a']['description'] == 'First (:thing) number to add'
+  assert tool['function']['parameters']['properties']['b']['description'] == 'Second number to add'