main.py

import google.generativeai as palm
import pandas as pd
import os
import io
from flask import Flask, request, jsonify
from flask_cors import CORS, cross_origin
import pandas as pd
import firebase_admin
from firebase_admin import credentials, firestore, auth
import requests
import pandas as pd
from datetime import datetime
import os
from pandasai.llm import GoogleGemini
from pandasai import SmartDataframe, SmartDatalake

from pandasai.responses.response_parser import  ResponseParser
import matplotlib.pyplot as plt
from wordcloud import WordCloud
import random

from dotenv import load_dotenv
import json


from dotenv import load_dotenv

load_dotenv()


app = Flask(__name__)
cors = CORS(app)

class FlaskResponse(ResponseParser):
  def __init__(self, context) -> None:
      super().__init__(context)

  def format_dataframe(self, result):
      return result['value'].to_html()

  def format_plot(self, result):
        # Save the plot using savefig
      try:
          
          img_path = result['value']
      
        
      except ValueError:
          img_path = str(result['value'])
          print("value error!", img_path)
        
      print("response_class_path:", img_path)
      return img_path

  def format_other(self, result):
      return str(result['value'])

gemini_api_key = os.environ['Gemini']

@app.route("/", methods=["GET"])
def home():
   
    return "Hello Qx!"



llm = GoogleGemini(api_key=gemini_api_key)
        
# Initialize Firebase app
if not firebase_admin._apps:
   
    cred = credentials.Certificate("quant-app-99d09-firebase-adminsdk-6prb1-37f34e1c91.json")
    firebase_admin.initialize_app(cred)

db = firestore.client()




@app.route("/predict", methods=["POST"])
@cross_origin()
def bot():


    user_id = request.json.get("user_id")
    user_question = request.json.get("user_question")
    load_dotenv()
    inventory_ref = db.collection("test_users").document(user_id).collection('inventory')


    sales_ref = db.collection("test_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, "enable_cache": False, "save_logs":False})
    response = lake.chat(user_question)
    print(user_question)

    resp = str(response)
    
    return jsonify(resp)


if __name__ == "__main__":
    app.run(debug=True,host="0.0.0.0", port=7860)