RFP_Analyzer_Agent_backup / components /collection_manager.py
cryogenic22's picture
Update components/collection_manager.py
04d1254 verified
# components/collection_manager.py
import streamlit as st # Import Streamlit here
from utils.database import (
get_all_documents,
create_collection,
add_document_to_collection,
get_collections,
search_documents,
get_collection_documents,
)
st.session_state.reinitialize_chat = True
def show_collection_creation_dialog():
"""Show dialog for creating a new collection."""
with st.form("create_collection"):
st.subheader("Create New Collection")
name = st.text_input("Collection Name", key="new_col_name")
description = st.text_area("Description", key="new_col_desc")
# Document selection
existing_docs = get_all_documents(st.session_state.db_conn)
if existing_docs:
st.subheader("Add Existing Documents")
selected_docs = st.multiselect(
"Select documents to add",
options=[doc['id'] for doc in existing_docs],
format_func=lambda x: next(doc['name'] for doc in existing_docs if doc['id'] == x)
)
submitted = st.form_submit_button("Create Collection")
if submitted and name:
collection_id = create_collection(st.session_state.db_conn, name, description)
if collection_id and selected_docs:
for doc_id in selected_docs:
add_document_to_collection(st.session_state.db_conn, doc_id, collection_id)
st.success(f"Collection '{name}' created!")
st.session_state.show_collection_dialog = False
st.rerun()
def enhanced_document_preview(doc):
"""Enhanced document preview with actions."""
col1, col2 = st.columns([3, 1])
with col1:
st.text_area(
"Preview",
value=doc['content'][:1000] + "..." if len(doc['content']) > 1000 else doc['content'],
height=200,
disabled=True
)
with col2:
st.button("πŸ—‘οΈ Remove", key=f"remove_{doc['id']}", use_container_width=True)
st.button("πŸ“‹ Copy", key=f"copy_{doc['id']}", use_container_width=True)
st.button("πŸ’¬ Chat", key=f"chat_{doc['id']}", use_container_width=True)
def display_enhanced_collections():
"""Enhanced collections management interface."""
col1, col2 = st.columns([1, 2])
with col1:
st.subheader("Collections")
# Search collections
search_query = st.text_input("πŸ” Search collections", key="collection_search")
# Create new collection button
if st.button("+ New Collection", type="primary", use_container_width=True):
st.session_state.show_collection_dialog = True
# Collection list with stats
collections = get_collections(st.session_state.db_conn)
if search_query:
collections = [c for c in collections
if search_query.lower() in c['name'].lower()]
st.divider()
for collection in collections:
with st.container():
col_button = st.button(
f"πŸ“ {collection['name']}",
key=f"col_{collection['id']}",
use_container_width=True
)
# Show collection stats
stats_cols = st.columns(3)
with stats_cols[0]:
st.caption(f"πŸ“„ {collection['doc_count']}")
with stats_cols[1]:
st.caption(f"πŸ“… {collection['created_at'][:10]}")
with stats_cols[2]:
st.caption("βš™οΈ Manage")
if col_button:
st.session_state.selected_collection = collection
st.rerun()
with col2:
if st.session_state.get('selected_collection'):
display_collection_content(st.session_state.selected_collection)
else:
st.info("Select a collection to view its contents")
def display_collection_content(collection):
"""Display enhanced collection contents."""
st.subheader(collection['name'])
# Collection actions
col1, col2, col3, col4 = st.columns(4)
with col1:
st.button("πŸ“₯ Add Documents", use_container_width=True)
with col2:
st.button("✏️ Edit Collection", use_container_width=True)
with col3:
st.button("πŸ“€ Export", use_container_width=True)
with col4:
st.button("πŸ—‘οΈ Delete", use_container_width=True)
# Document search within collection
search_query = st.text_input("πŸ” Search documents in collection")
# Filters
with st.expander("Filters"):
filter_cols = st.columns(3)
with filter_cols[0]:
st.date_input("Date Range")
with filter_cols[1]:
st.selectbox("Sort By", ["Name", "Date", "Size"])
with filter_cols[2]:
st.selectbox("View", ["List", "Grid"])
# Document list with enhanced display
documents = get_collection_documents(st.session_state.db_conn, collection['id'])
if search_query:
documents = [doc for doc in documents
if search_query.lower() in doc['name'].lower() or
search_query.lower() in doc.get('content', '').lower()]
# Display documents in a clean, modern format
for doc in documents:
with st.container():
doc_cols = st.columns([0.1, 2, 0.5])
with doc_cols[0]:
st.checkbox("", key=f"select_{doc['id']}")
with doc_cols[1]:
st.markdown(f"**{doc['name']}**")
st.caption(f"Uploaded: {doc.get('upload_date', 'N/A')}")
with doc_cols[2]:
st.button("Preview", key=f"preview_{doc['id']}", use_container_width=True)
st.button("Remove", key=f"remove_{doc['id']}", use_container_width=True)
# Document preview expander
if st.session_state.get(f"preview_{doc['id']}", False):
with st.expander("Document Preview", expanded=True):
st.text_area(
"Content",
value=doc['content'][:1000] + "..." if len(doc['content']) > 1000 else doc['content'],
height=200,
disabled=True
)
def show_document_preview(document):
"""Show document preview in a modal."""
with st.expander(f"Preview: {document['name']}", expanded=True):
st.text_area(
"Content",
value=document['content'][:1000] + "..." if len(document['content']) > 1000 else document['content'],
height=300,
disabled=True
)