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,
@@ -293,6 +297,30 @@ def chat(
     """
     Create a chat response using the requested model.
 
+    Args:
+      tools:
+        A JSON schema as a dict, an Ollama Tool or a Python Function.
+        Python functions need to follow Google style docstrings to be converted to an Ollama Tool.
+        For more information, see: https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings
+      stream: Whether to stream the response.
+      format: The format of the response.
+
+    Example:
+      def add_two_numbers(a: int, b: int) -> int:
+        \"""
+        Add two numbers together.
+
+        Args:
+          a: First number to add
+          b: Second number to add
+
+        Returns:
+          int: The sum of a and b
+        \"""
+        return a + b
+
+      client.chat(model='llama3.1:8b', tools=[add_two_numbers], messages=[...])
+
     Raises `RequestError` if a model is not provided.
 
     Raises `ResponseError` if the request could not be fulfilled.
@@ -750,7 +778,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,
@@ -771,6 +799,30 @@ async def chat(
     """
     Create a chat response using the requested model.
 
+    Args:
+      tools:
+        A JSON schema as a dict, an Ollama Tool or a Python Function.
+        Python functions need to follow Google style docstrings to be converted to an Ollama Tool.
+        For more information, see: https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings
+      stream: Whether to stream the response.
+      format: The format of the response.
+
+    Example:
+      def add_two_numbers(a: int, b: int) -> int:
+        \"""
+        Add two numbers together.
+
+        Args:
+          a: First number to add
+          b: Second number to add
+
+        Returns:
+          int: The sum of a and b
+        \"""
+        return a + b
+
+      client.chat(model='llama3.1:8b', tools=[add_two_numbers], messages=[...])
+
     Raises `RequestError` if a model is not provided.
 
     Raises `ResponseError` if the request could not be fulfilled.
@@ -1075,9 +1127,12 @@ 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:
+    yield convert_function_to_tool(unprocessed_tool) if callable(unprocessed_tool) else Tool.model_validate(unprocessed_tool)
 
 
 def _as_path(s: Optional[Union[str, PathLike]]) -> Union[Path, None]:
@@ -1152,6 +1207,8 @@ def _parse_host(host: Optional[str]) -> str:
   'https://[0001:002:003:0004::1]:56789/path'
   >>> _parse_host('[0001:002:003:0004::1]:56789/path/')
   'http://[0001:002:003:0004::1]:56789/path'
+  >>> _parse_host('http://host.docker.internal:11434/path')
+  'http://host.docker.internal:11434/path'
   """
 
   host, port = host or '', 11434

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):
+    model: Optional[str] = None
     modified_at: Optional[datetime] = None
     digest: Optional[str] = None
     size: Optional[ByteSize] = None

ollama/_utils.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+from collections import defaultdict
+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 = defaultdict(str)
+  if not doc_string:
+    return parsed_docstring
+
+  lowered_doc_string = doc_string.lower()
+
+  key = hash(doc_string)
+  parsed_docstring[key] = ''
+  for line in lowered_doc_string.splitlines():
+    if line.startswith('args:'):
+      key = 'args'
+    elif line.startswith('returns:') or line.startswith('yields:') or line.startswith('raises:'):
+      key = '_'
+
+    else:
+      # maybe change to a list and join later
+      parsed_docstring[key] += f'{line.strip()}\n'
+
+  last_key = None
+  for line in parsed_docstring['args'].splitlines():
+    line = line.strip()
+    if ':' in line and not line.startswith('args'):
+      # Split on first occurrence of '(' or ':' to separate arg name from description
+      split_char = '(' if '(' in line else ':'
+      arg_name, rest = line.split(split_char, 1)
+
+      last_key = arg_name.strip()
+      # Get description after the colon
+      arg_description = rest.split(':', 1)[1].strip() if split_char == '(' else rest.strip()
+      parsed_docstring[last_key] = arg_description
+
+    elif last_key and line:
+      parsed_docstring[last_key] += ' ' + line
+
+  return parsed_docstring
+
+
+def convert_function_to_tool(func: Callable) -> Tool:
+  doc_string_hash = hash(inspect.getdoc(func))
+  parsed_docstring = _parse_docstring(inspect.getdoc(func))
+  schema = type(
+    func.__name__,
+    (pydantic.BaseModel,),
+    {
+      '__annotations__': {k: v.annotation if v.annotation != inspect._empty else str for k, v in inspect.signature(func).parameters.items()},
+      '__signature__': inspect.signature(func),
+      '__doc__': parsed_docstring[doc_string_hash],
+    },
+  ).model_json_schema()
+
+  for k, v in schema.get('properties', {}).items():
+    # If type is missing, the default is string
+    types = {t.get('type', 'string') for t in v.get('anyOf')} if 'anyOf' in v else {v.get('type', 'string')}
+    if 'null' in types:
+      schema['required'].remove(k)
+      types.discard('null')
+
+    schema['properties'][k] = {
+      'description': parsed_docstring[k],
+      'type': ', '.join(types),
+    }
+
+  tool = Tool(
+    function=Tool.Function(
+      name=func.__name__,
+      description=schema.get('description', ''),
+      parameters=Tool.Function.Parameters(**schema),
+    )
+  )
+
+  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]))