Update app.py

app.py

@@ -278,12 +278,8 @@ def display_header():
 def display_collection_management():
     """Display collection management interface."""
     st.header("📁 Collection Management")
-    
-    # Get existing collections
-    collections = get_collections(st.session_state.db_conn)
-    
     col1, col2 = st.columns([2, 1])
-    
+
     with col1:
         # Create new collection form
         with st.form("create_collection_form"):
@@ -291,68 +287,58 @@ def display_collection_management():
             name = st.text_input("Collection Name")
             description = st.text_area("Description")
             submit = st.form_submit_button("Create Collection", use_container_width=True)
-            
+
             if submit and name:
                 collection_id = create_collection(st.session_state.db_conn, name, description)
                 if collection_id:
                     st.success(f"Collection '{name}' created successfully!")
+                    st.session_state.current_collection_id = collection_id
                     st.rerun()
-    
-    with col2:
-        # Upload to existing collection
+
+        # Display existing collections
+        collections = get_collections(st.session_state.db_conn)
         if collections:
-            st.subheader("Add to Collection")
-            selected_collection = st.selectbox(
-                "Select Collection",
-                options=[(c["name"], c["id"]) for c in collections],
-                format_func=lambda x: x[0]
-            )
-            
-            uploaded_files = st.file_uploader(
-                "Upload Documents",
-                type=['pdf'],
-                accept_multiple_files=True,
-                key="collection_uploader"
-            )
-            
-            if uploaded_files:
-                collection_id = selected_collection[1]
-                with st.spinner("Processing documents..."):
-                    if handle_document_upload(uploaded_files, collection_id=collection_id):
-                        st.success("Documents added to collection successfully!")
-                        st.rerun()
-    
-    # Display existing collections
-    if collections:
-        st.markdown("### Existing Collections")
-        for collection in collections:
-            with st.expander(f"📁 {collection['name']} ({collection['doc_count']} documents)"):
-                col1, col2 = st.columns([3, 1])
-                
-                with col1:
-                    st.write(f"**Description:** {collection.get('description', 'No description')}")
-                    st.write(f"**Created:** {collection['created_at']}")
-                    
-                    # Display documents in collection
-                    docs = get_collection_documents(st.session_state.db_conn, collection['id'])
-                    if docs:
-                        st.write("**Documents:**")
-                        for doc in docs:
-                            st.write(f"- {doc['name']}")
-                            tags = get_document_tags(st.session_state.db_conn, doc['id'])
-                            if tags:
-                                st.write(f"  Tags: {', '.join(tags)}")
-                
-                with col2:
-                    st.button("Start Chat", key=f"chat_{collection['id']}", use_container_width=True)
-                    st.button("Upload Files", key=f"upload_{collection['id']}", use_container_width=True)
-                    if st.button("Delete", key=f"delete_{collection['id']}", use_container_width=True):
-                        if st.warning("Are you sure you want to delete this collection?"):
-                            if delete_collection(st.session_state.db_conn, collection['id']):
-                                st.success("Collection deleted successfully!")
+            st.markdown("### Existing Collections")
+            for collection in collections:
+                with st.expander(f"📁 {collection['name']} ({collection['doc_count']} documents)"):
+                    col1, col2 = st.columns([3, 1])
+                    with col1:
+                        st.write(f"**Description:** {collection.get('description', 'No description')}")
+                        st.write(f"**Created:** {collection['created_at']}")
+
+                        # Display documents in collection
+                        docs = get_collection_documents(st.session_state.db_conn, collection['id'])
+                        if docs:
+                            st.write("**Documents:**")
+                            for doc in docs:
+                                st.write(f"- {doc['name']}")
+                                tags = get_document_tags(st.session_state.db_conn, doc['id'])
+                                if tags:
+                                    st.write(f" Tags: {', '.join(tags)}")
+
+                    with col2:
+                        # Add documents to collection
+                        uploaded_files = st.file_uploader(
+                            "Add Documents",
+                            type=['pdf'],
+                            accept_multiple_files=True,
+                            key=f"collection_upload_{collection['id']}"
+                        )
+                        if uploaded_files:
+                            if handle_document_upload(uploaded_files, collection_id=collection['id']):
+                                st.success("Documents added successfully!")
                                 st.rerun()
 
+                        if st.button("Start Chat", key=f"chat_{collection['id']}", use_container_width=True):
+                            st.session_state.selected_collection = collection
+                            initialize_chat_system(collection['id'])
+                            st.rerun()
 
+                        if st.button("Delete Collection", key=f"delete_{collection['id']}", use_container_width=True):
+                            if st.warning("Are you sure you want to delete this collection?"):
+                                if delete_collection(st.session_state.db_conn, collection['id']):
+                                    st.success("Collection deleted successfully!")
+                                    st.rerun()
 def display_chat_interface():
     """Display the main chat interface with persistent storage."""
     st.header("💬 Ask your documents")