DH-4779/handling the bug with malformed tool input

Dataherald · Oct 3, 2023 · c66cc14 · c66cc14
1 parent e951afa
commit c66cc14
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/dataherald/sql_generator/dataherald_sqlagent.py b/dataherald/sql_generator/dataherald_sqlagent.py
@@ -313,8 +313,11 @@ def _run(
         tool_input: str,
         run_manager: CallbackManagerForToolRun | None = None,  # noqa: ARG002
     ) -> str:
-        schema, entity = tool_input.split(",")
-        table_name, column_name = schema.split("->")
+        try:
+            schema, entity = tool_input.split(",")
+            table_name, column_name = schema.split("->")
+        except ValueError:
+            return "Invalid input format, use following format: table_name -> column_name, entity (entity should be a string without ',')"
         search_pattern = f"%{entity.strip().lower()}%"
         meta = MetaData(bind=self.db.engine)
         table = sqlalchemy.Table(table_name.strip(), meta, autoload=True)