diff --git a/ollama/_utils.py b/ollama/_utils.py index 15f1cc0c..75fd8a45 100644 --- a/ollama/_utils.py +++ b/ollama/_utils.py @@ -3,24 +3,30 @@ import inspect import re from collections import defaultdict -from typing import Callable, Union +from typing import Callable, Container, Union import pydantic from ollama._types import Tool -def _parse_docstring(doc_string: Union[str, None]) -> dict[str, str]: +def _parse_docstring(doc_string: Union[str, None], parameters: Container[str] = ()) -> dict[str, str]: parsed_docstring = defaultdict(str) if not doc_string: return parsed_docstring key = str(hash(doc_string)) + args_indent = 0 for line in doc_string.splitlines(): lowered_line = line.lower().strip() + indent = len(line) - len(line.lstrip()) + name, _, text = line.strip().partition(':') + # A deeper line under Args like `raises: ...` documents the parameter of that exact name; bare or other headers still end Args + is_parameter = key == 'args' and indent > args_indent and name in parameters and text.strip() != '' if lowered_line.startswith('args:'): key = 'args' - elif lowered_line.startswith(('returns:', 'yields:', 'raises:')): + args_indent = indent + elif lowered_line.startswith(('returns:', 'yields:', 'raises:')) and not is_parameter: key = '_' else: @@ -55,7 +61,7 @@ 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)) + parsed_docstring = _parse_docstring(inspect.getdoc(func), inspect.signature(func).parameters) schema = type( func.__name__, (pydantic.BaseModel,), diff --git a/tests/test_utils.py b/tests/test_utils.py index cb9e0d4f..3bb77709 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -256,3 +256,72 @@ def func_with_parentheses_and_args(a: int, b: int): tool = convert_function_to_tool(func_with_parentheses_and_args).model_dump() 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_function_with_parameters_named_like_sections(): + def fetch(url: str, raises: bool, timeout: int) -> str: + """ + Fetch a URL. + + Args: + url: The URL to fetch + raises: Whether to raise on HTTP errors + timeout: Seconds to wait + """ + + def lookup(key: str, returns: str, limit: int) -> str: + """ + Look up a record. + + Args: + key: The record key + returns: Which field to return + limit: Maximum number of records + + Returns: + The selected field + """ + + tool = convert_function_to_tool(fetch).model_dump() + assert tool['function']['description'] == 'Fetch a URL.' + assert tool['function']['parameters']['properties']['url']['description'] == 'The URL to fetch' + assert tool['function']['parameters']['properties']['raises']['description'] == 'Whether to raise on HTTP errors' + assert tool['function']['parameters']['properties']['timeout']['description'] == 'Seconds to wait' + + tool = convert_function_to_tool(lookup).model_dump() + assert tool['function']['description'] == 'Look up a record.' + assert tool['function']['parameters']['properties']['key']['description'] == 'The record key' + assert tool['function']['parameters']['properties']['returns']['description'] == 'Which field to return' + assert tool['function']['parameters']['properties']['limit']['description'] == 'Maximum number of records' + + +def test_overindented_sections_still_end_args(): + def fetch(url: str, timeout: int) -> str: + """ + Fetch a URL. + + Args: + url: The URL to fetch + timeout: Seconds to wait + Raises: + timeout: Not a parameter description + """ + + def lookup(key: str, returns: str) -> str: + """ + Look up a record. + + Args: + key: The record key + returns: Which field to return + Returns: The selected field + key: Not a parameter description + """ + + tool = convert_function_to_tool(fetch).model_dump() + assert tool['function']['parameters']['properties']['url']['description'] == 'The URL to fetch' + assert tool['function']['parameters']['properties']['timeout']['description'] == 'Seconds to wait' + + tool = convert_function_to_tool(lookup).model_dump() + assert tool['function']['parameters']['properties']['key']['description'] == 'The record key' + assert tool['function']['parameters']['properties']['returns']['description'] == 'Which field to return'