mirror of
https://github.com/ollama/ollama-python.git
synced 2026-01-13 21:57:16 +08:00
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:
parent
c41263d8f7
commit
49ed36bf47
@ -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:
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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
42
ollama/_tools.py
Normal 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
|
||||
@ -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
97
tests/test_tools.py
Normal 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"]
|
||||
@ -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"},
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user