Update main.py

main.py

@@ -12,7 +12,8 @@ import pandas as pd
 from datetime import datetime
 import os
 from pandasai.llm import GoogleGemini
-from pandasai import SmartDataframe
+from pandasai import SmartDataframe, SmartDatalake
+
 from pandasai.responses.response_parser import  ResponseParser
 import matplotlib.pyplot as plt
 from wordcloud import WordCloud
@@ -30,7 +31,7 @@ load_dotenv()
 app = Flask(__name__)
 cors = CORS(app)
 
-class StreamLitResponse(ResponseParser):
+class FlaskResponse(ResponseParser):
         def __init__(self,context) -> None:
               super().__init__(context)
         def format_dataframe(self,result):
@@ -51,12 +52,9 @@ def home():
     return "Hello Qx!"
 
 
-def generateResponse(dataFrame,prompt):
-        llm = GoogleGemini(api_key=gemini_api_key)
-        pandas_agent = SmartDataframe(dataFrame,config={"llm":llm, "response_parser":StreamLitResponse})
-        answer = pandas_agent.chat(prompt)
-        return answer
 
+llm = GoogleGemini(api_key=gemini_api_key)
+        
 # Initialize Firebase app
 if not firebase_admin._apps:
    
@@ -65,43 +63,41 @@ if not firebase_admin._apps:
 
 db = firestore.client()
 
-inventory_ref = db.collection('inventory')
 
 
-sales_ref = db.collection('sales')
 
-inventory_list = []
-for doc in inventory_ref.stream():
-    a = doc.to_dict()
-    inventory_list.append(a)
+@app.route("/predict", methods=["POST"])
+@cross_origin()
+def bot():
 
-sales_list = []
-for doc in sales_ref.stream():
-    a = doc.to_dict()
-    sales_list.append(a)
 
-inventory_df = pd.DataFrame(inventory_list)
-sales_df = pd.DataFrame(sales_list)
+    user_id = request.json.get("user_id")
+    user_question = request.json.get("user_question")
+    load_dotenv()
+    inventory_ref = db.collection("users").document(user_id).collection('inventory')
 
 
+    sales_ref = db.collection("users").document(user_id).collection('sales')
+
+    inventory_list = []
+    for doc in inventory_ref.stream():
+        a = doc.to_dict()
+        inventory_list.append(a)
+
+    sales_list = []
+    for doc in sales_ref.stream():
+        a = doc.to_dict()
+        sales_list.append(a)
+
+    inventory_df = pd.DataFrame(inventory_list)
+    sales_df = pd.DataFrame(sales_list)
+
+    lake = SmartDatalake([inventory_df, sales_df], config={"llm":llm, "response_parser":FlaskResponse})
+    response = lake.chat(user_question)
 
-@app.route("/predict", methods=["POST"])
-@cross_origin()
-def bot():
-    load_dotenv()
-    #
-    json_table = request.json.get("json_table")
-    user_question = request.json.get("user_question")
-    #data = request.get_json(force=True)TRye
-    #print(req_body)
-    #data = eval(req_body)
-    #json_table = data["json_table"]
-    #user_question = data["user_question"]
-    #print(json_table)
     print(user_question)
-    data = eval(str(json_table))
-    df = pd.DataFrame(data)
-    print(list(df))
+    
+   
     return jsonify(response)