Spaces:
Build error
Build error
| # pages/document_store.py | |
| import streamlit as st | |
| from utils.database import ( | |
| create_collection, | |
| get_collections, | |
| get_collection_documents, | |
| delete_collection, | |
| add_document_to_collection, | |
| process_and_store_document | |
| ) | |
| 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", key="new_collection_name") | |
| description = st.text_area("Description", key="new_collection_desc") | |
| with col2: | |
| if st.button("Create Collection", key="create_collection_btn") 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: | |
| # Process and store the document | |
| 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 | |
| # Add this if you need it for the page | |
| def main(): | |
| display_collection_manager() | |
| if __name__ == "__main__": | |
| main() |