File size: 3,804 Bytes
cb88b52 c9ae383 cb88b52 c9ae383 cb88b52 c9ae383 cb88b52 c9ae383 cb88b52 c9ae383 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
# 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() |