Update app.py via AI Editor

app.py

@@ -6,7 +6,7 @@ from docx import Document
 from io import BytesIO
 import dash
 import dash_bootstrap_components as dbc
-from dash import html, dcc, Input, Output, State, callback_context, ALL
+from dash import html, dcc, Input, Output, State, callback_context, ALL, MATCH
 from dash.dash_table import DataTable
 from PyPDF2 import PdfReader
 import logging
@@ -229,7 +229,7 @@ def get_left_col_content():
         chat_card
     ]
 
-def get_right_col_content(selected_type, shred_doc, pink_doc, pink_review_doc, red_doc, red_review_doc, gold_doc):
+def get_right_col_content(selected_type, store_data):
     controls = []
     controls.append(html.Div([
         html.Div(className="blinking-dot", style={'margin':'0 auto','width':'16px','height':'16px'}),
@@ -239,11 +239,13 @@ def get_right_col_content(selected_type, shred_doc, pink_doc, pink_review_doc, r
         type="dot",
         children=[html.Div(id="loading-output")]
     ))
+
+    # For Shred, show upload
     if selected_type == "Shred":
         controls.append(
             html.Div([
                 dcc.Upload(
-                    id='upload-document',
+                    id={'type': 'upload-doc-type', 'subtype': 'shred', 'index': selected_type},
                     children=html.Div([
                         'Drag and Drop or ',
                         html.A('Select Files')
@@ -263,40 +265,35 @@ def get_right_col_content(selected_type, shred_doc, pink_doc, pink_review_doc, r
                 file_list_component()
             ])
         )
+    # For others, show dependency uploads and previews
     elif selected_type in doc_dependencies:
         sources = doc_dependencies[selected_type]["source"]
         for src in sources:
             label = ""
             loaded_preview = None
-            store_var = None
+            store_var = store_data.get(src)
             if src == "shred":
                 label = "Shred (Requirements)"
-                store_var = shred_doc
-                loaded_preview = markdown_table_preview(shred_doc)
+                loaded_preview = markdown_table_preview(store_var)
             elif src == "pink":
                 label = "Pink Document"
-                store_var = pink_doc
-                loaded_preview = markdown_narrative_preview(pink_doc)
+                loaded_preview = markdown_narrative_preview(store_var)
             elif src == "pink_review":
                 label = "Pink Review"
-                store_var = pink_review_doc
-                loaded_preview = markdown_table_preview(pink_review_doc)
+                loaded_preview = markdown_table_preview(store_var)
             elif src == "red":
                 label = "Red Document"
-                store_var = red_doc
-                loaded_preview = markdown_narrative_preview(red_doc)
+                loaded_preview = markdown_narrative_preview(store_var)
             elif src == "red_review":
                 label = "Red Review"
-                store_var = red_review_doc
-                loaded_preview = markdown_table_preview(red_review_doc)
+                loaded_preview = markdown_table_preview(store_var)
             elif src == "gold":
                 label = "Gold Document"
-                store_var = gold_doc
-                loaded_preview = markdown_narrative_preview(gold_doc)
+                loaded_preview = markdown_narrative_preview(store_var)
             controls.append(html.Div([
                 html.Label(f"Upload {label}"),
                 dcc.Upload(
-                    id={'type': f'upload-doc-type-{src}', 'index': selected_type},
+                    id={'type': 'upload-doc-type', 'subtype': src, 'index': selected_type},
                     children=html.Div(['Drag and Drop or ', html.A('Select File')]),
                     style={
                         'width': '100%',
@@ -310,14 +307,13 @@ def get_right_col_content(selected_type, shred_doc, pink_doc, pink_review_doc, r
                     },
                     multiple=False
                 ),
-                html.Div(id={'type': f'uploaded-doc-name-{src}', 'index': selected_type}),
+                html.Div(id={'type': 'uploaded-doc-name', 'subtype': src, 'index': selected_type}),
                 dbc.RadioItems(
-                    id={'type': f'radio-doc-source-{src}', 'index': selected_type},
+                    id={'type': 'radio-doc-source', 'subtype': src, 'index': selected_type},
                     options=[
-                        {'label': 'Loaded Document', 'value': 'loaded'},
-                        {'label': 'Uploaded Document', 'value': 'uploaded'}
-                    ] if store_var else [{'label': 'Uploaded Document', 'value': 'uploaded'}],
-                    value='loaded' if store_var else 'uploaded',
+                        {'label': 'Loaded Document', 'value': 'loaded'}
+                    ] if store_var else [],
+                    value='loaded' if store_var else None,
                     inline=True,
                     className="mb-2"
                 ),
@@ -438,75 +434,143 @@ def update_selected_doc_type(n_clicks_list, btn_ids):
     State('store-red', 'data'),
     State('store-red-review', 'data'),
     State('store-gold', 'data'),
-)
-def update_right_col(selected_type, shred_doc, pink_doc, pink_review_doc, red_doc, red_review_doc, gold_doc):
-    return get_right_col_content(selected_type, shred_doc, pink_doc, pink_review_doc, red_doc, red_review_doc, gold_doc)
-
-@app.callback(
-    Output('file-list', 'children'),
-    Output('store-shred', 'data'),
-    Output('document-preview', 'children'),
-    Output('loading-output', 'children'),
-    Output('store-generated-doc', 'data'),
-    Input('upload-document', 'contents'),
-    State('upload-document', 'filename'),
-    State('file-list', 'children'),
-    State('store-shred', 'data'),
-    Input({'type': 'btn-generate-doc', 'index': ALL}, 'n_clicks'),
-    State('selected-doc-type', '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 combined_callback(
-    upload_contents, upload_filenames, existing_files, current_shred,
-    n_clicks, selected_type, shred, pink, pink_review, red, red_review, gold, gold_review, loe, virtual_board
-):
-    ctx = callback_context
-    trigger = ctx.triggered[0]['prop_id'].split('.')[0] if ctx.triggered else ""
-    files_display = existing_files if existing_files else []
-    new_shred = current_shred
-    preview = dash.no_update
-    loading_msg = ""
-    generated_doc = dash.no_update
+def update_right_col(selected_type, shred, pink, pink_review, red, red_review, gold, gold_review, loe, virtual_board):
+    store_data = {
+        "shred": shred,
+        "pink": pink,
+        "pink_review": pink_review,
+        "red": red,
+        "red_review": red_review,
+        "gold": gold,
+        "gold_review": gold_review,
+        "loe": loe,
+        "virtual_board": virtual_board
+    }
+    return get_right_col_content(selected_type, store_data)
 
-    # Only handle upload if upload-document is present in layout
-    if trigger == "upload-document" and upload_contents is not None and upload_filenames is not None:
-        for content, name in zip(upload_contents, upload_filenames):
-            file_content = process_document(content, name)
-            uploaded_files[name] = file_content
-            new_shred = file_content
-            files_display.append(
-                dbc.Row([
-                    dbc.Col(
-                        html.Button('×', id={'type': 'remove-file', 'index': name}, style={'marginRight': '5px', 'fontSize': '10px'}),
-                        width=1
-                    ),
-                    dbc.Col(
-                        html.Span(name, style={'wordBreak': 'break-all'}),
-                        width=11
-                    )
-                ], id={'type': 'file-row', 'index': name}, align="center", className="mb-1")
-            )
-        logging.info("Shred document uploaded and stored.")
-        preview = markdown_table_preview(new_shred)
-        generated_doc = dash.no_update
-        return files_display, new_shred, preview, loading_msg, generated_doc
+@app.callback(
+    [
+        Output('store-shred', 'data'),
+        Output('file-list', 'children'),
+        Output({'type': 'uploaded-doc-name', 'subtype': 'shred', 'index': 'Shred'}, 'children')
+    ],
+    [
+        Input({'type': 'upload-doc-type', 'subtype': 'shred', 'index': 'Shred'}, 'contents')
+    ],
+    [
+        State({'type': 'upload-doc-type', 'subtype': 'shred', 'index': 'Shred'}, 'filename'),
+        State('store-shred', 'data'),
+        State('file-list', 'children')
+    ],
+    prevent_initial_call=True
+)
+def handle_shred_upload(contents, filenames, current_shred, file_list):
+    if not contents or not filenames:
+        raise dash.exceptions.PreventUpdate
+    if isinstance(contents, str):
+        contents = [contents]
+        filenames = [filenames]
+    file_previews = []
+    latest_text = current_shred
+    for content, name in zip(contents, filenames):
+        file_text = process_document(content, name)
+        uploaded_files[name] = file_text
+        latest_text = file_text
+        file_previews.append(
+            dbc.Row([
+                dbc.Col(
+                    html.Button('×', id={'type': 'remove-file', 'index': name}, style={'marginRight': '5px', 'fontSize': '10px'}),
+                    width=1
+                ),
+                dbc.Col(
+                    html.Span(name, style={'wordBreak': 'break-all'}),
+                    width=11
+                )
+            ], id={'type': 'file-row', 'index': name}, align="center", className="mb-1")
+        )
+    logging.info("Shred document uploaded and stored.")
+    return latest_text, file_previews, filenames[-1]
 
-    # Always show preview for current shred if nothing else triggered, for Shred only
-    if selected_type == "Shred" and not any(n_clicks):
-        preview = markdown_table_preview(new_shred)
-        return files_display, new_shred, preview, loading_msg, generated_doc
+# Pattern-matching callback for all dependency uploads
+@app.callback(
+    [
+        Output('store-pink', 'data'),
+        Output('store-pink-review', 'data'),
+        Output('store-red', 'data'),
+        Output('store-red-review', 'data'),
+        Output('store-gold', 'data'),
+        Output('store-gold-review', 'data'),
+        Output('store-loe', 'data'),
+        Output('store-virtual-board', 'data'),
+        Output({'type': 'uploaded-doc-name', 'subtype': MATCH, 'index': MATCH}, 'children')
+    ],
+    [
+        Input({'type': 'upload-doc-type', 'subtype': ALL, 'index': ALL}, 'contents')
+    ],
+    [
+        State({'type': 'upload-doc-type', 'subtype': ALL, 'index': ALL}, 'filename')
+    ],
+    prevent_initial_call=True
+)
+def handle_dep_upload(contents_list, filenames_list):
+    stores = [dash.no_update]*8
+    upload_name = ""
+    if not contents_list or not filenames_list:
+        return [dash.no_update]*8 + [upload_name]
+    for i, (contents, filenames) in enumerate(zip(contents_list, filenames_list)):
+        if contents and filenames:
+            if isinstance(contents, str):
+                contents = [contents]
+                filenames = [filenames]
+            for content, name in zip(contents, filenames):
+                file_text = process_document(content, name)
+                if 'pink' in name.lower():
+                    stores[0] = file_text
+                elif 'pink_review' in name.lower():
+                    stores[1] = file_text
+                elif 'red' in name.lower() and 'review' not in name.lower():
+                    stores[2] = file_text
+                elif 'red_review' in name.lower():
+                    stores[3] = file_text
+                elif 'gold' in name.lower() and 'review' not in name.lower():
+                    stores[4] = file_text
+                elif 'gold_review' in name.lower():
+                    stores[5] = file_text
+                elif 'loe' in name.lower():
+                    stores[6] = file_text
+                elif 'virtual' in name.lower():
+                    stores[7] = file_text
+                upload_name = name
+    return stores + [upload_name]
 
-    # Handle generate document
-    if any(n_clicks):
+@app.callback(
+    Output('document-preview', 'children'),
+    [
+        Input({'type': 'btn-generate-doc', 'index': ALL}, 'n_clicks'),
+        Input('selected-doc-type', 'data'),
+        Input('store-shred', 'data'),
+        Input('store-pink', 'data'),
+        Input('store-pink-review', 'data'),
+        Input('store-red', 'data'),
+        Input('store-red-review', 'data'),
+        Input('store-gold', 'data'),
+        Input('store-gold-review', 'data'),
+        Input('store-loe', 'data'),
+        Input('store-virtual-board', 'data'),
+        State('store-generated-doc', 'data')
+    ],
+    prevent_initial_call=True
+)
+def preview_or_generate_doc(n_clicks_list, 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 trigger: {trigger}")
+    # If generate button pressed
+    if any(n_clicks_list):
         logging.info(f"Generating document for type: {selected_type}")
         inputs = {
             "shred": shred,
@@ -541,46 +605,115 @@ def combined_callback(
         t.start()
         t.join(timeout=60)
         generated_doc = result_holder.get('result', 'Error: No document generated.')
-        # Always update preview for all doc types
+        app.server.config['last_generated_doc'] = generated_doc
         if selected_type in spreadsheet_types:
             preview = markdown_table_preview(generated_doc)
         else:
             preview = markdown_narrative_preview(generated_doc)
         logging.info("Document preview updated.")
-        return files_display, new_shred, preview, "", generated_doc
-
-    # For all other doc types, show preview of generated doc if present
-    if selected_type != "Shred":
-        if generated_doc not in (None, "", dash.no_update):
-            if selected_type in spreadsheet_types:
-                preview = markdown_table_preview(generated_doc)
-            else:
-                preview = markdown_narrative_preview(generated_doc)
-            return files_display, new_shred, preview, loading_msg, generated_doc
-
-    # Default: clear preview for non-shred if nothing else
-    if selected_type != "Shred":
-        return files_display, new_shred, html.Div(), loading_msg, generated_doc
+        return preview
+    # If just switching doc type, show preview of loaded doc (from store)
+    else:
+        # For Shred, show table preview
+        if selected_type == "Shred" and shred:
+            return markdown_table_preview(shred)
+        # For spreadsheet types, preview from store
+        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)
+        # For narrative types, preview from store
+        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)
+        return html.Div("No document loaded.")
 
-    return files_display, new_shred, markdown_table_preview(new_shred), loading_msg, generated_doc
+@app.callback(
+    Output('store-generated-doc', 'data'),
+    [
+        Input({'type': 'btn-generate-doc', 'index': ALL}, 'n_clicks'),
+        Input('selected-doc-type', 'data'),
+        Input('document-preview', 'children')
+    ],
+    [
+        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'),
+        State('store-generated-doc', 'data')
+    ],
+    prevent_initial_call=True
+)
+def update_generated_doc(n_clicks_list, selected_type, preview_content, 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 ""
+    if any(n_clicks_list):
+        return app.server.config.get('last_generated_doc', prev_generated)
+    return prev_generated
 
 @app.callback(
     Output("download-document", "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):
-    if not generated_doc:
+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 generated, try to get from store for preview/download
+    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
     if selected_type in spreadsheet_types:
-        xlsx_io = markdown_tables_to_xlsx(generated_doc)
+        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")
     else:
-        content = generated_doc.encode('utf-8')
+        content = doc.encode('utf-8')
         logging.info("Narrative document prepared for download.")
         return dcc.send_bytes(content, filename=f"{selected_type}_output.md")