Update app.py

app.py

@@ -46,10 +46,21 @@ if not GOOGLE_API_KEY:
 
 
 def generateResponse(prompt, dfs):
-        llm = GoogleGemini(api_key=GOOGLE_API_KEY)
-        pandas_agent = SmartDatalake(dfs,config={"llm":llm, "response_parser":StreamLitResponse})
-        answer = pandas_agent.chat(prompt)
-        return answer
+    llm = GoogleGemini(api_key=GOOGLE_API_KEY)
+    response_parser = StreamLitResponse({})
+    pandas_agent = SmartDatalake(dfs, config={"llm": llm, "response_parser": response_parser})
+    try:
+        pandas_agent.chat(prompt)
+        results = response_parser.get_results()
+        if not results:  # If results is empty
+            return {"type": "text", "value": "No results found for your query."}  # Return a default message as a dictionary
+        elif len(results) == 1:  # If results contains only one object, return the object directly
+            return results[0]
+        else:
+            return results  # Return the list of results
+    except Exception as e:  # Catch any exceptions during PandasAI processing
+        st.error(f"Error in PandasAI: {e}")  # Display the error in Streamlit
+        return {"type": "text", "value": f"An error occurred: {e}"} 
 
 
 # Processing pdfs
@@ -92,28 +103,27 @@ def handle_userinput(question, pdf_vectorstore, dfs):
 
         st.rerun()
 
-    elif dfs:  # PandasAI branch (handling DataFrames and plots)
-        results = generateResponse(question, dfs)
+    elif dfs:  # PandasAI
+        assistant_response = generateResponse(question, dfs)  # Get the single response
+
         st.session_state.chat_history.append({"role": "user", "content": question})
 
-        if results is not None:  # Check if results is NOT None
-            for assistant_response in results:
-                if assistant_response:  # Check if not None or empty
-                    if isinstance(assistant_response, dict) and 'value' in assistant_response:
-                        content_type = assistant_response.get('type')
-                        content_value = assistant_response.get('value')
-
-                        if content_type == "dataframe":
-                            st.session_state.chat_history.append({"role": "assistant", "content": "DataFrame"})
-                            st.session_state.chat_history.append({"role": "assistant", "dataframe": content_value})
-                        elif content_type == "plot":
-                            st.session_state.chat_history.append({"role": "assistant", "content": "Plot"})
-                            st.session_state.chat_history.append({"role": "assistant", "plot": content_value})
-                        else:  # Text or other
-                            st.session_state.chat_history.append({"role": "assistant", "content": assistant_response})
-
-                    else:  # Text or other
-                        st.session_state.chat_history.append({"role": "assistant", "content": assistant_response})
+        if assistant_response:  # Check if assistant_response is not None or empty
+            if isinstance(assistant_response, dict) and 'value' in assistant_response:
+                content_type = assistant_response.get('type')
+                content_value = assistant_response.get('value')
+
+                if content_type == "dataframe":
+                    st.session_state.chat_history.append({"role": "assistant", "content": "DataFrame"})
+                    st.session_state.chat_history.append({"role": "assistant", "dataframe": content_value})
+                elif content_type == "plot":
+                    st.session_state.chat_history.append({"role": "assistant", "content": "Plot"})
+                    st.session_state.chat_history.append({"role": "assistant", "plot": content_value})
+                else:  # Text or other
+                    st.session_state.chat_history.append({"role": "assistant", "content": assistant_response})
+
+            else:  # Text or other (including None if that's what it is)
+                st.session_state.chat_history.append({"role": "assistant", "content": assistant_response})
 
         st.rerun()
         return  # Exit early after PandasAI processing