From 49ed36bf4789c754102fc05d2f911bbec5ea9cc6 Mon Sep 17 00:00:00 2001 From: Caua Ramos Date: Tue, 20 May 2025 23:39:17 -0300 Subject: [PATCH] feat(tools): Refactor tools system and improve AI context handling - Create dedicated _tools.py module for better organization - Move tool-related functions from _utils.py to _tools.py (not including create_function_tool) - Rename tool getter functions for clarity: * get_ollama_tools_name * get_ollama_tools * get_ollama_name_async_tools - Add get_ollama_tool_description for enhanced AI context extraction (it returns a list with name and description, useful for adding to the "content" of the system) - Fix wrapper function metadata preservation using functools.wraps - Fix tool registration and lookup bugs found in tests - Add comprehensive test suite in test_tools.py - Improve code organization and modularity - Add proper docstring handling for tool descriptions This change improves the overall architecture of the tools system and adds better support for AI context understanding through proper function metadata and descriptions. --- examples/async-tools-decorators.py | 20 ++++-- examples/async-tools.py | 8 ++- examples/tools-decorators.py | 20 ++++-- examples/tools.py | 8 ++- ollama/__init__.py | 11 ++-- ollama/_tools.py | 42 +++++++++++++ ollama/_utils.py | 28 --------- tests/test_tools.py | 97 ++++++++++++++++++++++++++++++ tests/test_utils.py | 66 ++++---------------- 9 files changed, 194 insertions(+), 106 deletions(-) create mode 100644 ollama/_tools.py create mode 100644 tests/test_tools.py diff --git a/examples/async-tools-decorators.py b/examples/async-tools-decorators.py index e32887e..1dbf43e 100644 --- a/examples/async-tools-decorators.py +++ b/examples/async-tools-decorators.py @@ -1,7 +1,13 @@ import asyncio import ollama from ollama import ChatResponse -from ollama import ollama_tool, ollama_async_tool, get_tools, get_name_async_tools, get_tools_name +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 @@ -39,13 +45,15 @@ async def web_search(query: str) -> str: """ return f"Searching the web for {query}" -available_functions = get_tools_name() # this is a dictionary of tools +available_functions = get_ollama_tools_name() # this is a dictionary of tools # tools are treated differently in synchronous code -async_available_functions = get_name_async_tools() +async_available_functions = get_ollama_name_async_tools() -messages = [{'role': 'user', 'content': 'What is three plus one? and Search the web for what is ollama'}] -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']) async def main(): client = ollama.AsyncClient() @@ -53,7 +61,7 @@ async def main(): response: ChatResponse = await client.chat( 'llama3.1', messages=messages, - tools=get_tools(), + tools=get_ollama_tools(), ) if response.message.tool_calls: 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 @@ subtract_two_numbers_tool = { }, } -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 index a6536f4..c67e9e8 100644 --- a/examples/tools-decorators.py +++ b/examples/tools-decorators.py @@ -1,6 +1,12 @@ import asyncio from ollama import ChatResponse, chat -from ollama import ollama_tool, ollama_async_tool, get_tools, get_name_async_tools, get_tools_name +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: @@ -37,18 +43,20 @@ async def web_search(query: str) -> str: """ return f"Searching the web for {query}" -available_functions = get_tools_name() # this is a dictionary of tools +available_functions = get_ollama_tools_name() # this is a dictionary of tools # tools are treated differently in synchronous code -async_available_functions = get_name_async_tools() +async_available_functions = get_ollama_name_async_tools() -messages = [{'role': 'user', 'content': 'What is three plus one? and Search the web for what is ollama'}] -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']) response: ChatResponse = chat( 'llama3.1', messages=messages, - tools=get_tools(), # this is the list of tools using decorators + tools=get_ollama_tools(), # this is the list of tools using decorators ) if response.message.tool_calls: diff --git a/examples/tools.py b/examples/tools.py index 72727ca..92b79e1 100644 --- a/examples/tools.py +++ b/examples/tools.py @@ -1,5 +1,5 @@ from ollama import ChatResponse, chat, create_function_tool - +from ollama import get_ollama_tool_description def add_two_numbers(a: int, b: int) -> int: """ @@ -56,8 +56,10 @@ multiply_two_numbers_tool = create_function_tool(tool_name="multiply_two_numbers "b": {"type": "integer", "description": "The second number"}}], required_parameters=["a", "b"]) -messages = [{'role': 'user', 'content': 'What is three plus one? And what is three times two?'}] -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/ollama/__init__.py b/ollama/__init__.py index efdede7..acae311 100644 --- a/ollama/__init__.py +++ b/ollama/__init__.py @@ -1,11 +1,12 @@ from ollama._client import AsyncClient, Client -from ollama._utils import ( +from ollama._utils import create_function_tool +from ollama._tools import ( ollama_tool, ollama_async_tool, - get_tools, get_tools_name, - get_name_async_tools, - create_function_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 4cc7b70..dcae6b3 100644 --- a/ollama/_utils.py +++ b/ollama/_utils.py @@ -114,31 +114,3 @@ def create_function_tool(tool_name: str, description: str, parameter_list: list, } } return tool_definition - -list_tools = [] -async_list_tools = [] - -def ollama_async_tool(func): - async_list_tools.append(func) - def wrapper(*args, **kwargs): - return func(*args, **kwargs) - return wrapper - -def ollama_tool(func): - list_tools.append(func) - def wrapper(*args, **kwargs): - return func(*args, **kwargs) - return wrapper - -def get_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_tools(): - return list_tools + async_list_tools - -def get_name_async_tools(): - return {f"{func.__name__}" for func in async_list_tools} \ No newline at end of file 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 bb93ad0..7415604 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -257,19 +257,6 @@ def test_function_with_parentheses(): 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__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"}, - } - def test_create_function_tool(): from ollama._utils import create_function_tool tool = create_function_tool( @@ -285,45 +272,14 @@ def test_create_function_tool(): assert tool["function"]["parameters"]["properties"]["foo"]["description"] == "bar" assert tool["function"]["parameters"]["required"] == ["foo"] -def test_tool_and_async_tool_registration(): - import types - from ollama import _utils - # Limpar listas para evitar interferência - _utils.list_tools.clear() - _utils.async_list_tools.clear() - - @(_utils.ollama_tool) - def t1(): - return "ok" - - @(_utils.ollama_async_tool) - async def t2(): - return "ok" - - assert t1 in _utils.list_tools - assert t2 in _utils.async_list_tools - # Testa wrappers - assert t1() == "ok" - import asyncio - assert asyncio.run(t2()) == "ok" - -def test_get_tools_name_and_get_tools(): - from ollama import _utils - _utils.list_tools.clear() - _utils.async_list_tools.clear() - - @(_utils.ollama_tool) - def t3(): - return 1 - @(_utils.ollama_async_tool) - async def t4(): - return 2 - - names = _utils.get_tools_name() - assert "t3" in names - assert "t4" in names - assert callable(names["t3"]) - assert callable(names["t4"]) - tools = _utils.get_tools() - assert t3 in tools - assert t4 in tools +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