From c41263d8f7b542b2ed97978ae4f340978f928dc0 Mon Sep 17 00:00:00 2001 From: Caua Ramos Date: Sun, 18 May 2025 12:55:32 -0300 Subject: [PATCH 1/2] feat: Adds decorators to simplify tool registration This commit introduces decorators and helper functions to simplify the process of defining and registering synchronous and asynchronous tools in the ollama-python library. Main changes: - Added `@ollama_tool` and `@ollama_async_tool` decorators to `ollama/_utils.py` to automatically register functions as tools. New helper functions in `ollama/_utils.py`: - `get_tools()`: Returns a list of all registered tools (synchronous and asynchronous). - `get_tools_name()`: Returns a dictionary mapping tool names to their respective functions. - `get_name_async_tools()`: Returns an array of asynchronous tool names. - `ollama/__init__.py` has been updated to export the new decorators and functions. - `examples/tools.py` has been updated to reflect the changes and simplify the example usage of manually defined tools. - New example files have been created: - `examples/tools-decorators.py`: Demonstrates the use of decorators for synchronous and asynchronous tools in a synchronous context. - `examples/async-tools-decorators.py`: Demonstrates the use of decorators in an asynchronous context with `asyncio`. - Tests have been added to `tests/test_utils.py` to cover the new decorators and helper functions (`test_tool_and_async_tool_registration`, `test_get_tools_name_and_get_tools`). These changes aim to make it easier for developers to integrate custom functionality with ollama, reducing the need for boilerplate to define and manage tools. --- examples/async-tools-decorators.py | 93 ++++++++++++++++++++++++++++++ examples/tools-decorators.py | 81 ++++++++++++++++++++++++++ examples/tools.py | 19 +++++- ollama/__init__.py | 7 +++ ollama/_utils.py | 55 ++++++++++++++++++ tests/test_utils.py | 71 +++++++++++++++++++++++ 6 files changed, 323 insertions(+), 3 deletions(-) create mode 100644 examples/async-tools-decorators.py create mode 100644 examples/tools-decorators.py diff --git a/examples/async-tools-decorators.py b/examples/async-tools-decorators.py new file mode 100644 index 00000000..e32887eb --- /dev/null +++ b/examples/async-tools-decorators.py @@ -0,0 +1,93 @@ +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 + + +@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_tools_name() # this is a dictionary of tools + +# tools are treated differently in synchronous code +async_available_functions = get_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']) + +async def main(): + client = ollama.AsyncClient() + + response: ChatResponse = await client.chat( + 'llama3.1', + messages=messages, + tools=get_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/tools-decorators.py b/examples/tools-decorators.py new file mode 100644 index 00000000..a6536f4f --- /dev/null +++ b/examples/tools-decorators.py @@ -0,0 +1,81 @@ +import asyncio +from ollama import ChatResponse, chat +from ollama import ollama_tool, ollama_async_tool, get_tools, get_name_async_tools, get_tools_name + +@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_tools_name() # this is a dictionary of tools + +# tools are treated differently in synchronous code +async_available_functions = get_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']) + +response: ChatResponse = chat( + 'llama3.1', + messages=messages, + tools=get_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 d6f3fcf2..72727ca4 100644 --- a/examples/tools.py +++ b/examples/tools.py @@ -1,4 +1,4 @@ -from ollama import ChatResponse, chat +from ollama import ChatResponse, chat, create_function_tool 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,26 @@ def subtract_two_numbers(a: int, b: int) -> int: }, } -messages = [{'role': 'user', 'content': 'What is three plus one?'}] +# 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': 'user', 'content': 'What is three plus one? And what is three times two?'}] print('Prompt:', messages[0]['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 afe8ce71..efdede7a 100644 --- a/ollama/__init__.py +++ b/ollama/__init__.py @@ -1,4 +1,11 @@ from ollama._client import AsyncClient, Client +from ollama._utils import ( + ollama_tool, + ollama_async_tool, + get_tools, get_tools_name, + get_name_async_tools, + create_function_tool, + ) from ollama._types import ( ChatResponse, EmbeddingsResponse, diff --git a/ollama/_utils.py b/ollama/_utils.py index 653a04c7..4cc7b707 100644 --- a/ollama/_utils.py +++ b/ollama/_utils.py @@ -87,3 +87,58 @@ 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 + +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_utils.py b/tests/test_utils.py index cb9e0d4f..bb93ad0f 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -256,3 +256,74 @@ 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__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( + 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_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 From 49ed36bf4789c754102fc05d2f911bbec5ea9cc6 Mon Sep 17 00:00:00 2001 From: Caua Ramos Date: Tue, 20 May 2025 23:39:17 -0300 Subject: [PATCH 2/2] 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 e32887eb..1dbf43e1 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 55782298..d85b7e82 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 index a6536f4f..c67e9e89 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 72727ca4..92b79e17 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 @@ def multiply_two_numbers(a: int, b: int) -> int: "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 efdede7a..acae3115 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 00000000..e4b646b9 --- /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 4cc7b707..dcae6b39 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 00000000..dcd2f583 --- /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 bb93ad0f..74156045 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -257,19 +257,6 @@ def func_with_parentheses_and_args(a: int, b: int): 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