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.
This commit is contained in:
Caua Ramos 2025-05-20 23:39:17 -03:00
parent c41263d8f7
commit 49ed36bf47
9 changed files with 194 additions and 106 deletions

View File

@ -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:

View File

@ -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,

View File

@ -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:

View File

@ -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,

View File

@ -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,

42
ollama/_tools.py Normal file
View File

@ -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

View File

@ -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}

97
tests/test_tools.py Normal file
View File

@ -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"]

View File

@ -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"},
}