feat: Adds decorators to simplify tool registration

examples/async-tools-decorators.py

@@ -0,0 +1,101 @@
+import asyncio
+import ollama
+from ollama import ChatResponse
+from ollama import (
+  ollama_tool, 
+  ollama_async_tool, 
+  get_ollama_tools, 
+  get_ollama_name_async_tools, 
+  get_ollama_tools_name, 
+  get_ollama_tool_description)
+
+
+@ollama_tool
+def add_two_numbers(a: int, b: int) -> int:
+  """
+  Add two numbers
+  Args:
+    a (int): The first number
+    b (int): The second number
+  Returns:
+    int: The sum of the two numbers
+  """
+  return a + b
+
+@ollama_tool
+def subtract_two_numbers(a: int, b: int) -> int:
+  """
+  Subtract two numbers
+  Args:
+    a (int): The first number
+    b (int): The second number
+  Returns:
+    int: The difference of the two numbers
+  """
+  return a - b
+
+@ollama_async_tool
+async def web_search(query: str) -> str:
+  """
+  Search the web for information,
+  Args:
+    query (str): The query to search the web for
+  Returns:
+    str: The result of the web search
+  """
+  return f"Searching the web for {query}"
+
+available_functions = get_ollama_tools_name() # this is a dictionary of tools
+
+# tools are treated differently in synchronous code
+async_available_functions = get_ollama_name_async_tools()
+
+messages = [
+  {'role': 'system', 'content': f'You are a helpful assistant, with access to these tools: {get_ollama_tool_description()}'}, #usage example for the get_ollama_tool_description function
+  {'role': 'user', 'content': 'What is three plus one? and Search the web for what is ollama'}]
+print('Prompt:', messages[1]['content'])
+
+async def main():
+  client = ollama.AsyncClient()
+
+  response: ChatResponse = await client.chat(
+    'llama3.1',
+    messages=messages,
+    tools=get_ollama_tools(),
+  )
+
+  if response.message.tool_calls:
+    # There may be multiple tool calls in the response
+    for tool in response.message.tool_calls:
+      # Ensure the function is available, and then call it
+      if function_to_call := available_functions.get(tool.function.name):
+        print('Calling function:', tool.function.name)
+        print('Arguments:', tool.function.arguments)
+        # if the function is in the list of asynchronous functions it is executed with asyncio.run()
+        if tool.function.name in async_available_functions:
+          output = await function_to_call(**tool.function.arguments)
+        else:
+          output = function_to_call(**tool.function.arguments)
+        print('Function output:', output)
+      else:
+        print('Function', tool.function.name, 'not found')
+
+  # Only needed to chat with the model using the tool call results
+  if response.message.tool_calls:
+    # Add the function response to messages for the model to use
+    messages.append(response.message)
+    messages.append({'role': 'tool', 'content': str(output), 'name': tool.function.name})
+
+    # Get final response from model with function outputs
+    final_response = await client.chat('llama3.1', messages=messages)
+    print('Final response:', final_response.message.content)
+
+  else:
+    print('No tool calls returned from model')
+
+
+if __name__ == '__main__':
+  try:
+    asyncio.run(main())
+  except KeyboardInterrupt:
+    print('\nGoodbye!')

examples/async-tools.py

@@ -1,7 +1,7 @@
 import asyncio
 
 import ollama
-from ollama import ChatResponse
+from ollama import ChatResponse, get_ollama_tool_description
 
 
 def add_two_numbers(a: int, b: int) -> int:
@@ -42,8 +42,10 @@ def subtract_two_numbers(a: int, b: int) -> int:
   },
 }
 
