Update app.py

app.py

@@ -6,8 +6,36 @@ import pandas as pd
 from datetime import datetime
 from st_on_hover_tabs import on_hover_tabs
 import streamlit as st
+import os
+from pandasai.llm import GoogleGemini
+from pandasai import SmartDataframe
+from pandasai.responses.response_parser import  ResponseParser
+
 st.set_page_config(layout="wide")
 
+class StreamLitResponse(ResponseParser):
+        def __init__(self,context) -> None:
+              super().__init__(context)
+        def format_dataframe(self,result):
+               st.dataframe(result['value'])
+               return
+        def format_plot(self,result):
+               st.image(result['value'])
+               return
+        def format_other(self, result):
+               st.write(result['value'])
+               return
+
+gemini_api_key = os.environ['Gemini']
+
+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
+
+
+
 # Initialize Firebase app
 if not firebase_admin._apps:
     cred = credentials.Certificate("ecomplaintbook-firebase-adminsdk-4q5bo-1f83312d02.json")
@@ -15,6 +43,14 @@ if not firebase_admin._apps:
 
 db = firestore.client()
 
+complaints_ref = db.collection('complaints')
+
+complaints_list = []
+for doc in complaints_ref.stream():
+    a = doc.to_dict()
+    complaints_list.append(a)
+
+complaints_df = pd.DataFrame(complaints_list)
 
 # Function to fetch user profile from Firebase
 def fetch_user_profile_from_firebase(user_id):
@@ -82,14 +118,7 @@ def main():
         st.title("Navigation Bar")
         st.header("Complaints")
     # Then query to list all users
-        complaints_ref = db.collection('complaints')
-
-        complaints_list = []
-        for doc in complaints_ref.stream():
-            a = doc.to_dict()
-            complaints_list.append(a)
 
-        complaints_df = pd.DataFrame(complaints_list)
 
         st.dataframe(complaints_df)
 
@@ -98,8 +127,19 @@ def main():
         st.write('Name of option is {}'.format(tabs))
 
     elif tabs == 'Chat':
-        st.title("Chat Powered by Gemini")
+        st.title("eThekwini Chat Powered by Gemini")
         st.write('Name of option is {}'.format(tabs))
+    
+
+        # Display the data
+        with st.expander("Preview"):
+            st.write(df.head())
+
+        # Plot the data
+        user_input = st.text_input("Type your message here",placeholder="Ask me about complaints from eThekwini residents")
+        if user_input:
+            answer = generateResponse(dataFrame=complaints_df,prompt=user_input)
+            st.write(answer)