Update app.py via AI Editor

app.py

@@ -472,7 +472,6 @@ def update_selected_doc_type(n_clicks_list, btn_ids):
     Input('selected-doc-type', 'data'),
 )
 def update_right_col(selected_type):
-    # Load all doc stores from sqlite
     store_data = {}
     for store_key in ["shred","pink","pink_review","red","red_review","gold","gold_review","loe","virtual_board"]:
         doc_content, _ = get_uploaded_doc(store_key)
@@ -530,7 +529,6 @@ def handle_doc_upload(contents_lists, filenames_lists,
     if not triggered or not contents_lists:
         return [dash.no_update]*19 + [dash.no_update]
 
-    # Map subtype to store index and doc_name index
     subtype_map = {
         "shred": (0, 0),
         "pink": (1, 1),
@@ -543,7 +541,6 @@ def handle_doc_upload(contents_lists, filenames_lists,
         "virtual_board": (8, 8)
     }
 
-    # Find which upload triggered
     for i, (contents, filenames) in enumerate(zip(contents_lists, filenames_lists)):
         if contents and filenames:
             id_dict = callback_context.inputs_list[0][i]['id']
@@ -552,7 +549,6 @@ def handle_doc_upload(contents_lists, filenames_lists,
             if isinstance(contents, str):
                 contents = [contents]
                 filenames = [filenames]
-            # Only take the last uploaded file for that slot
             content = contents[-1]
             name = filenames[-1]
             file_text = process_document(content, name)
@@ -560,7 +556,6 @@ def handle_doc_upload(contents_lists, filenames_lists,
             stores[idx] = file_text
             doc_names[doc_name_idx] = name
             if subtype == "shred":
-                # For shred, build file list
                 file_list = [
                     dbc.Row([
                         dbc.Col(
@@ -582,10 +577,12 @@ def handle_doc_upload(contents_lists, filenames_lists,
     [
         Output('document-preview', 'children'),
         Output('store-generated-doc', 'data'),
+        Output('download-document', 'data'),
         Output('progress-dot-store', 'data')
     ],
     [
         Input({'type': 'btn-generate-doc', 'index': ALL}, 'n_clicks'),
+        Input("btn-download", "n_clicks"),
         Input('selected-doc-type', 'data')
     ],
     [
@@ -602,14 +599,50 @@ def handle_doc_upload(contents_lists, filenames_lists,
     ],
     prevent_initial_call=True
 )
-def handle_preview_and_generate(n_clicks_list, selected_type,
-                               shred, pink, pink_review, red, red_review, gold, gold_review, loe, virtual_board, prev_generated):
+def handle_preview_generate_download(n_clicks_generate_list, n_clicks_download, selected_type,
+                                     shred, pink, pink_review, red, red_review, gold, gold_review, loe, virtual_board, prev_generated):
     ctx = callback_context
     trigger = ctx.triggered[0]['prop_id'] if ctx.triggered else ""
-    logging.info(f"Document preview/generate trigger: {trigger}")
+    logging.info(f"Unified preview/generate/download trigger: {trigger}")
 
-    # If generate button pressed
-    if any(n_clicks_list):
+    # Download logic
+    if "btn-download" in trigger:
+        doc = prev_generated
+        if not doc:
+            doc = get_generated_doc(selected_type)
+        if not doc:
+            if selected_type == "Shred":
+                doc = shred
+            elif selected_type == "Pink":
+                doc = pink
+            elif selected_type == "Pink Review":
+                doc = pink_review
+            elif selected_type == "Red":
+                doc = red
+            elif selected_type == "Red Review":
+                doc = red_review
+            elif selected_type == "Gold":
+                doc = gold
+            elif selected_type == "Gold Review":
+                doc = gold_review
+            elif selected_type == "LOE":
+                doc = loe
+            elif selected_type == "Virtual Board":
+                doc = virtual_board
+        if not doc:
+            logging.warning("No generated document to download.")
+            return dash.no_update, dash.no_update, dash.no_update, dash.no_update
+        if selected_type in spreadsheet_types:
+            xlsx_io = markdown_tables_to_xlsx(doc)
+            logging.info("Spreadsheet document prepared for download.")
+            return dash.no_update, dash.no_update, dcc.send_bytes(xlsx_io.getvalue(), filename=f"{selected_type}_output.xlsx"), "progress"
+        else:
+            content = doc.encode('utf-8')
+            logging.info("Narrative document prepared for download.")
+            return dash.no_update, dash.no_update, dcc.send_bytes(content, filename=f"{selected_type}_output.md"), "progress"
+
+    # Generate logic
+    if any(n_clicks_generate_list):
         logging.info(f"Generating document for type: {selected_type}")
         inputs = {
             "shred": shred,
@@ -650,89 +683,36 @@ def handle_preview_and_generate(n_clicks_list, selected_type,
         else:
             preview = markdown_narrative_preview(generated_doc)
         logging.info("Document preview updated.")
-        return preview, generated_doc, "progress"
-    else:
-        # Just preview from store or sqlite
-        doc = get_generated_doc(selected_type)
-        if not doc:
-            # Fallback to uploaded doc for preview
-            if selected_type == "Shred" and shred:
-                return markdown_table_preview(shred), dash.no_update, dash.no_update
-            elif selected_type in spreadsheet_types:
-                doc_store = {
-                    "Pink Review": pink_review,
-                    "Red Review": red_review,
-                    "Gold Review": gold_review,
-                    "Virtual Board": virtual_board,
-                    "LOE": loe
-                }
-                doc = doc_store.get(selected_type, "")
-                return markdown_table_preview(doc), dash.no_update, dash.no_update
-            elif selected_type in narrative_types:
-                doc_store = {
-                    "Pink": pink,
-                    "Red": red,
-                    "Gold": gold
-                }
-                doc = doc_store.get(selected_type, "")
-                return markdown_narrative_preview(doc), dash.no_update, dash.no_update
-            return html.Div("No document loaded."), dash.no_update, dash.no_update
-        if selected_type in spreadsheet_types:
-            return markdown_table_preview(doc), doc, dash.no_update
-        else:
-            return markdown_narrative_preview(doc), doc, dash.no_update
+        return preview, generated_doc, dash.no_update, "progress"
 
-@app.callback(
-    Output("download-document", "data"),
-    Output('progress-dot-store', 'data'),
-    Input("btn-download", "n_clicks"),
-    State('selected-doc-type', 'data'),
-    State('store-generated-doc', 'data'),
-    State('store-shred', 'data'),
-    State('store-pink', 'data'),
-    State('store-pink-review', 'data'),
-    State('store-red', 'data'),
-    State('store-red-review', 'data'),
-    State('store-gold', 'data'),
-    State('store-gold-review', 'data'),
-    State('store-loe', 'data'),
-    State('store-virtual-board', 'data'),
-    prevent_initial_call=True
-)
-def download_document(n_clicks, selected_type, generated_doc, shred, pink, pink_review, red, red_review, gold, gold_review, loe, virtual_board):
-    doc = generated_doc
-    if not doc:
-        doc = get_generated_doc(selected_type)
-    if not doc:
-        if selected_type == "Shred":
-            doc = shred
-        elif selected_type == "Pink":
-            doc = pink
-        elif selected_type == "Pink Review":
-            doc = pink_review
-        elif selected_type == "Red":
-            doc = red
-        elif selected_type == "Red Review":
-            doc = red_review
-        elif selected_type == "Gold":
-            doc = gold
-        elif selected_type == "Gold Review":
-            doc = gold_review
-        elif selected_type == "LOE":
-            doc = loe
-        elif selected_type == "Virtual Board":
-            doc = virtual_board
+    # Preview logic (selected type changed)
+    doc = get_generated_doc(selected_type)
     if not doc:
-        logging.warning("No generated document to download.")
-        return dash.no_update, dash.no_update
+        if selected_type == "Shred" and shred:
+            return markdown_table_preview(shred), dash.no_update, dash.no_update, dash.no_update
+        elif selected_type in spreadsheet_types:
+            doc_store = {
+                "Pink Review": pink_review,
+                "Red Review": red_review,
+                "Gold Review": gold_review,
+                "Virtual Board": virtual_board,
+                "LOE": loe
+            }
+            doc = doc_store.get(selected_type, "")
+            return markdown_table_preview(doc), dash.no_update, dash.no_update, dash.no_update
+        elif selected_type in narrative_types:
+            doc_store = {
+                "Pink": pink,
+                "Red": red,
+                "Gold": gold
+            }
+            doc = doc_store.get(selected_type, "")
+            return markdown_narrative_preview(doc), dash.no_update, dash.no_update, dash.no_update
+        return html.Div("No document loaded."), dash.no_update, dash.no_update, dash.no_update
     if selected_type in spreadsheet_types:
-        xlsx_io = markdown_tables_to_xlsx(doc)
-        logging.info("Spreadsheet document prepared for download.")
-        return dcc.send_bytes(xlsx_io.getvalue(), filename=f"{selected_type}_output.xlsx"), "progress"
+        return markdown_table_preview(doc), doc, dash.no_update, dash.no_update
     else:
-        content = doc.encode('utf-8')
-        logging.info("Narrative document prepared for download.")
-        return dcc.send_bytes(content, filename=f"{selected_type}_output.md"), "progress"
+        return markdown_narrative_preview(doc), doc, dash.no_update, dash.no_update
 
 @app.callback(
     Output("main-progress-dot-output", "children"),