Update app.py

app.py

@@ -17,24 +17,17 @@ import io
 import base64
 
 class StreamLitResponse(ResponseParser):
-    def __init__(self, context):
-        super().__init__(context)
-        self.results = []
-
-    def format_dataframe(self, result):
-        self.results.append({"type": "dataframe", "value": result['value']})
-        return
-
-    def format_plot(self, result):
-        self.results.append({"type": "plot", "value": result['value']})
-        return
-
-    def format_other(self, result):
-        self.results.append({"type": "text", "value": result['value']})
-        return
-
-    def get_results(self):
-        return self.results
+        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
 
 
 load_dotenv()  # Load environment variables at the beginning
@@ -47,8 +40,7 @@ if not GOOGLE_API_KEY:
 
 def generateResponse(prompt, dfs):
     llm = GoogleGemini(api_key=GOOGLE_API_KEY)
-    response_parser = StreamLitResponse({})
-    pandas_agent = SmartDatalake(dfs, config={"llm": llm, "response_parser": response_parser})
+    pandas_agent = SmartDataframe(dataFrame,config={"llm":llm, "response_parser":StreamLitResponse}
     answer = pandas_agent.chat(prompt)
     return answer