diff --git a/ollama/_client.py b/ollama/_client.py index 3654e3c0..f08a20b3 100644 --- a/ollama/_client.py +++ b/ollama/_client.py @@ -362,6 +362,7 @@ def chat( tools: A JSON schema as a dict, an Ollama Tool or a Python Function. Python functions need to follow Google style docstrings to be converted to an Ollama Tool. + Use functools.partial to bind positional arguments that should not be exposed as tool parameters. For more information, see: https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings stream: Whether to stream the response. format: The format of the response. @@ -1021,6 +1022,7 @@ async def chat( tools: A JSON schema as a dict, an Ollama Tool or a Python Function. Python functions need to follow Google style docstrings to be converted to an Ollama Tool. + Use functools.partial to bind positional arguments that should not be exposed as tool parameters. For more information, see: https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings stream: Whether to stream the response. format: The format of the response. diff --git a/ollama/_utils.py b/ollama/_utils.py index 15f1cc0c..7b2cabd4 100644 --- a/ollama/_utils.py +++ b/ollama/_utils.py @@ -3,6 +3,7 @@ import inspect import re from collections import defaultdict +from functools import partial from typing import Callable, Union import pydantic @@ -54,14 +55,31 @@ def _parse_docstring(doc_string: Union[str, None]) -> dict[str, str]: def convert_function_to_tool(func: Callable) -> Tool: - doc_string_hash = str(hash(inspect.getdoc(func))) - parsed_docstring = _parse_docstring(inspect.getdoc(func)) + name_source = doc_source = func + while isinstance(name_source, partial) and not hasattr(name_source, '__name__'): + name_source = name_source.func + while isinstance(doc_source, partial) and '__doc__' not in vars(doc_source): + doc_source = doc_source.func + + # A partial's __wrapped__ metadata must not hide its bound arguments from inspect. + signature_source = func + bound_partials = [] + while isinstance(signature_source, partial) and getattr(signature_source, '__signature__', None) is None: + bound_partials.append(signature_source) + signature_source = signature_source.func + for bound in reversed(bound_partials): + signature_source = partial(signature_source, *bound.args, **bound.keywords) + signature = inspect.signature(signature_source) + + doc_string = inspect.getdoc(doc_source) + doc_string_hash = str(hash(doc_string)) + parsed_docstring = _parse_docstring(doc_string) schema = type( - func.__name__, + name_source.__name__, (pydantic.BaseModel,), { - '__annotations__': {k: v.annotation if v.annotation != inspect._empty else str for k, v in inspect.signature(func).parameters.items()}, - '__signature__': inspect.signature(func), + '__annotations__': {k: v.annotation if v.annotation != inspect._empty else str for k, v in signature.parameters.items()}, + '__signature__': signature, '__doc__': parsed_docstring[doc_string_hash], }, ).model_json_schema() @@ -81,7 +99,7 @@ def convert_function_to_tool(func: Callable) -> Tool: tool = Tool( type='function', function=Tool.Function( - name=func.__name__, + name=name_source.__name__, description=schema.get('description', ''), parameters=Tool.Function.Parameters(**schema), ), diff --git a/tests/test_partial_tools.py b/tests/test_partial_tools.py new file mode 100644 index 00000000..20776c6d --- /dev/null +++ b/tests/test_partial_tools.py @@ -0,0 +1,130 @@ +import asyncio +import inspect +import json +from functools import partial, update_wrapper, wraps +from typing import Optional + +import httpx +import pytest + +from ollama import AsyncClient, Client +from ollama._utils import convert_function_to_tool + + +def scoped_lookup(tenant: str, query: str) -> str: + """Look up a record. + + Args: + tenant: The application tenant. + query: The record to retrieve. + """ + return tenant + ':' + query + + +def make_tool(case): + tool = partial(scoped_lookup, 'bound-tenant-fixture') + name, description = 'scoped_lookup', 'Look up a record.' + parameters = ['query'] + if case == 'custom_name': + tool.__name__ = name = 'tenant_lookup' + elif case == 'custom_doc': + tool.__doc__ = 'Search this tenant.\n\nArgs:\n query: The record to retrieve.' + description = 'Search this tenant.' + elif case == 'wrapped': + update_wrapper(tool, scoped_lookup) + elif case == 'nested_wrapped': + update_wrapper(tool, scoped_lookup) + tool = partial(tool, 'bound-query-fixture') + parameters = [] + elif case == 'decorated_function': + + @wraps(scoped_lookup) + def decorated(*args, **kwargs): + return scoped_lookup(*args, **kwargs) + + tool = partial(decorated, 'bound-tenant-fixture') + elif case == 'explicit_signature': + tool.__signature__ = inspect.Signature([inspect.Parameter('query', inspect.Parameter.KEYWORD_ONLY, annotation=str)]) + return tool, name, description, parameters + + +def chat_with_tool(tool, asynchronous): + requests = [] + + def handle(request): + requests.append(json.loads(request.content)) + return httpx.Response(200, json={'model': 'fixture', 'message': {'role': 'assistant', 'content': 'ok'}, 'done': True}) + + async def async_chat(): + async with AsyncClient(host='http://offline.invalid', transport=httpx.MockTransport(handle), trust_env=False) as client: + response = await client.chat('fixture', tools=[tool]) + assert response.message.content == 'ok' + + if asynchronous: + asyncio.run(async_chat()) + else: + with Client(host='http://offline.invalid', transport=httpx.MockTransport(handle), trust_env=False) as client: + response = client.chat('fixture', tools=[tool]) + assert response.message.content == 'ok' + + assert len(requests) == 1 + return requests[0]['tools'][0]['function'] + + +@pytest.mark.parametrize('asynchronous', [False, True]) +@pytest.mark.parametrize('case', ['plain', 'custom_name', 'custom_doc', 'wrapped', 'nested_wrapped', 'decorated_function', 'explicit_signature']) +def test_chat_partial_tool(case, asynchronous): + tool, name, description, parameters = make_tool(case) + before = vars(tool).copy() + function = chat_with_tool(tool, asynchronous) + + assert function['name'] == name + assert function['description'].strip() == description + assert set(function['parameters']['properties']) == set(parameters) + assert function['parameters'].get('required', []) == parameters + if parameters: + assert function['parameters']['properties']['query'] == {'type': 'string', 'description': 'The record to retrieve.'} + assert 'bound-tenant-fixture' not in json.dumps(function) + assert 'bound-query-fixture' not in json.dumps(function) + assert vars(tool) == before + + +@pytest.mark.parametrize('asynchronous', [False, True]) +def test_partial_keyword_bindings_keep_existing_function_schema(asynchronous): + def lookup(tenant: str, query: str, limit: int = 5, label: Optional[str] = None) -> str: + """Look up a record. + + Args: + tenant: The application tenant. + query: The record to retrieve. + limit: The result limit. + label: An optional label. + """ + return tenant + ':' + query + + def ordinary(query: str, *, limit: int = 5, label: Optional[str] = None) -> str: + """Look up a record. + + Args: + query: The record to retrieve. + limit: The result limit. + label: An optional label. + """ + return query + + tool = partial(lookup, 'bound-tenant-fixture', limit=5, label=None) + function = chat_with_tool(tool, asynchronous) + ordinary_schema = chat_with_tool(ordinary, asynchronous) + assert function['parameters'] == ordinary_schema['parameters'] + assert set(function['parameters']['properties']) == {'query', 'limit', 'label'} + assert 'tenant' not in function['parameters']['properties'] + assert 'label' not in function['parameters'].get('required', []) + assert tool(query='query-fixture', limit=10) == 'bound-tenant-fixture:query-fixture' + + +def test_plain_decorated_functions_keep_their_signature(): + @wraps(scoped_lookup) + def decorated(*args, **kwargs): + return scoped_lookup(*args, **kwargs) + + assert convert_function_to_tool(decorated) == convert_function_to_tool(scoped_lookup)