fix(utils): preserve header-named parameter descriptions

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

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'