Update app.py

app.py

@@ -1,179 +1,262 @@
 import streamlit as st
 import backend
 from langchain.schema import HumanMessage, AIMessage
+import time
+from datetime import datetime
 
-# Streamlit App Interface
-def main():
-    st.title("**SYNAPTYX - RFP Analysis Agent**")
-    st.markdown(
-        "<h3 style='color: #1E3A8A;'>Upload RFP documents, provide a URL, search, and get intelligent answers.</h3>",
-        unsafe_allow_html=True,
-    )
+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
 
-    # Database Initialization
-    database = "rfp_agent.db"
-    conn = backend.create_connection(database)
-    if conn is not None:
-        backend.create_tables(conn)  # Use backend.create_tables(conn)
-    else:
-        st.error("Error! Cannot create the database connection.")
-
-    # Dashboard Overview Tab
-    st.sidebar.markdown(
-        "<h2 style='color: #1E3A8A;'>Dashboard Overview</h2>",
-        unsafe_allow_html=True,
+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}")
+    
+    # Chat input
+    user_query = st.text_input(
+        "Ask about the RFPs:",
+        placeholder="e.g., What are the key requirements?",
+        key="chat_input"
     )
-    if conn is not None:
+    
+    if user_query:
+        # Add user message to chat history
+        st.session_state.chat_history.append(HumanMessage(content=user_query))
+        
         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]
-
-            st.sidebar.write(f"Total Documents: {total_documents}")
-            st.sidebar.write(f"Total Queries: {total_queries}")
+            with st.spinner("Analyzing documents..."):
+                response = st.session_state.qa_system.invoke({
+                    "input": user_query,
+                    "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.experimental_rerun()
+            
         except Exception as e:
-            st.error(f"Error retrieving dashboard data: {e}")
+            st.error(f"Error generating response: {e}")
 
-    # Sidebar Knowledge Base Tab
-    st.sidebar.markdown(
-        "<h2 style='color: #1E3A8A;'>Knowledge Base</h2>", unsafe_allow_html=True
-    )
-    st.sidebar.markdown(
-        "<p style='color: #1E3A8A;'>View and select documents for search.</p>",
-        unsafe_allow_html=True,
-    )
+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}")
 
-    # Retrieve Documents from Database
+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.experimental_rerun()  # Refresh the page to update the knowledge base
+                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:
         try:
             cursor = conn.cursor()
-            cursor.execute("SELECT id, name FROM documents")
+            cursor.execute("SELECT id, name, upload_date FROM documents ORDER BY upload_date DESC")
             documents_in_db = cursor.fetchall()
-
+            
             if documents_in_db:
-                # Use st.multiselect instead of st.selectbox
-                selected_doc_ids = st.sidebar.multiselect(
-                    "Select documents to include in the search:",
-                    options=[doc[0] for doc in documents_in_db],
-                    format_func=lambda doc_id: next(
-                        doc[1] for doc in documents_in_db if doc[0] == doc_id
-                    ),
-                    default=[
-                        doc[0] for doc in documents_in_db
-                    ],  # Select all documents by default
-                )
-
-                if selected_doc_ids:
-                    selected_documents = []
-                    selected_doc_names = []  # Also keep track of the document names
-                    for doc_id in selected_doc_ids:
-                        cursor.execute(
-                            "SELECT content, name FROM documents WHERE id = ?",
-                            (doc_id,),
-                        )
-                        result = cursor.fetchone()
-                        selected_documents.append(result[0])
-                        selected_doc_names.append(result[1])  # Add the name to the list
-
-                    # Initialize FAISS and Store Embeddings for Selected Documents
-                    embeddings = backend.get_embeddings_model()
-                    if embeddings:
-                        st.sidebar.write(
-                            "Creating embeddings..."
-                        )  # Indicate embedding creation
-                        vector_store = backend.initialize_faiss(
-                            embeddings,
-                            selected_documents,
-                            selected_doc_names,
-                        )  # Use selected_doc_names here
-                        if vector_store:
-                            st.sidebar.success(
-                                "Embeddings for selected documents stored successfully.",
-                                icon="📁",
+                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])
+                        
+                        # Document name and checkbox
+                        with col1:
+                            selected = st.checkbox(
+                                name,
+                                value=doc_id in st.session_state.selected_docs,
+                                key=f"doc_{doc_id}"
                             )
