Multi-line docstrings and exhaustive tests

Author: ParthSareen

ollama/_utils.py

@@ -6,7 +6,7 @@
 from typing import Dict, Set as TypeSet
 import sys
 
-# Type compatibility layer
+# Type compatibility layer for Union and UnionType
 if sys.version_info >= (3, 10):
   from types import UnionType
 
@@ -30,6 +30,7 @@ def is_union(tp: Any) -> bool:
   bool: 'boolean',
   'bool': 'boolean',
   type(None): 'null',
+  None: 'null',
   'None': 'null',
   # Collection types
   list: 'array',
@@ -83,16 +84,16 @@ def _get_json_type(python_type: Any) -> str | List[str]:
       if issubclass(get_origin(python_type), (dict, Mapping)):
         return 'object'
 
-  # Handle both type objects and type references
+  # Handle both type objects and type references (older Python versions)
   type_key = python_type
   if isinstance(python_type, type):
     type_key = python_type
   elif isinstance(python_type, str):
-    type_key = python_type.lower()
+    type_key = python_type
 
   # If type not found in map, try to get the type name
   if type_key not in TYPE_MAP and hasattr(python_type, '__name__'):
-    type_key = python_type.__name__.lower()
+    type_key = python_type.__name__
 
   if type_key in TYPE_MAP:
     return TYPE_MAP[type_key]
@@ -138,23 +139,47 @@ def convert_function_to_tool(func: Callable) -> Tool:
     if param_name == 'return':
       continue
 
-    param_desc = None
+    param_desc_lines = []
+    found_param = False
+    indent_level = None
+
+    # Process docstring lines
     for line in args_section.split('\n'):
-      line = line.strip()
-      # Check for parameter name with or without colon, space, or parentheses to mitigate formatting issues
-      if line.startswith(param_name + ':') or line.startswith(param_name + ' ') or line.startswith(param_name + '('):
-        param_desc = line.split(':', 1)[1].strip()
+      stripped_line = line.strip()
+      if not stripped_line:
+        continue
+
+      # Check for parameter start
+      if stripped_line.startswith(f'{param_name}:') or stripped_line.startswith(f'{param_name} ') or stripped_line.startswith(f'{param_name}('):
+        found_param = True
+        # Get the description part after the parameter name
+        desc_part = stripped_line.split(':', 1)[1].strip() if ':' in stripped_line else ''
+        if desc_part:
+          param_desc_lines.append(desc_part)
+        # Get the indentation level for continuation lines
+        indent_level = len(line) - len(line.lstrip())
+        continue
+
+      # Handle continuation lines
+      if found_param and line.startswith(' ' * (indent_level + 4)):
+        # Add continuation line, stripped of extra indentation
+        param_desc_lines.append(stripped_line)
+      elif found_param and stripped_line:
+        # If we hit a line with different indentation, we're done with this parameter
         break
 
-    if not param_desc:
+    if not found_param:
       raise ValueError(f'Parameter {param_name} must have a description in the Args section')
 
+    # Join all lines with spaces
+    param_desc = ' '.join(param_desc_lines).strip()
+
     parameters['properties'][param_name] = {
       'type': _get_json_type(param_type),
       'description': param_desc,
     }
 
-    # Only add to required if not optional - could capture and map earlier to save this call
+    # Only add to required if not optional
     if not _is_optional_type(param_type):
       parameters['required'].append(param_name)
 

tests/test_utils.py

@@ -98,14 +98,14 @@ def test_process_tools():
   def func1(x: int) -> str:
     """Simple function 1.
     Args:
-        x: A number
+        x (integer): A number
     """
     pass
 
   def func2(y: str) -> int:
     """Simple function 2.
     Args:
-        y: A string
+        y (string): A string
     """
     pass
 
@@ -126,3 +126,81 @@ def func2(y: str) -> int:
   assert len(tools) == 2
   assert tools[0].function.name == 'func1'
   assert tools[1].function.name == 'test'
