Update app.py

app.py

@@ -1,114 +1,52 @@
 import streamlit as st
 import backend
 from langchain.schema import HumanMessage, AIMessage
-import time
 from datetime import datetime
 
-def initialize_session_state():
-    """Initialize session state variables"""
-    if 'chat_history' not in st.session_state:
-        st.session_state.chat_history = []
-    if 'current_chat' not in st.session_state:
-        st.session_state.current_chat = False
-    if 'selected_docs' not in st.session_state:
-        st.session_state.selected_docs = []
-    if 'qa_system' not in st.session_state:
-        st.session_state.qa_system = None
-    if 'documents_initialized' not in st.session_state:
-        st.session_state.documents_initialized = False
-    if 'should_rerun' not in st.session_state:
-        st.session_state.should_rerun = False
-    if 'upload_complete' not in st.session_state:
-        st.session_state.upload_complete = False
+# Initialize session state
+if 'chat_history' not in st.session_state:
+    st.session_state.chat_history = []
+if 'current_chat' not in st.session_state:
+    st.session_state.current_chat = False
+if 'selected_docs' not in st.session_state:
+    st.session_state.selected_docs = []
+if 'qa_system' not in st.session_state:
+    st.session_state.qa_system = None
+if 'documents_initialized' not in st.session_state:
+    st.session_state.documents_initialized = False
+
+def handle_doc_select():
+    st.session_state.current_chat = True
+    st.session_state.chat_history = []
+
+def handle_start_chat():
+    st.session_state.current_chat = True
+    st.session_state.chat_history = []
 
 def display_chat_interface():
-    """Display the chat interface"""
     st.markdown("### 💬 RFP Analysis Chat")
     
     # Display chat history
-    chat_container = st.container()
-    with chat_container:
-        for message in st.session_state.chat_history:
-            if isinstance(message, HumanMessage):
-                st.markdown(f"🧑‍💼 **You**: {message.content}")
-            else:
-                st.markdown(f"🤖 **Assistant**: {message.content}")
+    for message in st.session_state.chat_history:
+        if isinstance(message, HumanMessage):
+            st.markdown(f"🧑‍💼 **You**: {message.content}")
+        else:
+            st.markdown(f"🤖 **Assistant**: {message.content}")
     
     # Chat input
-    user_query = st.text_input(
-        "Ask about the RFPs:",
-        placeholder="e.g., What are the key requirements?",
-        key="chat_input"
-    )
-    
-    if user_query:
-        # Add user message to chat history
-        st.session_state.chat_history.append(HumanMessage(content=user_query))
-        
+    if prompt := st.chat_input("Ask about the RFPs..."):
+        st.session_state.chat_history.append(HumanMessage(content=prompt))
         try:
             with st.spinner("Analyzing documents..."):
                 response = st.session_state.qa_system.invoke({
-                    "input": user_query,
+                    "input": prompt,
                     "chat_history": st.session_state.chat_history
                 })
-            
-            # Add AI message to chat history
             st.session_state.chat_history.append(AIMessage(content=response["output"]))
-            
-            # Force refresh to show new messages
-            st.rerun()
-            
         except Exception as e:
             st.error(f"Error generating response: {e}")
 
-def display_document_stats(conn):
-    """Display document statistics"""
-    try:
-        cursor = conn.cursor()
-        cursor.execute("SELECT COUNT(*) FROM documents")
-        total_documents = cursor.fetchone()[0]
-        
-        cursor.execute("SELECT COUNT(*) FROM queries")
-        total_queries = cursor.fetchone()[0]
-        
-        cursor.execute("SELECT COUNT(*) FROM annotations")
-        total_annotations = cursor.fetchone()[0]
-        
-        # Create three columns for metrics
-        col1, col2, col3 = st.columns(3)
-        
-        with col1:
-            st.metric("📚 Documents", total_documents)
-        with col2:
-            st.metric("❓ Queries", total_queries)
-        with col3:
-            st.metric("📝 Annotations", total_annotations)
-            
-    except Exception as e:
-        st.error(f"Error retrieving statistics: {e}")
-
-def handle_document_upload(conn, uploaded_documents):
-    """Handle document upload process"""
-    if uploaded_documents:
-        with st.spinner(f"Processing {len(uploaded_documents)} documents..."):
-            all_texts, document_names, document_pages = backend.upload_and_parse_documents(uploaded_documents)
-            
-            if all_texts:
-                try:
-                    with conn:
-                        for doc, doc_name in zip(all_texts, document_names):
-                            conn.execute(
-                                "INSERT INTO documents (name, content) VALUES (?, ?)",
-                                (doc_name, doc)
-                            )
-                    st.success(f"Successfully uploaded {len(uploaded_documents)} documents!", icon="✅")
-                    time.sleep(1)  # Give users time to see the success message
-                    st.rerun()  # Updated from experimental_rerun
-                except Exception as e:
-                    st.error(f"Error saving documents to database: {e}")
-
 def display_knowledge_base(conn):
-    """Display and manage knowledge base"""
     st.markdown("### 📚 Knowledge Base")
     
     if conn is not None:
@@ -120,34 +58,26 @@ def display_knowledge_base(conn):
             if documents_in_db:
                 st.markdown("#### Available Documents")
                 
-                # Create a container for the document list
-                doc_container = st.container()
-                with doc_container:
-                    for doc_id, name, upload_date in documents_in_db:
-                        col1, col2, col3 = st.columns([3, 2, 1])
+                for doc_id, name, upload_date in documents_in_db:
+                    col1, col2 = st.columns([3, 1])
+                    
+                    with col1:
+                        selected = st.checkbox(
+                            name,
+                            value=doc_id in st.session_state.selected_docs,
+                            key=f"doc_{doc_id}",
+                            on_change=handle_doc_select
+                        )
                         
