Merge pull request #744 from madeindjs/fix-kg-description

feat(writer): include description in `GraphTool`
writer · Feb 20, 2025 · 29ad187 · 29ad187
2 parents f72eccd + 1dfd8b5
commit 29ad187
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/src/writer/ai.py b/src/writer/ai.py
@@ -88,6 +88,7 @@ class Tool(TypedDict, total=False):
 class GraphTool(Tool):
     graph_ids: List[str]
     subqueries: bool
+    description: Optional[str]
 
 
 class FunctionToolParameterMeta(TypedDict):
@@ -1419,7 +1420,8 @@ def validate_graph_ids(graph_ids: List[str]) -> bool:
                                 "graph_ids": tool_instance["graph_ids"],
                                 "subqueries": tool_instance.get(
                                     "subqueries", False
-                                    )
+                                    ),
+                                "description": tool_instance.get("description", None)
                             }
                         }
                     )

diff --git a/src/writer/blocks/writerchat.py b/src/writer/blocks/writerchat.py
@@ -90,10 +90,12 @@ def run(self):
                         parameters=tool_raw.get("parameters")
                     )
                 elif tool_type == "graph":
-                    tool = {
-                        "type": "graph",
-                        "graph_ids": tool_raw.get("graph_ids")
-                    }
+                    tool = writer.ai.GraphTool(
+                        type="graph",
+                        graph_ids=tool_raw.get("graph_ids"),
+                        subqueries=False,
+                        description=tool_name,
+                    )
                 else:
                     continue
                 tools.append(tool)