Update app.py via AI Editor

app.py

@@ -13,6 +13,8 @@ import logging
 import threading
 import re
 import google.generativeai as genai
+import sqlite3
+import time
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
 
@@ -20,10 +22,49 @@ GEMINI_KEY = os.environ.get("GEMINI_KEY", "")
 genai.configure(api_key=GEMINI_KEY)
 GEMINI_MODEL = "gemini-2.5-pro-preview-03-25"
 
-app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP], suppress_callback_exceptions=True)
-
-uploaded_files = {}
-uploaded_doc_contents = {}
+DB_PATH = "maiko_dash_docs.sqlite"
+
+def init_db():
+    with sqlite3.connect(DB_PATH) as conn:
+        c = conn.cursor()
+        c.execute('''CREATE TABLE IF NOT EXISTS uploaded_docs
+                     (doc_type TEXT PRIMARY KEY, doc_content TEXT, file_name TEXT, last_update TIMESTAMP)''')
+        c.execute('''CREATE TABLE IF NOT EXISTS generated_docs
+                     (doc_type TEXT PRIMARY KEY, doc_content TEXT, last_update TIMESTAMP)''')
+        conn.commit()
+init_db()
+
+def set_uploaded_doc(doc_type, doc_content, file_name):
+    with sqlite3.connect(DB_PATH) as conn:
+        c = conn.cursor()
+        c.execute('''INSERT OR REPLACE INTO uploaded_docs (doc_type, doc_content, file_name, last_update)
+                     VALUES (?, ?, ?, ?)''', (doc_type, doc_content, file_name, time.time()))
+        conn.commit()
+
+def get_uploaded_doc(doc_type):
+    with sqlite3.connect(DB_PATH) as conn:
+        c = conn.cursor()
+        c.execute('''SELECT doc_content, file_name FROM uploaded_docs WHERE doc_type=?''', (doc_type,))
+        row = c.fetchone()
+        if row:
+            return row[0], row[1]
+        return None, None
+
+def set_generated_doc(doc_type, doc_content):
+    with sqlite3.connect(DB_PATH) as conn:
+        c = conn.cursor()
+        c.execute('''INSERT OR REPLACE INTO generated_docs (doc_type, doc_content, last_update)
+                     VALUES (?, ?, ?)''', (doc_type, doc_content, time.time()))
+        conn.commit()
+
+def get_generated_doc(doc_type):
+    with sqlite3.connect(DB_PATH) as conn:
+        c = conn.cursor()
+        c.execute('''SELECT doc_content FROM generated_docs WHERE doc_type=?''', (doc_type,))
+        row = c.fetchone()
+        if row:
+            return row[0]
+        return None
 
 spreadsheet_types = ["Shred", "Pink Review", "Red Review", "Gold Review", "Virtual Board", "LOE"]
 narrative_types = ["Pink", "Red", "Gold"]
@@ -333,6 +374,8 @@ def get_right_col_content(selected_type, store_data):
     controls.append(html.Div(id='document-preview', className="border p-3 mb-3"))
     return dbc.Card(dbc.CardBody(controls))
 
