fix

Canner · Aug 22, 2024 · ea9d190 · ea9d190
1 parent 5248794
commit ea9d190
Showing 1 changed file with 14 additions and 13 deletions.
diff --git a/wren-ai-service/src/pipelines/sql_explanation/generation.py b/wren-ai-service/src/pipelines/sql_explanation/generation.py
@@ -114,11 +114,12 @@ def _collect_relations(relation, result, cte_names, top_level: bool = True):
                 {
                     "values": {
                         "type": relation["type"],
-                        "criteria": relation["criteria"],
+                        "criteria": relation["criteria"]["expression"],
                         "exprSources": [
                             {
                                 "expression": expr_source["expression"],
                                 "sourceDataset": expr_source["sourceDataset"],
+                                "sourceColumn": expr_source["sourceColumn"],
                             }
                             for expr_source in relation["exprSources"]
                         ],
@@ -360,12 +361,12 @@ def run(
                                 },
                             }
                             for select_item, sql_explanation in zip(
-                                preprocessed_sql_analysis_results["selectItems"][
-                                    "withFunctionCallOrMathematicalOperation"
-                                ],
-                                sql_explanation_results["selectItems"][
-                                    "withFunctionCallOrMathematicalOperation"
-                                ],
+                                preprocessed_sql_analysis_results.get(
+                                    "selectItems", {}
+                                ).get("withFunctionCallOrMathematicalOperation", []),
+                                sql_explanation_results.get("selectItems", {}).get(
+                                    "withFunctionCallOrMathematicalOperation", []
+                                ),
                             )
                         ] + [
                             {
@@ -378,12 +379,12 @@ def run(
                                 },
                             }
                             for select_item, sql_explanation in zip(
-                                preprocessed_sql_analysis_results["selectItems"][
-                                    "withoutFunctionCallOrMathematicalOperation"
-                                ],
-                                sql_explanation_results["selectItems"][
-                                    "withoutFunctionCallOrMathematicalOperation"
-                                ],
+                                preprocessed_sql_analysis_results.get(
+                                    "selectItems", {}
+                                ).get("withoutFunctionCallOrMathematicalOperation", []),
+                                sql_explanation_results.get("selectItems", {}).get(
+                                    "withoutFunctionCallOrMathematicalOperation", []
+                                ),
                             )
                         ]