Update utility.py

utility.py

@@ -4,6 +4,13 @@ from datetime import datetime
 import openai
 from google.cloud import firestore
 from dotenv import load_dotenv
+from pandasai import SmartDatalake
+from pandasai import Agent
+from pandasai.responses.response_parser import  ResponseParser
+from langchain_experimental.agents import create_pandas_dataframe_agent
+import pandas as pd
+from pandasai.llm import OpenAI
+
 
 # Make API connection
 load_dotenv()
@@ -19,6 +26,36 @@ client = openai.OpenAI(
     base_url="https://api.sambanova.ai/v1",
 )
 
+llm = OpenAI(
+    api_key=os.environ.get("sambanova_api_key"),
+    base_url="https://api.sambanova.ai/v1",
+)
+
+
+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'])
+      
 # Generate AI response from user input
 def generateResponse(prompt,model='Meta-Llama-3.1-70B-Instruct'):
   #----- Call API to classify and extract relevant transaction information
@@ -218,6 +255,23 @@ 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')
+
+    inventory_list = [doc.to_dict() for doc in inventory_ref.stream()]
+    sales_list = [doc.to_dict() for doc in sales_ref.stream()]
+
+    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)
+
+    return response
+
 def create_inventory(user_phone, transaction_data):
     for transaction in transaction_data:
         item_name = transaction['details']['item'] # assumes unique item name
@@ -289,6 +343,8 @@ def delete_transaction(user_phone, transaction_id):
     print("Transaction deleted successfully!")
 
 
+
+
 # Example usage
 # response_text = """
 # The information provided indicates that you want to **create/record** a new transaction.