fix

Canner · Oct 2, 2024 · 4eda8c9 · 4eda8c9
1 parent 0b2f61c
commit 4eda8c9
Showing 1 changed file with 14 additions and 13 deletions.
diff --git a/wren-ai-service/src/pipelines/generation/sql_explanation.py b/wren-ai-service/src/pipelines/generation/sql_explanation.py
@@ -197,11 +197,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"]
                         ],
@@ -443,12 +444,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", []
+                                ),
                             )
                         ] + [
                             {
@@ -461,12 +462,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", []
+                                ),
                             )
                         ]