From ea9d19049c07ce9741c51f588fab94174054c102 Mon Sep 17 00:00:00 2001 From: ChihYu Yeh Date: Tue, 30 Jul 2024 21:16:05 +0800 Subject: [PATCH] fix --- .../pipelines/sql_explanation/generation.py | 27 ++++++++++--------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/wren-ai-service/src/pipelines/sql_explanation/generation.py b/wren-ai-service/src/pipelines/sql_explanation/generation.py index ba65d3131..dd1ba73bf 100644 --- 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", [] + ), ) ]