+app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP], suppress_callback_exceptions=True)
+
 app.layout = dbc.Container([
     dcc.Store(id='selected-doc-type', data="Shred"),
     dcc.Store(id='store-shred'),
@@ -427,265 +470,145 @@ def update_selected_doc_type(n_clicks_list, btn_ids):
 @app.callback(
     Output('right-col-content', 'children'),
     Input('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'),
 )
-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
-    }
+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)
+        store_data[store_key] = doc_content
     return get_right_col_content(selected_type, store_data)
 
 @app.callback(
     [
         Output('store-shred', 'data'),
+        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('file-list', 'children'),
         Output('uploaded-doc-name-shred', 'children'),
+        Output('uploaded-doc-name-pink', 'children'),
+        Output('uploaded-doc-name-pink_review', 'children'),
+        Output('uploaded-doc-name-red', 'children'),
+        Output('uploaded-doc-name-red_review', 'children'),
+        Output('uploaded-doc-name-gold', 'children'),
+        Output('uploaded-doc-name-gold_review', 'children'),
+        Output('uploaded-doc-name-loe', 'children'),
+        Output('uploaded-doc-name-virtual_board', 'children'),
         Output('progress-dot-store', 'data')
     ],
     [
-        Input({'type': 'upload-doc-type', 'subtype': 'shred', 'index': 'Shred'}, 'contents')
+        Input({'type': 'upload-doc-type', 'subtype': ALL, 'index': ALL}, 'contents'),
     ],
     [
-        State({'type': 'upload-doc-type', 'subtype': 'shred', 'index': 'Shred'}, 'filename'),
+        State({'type': 'upload-doc-type', 'subtype': ALL, 'index': ALL}, 'filename'),
         State('store-shred', 'data'),
-        State('file-list', 'children')
+        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 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], "progress"
-
-@app.callback(
-    Output('store-pink', 'data'),
-    Output('uploaded-doc-name-pink', 'children'),
-    Output('progress-dot-store', 'data'),
-    Input({'type': 'upload-doc-type', 'subtype': 'pink', 'index': ALL}, 'contents'),
-    State({'type': 'upload-doc-type', 'subtype': 'pink', 'index': ALL}, 'filename'),
-    prevent_initial_call=True
-)
-def handle_pink_upload(contents_list, filenames_list):
-    if not contents_list or not filenames_list:
-        return dash.no_update, dash.no_update, dash.no_update
-    for contents, filenames in 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)
-                return file_text, name, "progress"
-    return dash.no_update, dash.no_update, dash.no_update
-
-@app.callback(
-    Output('store-pink-review', 'data'),
-    Output('uploaded-doc-name-pink_review', 'children'),
-    Output('progress-dot-store', 'data'),
-    Input({'type': 'upload-doc-type', 'subtype': 'pink_review', 'index': ALL}, 'contents'),
-    State({'type': 'upload-doc-type', 'subtype': 'pink_review', 'index': ALL}, 'filename'),
-    prevent_initial_call=True
-)
-def handle_pink_review_upload(contents_list, filenames_list):
-    if not contents_list or not filenames_list:
-        return dash.no_update, dash.no_update, dash.no_update
-    for contents, filenames in 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)
-                return file_text, name, "progress"
-    return dash.no_update, dash.no_update, dash.no_update
-
-@app.callback(
-    Output('store-red', 'data'),
-    Output('uploaded-doc-name-red', 'children'),
-    Output('progress-dot-store', 'data'),
-    Input({'type': 'upload-doc-type', 'subtype': 'red', 'index': ALL}, 'contents'),
-    State({'type': 'upload-doc-type', 'subtype': 'red', 'index': ALL}, 'filename'),
-    prevent_initial_call=True
-)
-def handle_red_upload(contents_list, filenames_list):
-    if not contents_list or not filenames_list:
-        return dash.no_update, dash.no_update, dash.no_update
-    for contents, filenames in 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)
-                return file_text, name, "progress"
-    return dash.no_update, dash.no_update, dash.no_update
-
-@app.callback(
-    Output('store-red-review', 'data'),
-    Output('uploaded-doc-name-red_review', 'children'),
-    Output('progress-dot-store', 'data'),
-    Input({'type': 'upload-doc-type', 'subtype': 'red_review', 'index': ALL}, 'contents'),
-    State({'type': 'upload-doc-type', 'subtype': 'red_review', 'index': ALL}, 'filename'),
-    prevent_initial_call=True
-)
-def handle_red_review_upload(contents_list, filenames_list):
-    if not contents_list or not filenames_list:
-        return dash.no_update, dash.no_update, dash.no_update
-    for contents, filenames in 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)
-                return file_text, name, "progress"
-    return dash.no_update, dash.no_update, dash.no_update
-
-@app.callback(
-    Output('store-gold', 'data'),
-    Output('uploaded-doc-name-gold', 'children'),
-    Output('progress-dot-store', 'data'),
-    Input({'type': 'upload-doc-type', 'subtype': 'gold', 'index': ALL}, 'contents'),
-    State({'type': 'upload-doc-type', 'subtype': 'gold', 'index': ALL}, 'filename'),
-    prevent_initial_call=True
-)
-def handle_gold_upload(contents_list, filenames_list):
-    if not contents_list or not filenames_list:
-        return dash.no_update, dash.no_update, dash.no_update
-    for contents, filenames in 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)
-                return file_text, name, "progress"
-    return dash.no_update, dash.no_update, dash.no_update
+def handle_doc_upload(contents_lists, filenames_lists,
+                     shred, pink, pink_review, red, red_review, gold, gold_review, loe, virtual_board):
+    stores = [shred, pink, pink_review, red, red_review, gold, gold_review, loe, virtual_board]
+    file_list = []
+    doc_names = [None]*9
+    progress = dash.no_update
 
-@app.callback(
-    Output('store-gold-review', 'data'),
-    Output('uploaded-doc-name-gold_review', 'children'),
-    Output('progress-dot-store', 'data'),
-    Input({'type': 'upload-doc-type', 'subtype': 'gold_review', 'index': ALL}, 'contents'),
-    State({'type': 'upload-doc-type', 'subtype': 'gold_review', 'index': ALL}, 'filename'),
-    prevent_initial_call=True
-)
-def handle_gold_review_upload(contents_list, filenames_list):
-    if not contents_list or not filenames_list:
-        return dash.no_update, dash.no_update, dash.no_update
-    for contents, filenames in 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)
-                return file_text, name, "progress"
-    return dash.no_update, dash.no_update, dash.no_update
+    triggered = callback_context.triggered
+    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),
+        "pink_review": (2, 2),
+        "red": (3, 3),
+        "red_review": (4, 4),
+        "gold": (5, 5),
+        "gold_review": (6, 6),
+        "loe": (7, 7),
+        "virtual_board": (8, 8)
+    }
 
