diff --git a/src/writer/ai.py b/src/writer/ai.py
index 71af28870..557f6003d 100644
--- a/src/writer/ai.py
+++ b/src/writer/ai.py
@@ -1210,7 +1210,7 @@ def _register_callable(
             self,
             callable_to_register: Callable,
             name: str,
-            parameters: Dict[str, Dict[str, str]]
+            parameters: Dict[str, FunctionToolParameterMeta]
     ):
         """
         Internal helper function to store a provided callable
@@ -1280,7 +1280,9 @@ def _prepare_tool(
         Internal helper function to process a tool instance
         into the required format.
         """
-        def validate_parameters(parameters: Dict[str, Dict[str, str]]) -> bool:
+        def validate_parameters(
+                parameters: Dict[str, FunctionToolParameterMeta]
+        ) -> bool:
             """
             Validates the `parameters` dictionary to ensure that each key
             is a parameter name, and each value is a dictionary containing