-
-                            # Initialize QA System for Selected Documents
-                            qa_system = backend.initialize_qa_system(
-                                vector_store
+                            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")
+                
+                # Initialize selected documents
+                if st.session_state.selected_docs and not st.session_state.documents_initialized:
+                    with st.spinner("Initializing document analysis..."):
+                        selected_documents = []
+                        selected_doc_names = []
+                        
+                        for doc_id in st.session_state.selected_docs:
+                            cursor.execute(
+                                "SELECT content, name FROM documents WHERE id = ?",
+                                (doc_id,)
                             )
-                            if qa_system:
-                                chat_history = []  # Initialize chat history
-                                user_query = st.text_input(
-                                    "Enter your query about the RFPs:",
-                                    placeholder="e.g., What are the evaluation criteria?",
-                                    label_visibility="visible",
-                                )
-                                if user_query:
-                                    st.markdown(
-                                        "<p style='color: #1E3A8A;'>Retrieving answer...</p>",
-                                        unsafe_allow_html=True,
-                                    )
-                                    try:
-                                        # Add user message to chat history
-                                        chat_history.append(
-                                            HumanMessage(content=user_query)
-                                        )
-
-                                        # Use agent_executor.invoke() to run the agent with chat history
-                                        response = qa_system.invoke(
-                                            {
-                                                "input": user_query,
-                                                "chat_history": chat_history,
-                                            }
-                                        )
-
-                                        # Add AI message to chat history
-                                        chat_history.append(
-                                            AIMessage(content=response["output"])
-                                        )
-
-                                        st.markdown(
-                                            "<h4 style='color: #1E3A8A;'>Answer:</h4>",
-                                            unsafe_allow_html=True,
-                                        )
-                                        st.write(
-                                            response["output"]
-                                        )  # Access the answer text
-
-                                    except Exception as e:
-                                        st.error(f"Error generating response: {e}")
+                            result = cursor.fetchone()
+                            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(
+                                embeddings,
+                                selected_documents,
+                                selected_doc_names
+                            )
+                            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.experimental_rerun()
+            else:
+                st.info("No documents in the knowledge base. Upload some documents to get started!")
+                
         except Exception as e:
-            st.error(f"Error retrieving documents from database: {e}")
+            st.error(f"Error accessing knowledge base: {e}")
 
-    # Document Upload Section
-    st.markdown(
-        "<h2 style='color: #1E3A8A;'>Upload RFP Documents</h2>",
-        unsafe_allow_html=True,
+def main():
+    # Page configuration
+    st.set_page_config(
+        page_title="SYNAPTYX - RFP Analysis Agent",
+        page_icon="🤖",
+        layout="wide"
     )
-    uploaded_documents = st.file_uploader(
-        "Upload PDF documents", type="pdf", accept_multiple_files=True
+    
+    # Initialize session state
+    initialize_session_state()
+    
+    # Main title and description
+    st.title("🤖 SYNAPTYX - RFP Analysis Agent")
+    st.markdown(
+        """
+        <p style='color: #1E3A8A;'>
+        Upload RFP documents, analyze requirements, and get intelligent answers powered by AI.
+        </p>
+        """,
+        unsafe_allow_html=True
     )
-    if uploaded_documents:
-        st.write(f"Uploaded {len(uploaded_documents)} documents.")
-        (
-            all_texts,
-            document_names,
-            document_pages,
-        ) = backend.upload_and_parse_documents(uploaded_documents)
-        if all_texts:
-            # Store Documents in Database
-            if conn is not None:
-                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(
-                        "Documents uploaded and parsed successfully.", icon="✅"
-                    )
-                except Exception as e:
-                    st.error(f"Error saving documents to database: {e}")
-
+    
+    # Database initialization
+    database = "rfp_agent.db"
+    conn = backend.create_connection(database)
+    if conn is not None:
+        backend.create_tables(conn)
+    else:
+        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(
+            "Upload PDF documents",
+            type="pdf",
+            accept_multiple_files=True,
+            help="Select one or more PDF files to upload"
+        )
+        handle_document_upload(conn, uploaded_documents)
+        
+        # Display document statistics
+        st.markdown("### 📊 Statistics")
+        display_document_stats(conn)
+        
+        # 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:
+            st.markdown(
+                """
+                ### 👋 Welcome to SYNAPTYX!
+                
+                To get started:
+                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
+                """
+            )
 
 if __name__ == "__main__":
     main()