Create pages/document_store.py

pages/document_store.py

@@ -0,0 +1,82 @@
+# pages/document_store.py
+
+import streamlit as st
+from utils.collections import *
+from utils.database import create_connection
+
+def display_collection_manager():
+    """Display the collection management interface."""
+    st.header("Collection Management")
+    
+    # Get database connection from session state
+    conn = st.session_state.db_conn
+    
+    # Create new collection section
+    with st.expander("Create New Collection", expanded=False):
+        col1, col2 = st.columns([3, 1])
+        with col1:
+            name = st.text_input("Collection Name")
+            description = st.text_area("Description")
+        with col2:
+            if st.button("Create Collection") and name:
+                if create_collection(conn, name, description):
+                    st.success(f"Collection '{name}' created!")
+                    st.rerun()
+
+    # Display existing collections
+    collections = get_collections(conn)
+    if collections:
+        st.subheader("Existing Collections")
+        
+        for collection in collections:
+            with st.expander(f"📁 {collection['name']} ({collection['doc_count']} documents)"):
+                # Collection details
+                st.write(f"**Description:** {collection['description']}")
+                st.write(f"**Created:** {collection['created_at']}")
+                
+                # Get documents in this collection
+                documents = get_collection_documents(conn, collection['id'])
+                
+                # Display documents
+                if documents:
+                    st.write("**Documents:**")
+                    for doc in documents:
+                        st.write(f"• {doc['name']}")
+                
+                # Collection management buttons
+                col1, col2, col3 = st.columns([1, 1, 1])
+                with col1:
+                    if st.button("Add Documents", key=f"add_{collection['id']}"):
+                        st.session_state.selected_collection = collection['id']
+                        st.session_state.show_upload = True
+                with col2:
+                    if st.button("Edit", key=f"edit_{collection['id']}"):
+                        st.session_state.editing_collection = collection['id']
+                with col3:
+                    if st.button("Delete", key=f"del_{collection['id']}"):
+                        if delete_collection(conn, collection['id']):
+                            st.success(f"Collection '{collection['name']}' deleted!")
+                            st.rerun()
+
+def handle_document_upload_with_collection(uploaded_files, collection_id):
+    """Handle document upload with collection assignment."""
+    conn = st.session_state.db_conn
+    
+    for uploaded_file in uploaded_files:
+        # Your existing document processing code here
+        document_id = process_and_store_document(uploaded_file)
+        
+        if document_id:
+            # Add to collection
+            if add_document_to_collection(conn, document_id, collection_id):
+                st.success(f"Added {uploaded_file.name} to collection!")
+            else:
+                st.error(f"Failed to add {uploaded_file.name} to collection!")
+
+# Initialize required session state
+if 'show_upload' not in st.session_state:
+    st.session_state.show_upload = False
+if 'selected_collection' not in st.session_state:
+    st.session_state.selected_collection = None
+if 'editing_collection' not in st.session_state:
+    st.session_state.editing_collection = None