Update app.py

app.py

@@ -1,19 +1,17 @@
 import os
 import streamlit as st
 import tempfile
+from pymongo import MongoClient
 from datetime import datetime
 from pathlib import Path
-from pymongo import MongoClient
-from urllib.parse import quote_plus
 from document_chunker import DocumentChunker
-import time
+from urllib.parse import quote_plus
 
 # === MongoDB connection via Hugging Face secrets ===
 user = quote_plus(os.getenv("MONGO_USER"))
 password = quote_plus(os.getenv("MONGO_PASS"))
 cluster = os.getenv("MONGO_CLUSTER")
 db_name = os.environ.get("MONGO_DB", "grant_docs")
-
 mongo_uri = f"mongodb+srv://{user}:{password}@{cluster}/{db_name}?retryWrites=true&w=majority&tls=true&tlsAllowInvalidCertificates=true"
 client = MongoClient(mongo_uri, tls=True, tlsAllowInvalidCertificates=True, serverSelectionTimeoutMS=20000)
 db = client[db_name]
@@ -24,14 +22,15 @@ st.title("📄 Document Chunker & Uploader")
 
 with st.sidebar:
     st.header("Settings")
+
+    # Fetch collection names for dropdown
     try:
         existing_collections = db.list_collection_names()
         existing_collections.append("Create New Collection")
-        default_index = existing_collections.index("doc_chunks_cat") if "doc_chunks_cat" in existing_collections else 0
-        selected_collection = st.selectbox("Choose MongoDB Collection", existing_collections, index=default_index)
+        selected_collection = st.selectbox("Choose MongoDB Collection", existing_collections, index=existing_collections.index("default_collection") if "default_collection" in existing_collections else 0)
     except Exception as e:
         st.error(f"Failed to list collections: {e}")
-        selected_collection = "doc_chunks_cat"
+        selected_collection = "default_collection"
 
     if selected_collection == "Create New Collection":
         selected_collection = st.sidebar.text_input("Enter Collection Name:")
@@ -43,47 +42,35 @@ with st.sidebar:
 
 uploaded_file = st.file_uploader("Upload a DOCX, TXT, or PDF file", type=["docx", "txt", "pdf"])
 
-# === Store session state after upload ===
-if uploaded_file and "ready_to_process" not in st.session_state:
+if uploaded_file:
     temp_path = Path(tempfile.gettempdir()) / uploaded_file.name
     with open(temp_path, "wb") as f:
         f.write(uploaded_file.getbuffer())
 
-    st.session_state["uploaded_file_name"] = uploaded_file.name
-    st.session_state["collection_name"] = selected_collection
-    st.session_state["is_grant_app"] = is_grant_app
-    st.session_state["temp_path"] = str(temp_path)
-    st.session_state["ready_to_process"] = True
-    st.rerun()
-
-# === Process document ===
-if st.session_state.get("ready_to_process"):
-    file_name = st.session_state["uploaded_file_name"]
-    collection_name = st.session_state["collection_name"]
-    is_grant_app = st.session_state["is_grant_app"]
-    temp_path = st.session_state["temp_path"]
+    st.success(f"Uploaded `{uploaded_file.name}`")
 
-    st.success(f"Uploaded `{file_name}`")
-    collection = db[collection_name]
+    modified_time = datetime.now().isoformat()
+    collection = db[selected_collection]
 
-    if collection.find_one({"metadata.title": file_name}):
+    if collection.find_one({"metadata.title": uploaded_file.name}):
         st.warning("⚠️ This file already exists in the collection. Skipping...")
     else:
         st.write("⏳ Processing with DocumentChunker...")
         chunker = DocumentChunker()
-        chunks = chunker.process_document(temp_path)
+        chunks = chunker.process_document(str(temp_path))
 
         if chunks:
             for chunk in chunks:
                 chunk['metadata'].update({
-                    "title": file_name,
-                    "uploaded_at": datetime.now().isoformat(),
+                    "title": uploaded_file.name,
+                    "uploaded_at": modified_time,
                     "is_grant_app": is_grant_app,
                 })
                 collection.insert_one(chunk)
 
-            st.success(f"✅ {len(chunks)} chunks inserted into `{collection_name}`")
+            st.success(f"✅ {len(chunks)} chunks inserted into `{selected_collection}`")
 
+            # Show a few previews
             for i, c in enumerate(chunks[:3]):
                 st.subheader(f"Chunk {i+1}: {c['metadata'].get('header') or 'No Header'}")
                 st.markdown(c['text'][:400] + "...")
@@ -92,16 +79,11 @@ if st.session_state.get("ready_to_process"):
 
             if len(chunks) > 3:
                 st.info(f"... and {len(chunks)-3} more chunks processed.")
+
         else:
             st.warning("⚠️ No chunks were generated.")
 
-    # Clean up
-    try:
-        os.remove(temp_path)
-    except Exception as e:
-        st.warning(f"⚠️ Could not delete temp file: {e}")
-    time.sleep(2)
-    # Reset session
-    for key in ["uploaded_file_name", "collection_name", "is_grant_app", "temp_path", "ready_to_process"]:
-        st.session_state.pop(key, None)
-    st.rerun()
+    # try:
+    #     os.remove(temp_path)
+    # except Exception as e:
+    #     st.warning(f"⚠️ Could not delete temp file: {e}")