+
+
+def test_advanced_json_type_conversion():
+  from typing import Optional, Union, List, Dict, Sequence, Mapping, Set, Tuple, Any
+
+  # Test nested collections
+  assert _get_json_type(List[List[int]]) == 'array'
+  assert _get_json_type(Dict[str, List[int]]) == 'object'
+
+  # Test multiple unions
+  assert set(_get_json_type(Union[int, str, float])) == {'integer', 'string', 'number'}
+
+  # Test collections.abc types
+  assert _get_json_type(Sequence[int]) == 'array'
+  assert _get_json_type(Mapping[str, int]) == 'object'
+  assert _get_json_type(Set[int]) == 'array'
+  assert _get_json_type(Tuple[int, str]) == 'array'
+
+  # Test nested optionals
+  assert _get_json_type(Optional[List[Optional[int]]]) == 'array'
+
+  # Test edge cases
+  assert _get_json_type(Any) == 'string'  # or however you want to handle Any
+  assert _get_json_type(None) == 'null'
+  assert _get_json_type(type(None)) == 'null'
+
+  # Test complex nested types
+  complex_type = Dict[str, Union[List[int], Optional[str], Dict[str, bool]]]
+  assert _get_json_type(complex_type) == 'object'
+
+
+def test_invalid_types():
+  # Test that invalid types raise appropriate errors
+  with pytest.raises(ValueError):
+    _get_json_type(lambda x: x)  # Function type
+
+  with pytest.raises(ValueError):
+    _get_json_type(type)  # metaclass
+
+
+def test_function_docstring_parsing():
+  from typing import List, Dict, Any
+
+  def func_with_complex_docs(x: int, y: List[str]) -> Dict[str, Any]:
+    """
+    Test function with complex docstring.
+
+    Args:
+        x (integer): A number
+            with multiple lines
+        y (array of string): A list
+            with multiple lines
+
+    Returns:
+        object: A dictionary
+            with multiple lines
+    """
+    pass
+
+  tool = convert_function_to_tool(func_with_complex_docs)
+  assert tool['function']['description'] == 'Test function with complex docstring.'
+  assert tool['function']['parameters']['properties']['x']['description'] == 'A number with multiple lines'
+  assert tool['function']['parameters']['properties']['y']['description'] == 'A list with multiple lines'
+
+
+def test_tool_validation():
+  # Test that malformed tool dictionaries are rejected
+  invalid_tool = {'type': 'invalid_type', 'function': {'name': 'test'}}
+  with pytest.raises(ValueError):
+    process_tools([invalid_tool])
+
+  # Test missing required fields
+  incomplete_tool = {
+    'type': 'function',
+    'function': {'name': 'test'},  # missing description and parameters
+  }
+  with pytest.raises(ValueError):
+    process_tools([incomplete_tool])

tests/test_utils_legacy.py

@@ -1,21 +1,16 @@
-from typing import Dict, List, Mapping, Optional, Sequence, Set, Tuple, Union
+from typing import Dict, List, Mapping, Sequence, Set, Tuple, Union
 
 from ollama._utils import _get_json_type, convert_function_to_tool
 
 
 def test_json_type_conversion():
   # Test basic types
-  assert _get_json_type(str) == 'string'
-  assert _get_json_type(int) == 'integer'
   assert _get_json_type(List) == 'array'
   assert _get_json_type(Dict) == 'object'
 
-  # Test Optional
-  assert _get_json_type(Optional[str]) == 'string'
-
 
 def test_function_to_tool_conversion():
-  def add_numbers(x: int, y: Union[int, None] = None) -> int:  # Changed Optional to Union
+  def add_numbers(x: int, y: Union[int, None] = None) -> int:
     """Add two numbers together.
     Args:
         x (integer): The first number
@@ -41,13 +36,13 @@ def test_function_with_all_typing_types():
   def all_types(
     x: int,
     y: str,
-    z: Sequence[int],
+    z: Sequence,
     w: Mapping[str, int],
     d: Dict[str, int],
     s: Set[int],
     t: Tuple[int, str],
     l: List[int],  # noqa: E741
-    o: Union[int, None],  # Changed Optional to Union
+    o: Union[int, None],
   ) -> Union[Mapping[str, int], str, None]:
     """
     A function with all types.