Update streamlit_app.py

streamlit_app.py

@@ -13,27 +13,22 @@ from src.graph_builder.graph_builder import GraphBuilder
 
 # Page configuration
 st.set_page_config(
-    page_title="🤖 PDF Agentic RAG",
-    page_icon="📄",
+    page_title="🤖 PDF Agentic Chat",
+    page_icon="💬",
     layout="wide"
 )
 
-# Custom CSS for a cleaner look
-st.markdown("""
-    <style>
-    .stAlert { margin-top: 1rem; }
-    .stButton > button { width: 100%; border-radius: 5px; height: 3em; }
-    </style>
-""", unsafe_allow_html=True)
-
 def init_session_state():
-    """Initialize session state variables"""
+    """Initialize session state variables for chat history and system state"""
     if 'rag_system' not in st.session_state:
         st.session_state.rag_system = None
+    if 'messages' not in st.session_state:
+        # Initialize with a greeting
+        st.session_state.messages = [
+            {"role": "assistant", "content": "Hi! Upload some PDFs in the sidebar, and I'll help you analyze them."}
+        ]
     if 'processed_files' not in st.session_state:
         st.session_state.processed_files = []
-    if 'history' not in st.session_state:
-        st.session_state.history = []
 
 def process_new_documents(uploaded_files):
     """Processes uploaded PDFs and initializes/updates the RAG system"""
@@ -43,28 +38,24 @@ def process_new_documents(uploaded_files):
             chunk_overlap=Config.CHUNK_OVERLAP
         )
         
-        # 1. Process PDFs into chunks
-        # Assuming your DocumentProcessor has a method for uploaded files or local paths
-        # If not, you may need to save them to a temp directory first
         all_docs = []
         for uploaded_file in uploaded_files:
-            # Save temp file
+            # Save temp file for processing
             temp_path = Path(f"temp_{uploaded_file.name}")
             with open(temp_path, "wb") as f:
                 f.write(uploaded_file.getvalue())
             
-            # Process (Update this call based on your DocumentProcessor's actual method)
+            # Use the processor to extract text and chunk
             docs = doc_processor.process_pdf(str(temp_path)) 
             all_docs.extend(docs)
             
             # Cleanup temp file
             temp_path.unlink()
 
-        # 2. Initialize components
+        # Build the vector store and graph
         vector_store = VectorStore()
         vector_store.create_vectorstore(all_docs)
         
-        # 3. Build Graph
         graph_builder = GraphBuilder(
             retriever=vector_store.get_retriever(),
             llm=Config.get_llm()
@@ -81,75 +72,77 @@ def main():
     
     # --- Sidebar: Document Upload ---
     with st.sidebar:
-        st.title("📁 Document Management")
+        st.title("📁 Document Portal")
         uploaded_files = st.file_uploader(
-            "Upload PDF documents", 
+            "Upload PDFs", 
             type="pdf", 
             accept_multiple_files=True
         )
         
-        process_btn = st.button("🚀 Process Documents")
-        
-        if process_btn and uploaded_files:
-            with st.spinner("Analyzing PDFs and building index..."):
+        if st.button("🚀 Index Documents") and uploaded_files:
+            with st.spinner("Processing documents..."):
                 rag_system, num_chunks = process_new_documents(uploaded_files)
                 if rag_system:
                     st.session_state.rag_system = rag_system
                     st.session_state.processed_files = [f.name for f in uploaded_files]
-                    st.success(f"Indexed {len(uploaded_files)} files ({num_chunks} chunks)")
+                    st.success(f"Successfully indexed {num_chunks} chunks.")
+                    # Add a status message to chat
+                    st.session_state.messages.append({
+                        "role": "assistant", 
+                        "content": f"I've finished reading: {', '.join(st.session_state.processed_files)}. What would you like to know?"
+                    })
         
         if st.session_state.processed_files:
             st.markdown("---")
-            st.markdown("**Currently Loaded:**")
+            st.markdown("**Active Documents:**")
             for f in st.session_state.processed_files:
-                st.caption(f"✅ {f}")
+                st.caption(f"📄 {f}")
+            if st.button("Clear Chat History"):
+                st.session_state.messages = [{"role": "assistant", "content": "Chat history cleared. How can I help?"}]
+                st.rerun()
 
-    # --- Main UI: Search ---
-    st.title("🔍 Agentic RAG Search")
-    
-    if not st.session_state.rag_system:
-        st.info("👈 Please upload and process PDF documents in the sidebar to start searching.")
-        return
+    # --- Main Chat Interface ---
+    st.title("💬 PDF AI Assistant")
 
-    # Search interface
-    with st.container():
-        question = st.text_input("Ask a question about your documents:")
-        search_cols = st.columns([1, 4])
-        submit = search_cols[0].button("Search")
+    # Display chat messages from history on app rerun
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
 
-    if (submit or question) and question:
-        with st.spinner("Agent is thinking..."):
-            start_time = time.time()
-            
-            # Execute RAG pipeline
-            result = st.session_state.rag_system.run(question)
-            
-            elapsed_time = time.time() - start_time
-            
-            # Update History
-            st.session_state.history.append({
-                'question': question,
-                'answer': result['answer'],
-                'time': elapsed_time
-            })
-            
-            # Display results
-            st.markdown("### 💡 Answer")
-            st.write(result['answer'])
-            
-            with st.expander("📄 View Source Context"):
-                for i, doc in enumerate(result.get('retrieved_docs', []), 1):
-                    st.markdown(f"**Source {i}:**")
-                    st.info(doc.page_content)
+    # React to user input
+    if prompt := st.chat_input("Ask a question about your documents..."):
+        # Display user message in chat message container
+        st.chat_message("user").markdown(prompt)
+        # Add user message to chat history
+        st.session_state.messages.append({"role": "user", "content": prompt})
 
-    # --- History Section ---
-    if st.session_state.history:
-        st.markdown("---")
-        st.subheader("📜 Search History")
-        for item in reversed(st.session_state.history):
-            with st.expander(f"Q: {item['question']}"):
-                st.write(item['answer'])
-                st.caption(f"Response time: {item['time']:.2f}s")
+        # Generate response
+        if st.session_state.rag_system:
+            with st.chat_message("assistant"):
+                with st.spinner("Thinking..."):
+                    try:
+                        # Run the Agentic RAG pipeline
+                        result = st.session_state.rag_system.run(prompt)
+                        response = result['answer']
+                        
+                        # Display response
+                        st.markdown(response)
+                        
+                        # Optional: Show sources in an expader inside the bubble
+                        if result.get('retrieved_docs'):
+                            with st.expander("View Sources"):
+                                for i, doc in enumerate(result['retrieved_docs'], 1):
+                                    st.markdown(f"**Source {i}:**\n{doc.page_content[:500]}...")
+                        
+                        # Add assistant response to chat history
+                        st.session_state.messages.append({"role": "assistant", "content": response})
+                    except Exception as e:
+                        error_msg = f"I encountered an error: {str(e)}"
+                        st.error(error_msg)
+                        st.session_state.messages.append({"role": "assistant", "content": error_msg})
+        else:
+            with st.chat_message("assistant"):
+                st.warning("Please upload and index some PDFs in the sidebar first!")
 
 if __name__ == "__main__":
     main()