Update app.py

app.py

@@ -2,8 +2,6 @@ import streamlit as st
 import tempfile
 import base64
 import os
-import nltk
-nltk.download('punkt')
 from src.utils.ingest_text import create_vector_database
 from src.utils.ingest_image import extract_and_store_images
 from src.utils.text_qa import qa_bot
@@ -13,13 +11,20 @@ from langchain.memory import ConversationBufferWindowMemory
 from langchain_community.chat_message_histories import StreamlitChatMessageHistory
 from dotenv import load_dotenv
 
-
-
+# Setup
 nest_asyncio.apply()
 load_dotenv()
-
 st.set_page_config(layout='wide', page_title="InsightFusion Chat")
 
+# ---- Session State Initialization ----
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+if "chain" not in st.session_state:
+    st.session_state.chain = None
+if "image_vdb" not in st.session_state:
+    st.session_state.image_vdb = None
+
+# ---- Chat Memory Setup ----
 memory_storage = StreamlitChatMessageHistory(key="chat_messages")
 memory = ConversationBufferWindowMemory(
     memory_key="chat_history",
@@ -28,18 +33,20 @@ memory = ConversationBufferWindowMemory(
     k=3
 )
 
+# ---- Background Image Setup ----
 image_bg = r"data/pexels-andreea-ch-371539-1166644.jpg"
 
 def add_bg_from_local(image_file):
-    with open(image_file, "rb") as image_file:
-        encoded_string = base64.b64encode(image_file.read())
+    with open(image_file, "rb") as img:
+        encoded_string = base64.b64encode(img.read())
     st.markdown(f"""<style>.stApp {{
-        background-image: url(data:image/{"png"};base64,{encoded_string.decode()});
-        background-size: cover
+        background-image: url(data:image/png;base64,{encoded_string.decode()});
+        background-size: cover;
     }}</style>""", unsafe_allow_html=True)
 
 add_bg_from_local(image_bg)
 
+# ---- Title ----
 st.markdown("""
     <svg width="600" height="100">
         <text x="50%" y="50%" font-family="San serif" font-size="42px" fill="Black" text-anchor="middle" stroke="white"
@@ -48,6 +55,7 @@ st.markdown("""
     </svg>
 """, unsafe_allow_html=True)
 
+# ---- Utility ----
 def get_answer(query, chain):
     try:
         response = chain.invoke(query)
@@ -56,58 +64,56 @@ def get_answer(query, chain):
         st.error(f"Error in get_answer: {e}")
         return None
 
-uploaded_file = st.file_uploader("File upload", type="pdf")
+# ---- File Upload ----
 path = None
+uploaded_file = st.file_uploader("Upload a PDF file", type="pdf")
 
-# Handle uploaded file
 if uploaded_file is not None:
     temp_file_path = os.path.join("temp", uploaded_file.name)
     os.makedirs("temp", exist_ok=True)
     with open(temp_file_path, "wb") as f:
         f.write(uploaded_file.getbuffer())
     path = os.path.abspath(temp_file_path)
+    st.success("Document uploaded successfully!")
     st.write(f"File saved to: {path}")
-    st.write("Document uploaded successfully!")
 
-# Option to use a predefined demo PDF from pdf_resource folder
+# ---- Demo PDF Option ----
 st.markdown("### Or use a demo file:")
 if st.button("Use Demo PDF"):
-    demo_file_path = os.path.join("pdf_resource", "sample.pdf")  # Replace with actual demo file name
+    demo_file_path = os.path.join("pdf_resource", "sample.pdf")  # Replace with actual demo name
     if os.path.exists(demo_file_path):
         path = os.path.abspath(demo_file_path)
-        st.write(f"Using demo file: {path}")
-        st.success("Demo file loaded successfully!")
-
+        st.success(f"Using demo file: {path}")
         with st.spinner("Processing demo file..."):
             try:
                 client = create_vector_database(path)
                 image_vdb = extract_and_store_images(path)
                 chain = qa_bot(client)
-                st.session_state['chain'] = chain
-                st.session_state['image_vdb'] = image_vdb
-                st.success("Demo file processing complete.")
+                st.session_state.chain = chain
+                st.session_state.image_vdb = image_vdb
+                st.success("Demo file processed successfully.")
             except Exception as e:
                 st.error(f"Error processing demo PDF: {e}")
     else:
-        st.error("Demo file not found. Make sure 'pdf_resource/sample.pdf' exists.")
+        st.error("Demo PDF not found in pdf_resource/.")
 
-# Process uploaded file on button click
+# ---- Processing Button ----
 if st.button("Start Processing"):
     if path is not None:
-        with st.spinner("Processing"):
+        with st.spinner("Processing file..."):
             try:
                 client = create_vector_database(path)
                 image_vdb = extract_and_store_images(path)
                 chain = qa_bot(client)
-                st.session_state['chain'] = chain
-                st.session_state['image_vdb'] = image_vdb
-                st.success("Processing complete.")
+                st.session_state.chain = chain
+                st.session_state.image_vdb = image_vdb
+                st.success("File processed successfully.")
             except Exception as e:
                 st.error(f"Error during processing: {e}")
     else:
-        st.error("Please upload a file or use the demo before starting processing.")
+        st.error("Please upload a file or select a demo.")
 
-# Custom input background
+# ---- Style Customization ----
 st.markdown("""
     <style> 
     .stChatInputContainer > div {
@@ -116,48 +122,38 @@ st.markdown("""
     </style>
 """, unsafe_allow_html=True)
 
-# Chat logic
+# ---- Chat Interface ----
 if user_input := st.chat_input("User Input"):
-    if 'chain' in st.session_state and 'image_vdb' in st.session_state:
-        chain = st.session_state['chain']
-        image_vdb = st.session_state['image_vdb']
-
+    if st.session_state.chain and st.session_state.image_vdb:
         with st.chat_message("user"):
             st.markdown(user_input)
-        
+
         with st.spinner("Generating Response..."):
-            response = get_answer(user_input, chain)
+            response = get_answer(user_input, st.session_state.chain)
             if response:
                 with st.chat_message("assistant"):
                     st.markdown(response)
-                
-                memory.save_context(
-                    {"input": user_input},
-                    {"output": response}
-                )
+
+                memory.save_context({"input": user_input}, {"output": response})
 
                 st.session_state.messages.append({"role": "user", "content": user_input})
                 st.session_state.messages.append({"role": "assistant", "content": response})
 
                 try:
-                    query_and_print_results(image_vdb, user_input)
+                    query_and_print_results(st.session_state.image_vdb, user_input)
                 except Exception as e:
                     st.error(f"Error querying image database: {e}")
             else:
-                st.error("Failed to generate response.")
+                st.error("Failed to generate a response.")
     else:
-        st.error("Please start processing before entering user input.")
-
-# Initialize message state
-if "messages" not in st.session_state:
-    st.session_state.messages = []
+        st.error("Please process a file before chatting.")
 
-# Display message history
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.write(message["content"])
+# ---- Display Chat History ----
+for msg in st.session_state.messages:
+    with st.chat_message(msg["role"]):
+        st.write(msg["content"])
 
-# Display chat memory history (LangChain)
+# ---- Display Memory History (LangChain) ----
 for i, msg in enumerate(memory_storage.messages):
-    name = "user" if i % 2 == 0 else "assistant"
-    st.chat_message(name).markdown(msg.content)
+    role = "user" if i % 2 == 0 else "assistant"
+    st.chat_message(role).markdown(msg.content)