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