DH5610/fix the intermediate steps

Dataherald · Mar 16, 2024 · 56df513 · 56df513
1 parent 1205d8a
commit 56df513
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/dataherald/sql_generator/dataherald_finetuning_agent.py b/dataherald/sql_generator/dataherald_finetuning_agent.py
@@ -585,7 +585,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)

diff --git a/dataherald/sql_generator/dataherald_sqlagent.py b/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)