diff --git a/ollama/_utils.py b/ollama/_utils.py index 15f1cc0c..b1cffbc3 100644 --- a/ollama/_utils.py +++ b/ollama/_utils.py @@ -54,15 +54,17 @@ def _parse_docstring(doc_string: Union[str, None]) -> dict[str, str]: def convert_function_to_tool(func: Callable) -> Tool: + signature = inspect.signature(func) doc_string_hash = str(hash(inspect.getdoc(func))) parsed_docstring = _parse_docstring(inspect.getdoc(func)) schema = type( func.__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], + **{k: v.default for k, v in signature.parameters.items() if v.default is not inspect.Parameter.empty}, }, ).model_json_schema() @@ -70,7 +72,8 @@ def convert_function_to_tool(func: Callable) -> Tool: # If type is missing, the default is string types = {t.get('type', 'string') for t in v.get('anyOf')} if 'anyOf' in v else {v.get('type', 'string')} if 'null' in types: - schema['required'].remove(k) + if k in schema.get('required', []): + schema['required'].remove(k) types.discard('null') schema['properties'][k] = { diff --git a/tests/test_utils.py b/tests/test_utils.py index cb9e0d4f..1a24c214 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,10 +1,52 @@ import json import sys +from enum import IntEnum from typing import Dict, List, Mapping, Sequence, Set, Tuple, Union from ollama._utils import convert_function_to_tool +def test_function_defaults_are_not_required(): + def search(query: str, limit: int = 5): + return query, limit + + parameters = convert_function_to_tool(search).function.parameters + assert parameters.required == ['query'] + assert parameters.properties['limit'].type == 'integer' + + +def test_function_default_with_custom_comparison(): + class DefaultLimit(IntEnum): + FIVE = 5 + + def __ne__(self, other): + raise RuntimeError('Default value comparison is not supported') + + def search(query: str, limit: int = DefaultLimit.FIVE): + return query, limit + + parameters = convert_function_to_tool(search).function.parameters + assert parameters.required == ['query'] + assert parameters.properties['limit'].type == 'integer' + + +def test_nullable_default_with_no_required_arguments(): + def search(query: Union[str, None] = None): + return query + + parameters = convert_function_to_tool(search).function.parameters + assert not parameters.required + assert parameters.properties['query'].type == 'string' + + +def test_keyword_only_defaults_are_not_required(): + def search(query: str, *, limit: int = 5): + return query, limit + + parameters = convert_function_to_tool(search).function.parameters + assert parameters.required == ['query'] + + def test_function_to_tool_conversion(): def add_numbers(x: int, y: Union[int, None] = None) -> int: """Add two numbers together.