Update app.py

app.py

@@ -5,44 +5,124 @@ import os
 from pathlib import Path
 from components.chat import display_chat_interface
 from utils.database import (
-    display_vector_store_info,
-    handle_document_upload,
     create_connection,
     create_tables,
-    get_collections
+    handle_document_upload,
+    verify_vector_store
+)
+from utils.collections import (
+    create_collection_tables,
+    get_collections,
+    get_collection_documents,
+    create_collection,
+    add_document_to_collection
 )
 import time
 from threading import Lock
 
-# Create a lock for database operations
+# Create locks for thread-safe operations
 conn_lock = Lock()
+vectorstore_lock = Lock()
 
 def initialize_database():
-    """Initialize database connection and tables."""
-    # [Previous initialize_database code remains the same]
-    pass
+    """Initialize database connection and create necessary tables."""
+    try:
+        with conn_lock:
+            if 'db_conn' not in st.session_state:
+                data_dir = "data"
+                if not os.path.exists(data_dir):
+                    os.makedirs(data_dir)
+
+                db_path = os.path.join(data_dir, 'rfp_analysis.db')
+
+                try:
+                    with open(db_path, 'a') as f:
+                        pass
+                except IOError as e:
+                    st.error(f"Failed to access database: {e}")
+                    return False
+
+                conn = create_connection(db_path)
+
+                if conn is not None:
+                    # Create base tables
+                    create_tables(conn)
+                    # Create collection-specific tables
+                    create_collection_tables(conn)
+                    st.session_state.db_conn = conn
+                    return True
+                else:
+                    return False
+            else:
+                return True
+
+    except Exception as e:
+        st.error(f"Database initialization error: {e}")
+        return False
+
+def initialize_session_state():
+    """Initialize all session state variables."""
+    if 'chat_ready' not in st.session_state:
+        st.session_state.chat_ready = False
+    if 'current_collection' not in st.session_state:
+        st.session_state.current_collection = None
+    if 'processed_files' not in st.session_state:
+        st.session_state.processed_files = []
+    if 'show_upload_dialog' not in st.session_state:
+        st.session_state.show_upload_dialog = False
+    if 'vector_stores' not in st.session_state:
+        st.session_state.vector_stores = {}
+
+def display_top_bar():
+    """Display the application's top navigation bar."""
+    col1, col2, col3 = st.columns([1, 3, 1])
+    
+    with col1:
+        if os.path.exists("img/logo.png"):
+            st.image("img/logo.png", width=100)
+        else:
+            st.warning("Logo not found at img/logo.png")
+    
+    with col2:
+        st.title("Synaptyx.AI - RFP Analysis Agent")
+    
+    with col3:
+        if st.session_state.current_collection:
+            st.info(f"📁 Active Collection: {st.session_state.current_collection['name']}")
 
-def display_sidebar():
-    """Display the document management sidebar."""
+def display_collection_sidebar():
+    """Display the collection management sidebar."""
     with st.sidebar:
         st.title("📚 Document Manager")
+
+        # Collection Selection
+        st.header("Collections")
+        collections = get_collections(st.session_state.db_conn)
         
+        # Create New Collection button
+        if st.button("➕ Create New Collection"):
+            st.session_state.show_collection_dialog = True
+
         # Collection selector
-        st.subheader("Select Collection")
-        collections = get_collections()
-        collection_names = ["All Documents"] + [c['name'] for c in collections]
+        collection_options = ["All Documents"] + [c['name'] for c in collections]
         selected_collection = st.selectbox(
-            "Working Collection",
-            collection_names,
+            "Select Collection",
+            collection_options,
             index=0,
             key="collection_selector"
         )
-        
-        # Update session state
-        st.session_state.current_collection = selected_collection if selected_collection != "All Documents" else None
-        
+
+        # Update current collection in session state
+        if selected_collection != "All Documents":
+            st.session_state.current_collection = next(
+                (c for c in collections if c['name'] == selected_collection),
+                None
+            )
+        else:
+            st.session_state.current_collection = None
+
         # Upload Section
-        st.header("Upload Documents", anchor=False)
+        st.header("Upload Documents")
         uploaded_files = st.file_uploader(
             "Upload PDF documents",
             type=['pdf'],
@@ -50,115 +130,112 @@ def display_sidebar():
             help="Limit 200MB per file • PDF"
         )
 
-        # Process uploads
         if uploaded_files:
-            if 'processed_files' not in st.session_state or uploaded_files != st.session_state.processed_files:
+            if st.button("Process Documents"):
                 with st.spinner("Processing documents..."):
