Add prompt, response, and messages to Firestore -> Snowflake DAG

astronomer · Apr 3, 2024 · b1f3a90 · b1f3a90
1 parent e8ee422
commit b1f3a90
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/airflow/dags/metrcis/load_firestore_to_snowflake.py b/airflow/dags/metrcis/load_firestore_to_snowflake.py
@@ -62,7 +62,10 @@ def load_request_data_from_firestore() -> list[tuple[str, int, bool, int]]:
             status = doc_dict.get("status") == "complete"
             sent_at = datetime.fromtimestamp(doc_dict.get("sent_at"))
             client = doc_dict.get("client")
-            rows.append((uuid, score, status, sent_at, client))
+            prompt = doc_dict.get("prompt")
+            response = doc_dict.get("response")
+            messages = doc_dict.get("messages")
+            rows.append((uuid, score, status, sent_at, client, prompt, response, messages))
 
         return rows
 
@@ -78,9 +81,9 @@ def write_request_data_to_snowflake(rows: list[tuple[str, int, bool, int]]) -> N
 
         insert_sql = f"""
             INSERT INTO
-                {METRICS_SNOWFLAKE_DB_DATABASE}.{METRICS_SNOWFLAKE_DB_SCHEMA}.request(uuid, score, success, created_at, client)
+                {METRICS_SNOWFLAKE_DB_DATABASE}.{METRICS_SNOWFLAKE_DB_SCHEMA}.request(uuid, score, success, created_at, client, prompt, response, messages)
             VALUES
-                (?, ?, ?, ?, ?)
+                (?, ?, ?, ?, ?, ?, ?, ?)
         """
         conn.cursor().executemany(insert_sql, rows)