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/_json_type_map.py

@@ -0,0 +1,109 @@
+import sys
+from typing import Any, List, Mapping, Optional, Sequence, Union, get_origin, get_args
+from collections.abc import Set
+from typing import Dict, Set as TypeSet, TypeVar
+
+T = TypeVar('T')
+if sys.version_info >= (3, 10):
+  from types import UnionType
+
+  def is_union(tp: Any) -> bool:
+    return get_origin(tp) in (Union, UnionType)
+
+else:
+  UnionType = Union[T]
+
+  def is_union(tp: Any) -> bool:
+    return get_origin(tp) is Union
+
+
+# Python doesn't have a type serializer, so we need to map types to JSON types
+TYPE_MAP = {
+  # Basic types
+  int: 'integer',
+  'int': 'integer',
+  'integer': 'integer',
+  str: 'string',
+  'str': 'string',
+  'string': 'string',
+  float: 'number',
+  'float': 'number',
+  'number': 'number',
+  bool: 'boolean',
+  'bool': 'boolean',
+  'boolean': 'boolean',
+  type(None): 'null',
+  None: 'null',
+  'None': 'null',
+  'null': 'null',
+  # Collection types
+  list: 'array',
+  'list': 'array',
+  List: 'array',
+  'List': 'array',
+  Sequence: 'array',
+  'Sequence': 'array',
+  tuple: 'array',
+  'tuple': 'array',
+  set: 'array',
+  'set': 'array',
+  Set: 'array',
+  TypeSet: 'array',
+  'Set': 'array',
+  'array': 'array',
+  # Mapping types
+  dict: 'object',
+  'dict': 'object',
+  Dict: 'object',
+  'Dict': 'object',
+  Mapping: 'object',
+  'Mapping': 'object',
+  'object': 'object',
+  Any: 'string',
+  'Any': 'string',
+}
+
+
+def _map_type(python_type: Any) -> str:
+  # Handle generic types (List[int], Dict[str, int], etc.)
+  origin = get_origin(python_type)
+  if origin is not None:
+    # Get the base type (List, Dict, etc.)
+    base_type = TYPE_MAP.get(origin, None)
+    if base_type:
+      return base_type
+    # If it's a subclass of known abstract base classes, map to appropriate type
+    if isinstance(origin, type):
+      if issubclass(origin, (list, Sequence, tuple, set, Set)):
+        return 'array'
+      if issubclass(origin, (dict, Mapping)):
+        return 'object'
+
+  # Handle both type objects and type references (older Python versions)
+  type_key = python_type
+  if isinstance(python_type, type):
+    type_key = python_type
+  elif isinstance(python_type, str):
+    type_key = python_type
+
+  # If type not found in map, try to get the type name
+  if type_key not in TYPE_MAP and hasattr(python_type, '__name__'):
+    type_key = python_type.__name__
+
+  if type_key in TYPE_MAP:
+    return TYPE_MAP[type_key]
+
+  raise ValueError(f'Could not map Python type {python_type} to a valid JSON type')
+
+
+def get_json_type(python_type: Union[type, UnionType, Optional[T]]) -> Union[str, List[str]]:
+  # Handle Optional types (Union[type, None] and type | None)
+  if is_union(python_type):
+    args = get_args(python_type)
+    # Filter out None/NoneType from union args
+    if non_none_args := [arg for arg in args if arg not in (None, type(None))]:
+      if len(non_none_args) == 1:
+        return _map_type(non_none_args[0])
+      # For multiple return types (e.g., int | str | None), return stringified array of types -> "['integer', 'string', 'null']"
+      return str([_map_type(arg) for arg in non_none_args]).replace(' ', '')
+  return _map_type(python_type)

ollama/_types.py

@@ -2,25 +2,21 @@
 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 ollama._json_type_map import get_json_type, T, UnionType
+
+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 +225,19 @@ 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: Union[type, UnionType, Optional[T], str]
+        description: str
+
+        @model_serializer
+        def serialize_model(self) -> dict:
+          return {'type': get_json_type(self.type), 'description': self.description}
+
+      properties: Optional[Mapping[str, Property]] = None
 
     parameters: Parameters
 
@@ -335,6 +341,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,110 @@
+from __future__ import annotations
+from typing import Any, Callable, Union, get_args
+from ollama._json_type_map import is_union
+from ollama._types import Tool
+from typing import Dict
+
+
+def _parse_docstring(func: Callable, doc_string: Union[str, None]) -> tuple[str, Dict[str, str]]:
+  # Extract description from docstring - get all lines before Args:
+  if not doc_string:
+    return '', {}
+
+  description_lines = []
+  for line in doc_string.split('\n'):
+    line = line.strip()
+    if line.startswith('Args:'):
+      break
+    if line:
+      description_lines.append(line)
+
+  description = ' '.join(description_lines).strip()
+
+  if 'Args:' not in doc_string:
+    return description, {}
+
+  args_section = doc_string.split('Args:')[1]
+  if 'Returns:' in args_section:
+    args_section = args_section.split('Returns:')[0]
+
+  param_descriptions = {}
+  current_param = None
+  param_desc_lines = []
+  indent_level = None
+
+  for line in args_section.split('\n'):
+    stripped_line = line.strip()
+    if not stripped_line:
+      continue
+
+    # Check for new parameter
+    for param_name in func.__annotations__:
+      if param_name == 'return':
+        continue
+      if stripped_line.startswith(f'{param_name}:') or stripped_line.startswith(f'{param_name} ') or stripped_line.startswith(f'{param_name}('):
+        # Save previous parameter if exists
+        if current_param:
+          param_descriptions[current_param] = ' '.join(param_desc_lines).strip()
+          param_desc_lines = []
+
+        current_param = param_name
+        # Get description after parameter name
+        desc_part = stripped_line.split(':', 1)[1].strip() if ':' in stripped_line else ''
+        if desc_part:
+          param_desc_lines.append(desc_part)
+        indent_level = len(line) - len(line.lstrip())
+        break
+    else:
+      # Handle continuation lines
+      if current_param and line.startswith(' ' * (indent_level + 4 if indent_level else 0)):
+        param_desc_lines.append(stripped_line)
+      elif current_param and stripped_line:
+        # Different indentation means new parameter
+        param_descriptions[current_param] = ' '.join(param_desc_lines).strip()
+        param_desc_lines = []
+        current_param = None
+
+  if current_param:
+    param_descriptions[current_param] = ' '.join(param_desc_lines).strip()
+
+  # Verify all parameters have descriptions
+  for param_name in func.__annotations__:
+    if param_name == 'return':
+      continue
+    if param_name not in param_descriptions:
+      param_descriptions[param_name] = ''
+
+  return description, param_descriptions
+
+
+def is_optional_type(python_type: Any) -> bool:
+  if is_union(python_type):
+    args = get_args(python_type)
+    return any(arg in (None, type(None)) for arg in args)
+  return False
+
+
+def convert_function_to_tool(func: Callable) -> Tool:
+  doc_string = func.__doc__
+
+  description, param_descriptions = _parse_docstring(func, doc_string)
+
+  parameters = Tool.Function.Parameters(type='object', properties={}, required=[])
+
+  for param_name, param_type in func.__annotations__.items():
+    if param_name == 'return':
+      continue
+
+    parameters.properties[param_name] = Tool.Function.Parameters.Property(type=param_type, description=param_descriptions.get(param_name, ''))
+
+    if not is_optional_type(param_type):
+      parameters.required.append(param_name)
+
+  return Tool(
+    function=Tool.Function(
+      name=func.__name__,
+      description=description,
+      parameters=parameters,
+      return_type=None,
+    )
+  )