-                    # Get collection ID if a collection is selected
-                    collection_id = None
-                    if st.session_state.current_collection:
-                        collection_id = next(
-                            (c['id'] for c in collections if c['name'] == st.session_state.current_collection),
-                            None
-                        )
-                    
-                    handle_document_upload(uploaded_files, collection_id)
-                    st.session_state.processed_files = uploaded_files
-                    st.session_state.chat_ready = True
-                    time.sleep(1)
-                    st.rerun()
-
-        # Knowledge Base Status
-        if st.session_state.get('vector_store'):
-            st.success("✅ Documents ready for analysis")
-            display_vector_store_info()
+                    collection_id = st.session_state.current_collection['id'] if st.session_state.current_collection else None
+                    success = handle_document_upload(uploaded_files, collection_id)
+                    if success:
+                        st.success("Documents processed successfully!")
+                        st.session_state.chat_ready = True
+                        st.rerun()
+
+        # Display current collection info
+        if st.session_state.current_collection:
+            st.header("Collection Info")
+            collection = st.session_state.current_collection
+            documents = get_collection_documents(st.session_state.db_conn, collection['id'])
+            
+            st.markdown(f"""
+            **Collection:** {collection['name']}  
+            **Documents:** {len(documents)}  
+            **Created:** {collection['created_at']}
+            """)
+
+            if documents:
+                st.subheader("Documents in Collection")
+                for doc in documents:
+                    st.write(f"📄 {doc['name']}")
+
+def display_collection_dialog():
+    """Display the create collection dialog."""
+    with st.sidebar:
+        st.subheader("Create New Collection")
+        name = st.text_input("Collection Name")
+        description = st.text_area("Description")
         
-        # Document List
-        if uploaded_files:
-            st.subheader("📑 Uploaded Documents")
-            for doc in uploaded_files:
-                st.write(f"• {doc.name}")
-
-def display_top_navigation():
-    """Display the top navigation bar."""
-    with st.container():
-        cols = st.columns([1, 3, 1])
-        with cols[0]:
-            if os.path.exists("img/logo.png"):
-                st.image("img/logo.png", width=100)
-            else:
-                st.error("Logo not found")
+        col1, col2 = st.columns(2)
+        with col1:
+            if st.button("Create"):
+                if name:
+                    if create_collection(st.session_state.db_conn, name, description):
+                        st.success(f"Collection '{name}' created!")
+                        st.session_state.show_collection_dialog = False
+                        st.rerun()
+                else:
+                    st.error("Please enter a collection name")
         
-        with cols[1]:
-            st.title("Synaptyx.AI - RFP Analysis Agent")
+        with col2:
+            if st.button("Cancel"):
+                st.session_state.show_collection_dialog = False
+                st.rerun()
+
+def display_chat_area():
+    """Display the main chat interface."""
+    if not st.session_state.chat_ready:
+        st.title("🤖 Welcome to SYNAPTYX")
+        st.markdown("### Your AI-powered RFP Analysis Assistant")
         
-        with cols[2]:
-            if st.session_state.get('current_collection'):
-                st.write(f"📁 Collection: {st.session_state.current_collection}")
-
-def display_welcome_screen():
-    """Display the welcome screen for new users."""
-    st.title("🤖 SYNAPTYX - RFP Analysis Agent")
-    st.markdown("### Welcome to your AI-powered RFP analysis assistant!")
-    
-    col1, col2 = st.columns(2)
-    with col1:
-        st.markdown("""
-        #### Getting Started:
-        1. Upload your RFP documents using the sidebar
-        2. Wait for the processing to complete
-        3. Start chatting with your documents!
-        """)
-    
-    with col2:
-        st.markdown("#### Example Questions:")
-        examples = [
-            "📊 Summarize the main points of the document",
-            "📝 Draft a 'Why Us' section based on the document",
-            "🎯 Extract key success metrics and outcomes",
-            "💡 What are the innovative solutions mentioned?",
-            "🤝 Analyze the partnership benefits described"
-        ]
-        for example in examples:
-            st.markdown(f"• {example}")
+        col1, col2 = st.columns(2)
+        with col1:
+            st.markdown("""
+            #### Getting Started:
+            1. Select or create a collection in the sidebar
+            2. Upload your RFP documents
+            3. Start analyzing with AI!
+            """)
+        
+        with col2:
+            st.markdown("#### Example Questions:")
+            examples = [
+                "📊 Summarize the main points of the document",
+                "📝 Draft a 'Why Us' section based on the document",
+                "🎯 Extract key success metrics and outcomes",
+                "💡 What are the innovative solutions mentioned?",
+                "🤝 Analyze the partnership benefits described"
+            ]
+            for example in examples:
+                st.markdown(f"• {example}")
+    else:
+        display_chat_interface()
 
 def main():
-    # Set up the page configuration
+    # Configure the page
     st.set_page_config(
         layout="wide",
         page_title="SYNAPTYX - RFP Analysis Agent",
         initial_sidebar_state="expanded"
     )
 
-    # Initialize database
+    # Initialize database and session state
     if not initialize_database():
         st.error("Failed to initialize database. Please contact support.")
         return
 
-    # Initialize session state
-    if 'chat_ready' not in st.session_state:
-        st.session_state.chat_ready = False
-    if 'current_collection' not in st.session_state:
-        st.session_state.current_collection = None
+    initialize_session_state()
 
-    # Display sidebar (maintain existing functionality)
-    display_sidebar()
+    # Display the top navigation bar
+    display_top_bar()
 
-    # Main content area
-    display_top_navigation()
+    # Display the collection sidebar
+    display_collection_sidebar()
 
-    # Tabs for different sections
-    tab1, tab2 = st.tabs(["Chat", "Document Store"])
-    
-    with tab1:
-        if not st.session_state.chat_ready:
-            display_welcome_screen()
-        else:
-            display_chat_interface()
-    
-    with tab2:
-        from pages.document_store import display_document_store
-        display_document_store()
+    # Show collection creation dialog if triggered
+    if st.session_state.get('show_collection_dialog', False):
+        display_collection_dialog()
+
+    # Main chat area
+    display_chat_area()
 
 if __name__ == "__main__":
     main()