[DH-5597] Fix sql-generation

Author: jcjc712

dataherald/sql_generator/__init__.py

@@ -60,7 +60,7 @@ def remove_markdown(self, query: str) -> str:
         matches = re.findall(pattern, query, re.DOTALL)
         if matches:
             return matches[0].strip()
-        return ""
+        return query
 
     @classmethod
     def get_upper_bound_limit(cls) -> int:
@@ -110,15 +110,14 @@ def extract_query_from_intermediate_steps(
             action = step[0]
             if type(action) == AgentAction and action.tool == "SqlDbQuery":
                 sql_query = self.format_sql_query(action.tool_input)
-                if "```sql" in sql_query:
+                if "SELECT" in sql_query.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 "```sql" in sql_query:
+                if "SELECT" in sql_query.upper():
                     sql_query = self.remove_markdown(sql_query)
-
         return sql_query
 
     @abstractmethod

dataherald/sql_generator/dataherald_sqlagent.py

@@ -730,7 +730,7 @@ def generate_response(
             sql_query = self.remove_markdown(result["output"])
         else:
             sql_query = self.extract_query_from_intermediate_steps(
-                result["intermediate"]
+                result["intermediate_steps"]
             )
         logger.info(f"cost: {str(cb.total_cost)} tokens: {str(cb.total_tokens)}")
         response.sql = replace_unprocessable_characters(sql_query)