[DH-5619] Fix SQL generation, it was returning the intermediate steps

Dataherald · Mar 19, 2024 · 399c830 · 399c830
1 parent 7906f03
commit 399c830
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/dataherald/sql_generator/__init__.py b/dataherald/sql_generator/__init__.py
@@ -109,14 +109,12 @@ def extract_query_from_intermediate_steps(
         for step in intermediate_steps:
             action = step[0]
             if type(action) == AgentAction and action.tool == "SqlDbQuery":
-                sql_query = self.format_sql_query(action.tool_input)
-                if "SELECT" in sql_query.upper():
+                if "SELECT" in self.format_sql_query(action.tool_input).upper():
                     sql_query = self.remove_markdown(sql_query)
         if sql_query == "":
             for step in intermediate_steps:
                 action = step[0]
-                sql_query = action.tool_input
-                if "SELECT" in sql_query.upper():
+                if "SELECT" in action.tool_input.upper():
                     sql_query = self.remove_markdown(sql_query)
         return sql_query