qbapi / main.py
rairo's picture
Update main.py
30743bb verified
from langchain_google_genai import ChatGoogleGenerativeAI
import pandas as pd
import os
import io
from flask import Flask, request, jsonify
from flask_cors import CORS, cross_origin
import logging
from dotenv import load_dotenv
from pandasai import SmartDatalake
from pandasai import SmartDataframe
from pandasai.responses.response_parser import ResponseParser
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from datetime import datetime
import matplotlib.pyplot as plt
import google.generativeai as genai
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):
try:
img_path = result['value']
except ValueError:
img_path = str(result['value'])
print("ValueError:", img_path)
print("response_class_path:", img_path)
return img_path
def format_other(self, result):
return str(result['value'])
gemini_api_key = os.getenv('Gemini')
llm = ChatGoogleGenerativeAI(api_key=gemini_api_key, model='gemini-1.5-flash', temperature=0.1)
gemini_api_key = os.environ['Gemini']
genai.configure(api_key=gemini_api_key)
generation_config = {
"temperature": 0.2,
"top_p": 0.95,
"max_output_tokens": 5000,
}
model = genai.GenerativeModel(
model_name="gemini-2.0-flash-thinking-exp",
generation_config=generation_config,
)
# Endpoint for chat
@app.route("/chat", methods=["POST"])
@cross_origin()
def bot():
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))
pandas_agent = SmartDataframe(df,config={"llm":llm, "response_parser":FlaskResponse})
answer = pandas_agent.chat(user_question)
return jsonify(answer)
return answer
#df = df.rename(co
# Reports endpoint
@app.route("/report", methods=["POST"])
@cross_origin()
def marketing_rec():
json_data = request.json.get("json_data")
prompt = """
You are an expert business analyst. Analyze the following data and generate a comprehensive and insightful business report, including appropriate key perfomance indicators and recommendations.
data:
""" + str(json_data)
response = model.generate_content(prompt)
report = response.text
return jsonify(str(report))
if __name__ == "__main__":
app.run(debug=True, host="0.0.0.0", port=7860)