Update app.py via AI Editor

app.py

@@ -673,7 +673,6 @@ def unified_master_callback(
     chat_input, cancel_clicks, preview_window_state,
     session_id_state, session_id_store_input
 ):
-    # Determine session_id
     sid = get_session_id(session_id_state if session_id_state else session_id_store_input)
     sess_data = get_session_data(sid)
     ctx = callback_context
@@ -690,7 +689,6 @@ def unified_master_callback(
     doc_delete_clicks = safe_get_n_clicks(ctx, 8)
     proposal_delete_clicks = safe_get_n_clicks(ctx, 12)
 
-    # On first page load, just show the lists
     if not ctx.triggered or triggered_id in ['session-id-store']:
         doc_options = [{'label': truncate_filename(fn), 'value': fn} for fn in sess_data["uploaded_documents"].keys()]
         doc_value = next(iter(sess_data["uploaded_documents"]), None) if sess_data["uploaded_documents"] else None
@@ -759,25 +757,29 @@ def unified_master_callback(
         else:
             logging.error(f"[{sid}] Failed to decode uploaded proposal: {proposal_filename}")
 
+    # --- Fix for delete buttons below ---
     if triggered_id and isinstance(doc_delete_clicks, list):
         for i, n_click in enumerate(doc_delete_clicks):
             if n_click:
                 btn_id = ctx.inputs_list[8][i]['id']
                 del_filename = btn_id['index']
+                # Remove from uploaded_documents
                 if del_filename in sess_data["uploaded_documents"]:
                     del sess_data["uploaded_documents"][del_filename]
-                    if del_filename in sess_data["uploaded_documents_fileid"]:
-                        try:
-                            genai.delete_file(sess_data["uploaded_documents_fileid"][del_filename])
-                        except Exception as e:
-                            logging.warning(f"[{sid}] Failed to delete Gemini file {del_filename}: {e}")
-                        del sess_data["uploaded_documents_fileid"][del_filename]
-                    if del_filename in sess_data["uploaded_documents_bytes"]:
-                        del sess_data["uploaded_documents_bytes"][del_filename]
-                    logging.info(f"[{sid}] Document deleted: {del_filename}")
+                # Remove from uploaded_documents_fileid and delete Gemini file if exists
+                if del_filename in sess_data["uploaded_documents_fileid"]:
+                    try:
+                        genai.delete_file(sess_data["uploaded_documents_fileid"][del_filename])
+                    except Exception as e:
+                        logging.warning(f"[{sid}] Failed to delete Gemini file {del_filename}: {e}")
+                    del sess_data["uploaded_documents_fileid"][del_filename]
+                # Remove from uploaded_documents_bytes
+                if del_filename in sess_data["uploaded_documents_bytes"]:
+                    del sess_data["uploaded_documents_bytes"][del_filename]
+                # Remove from shredded_documents
                 if del_filename in sess_data["shredded_documents"]:
                     del sess_data["shredded_documents"][del_filename]
-                    logging.info(f"[{sid}] Shredded doc deleted: {del_filename}")
+                logging.info(f"[{sid}] Document deleted: {del_filename}")
                 if selected_doc == del_filename:
                     selected_doc = None
                 break
@@ -787,15 +789,17 @@ def unified_master_callback(
             if n_click:
                 btn_id = ctx.inputs_list[12][i]['id']
                 del_filename = btn_id['index']
+                # Remove from proposals
                 if del_filename in sess_data["proposals"]:
                     del sess_data["proposals"][del_filename]
-                    if del_filename in sess_data["proposals_fileid"]:
-                        try:
-                            genai.delete_file(sess_data["proposals_fileid"][del_filename])
-                        except Exception as e:
-                            logging.warning(f"[{sid}] Failed to delete Gemini proposal file {del_filename}: {e}")
-                        del sess_data["proposals_fileid"][del_filename]
-                    logging.info(f"[{sid}] Proposal deleted: {del_filename}")
+                # Remove from proposals_fileid and delete Gemini file if exists
+                if del_filename in sess_data["proposals_fileid"]:
+                    try:
+                        genai.delete_file(sess_data["proposals_fileid"][del_filename])
+                    except Exception as e:
+                        logging.warning(f"[{sid}] Failed to delete Gemini proposal file {del_filename}: {e}")
+                    del sess_data["proposals_fileid"][del_filename]
+                logging.info(f"[{sid}] Proposal deleted: {del_filename}")
                 if selected_proposal == del_filename:
                     selected_proposal = None
                 break