Update app.py

app.py

@@ -10,10 +10,9 @@ import pandas as pd
 import io
 import time
 
-from langchain.document_loaders import UnstructuredFileLoader
+from langchain_community.document_loaders import UnstructuredFileLoader
 from typing import List, Dict, Tuple
 from langchain_openai import AzureChatOpenAI,AzureOpenAIEmbeddings
-from langchain.vectorstores import FAISS
 from langchain.text_splitter import CharacterTextSplitter
 
 class PDFExtract:
@@ -261,11 +260,13 @@ if st.session_state.file_processed:
                 with st.spinner("Thinking..."):
                     response = chain.invoke({"input": question, 'lease_queries': lease_flow})
                     answer = response['answer']
+                    source = response['context']
                 
                 # Add AI response to chat
                 st.session_state.messages.append({"role": "assistant", "content": answer})
                 with st.chat_message("assistant"):
-                    st.write(answer)
+                    st.write(source)
+                    st.success(answer)
                 
                 if "GASB" in answer:
                     st.info("Flow stopped due to GASB answer.")
@@ -292,11 +293,13 @@ if st.session_state.file_processed:
                     with st.spinner("Thinking..."):
                         response = chain.invoke({"input": question, 'lease_queries': lease_flow})
                         answer = response['answer']
+                        source = response['context']
                     
                     # Add AI response to chat
                     st.session_state.messages.append({"role": "assistant", "content": answer})
                     with st.chat_message("assistant"):
-                        st.write(answer)
+                        st.write(source)
+                        st.success(answer)
                     
                     if "GASB" in answer:
                         st.info("Flow stopped due to GASB answer.")