Update app.py

app.py

@@ -7,10 +7,14 @@ from langchain.document_loaders import DataFrameLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
-from langchain.chains import RetrievalQA
+from langchain.chains import RetrievalQAWithSourcesChain
 from langchain import HuggingFacePipeline
 from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 
+# Custom avatars
+USER_AVATAR = "https://raw.githubusercontent.com/achilela/vila_fofoka_analysis/9904d9a0d445ab0488cf7395cb863cce7621d897/USER_AVATAR.png"
+BOT_AVATAR = "https://raw.githubusercontent.com/achilela/vila_fofoka_analysis/991f4c6e4e1dc7a8e24876ca5aae5228bcdb4dba/Ataliba_Avatar.jpg"
+
 def preprocess_excel(file_path: str) -> pd.DataFrame:
     df_raw = pd.read_excel(file_path, sheet_name='Data Base', header=None)
     df = df_raw.iloc[4:].copy()
@@ -48,15 +52,20 @@ def create_qa_pipeline(vectorstore):
     llm = HuggingFacePipeline(pipeline=gen_pipeline)
 
     retriever = vectorstore.as_retriever()
-    qa = RetrievalQA.from_chain_type(llm=llm, retriever=retriever, chain_type="stuff", return_source_documents=False)
+    qa = RetrievalQAWithSourcesChain.from_llm(llm=llm, retriever=retriever)
     return qa
 
+# Streamlit app layout
 st.set_page_config(page_title="Excel-Aware RAG Chatbot", layout="wide")
 st.title("📊 Excel-Aware RAG Chatbot (Professional QA)")
 
 with st.sidebar:
     uploaded_file = st.file_uploader("Upload your Excel file (.xlsx or .xlsm with 'Data Base' sheet)", type=["xlsx", "xlsm"])
 
+# Persistent chat history
+if "chat_history" not in st.session_state:
+    st.session_state.chat_history = []
+
 if uploaded_file is not None:
     with st.spinner("Processing and indexing your Excel sheet..."):
         with tempfile.NamedTemporaryFile(delete=False, suffix=".xlsm") as tmp_file:
@@ -68,28 +77,36 @@ if uploaded_file is not None:
             vectorstore = build_vectorstore_from_dataframe(cleaned_df)
             qa = create_qa_pipeline(vectorstore)
             st.success("✅ File processed and chatbot ready! Ask your questions below.")
+        except Exception as e:
+            st.error(f"❌ Error processing file: {e}")
+        finally:
+            os.remove(tmp_path)
 
-            if "chat_history" not in st.session_state:
-                st.session_state.chat_history = []
+    # Show previous messages
+    for message in st.session_state.chat_history:
+        st.chat_message(message["role"], avatar=USER_AVATAR if message["role"] == "user" else BOT_AVATAR).markdown(message["content"])
 
-            with st.chat_message("assistant"):
-                st.markdown("How can I help you with the inspection data?")
+    user_prompt = st.chat_input("Ask about inspections, delays, backlogs...")
 
-            user_prompt = st.chat_input("Ask a question like 'How many backlog items are marked Yes?' or 'List overdue inspections'.")
+    if user_prompt:
+        st.session_state.chat_history.append({"role": "user", "content": user_prompt})
+        st.chat_message("user", avatar=USER_AVATAR).markdown(user_prompt)
 
-            if user_prompt:
-                st.chat_message("user").markdown(user_prompt)
-                with st.chat_message("assistant"):
-                    with st.spinner("Thinking..."):
-                        try:
-                            answer = qa.run(user_prompt)
-                            st.markdown(f"**Answer:** {answer}")
-                            st.session_state.chat_history.append((user_prompt, answer))
-                        except Exception as e:
-                            st.error(f"Error: {e}")
-        except Exception as e:
-            st.error(f"Error processing file: {e}")
-        finally:
-            os.remove(tmp_path)
+        with st.chat_message("assistant", avatar=BOT_AVATAR):
+            with st.spinner("Searching and generating..."):
+                try:
+                    response = qa.run(user_prompt)
+                    final_response = response['answer']
+                    placeholder = st.empty()
+                    streamed = ""
+
+                    for word in final_response.split():
+                        streamed += word + " "
+                        placeholder.markdown(streamed + "▌")
+
+                    placeholder.markdown(f"**{final_response.strip()}**")
+                    st.session_state.chat_history.append({"role": "assistant", "content": final_response})
+                except Exception as e:
+                    st.error(f"❌ Error: {e}")
 else:
-    st.info("Upload a file to get started.")
+    st.info("Upload a file on the left to get started.")