Update app.py

app.py

@@ -135,35 +135,38 @@ def handle_userinput(question, pdf_vectorstore, dfs):
     if pdf_vectorstore and st.session_state.conversation:
         response = st.session_state.conversation({"question": question})
         st.session_state.chat_history.append({"role": "user", "content": question})
-        assistant_response = response.get('answer')  # Safe access
+        assistant_response = response.get('answer')
 
-    elif dfs:
+        if assistant_response:  # Check if assistant_response is not None or empty
+            st.session_state.chat_history.append({"role": "assistant", "content": assistant_response})  # Directly add string
+
+        st.rerun()
+
+    elif dfs:  # PandasAI branch (handling DataFrames and plots)
         results = generateResponse(question, dfs)
         st.session_state.chat_history.append({"role": "user", "content": question})
-        # PandasAI's results are already a list, no need to wrap
-
-        # Process each result from PandasAI (or Langchain) in a consistent way:
-        for assistant_response in results: # Changed to loop through results
-            # Assistant Response could be text, dataframe or plot
-            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
-                    st.session_state.chat_history.append({"role": "assistant", "content": assistant_response})
+        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})
 
         st.rerun()
-        return # Exit early since we've processed the results
+        return  # Exit early after PandasAI processing
 
     else:
         st.write("Please upload and process your documents/data first.")