-@app.callback(
-    Output('store-loe', 'data'),
-    Output('uploaded-doc-name-loe', 'children'),
-    Output('progress-dot-store', 'data'),
-    Input({'type': 'upload-doc-type', 'subtype': 'loe', 'index': ALL}, 'contents'),
-    State({'type': 'upload-doc-type', 'subtype': 'loe', 'index': ALL}, 'filename'),
-    prevent_initial_call=True
-)
-def handle_loe_upload(contents_list, filenames_list):
-    if not contents_list or not filenames_list:
-        return dash.no_update, dash.no_update, dash.no_update
-    for contents, filenames in zip(contents_list, filenames_list):
+    # 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']
+            subtype = id_dict['subtype']
+            idx, doc_name_idx = subtype_map[subtype]
             if isinstance(contents, str):
                 contents = [contents]
                 filenames = [filenames]
-            for content, name in zip(contents, filenames):
-                file_text = process_document(content, name)
-                return file_text, name, "progress"
-    return dash.no_update, dash.no_update, dash.no_update
+            # Only take the last uploaded file for that slot
+            content = contents[-1]
+            name = filenames[-1]
+            file_text = process_document(content, name)
+            set_uploaded_doc(subtype, file_text, name)
+            stores[idx] = file_text
+            doc_names[doc_name_idx] = name
+            if subtype == "shred":
+                # For shred, build file list
+                file_list = [
+                    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")
+                ]
+            progress = "progress"
+            break
 
-@app.callback(
-    Output('store-virtual-board', 'data'),
-    Output('uploaded-doc-name-virtual_board', 'children'),
-    Output('progress-dot-store', 'data'),
-    Input({'type': 'upload-doc-type', 'subtype': 'virtual_board', 'index': ALL}, 'contents'),
-    State({'type': 'upload-doc-type', 'subtype': 'virtual_board', 'index': ALL}, 'filename'),
-    prevent_initial_call=True
-)
-def handle_virtual_board_upload(contents_list, filenames_list):
-    if not contents_list or not filenames_list:
-        return dash.no_update, dash.no_update, dash.no_update
-    for contents, filenames in 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)
-                return file_text, name, "progress"
-    return dash.no_update, dash.no_update, dash.no_update
+    return stores + [file_list] + doc_names + [progress]
 
 @app.callback(
-    Output('document-preview', 'children'),
-    Output('progress-dot-store', 'data'),
+    [
+        Output('document-preview', 'children'),
+        Output('store-generated-doc', 'data'),
+        Output('progress-dot-store', 'data')
+    ],
     [
         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'),
+        Input('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'),
         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):
+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):
     ctx = callback_context
     trigger = ctx.triggered[0]['prop_id'] if ctx.triggered else ""
-    logging.info(f"Document preview trigger: {trigger}")
+    logging.info(f"Document preview/generate trigger: {trigger}")
+
+    # If generate button pressed
     if any(n_clicks_list):
         logging.info(f"Generating document for type: {selected_type}")
         inputs = {
@@ -721,64 +644,43 @@ def preview_or_generate_doc(n_clicks_list, selected_type, shred, pink, pink_revi
         t.start()
         t.join(timeout=60)
         generated_doc = result_holder.get('result', 'Error: No document generated.')
-        app.server.config['last_generated_doc'] = generated_doc
+        set_generated_doc(selected_type, 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 preview, "progress"
+        return preview, generated_doc, "progress"
     else:
-        if selected_type == "Shred" and shred:
-            return markdown_table_preview(shred), 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
-        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
-        return html.Div("No document loaded."), dash.no_update
-
-@app.callback(
-    Output('store-generated-doc', 'data'),
-    Output('progress-dot-store', '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), "progress"
-    return prev_generated, dash.no_update
+        # 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
 
 @app.callback(
     Output("download-document", "data"),
@@ -799,6 +701,8 @@ def update_generated_doc(n_clicks_list, selected_type, preview_content, shred, p
 )
 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
@@ -835,7 +739,6 @@ def download_document(n_clicks, selected_type, generated_doc, shred, pink, pink_
     Input('progress-dot-store', 'data')
 )
 def update_main_progress_dot(progress):
-    # Dummy output to trigger dcc.Loading
     return ""
 
 if __name__ == '__main__':