fix(tools): preserve full type information in tool signatures (#865)

TimeToBuildBob · web-flow · commit 10ddeb6cb06f · 2025-11-21T13:50:15.000+01:00
* fix(tools): preserve full type information in tool signatures - Add support for typing.Union (not just types.UnionType) - Handle generic types with arguments (list[int], dict[str, int], etc.) - Special case for NoneType to display as 'None' - Add comprehensive test cases for generic and union types Fixes #349 * fix(tests): use modern union syntax for type annotations Use list[int] | None instead of Optional[list[int]] to satisfy ruff UP007 linting rule. Co-authored-by: Bob <bob@superuserlabs.org>
diff --git a/gptme/tools/base.py b/gptme/tools/base.py
@@ -17,7 +17,9 @@
     Literal,
     Protocol,
     TypeAlias,
+    Union,
     cast,
+    get_args,
     get_origin,
 )
 
@@ -129,14 +131,31 @@ class Parameter:
 
 # TODO: there must be a better way?
 def derive_type(t) -> str:
-    if get_origin(t) == Literal:
-        v = ", ".join(f'"{a}"' for a in t.__args__)
+    origin = get_origin(t)
+
+    # Handle Literal types
+    if origin == Literal:
+        v = ", ".join(f'"{a}"' for a in get_args(t))
         return f"Literal[{v}]"
-    elif get_origin(t) == types.UnionType:
-        v = ", ".join(derive_type(a) for a in t.__args__)
+
+    # Handle Union types (both typing.Union and types.UnionType)
+    if origin == Union or origin == types.UnionType:
+        v = ", ".join(derive_type(a) for a in get_args(t))
         return f"Union[{v}]"
-    else:
-        return t.__name__
+
+    # Handle other generic types (list[int], dict[str, int], etc.)
+    if origin is not None:
+        args = get_args(t)
+        if args:
+            type_args = ", ".join(derive_type(arg) for arg in args)
+            return f"{origin.__name__}[{type_args}]"
+
+    # Special case for NoneType
+    if t is type(None):
+        return "None"
+
+    # Fallback to type name
+    return t.__name__
 
 
 def callable_signature(func: Callable) -> str:
diff --git a/tests/test_tools_python.py b/tests/test_tools_python.py
@@ -44,3 +44,26 @@ def h(a: TestType) -> str:
         return str(a)
 
     assert callable_signature(h) == 'h(a: Literal["a", "b"]) -> str'
+
+    # Test generic types
+
+    def i(a: list[int]) -> str:
+        return str(a)
+
+    assert callable_signature(i) == "i(a: list[int]) -> str"
+
+    def j(a: list[int] | None) -> str:
+        return str(a)
+
+    assert callable_signature(j) == "j(a: Union[list[int], None]) -> str"
+
+    def k(a: dict[str, int]) -> str:
+        return str(a)
+
+    assert callable_signature(k) == "k(a: dict[str, int]) -> str"
+
+    # Test union types with | syntax
+    def m(a: int | str) -> str:
+        return str(a)
+
+    assert callable_signature(m) == "m(a: Union[int, str]) -> str"