-messages = [{'role': 'user', 'content': 'What is three plus one?'}]
-print('Prompt:', messages[0]['content'])
+messages = [
+  {'role': 'system', 'content': f'You are a helpful assistant, with access to these tools: {get_ollama_tool_description()}'}, #usage example for the get_ollama_tool_description function
+  {'role': 'user', 'content': 'What is three plus one? and Search the web for what is ollama'}]
+print('Prompt:', messages[1]['content'])
 
 available_functions = {
   'add_two_numbers': add_two_numbers,

examples/tools-decorators.py

@@ -0,0 +1,89 @@
+import asyncio
+from ollama import ChatResponse, chat
+from ollama import (
+  ollama_tool, 
+  ollama_async_tool, 
+  get_ollama_tools, 
+  get_ollama_name_async_tools, 
+  get_ollama_tools_name, 
+  get_ollama_tool_description)
+
+@ollama_tool
+def add_two_numbers(a: int, b: int) -> int:
+  """
+  Add two numbers
+  Args:
+    a (int): The first number
+    b (int): The second number
+  Returns:
+    int: The sum of the two numbers
+  """
+  return a + b
+
+@ollama_tool
+def subtract_two_numbers(a: int, b: int) -> int:
+  """
+  Subtract two numbers
+  Args:
+    a (int): The first number
+    b (int): The second number
+  Returns:
+    int: The difference of the two numbers
+  """
+  return a - b
+
+@ollama_async_tool
+async def web_search(query: str) -> str:
+  """
+  Search the web for information,
+  Args:
+    query (str): The query to search the web for
+  Returns:
+    str: The result of the web search
+  """
+  return f"Searching the web for {query}"
+
+available_functions = get_ollama_tools_name() # this is a dictionary of tools
+
+# tools are treated differently in synchronous code
+async_available_functions = get_ollama_name_async_tools()
+
+messages = [
+  {'role': 'system', 'content': f'You are a helpful assistant, with access to these tools: {get_ollama_tool_description()}'}, #usage example for the get_ollama_tool_description function
+  {'role': 'user', 'content': 'What is three plus one? and Search the web for what is ollama'}]
+print('Prompt:', messages[1]['content'])
+
+response: ChatResponse = chat(
+  'llama3.1',
+  messages=messages,
+  tools=get_ollama_tools(), # this is the list of tools using decorators
+)
+
+if response.message.tool_calls:
+  # There may be multiple tool calls in the response
+  for tool in response.message.tool_calls:
+    # Ensure the function is available, and then call it
+    if function_to_call := available_functions.get(tool.function.name):
+      print('Calling function:', tool.function.name)
+      print('Arguments:', tool.function.arguments)
+      # if the function is in the list of asynchronous functions it is executed with asyncio.run()
+      if tool.function.name in async_available_functions:
+        output = asyncio.run(function_to_call(**tool.function.arguments))
+      else:
+        output = function_to_call(**tool.function.arguments)
+      print('Function output:', output)
+    else:
+      print('Function', tool.function.name, 'not found')
+
+# Only needed to chat with the model using the tool call results
+if response.message.tool_calls:
+  # Add the function response to messages for the model to use
+  messages.append(response.message)
+  messages.append({'role': 'tool', 'content': str(output), 'name': tool.function.name})
+
+  # Get final response from model with function outputs
+  final_response = chat('llama3.1', messages=messages)
+  print('Final response:', final_response.message.content)
+
+else:
+  print('No tool calls returned from model')

examples/tools.py

@@ -1,5 +1,5 @@
-from ollama import ChatResponse, chat
-
+from ollama import ChatResponse, chat, create_function_tool
+from ollama import get_ollama_tool_description
 
 def add_two_numbers(a: int, b: int) -> int:
   """
@@ -26,6 +26,11 @@ def subtract_two_numbers(a: int, b: int) -> int:
   # The cast is necessary as returned tool call arguments don't always conform exactly to schema
   return int(a) - int(b)
 
+def multiply_two_numbers(a: int, b: int) -> int:
+  """
+  Multiply two numbers
+  """
+  return int(a) * int(b)
 
 # Tools can still be manually defined and passed into chat
 subtract_two_numbers_tool = {
@@ -44,18 +49,28 @@ def subtract_two_numbers(a: int, b: int) -> int:
   },
 }
 
-messages = [{'role': 'user', 'content': 'What is three plus one?'}]
-print('Prompt:', messages[0]['content'])
+# A simple way to define tools manually, even though it seems long
+multiply_two_numbers_tool = create_function_tool(tool_name="multiply_two_numbers", 
+                                                 description="Multiply two numbers", 
+                                                 parameter_list=[{"a": {"type": "integer", "description": "The first number"}, 
+                                                                  "b": {"type": "integer", "description": "The second number"}}], 
+                                                 required_parameters=["a", "b"])
+
+messages = [
+  {'role': 'system', 'content': f'You are a helpful assistant, with access to these tools: {get_ollama_tool_description()}'}, #usage example for the get_ollama_tool_description function
+  {'role': 'user', 'content': 'What is three plus one? and Search the web for what is ollama'}]
+print('Prompt:', messages[1]['content'])
 
 available_functions = {
   'add_two_numbers': add_two_numbers,
   'subtract_two_numbers': subtract_two_numbers,
+  'multiply_two_numbers': multiply_two_numbers,
 }
 
 response: ChatResponse = chat(
   'llama3.1',
   messages=messages,
-  tools=[add_two_numbers, subtract_two_numbers_tool],
+  tools=[add_two_numbers, subtract_two_numbers_tool, multiply_two_numbers_tool],
 )
 
 if response.message.tool_calls:

ollama/__init__.py

@@ -1,4 +1,12 @@
 from ollama._client import AsyncClient, Client
+from ollama._utils import create_function_tool
+from ollama._tools import (
+  ollama_tool, 
+  ollama_async_tool, 
+  get_ollama_tools, 
+  get_ollama_name_async_tools, 
+  get_ollama_tools_name,
+  get_ollama_tool_description)
 from ollama._types import (
   ChatResponse,
   EmbeddingsResponse,

ollama/_tools.py

@@ -0,0 +1,42 @@
+from functools import wraps
+
+_list_tools = []
+_async_list_tools = []
+
+def ollama_async_tool(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        return func(*args, **kwargs)
+    _async_list_tools.append(wrapper)
+    return wrapper
+
+def ollama_tool(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        return func(*args, **kwargs)
+    _list_tools.append(wrapper)
+    return wrapper
+
+def get_ollama_tools_name():
+    list_name_tools = {} 
+    for func in _list_tools + _async_list_tools:
+        if func.__name__ not in list_name_tools:
+            list_name_tools[func.__name__] = func
+    return list_name_tools
+
+def get_ollama_tools():
+    return _list_tools + _async_list_tools
+
+def get_ollama_name_async_tools():
+    return {f"{func.__name__}" for func in _async_list_tools}
+
+def get_ollama_tool_description():
+    from ollama._utils import _parse_docstring
+    result = {}
+    for func in _list_tools + _async_list_tools:
+        if func.__doc__:
+            parsed_docstring = _parse_docstring(func.__doc__)
+            if parsed_docstring and str(hash(func.__doc__)) in parsed_docstring:
+                result[func.__name__] = parsed_docstring[str(hash(func.__doc__))].strip()
+
+    return result

ollama/_utils.py

@@ -87,3 +87,30 @@ def convert_function_to_tool(func: Callable) -> Tool:
   )
 
   return Tool.model_validate(tool)
+
+def _get_parameters(parameters: list):
+    properties_dict = {}
+    for param_item in parameters:
+        for key, value in param_item.items():
+            properties_dict[key] = {
+                "type": value.get("type"),
+                "description": value.get("description")
+            }
+    return properties_dict
+
+def create_function_tool(tool_name: str, description: str, parameter_list: list, required_parameters: list):
+    properties = _get_parameters(parameter_list)
+
+    tool_definition = {
+            'type': 'function',
+            'function': {
+                'name': tool_name,
+                'description': description,
+                'parameters': {
+                    'type': 'object',
+                    'properties': properties,
+                    'required': required_parameters
+                }
+            }
+        }
+    return tool_definition