-                        # Document name and checkbox
-                        with col1:
-                            selected = st.checkbox(
-                                name,
-                                value=doc_id in st.session_state.selected_docs,
-                                key=f"doc_{doc_id}"
-                            )
-                            if selected and doc_id not in st.session_state.selected_docs:
-                                st.session_state.selected_docs.append(doc_id)
-                            elif not selected and doc_id in st.session_state.selected_docs:
-                                st.session_state.selected_docs.remove(doc_id)
-                        
-                        # Upload date
-                        with col2:
-                            upload_date = datetime.strptime(upload_date, '%Y-%m-%d %H:%M:%S')
-                            st.text(upload_date.strftime('%Y-%m-%d %H:%M'))
-                        
-                        # Document actions
-                        with col3:
-                            st.button("📝", key=f"annotate_{doc_id}", help="Add annotations")
+                        if selected and doc_id not in st.session_state.selected_docs:
+                            st.session_state.selected_docs.append(doc_id)
+                        elif not selected and doc_id in st.session_state.selected_docs:
+                            st.session_state.selected_docs.remove(doc_id)
+                    
+                    with col2:
+                        upload_date = datetime.strptime(upload_date, '%Y-%m-%d %H:%M:%S')
+                        st.text(upload_date.strftime('%Y-%m-%d'))
                 
-                # Initialize selected documents
                 if st.session_state.selected_docs and not st.session_state.documents_initialized:
                     with st.spinner("Initializing document analysis..."):
                         selected_documents = []
@@ -159,10 +89,10 @@ def display_knowledge_base(conn):
                                 (doc_id,)
                             )
                             result = cursor.fetchone()
-                            selected_documents.append(result[0])
-                            selected_doc_names.append(result[1])
+                            if result:
+                                selected_documents.append(result[0])
+                                selected_doc_names.append(result[1])
                         
-                        # Initialize FAISS and QA system
                         embeddings = backend.get_embeddings_model()
                         if embeddings:
                             vector_store = backend.initialize_faiss(
@@ -173,14 +103,11 @@ def display_knowledge_base(conn):
                             if vector_store:
                                 st.session_state.qa_system = backend.initialize_qa_system(vector_store)
                                 st.session_state.documents_initialized = True
-                                st.success("Documents initialized successfully!")
-                                
-                # Start Chat button
+                
                 if st.session_state.selected_docs:
-                    if st.button("🚀 Start New Chat", use_container_width=True):
-                        st.session_state.current_chat = True
-                        st.session_state.chat_history = []
-                        st.rerun()  # Updated from experimental_rerun
+                    st.button("🚀 Start New Chat", 
+                             on_click=handle_start_chat,
+                             use_container_width=True)
             else:
                 st.info("No documents in the knowledge base. Upload some documents to get started!")
                 
@@ -188,17 +115,12 @@ def display_knowledge_base(conn):
             st.error(f"Error accessing knowledge base: {e}")
 
 def main():
-    # Page configuration
     st.set_page_config(
         page_title="SYNAPTYX - RFP Analysis Agent",
         page_icon="🤖",
         layout="wide"
     )
     
-    # Initialize session state
-    initialize_session_state()
-    
-    # Main title and description
     st.title("🤖 SYNAPTYX - RFP Analysis Agent")
     st.markdown(
         """
@@ -209,7 +131,6 @@ def main():
         unsafe_allow_html=True
     )
     
-    # Database initialization
     database = "rfp_agent.db"
     conn = backend.create_connection(database)
     if conn is not None:
@@ -218,29 +139,34 @@ def main():
         st.error("Error! Cannot create the database connection.")
         return
     
-    # Create two columns for layout
     col1, col2 = st.columns([1, 2])
     
     with col1:
-        # Document upload section
         st.markdown("### 📤 Upload Documents")
-        uploaded_documents = st.file_uploader(
+        uploaded_files = st.file_uploader(
             "Upload PDF documents",
             type="pdf",
-            accept_multiple_files=True,
-            help="Select one or more PDF files to upload"
+            accept_multiple_files=True
         )
-        handle_document_upload(conn, uploaded_documents)
         
-        # Display document statistics
-        st.markdown("### 📊 Statistics")
-        display_document_stats(conn)
+        if uploaded_files:
+            with st.spinner("Processing documents..."):
+                all_texts, document_names, document_pages = backend.upload_and_parse_documents(uploaded_files)
+                if all_texts:
+                    try:
+                        with conn:
+                            for doc, doc_name in zip(all_texts, document_names):
+                                conn.execute(
+                                    "INSERT INTO documents (name, content) VALUES (?, ?)",
+                                    (doc_name, doc)
+                                )
+                        st.success(f"Successfully uploaded {len(uploaded_files)} documents!")
+                    except Exception as e:
+                        st.error(f"Error saving documents to database: {e}")
         
-        # Knowledge base management
         display_knowledge_base(conn)
     
     with col2:
-        # Chat interface
         if st.session_state.current_chat and st.session_state.qa_system:
             display_chat_interface()
         else:
@@ -252,13 +178,6 @@ def main():
                 1. Upload your RFP documents using the upload section
                 2. Select the documents you want to analyze from the knowledge base
                 3. Click "Start New Chat" to begin asking questions
-                
-                Features:
-                - 📚 Document Management
-                - 💬 Intelligent Chat Interface
-                - 📝 Document Annotations
-                - 📊 Analytics Dashboard
-                - 🔍 Advanced Search Capabilities
                 """
             )