Update utility.py

utility.py

@@ -272,20 +272,31 @@ def parse_multiple_transactions(response_text):
 
     return transactions
 
-def read_datalake(user_phone, user_question):
-    inventory_ref = db.collection("users").document(user_phone).collection("inventory")
-
-    sales_ref = db.collection("users").document(user_phone).collection('sales')
+def read_datalake(user_phone, user_question, transaction_type="inventory"):
+    # transaction_type = transaction_data[0]['transaction_type'].lower()
+    if transaction_type in ('purchases', 'purchases', 'inventory', 'stock'):
+        doc_ref = db.collection("users").document(user_phone).collection("inventory")
+    elif transaction_type in ('sales', 'sale'):
+        doc_ref = db.collection("users").document(user_phone).collection('sales')
+    else: 
+        return None
+        
+    # inventory_ref = db.collection("users").document(user_phone).collection("inventory")
+    # sales_ref = db.collection("users").document(user_phone).collection('sales')
 
-    inventory_list = [doc.to_dict() for doc in inventory_ref.stream()]
-    sales_list = [doc.to_dict() for doc in sales_ref.stream()]
+    # inventory_list = [doc.to_dict() for doc in inventory_ref.stream()]
+    # sales_list = [doc.to_dict() for doc in sales_ref.stream()]
+    
+    doc_list = [doc.to_dict() for doc in doc_ref.stream()]
+    
+    # inventory_df = pd.DataFrame(inventory_list)
+    # sales_df = pd.DataFrame(sales_list)
     
-    inventory_df = pd.DataFrame(inventory_list)
-    # inventory_df = inventory_df[inventory_df.columns not in ['created_at', 'intent', 'transaction_type']]
-    sales_df = pd.DataFrame(sales_list)
+    df = pd.DataFrame(doc_list)
     # return f"{inventory_df} \n\n {sales_df}"
 
-    lake = SmartDatalake([inventory_df, sales_df], config={"llm": llm, "custom_whitelisted_dependencies":["ast"], "response_parser": FlaskResponse, "enable_cache": False, "save_logs": False})
+    # lake = SmartDatalake([inventory_df, sales_df], config={"llm": llm, "custom_whitelisted_dependencies":["ast"], "response_parser": FlaskResponse, "enable_cache": False, "save_logs": False})
+    lake = SmartDatalake([df], config={"llm": llm, "custom_whitelisted_dependencies":["ast"], "response_parser": FlaskResponse, "enable_cache": False, "save_logs": False})
     response = lake.chat(user_question)
     # response = inventory_df