[DH-5540] Update data report (Dataherald#487)

qqq-tech · May 7, 2024 · 89cd1a5 · 89cd1a5
1 parent d652130
commit 89cd1a5
Showing 1 changed file with 18 additions and 7 deletions.
diff --git a/apps/ai/server/database/scripts/data_report.py b/apps/ai/server/database/scripts/data_report.py
@@ -76,6 +76,10 @@ def create_csv_file(
                     "original_sql_query": original_response["sql"],
                     "original_sql_status": original_response["status"],
                     "latest_sql_status": final_response["status"],
+                    "source": prompt["metadata"]["dh_internal"].get("source"),
+                    "query_status": prompt["metadata"]["dh_internal"].get(
+                        "generation_status"
+                    ),
                     "final_confidence": (
                         None
                         if len(sql_generations) == 1
@@ -89,15 +93,22 @@ def create_csv_file(
                     ),
                     "was_the_original_correct": (
                         True
-                        if prompt["metadata"]["dh_internal"]["generation_status"]
-                        == "VERIFIED"
-                        and (
-                            len(sql_generations) == 1
-                            or (
-                                len(sql_generations) > 1
-                                and original_response["sql"] == final_response["sql"]
+                        if (
+                            prompt["metadata"]["dh_internal"]["generation_status"]
+                            == "VERIFIED"
+                            and (
+                                len(sql_generations) == 1
+                                or (
+                                    len(sql_generations) > 1
+                                    and original_response["sql"]
+                                    == final_response["sql"]
+                                )
                             )
                         )
+                        or (
+                            prompt["metadata"]["dh_internal"].get("source") == "SLACK"
+                            and len(sql_generations) == 1
+                        )
                         else False
                     ),
                     "question_asked_date": str(original_response["created_at"]),