update sql explanation results

Canner · Jul 10, 2024 · 6eca8a9 · 6eca8a9
1 parent d6a6f70
commit 6eca8a9
Showing 1 changed file with 17 additions and 8 deletions.
diff --git a/wren-ai-service/src/pipelines/sql_explanation/generation.py b/wren-ai-service/src/pipelines/sql_explanation/generation.py
@@ -269,14 +269,23 @@ def run(
                         "selectItems" in preprocessed_sql_analysis_results
                         and "selectItems" in sql_explanation_results
                     ):
-                        sql_analysis_result_for_select_items = (
-                            preprocessed_sql_analysis_results["selectItems"][
-                                "withFunctionCallOrMathematicalOperation"
-                            ]
-                            + preprocessed_sql_analysis_results["selectItems"][
-                                "withoutFunctionCallOrMathematicalOperation"
-                            ]
-                        )
+                        sql_analysis_result_for_select_items = [
+                            {
+                                **select_item,
+                                "type": "withFunctionCallOrMathematicalOperation",
+                            }
+                            for select_item in preprocessed_sql_analysis_results[
+                                "selectItems"
+                            ]["withFunctionCallOrMathematicalOperation"]
+                        ] + [
+                            {
+                                **select_item,
+                                "type": "withoutFunctionCallOrMathematicalOperation",
+                            }
+                            for select_item in preprocessed_sql_analysis_results[
+                                "selectItems"
+                            ]["withoutFunctionCallOrMathematicalOperation"]
+                        ]
 
                         for (
                             select_item,