diff --git a/examples/async-tools-decorators.py b/examples/async-tools-decorators.py new file mode 100644 index 0000000..1dbf43e --- /dev/null +++ b/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!') diff --git a/examples/async-tools.py b/examples/async-tools.py index 5578229..d85b7e8 100644 --- a/examples/async-tools.py +++ b/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, diff --git a/examples/tools-decorators.py b/examples/tools-decorators.py new file mode 100644 index 0000000..c67e9e8 --- /dev/null +++ b/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') \ No newline at end of file diff --git a/examples/tools.py b/examples/tools.py index d6f3fcf..92b79e1 100644 --- a/examples/tools.py +++ b/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: diff --git a/ollama/__init__.py b/ollama/__init__.py index afe8ce7..acae311 100644 --- a/ollama/__init__.py +++ b/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, diff --git a/ollama/_tools.py b/ollama/_tools.py new file mode 100644 index 0000000..e4b646b --- /dev/null +++ b/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 diff --git a/ollama/_utils.py b/ollama/_utils.py index 653a04c..dcae6b3 100644 --- a/ollama/_utils.py +++ b/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 diff --git a/tests/test_tools.py b/tests/test_tools.py new file mode 100644 index 0000000..dcd2f58 --- /dev/null +++ b/tests/test_tools.py @@ -0,0 +1,97 @@ +def test_tool_and_async_tool_registration(): + import types + from ollama import _tools + _tools._list_tools.clear() + _tools._async_list_tools.clear() + + @(_tools.ollama_tool) + def t1(): + return "ok" + + @(_tools.ollama_async_tool) + async def t2(): + return "ok" + + assert t1 in _tools._list_tools + assert t2 in _tools._async_list_tools + assert t1() == "ok" + import asyncio + assert asyncio.run(t2()) == "ok" + +def test_get_tools_name_and_get_tools(): + from ollama import _tools + _tools._list_tools.clear() + _tools._async_list_tools.clear() + + @(_tools.ollama_tool) + def t3(): + return 1 + @(_tools.ollama_async_tool) + async def t4(): + return 2 + + names = _tools.get_ollama_tools_name() + assert "t3" in names + assert "t4" in names + assert callable(names["t3"]) + assert callable(names["t4"]) + tools = _tools.get_ollama_tools() + assert t3 in tools + assert t4 in tools + +def test_get_ollama_name_async_tools(): + from ollama import _tools + _tools._list_tools.clear() + _tools._async_list_tools.clear() + + @(_tools.ollama_tool) + def sync_tool(): + return 1 + + @(_tools.ollama_async_tool) + async def async_tool1(): + return 2 + + @(_tools.ollama_async_tool) + async def async_tool2(): + return 3 + + async_names = _tools.get_ollama_name_async_tools() + + assert "async_tool1" in async_names + assert "async_tool2" in async_names + assert "sync_tool" not in async_names + assert len(async_names) == 2 + +def test_get_ollama_tool_description(): + from ollama import _tools + _tools._list_tools.clear() + _tools._async_list_tools.clear() + + @(_tools.ollama_tool) + def tool_with_doc(): + """ + Test description for sync tool. + """ + return 1 + + @(_tools.ollama_async_tool) + async def async_tool_with_doc(): + """ + Test description for async tool. + """ + return 2 + + @(_tools.ollama_tool) + def tool_without_doc(): + return 3 + + descriptions = _tools.get_ollama_tool_description() + + + assert "tool_with_doc" in descriptions + assert "async_tool_with_doc" in descriptions + assert "tool_without_doc" not in descriptions + + assert "Test description for sync tool" in descriptions["tool_with_doc"] + assert "Test description for async tool" in descriptions["async_tool_with_doc"] diff --git a/tests/test_utils.py b/tests/test_utils.py index cb9e0d4..7415604 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -256,3 +256,30 @@ def func_with_parentheses_and_args(a: int, b: int): 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' + +def test_create_function_tool(): + from ollama._utils import create_function_tool + tool = create_function_tool( + tool_name="my_tool", + description="desc", + parameter_list=[{"foo": {"type": "string", "description": "bar"}}], + required_parameters=["foo"] + ) + assert tool["type"] == "function" + assert tool["function"]["name"] == "my_tool" + assert tool["function"]["description"] == "desc" + assert tool["function"]["parameters"]["properties"]["foo"]["type"] == "string" + assert tool["function"]["parameters"]["properties"]["foo"]["description"] == "bar" + assert tool["function"]["parameters"]["required"] == ["foo"] + +def test_get_parameters(): + params = [ + {"param1": {"type": "string", "description": "desc1"}}, + {"param2": {"type": "integer", "description": "desc2"}}, + ] + from ollama._utils import _get_parameters + result = _get_parameters(params) + assert result == { + "param1": {"type": "string", "description": "desc1"}, + "param2": {"type": "integer", "description": "desc2"}, + } \ No newline at end of file