From 4e122e62295d1c98aa1f27c0b29ea0cd682ca4df Mon Sep 17 00:00:00 2001
From: mmikita95 <hello@mmikita.me>
Date: Thu, 5 Dec 2024 12:31:10 +0300
Subject: [PATCH] chore: add a TypedDict to parameters definition

---
 src/writer/ai.py | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/src/writer/ai.py b/src/writer/ai.py
index bbdbbfbb8..71af28870 100644
--- a/src/writer/ai.py
+++ b/src/writer/ai.py
@@ -90,17 +90,31 @@ class GraphTool(Tool):
     subqueries: bool
 
 
+class FunctionToolParameterMeta(TypedDict):
+    type: Union[
+        Literal["string"],
+        Literal["number"],
+        Literal["integer"],
+        Literal["float"],
+        Literal["boolean"],
+        Literal["array"],
+        Literal["object"],
+        Literal["null"]
+        ]
+    description: str
+
+
 class FunctionTool(Tool):
     callable: Callable
     name: str
     description: Optional[str]
-    parameters: Dict[str, Dict[str, str]]
+    parameters: Dict[str, FunctionToolParameterMeta]
 
 
 def create_function_tool(
     callable: Callable,
     name: str,
-    parameters: Optional[Dict[str, Dict[str, str]]] = None,
+    parameters: Optional[Dict[str, FunctionToolParameterMeta]] = None,
     description: Optional[str] = None
 ) -> FunctionTool:
     parameters = parameters or {}