wip

Author: ParthSareen

ollama/_utils.py

@@ -1,19 +1,24 @@
-from __future__ import annotations
+# from __future__ import annotations
 
-from types import UnionType
-from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Union, get_args, get_origin
+from typing import Any, Callable, List, Mapping, Optional, Sequence, Union, get_args, get_origin
 from ollama._types import Tool
 
+# Map both the type and the type reference to the same JSON type
 PYTHON_TO_JSON_TYPES = {
-  str: 'string',
   int: 'integer',
+  'int': 'integer',
+  str: 'string',
+  'str': 'string',
   float: 'number',
+  'float': 'number',
   bool: 'boolean',
+  'bool': 'boolean',
   list: 'array',
+  'list': 'array',
   dict: 'object',
-  List: 'array',
-  Dict: 'object',
-  None: 'null',
+  'dict': 'object',
+  type(None): 'null',
+  'None': 'null',
 }
 
 """
@@ -26,36 +31,39 @@
 
 
 def _get_json_type(python_type: Any) -> str | List[str]:
+  # Print for debugging
+
   # Handle Optional types (Union[type, None] and type | None)
   origin = get_origin(python_type)
-  if origin is UnionType or origin is Union:
+  if origin is type(int | str) or origin is Union:
     args = get_args(python_type)
     # Filter out None/NoneType from union args
     non_none_args = [arg for arg in args if arg not in (None, type(None))]
-    non_none_args_types = [PYTHON_TO_JSON_TYPES[arg] for arg in non_none_args]
     if non_none_args:
       if len(non_none_args) == 1:
         return _get_json_type(non_none_args[0])
       else:
-        return non_none_args_types
-
+        return [_get_json_type(arg) for arg in non_none_args]
     return 'null'
 
-  # Get basic type mapping
+  # Direct type check first (like int, str, etc.)
   if python_type in PYTHON_TO_JSON_TYPES:
     return PYTHON_TO_JSON_TYPES[python_type]
 
   # Handle typing.List, typing.Dict etc.
-  if origin in PYTHON_TO_JSON_TYPES:
-    return PYTHON_TO_JSON_TYPES[origin]
+  if origin is not None:
+    return PYTHON_TO_JSON_TYPES.get(origin, 'string')
+
+  # Default to string if type is unknown
+  raise ValueError(f'Unknown type: {python_type}')
+
 
-  # TODO: Default to string if type is unknown - define beahvior
-  return 'string'
+#   return 'string'
 
 
 def _is_optional_type(python_type: Any) -> bool:
   origin = get_origin(python_type)
-  if origin is UnionType or origin is Union:
+  if origin is type(int | str) or origin is Union:
     args = get_args(python_type)
     return any(arg in (None, type(None)) for arg in args)
   return False