Update app.py via AI Editor

app.py

@@ -6,7 +6,7 @@ from docx import Document
 from io import BytesIO, StringIO
 import dash
 import dash_bootstrap_components as dbc
-from dash import html, dcc, Input, Output, State, callback_context
+from dash import html, dcc, Input, Output, State, callback_context, MATCH, ALL
 from docx.shared import Pt
 from docx.enum.style import WD_STYLE_TYPE
 from PyPDF2 import PdfReader
@@ -14,23 +14,19 @@ import openai
 import logging
 import threading
 
-# Logging configuration
 logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
 
-# Initialize Dash app
 app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])
 
-# Configure OpenAI
 openai.api_key = os.environ.get("OPENAI_API_KEY", "")
 
-# Global variables
 uploaded_files = {}
 current_document = None
 document_type = None
 shredded_document = None
 pink_review_document = None
+uploaded_doc_contents = {}
 
-# Document types and their descriptions
 document_types = {
     "Shred": "Generate a requirements spreadsheet of the Project Work Statement (PWS) identified by action words like shall, will, perform etc. by pws section, requirement.  Do not write as if you're responding to the proposal.  Its a spreadsheet to distill the requirements, not microhealth's approach",
     "Pink": "Create a Pink Team document based on the PWS outline.  Your goal is to be compliant and compelling.",
@@ -43,44 +39,34 @@ document_types = {
     "LOE": "Generate a Level of Effort (LOE) breakdown as a spreadsheet"
 }
 
-app.layout = dbc.Container([
-    dbc.Row([
-        dbc.Col([
-            html.H4("Proposal Documents", className="mt-3 mb-4"),
+def get_right_col_content(selected_type):
+    if selected_type == "Shred":
+        return [
             html.Div([
                 html.Div(className="blinking-dot", style={'margin':'0 auto','width':'16px','height':'16px'}),
             ], style={'textAlign':'center', 'marginBottom':'10px'}),
-            dcc.Upload(
-                id='upload-document',
-                children=html.Div([
-                    'Drag and Drop or ',
-                    html.A('Select Files')
-                ]),
-                style={
-                    'width': '100%',
-                    'height': '60px',
-                    'lineHeight': '60px',
-                    'borderWidth': '1px',
-                    'borderStyle': 'dashed',
-                    'borderRadius': '5px',
-                    'textAlign': 'center',
-                    'margin': '10px 0'
-                },
-                multiple=True
+            html.Div(id='status-bar', className="alert alert-info", style={'marginBottom': '20px'}),
+            dcc.Loading(
+                id="loading-indicator",
+                type="dot",
+                children=[html.Div(id="loading-output")]
             ),
-            html.Div(id='file-list'),
+            html.Div(id='document-preview', className="border p-3 mb-3"),
+            dbc.Button("Download Document", id="btn-download", color="success", className="mt-3"),
+            dcc.Download(id="download-document"),
             html.Hr(),
-            html.Div([
-                dbc.Button(
-                    doc_type,
-                    id=f'btn-{doc_type.lower().replace("_", "-")}',
-                    color="link",
-                    className="mb-2 w-100 text-left custom-button",
-                    style={'overflow': 'hidden', 'text-overflow': 'ellipsis', 'white-space': 'nowrap'}
-                ) for doc_type in document_types.keys()
-            ])
-        ], width=3),
-        dbc.Col([
+            dcc.Loading(
+                id="chat-loading",
+                type="dot",
+                children=[
+                    dbc.Input(id="chat-input", type="text", placeholder="Chat with AI to update document...", className="mb-2", style={'whiteSpace':'pre-wrap'}),
+                    dbc.Button("Send", id="btn-send-chat", color="primary", className="mb-3"),
+                    html.Div(id="chat-output")
+                ]
+            )
+        ]
+    else:
+        return [
             html.Div([
                 html.Div(className="blinking-dot", style={'margin':'0 auto','width':'16px','height':'16px'}),
             ], style={'textAlign':'center', 'marginBottom':'10px'}),
@@ -94,10 +80,11 @@ app.layout = dbc.Container([
             dbc.Button("Download Document", id="btn-download", color="success", className="mt-3"),
             dcc.Download(id="download-document"),
             html.Hr(),
-            html.Div(id='pink-review-upload', style={'display': 'none'}, children=[
+            html.Div([
+                html.Label(f"Upload {selected_type} Document"),
                 dcc.Upload(
-                    id='upload-pink-review',
-                    children=html.Div(['Drag and Drop or ', html.A('Select Pink Review File')]),
+                    id={'type': 'upload-doc-type', 'index': selected_type},
+                    children=html.Div(['Drag and Drop or ', html.A('Select File')]),
                     style={
                         'width': '100%',
                         'height': '60px',
@@ -110,8 +97,19 @@ app.layout = dbc.Container([
                     },
                     multiple=False
                 ),
-                html.Div(id='pink-review-file-name')
-            ]),
+                html.Div(id={'type': 'uploaded-doc-name', 'index': selected_type}),
+                dbc.RadioItems(
+                    id={'type': 'radio-doc-source', 'index': selected_type},
+                    options=[
+                        {'label': 'Loaded Document', 'value': 'loaded'},
+                        {'label': 'Uploaded Document', 'value': 'uploaded'}
+                    ],
+                    value='loaded',
+                    inline=True,
+                    className="mb-2"
+                ),
+                dbc.Button("Generate Document", id={'type': 'btn-generate-doc', 'index': selected_type}, color="primary", className="mb-3"),
+            ], id={'type': 'doc-type-controls', 'index': selected_type}),
             dcc.Loading(
                 id="chat-loading",
                 type="dot",
@@ -121,6 +119,47 @@ app.layout = dbc.Container([
                     html.Div(id="chat-output")
                 ]
             )
+        ]
+
+app.layout = dbc.Container([
+    dbc.Row([
+        dbc.Col([
+            html.H4("Proposal Documents", className="mt-3 mb-4"),
+            html.Div([
+                html.Div(className="blinking-dot", style={'margin':'0 auto','width':'16px','height':'16px'}),
+            ], style={'textAlign':'center', 'marginBottom':'10px'}),
+            dcc.Upload(
+                id='upload-document',
+                children=html.Div([
+                    'Drag and Drop or ',
+                    html.A('Select Files')
+                ]),
+                style={
+                    'width': '100%',
+                    'height': '60px',
+                    'lineHeight': '60px',
+                    'borderWidth': '1px',
+                    'borderStyle': 'dashed',
+                    'borderRadius': '5px',
+                    'textAlign': 'center',
+                    'margin': '10px 0'
+                },
+                multiple=True
+            ),
+            html.Div(id='file-list'),
+            html.Hr(),
+            html.Div([
+                dbc.Button(
+                    doc_type,
+                    id={'type': 'btn-doc-type', 'index': doc_type},
+                    color="link",
+                    className="mb-2 w-100 text-left custom-button",
+                    style={'overflow': 'hidden', 'text-overflow': 'ellipsis', 'white-space': 'nowrap'}
+                ) for doc_type in document_types.keys()
+            ])
+        ], width=3),
+        dbc.Col([
+            html.Div(id='right-col-content')
         ], width=9)
     ])
 ], fluid=True)
@@ -175,7 +214,7 @@ def update_output(list_of_contents, list_of_names, existing_files):
 @app.callback(
     Output('file-list', 'children', allow_duplicate=True),
     Output('status-bar', 'children', allow_duplicate=True),
-    Input({'type': 'remove-file', 'index': dash.ALL}, 'n_clicks'),
+    Input({'type': 'remove-file', 'index': ALL}, 'n_clicks'),
     State('file-list', 'children'),
     prevent_initial_call=True
 )
@@ -190,6 +229,23 @@ def remove_file(n_clicks, existing_files):
     logging.info(f"Removed file: {removed_file}")
     return [file for file in existing_files if file['props']['children'][1]['props']['children'] != removed_file], "Document removed. Please upload a document and click 'Shred' to begin."
 
+@app.callback(
+    Output('right-col-content', 'children'),
+    [Input({'type': 'btn-doc-type', 'index': ALL}, 'n_clicks')],
+    [State({'type': 'btn-doc-type', 'index': ALL}, 'id')]
+)
+def update_right_col(n_clicks_list, btn_ids):
+    triggered = callback_context.triggered
+    if not triggered or all(x is None for x in n_clicks_list):
+        selected_type = "Shred"
+    else:
+        idx = [i for i, x in enumerate(n_clicks_list) if x]
+        if idx:
+            selected_type = btn_ids[idx[-1]]['index']
+        else:
+            selected_type = "Shred"
+    return get_right_col_content(selected_type)
+
 def generate_document(document_type, file_contents):
     prompt = f"""Generate a {document_type} based on the following project artifacts:
 {' '.join(file_contents)}
@@ -243,70 +299,106 @@ Now, generate the {document_type}:
     Output('document-preview', 'children'),
     Output('loading-output', 'children'),
     Output('status-bar', 'children', allow_duplicate=True),
-    Output('pink-review-upload', 'style'),
-    [Input(f'btn-{doc_type.lower().replace("_", "-")}', 'n_clicks') for doc_type in document_types.keys()],
-    State('pink-review-file-name', 'children'),
+    Input({'type': 'btn-doc-type', 'index': 'Shred'}, 'n_clicks'),
+    prevent_initial_call=True
+)
+def generate_shred_doc(n_clicks):
+    global current_document, document_type, shredded_document
+    if not uploaded_files:
+        return html.Div("Please upload a document before shredding."), "", "Please upload a document before shredding."
+    file_contents = list(uploaded_files.values())
+    try:
+        shredded_document = generate_document("Shred", file_contents)
+        current_document = shredded_document
+        return dcc.Markdown(shredded_document), "Shred generated", "Document shredded. You can now proceed with other operations."
+    except Exception as e:
+        logging.error(f"Error generating document: {str(e)}")
+        return html.Div(f"Error generating document: {str(e)}"), "Error", "An error occurred while shredding the document."
+
+@app.callback(
+    Output({'type': 'uploaded-doc-name', 'index': MATCH}, 'children'),
+    Output({'type': 'upload-doc-type', 'index': MATCH}, 'contents'),
+    Input({'type': 'upload-doc-type', 'index': MATCH}, 'contents'),
+    State({'type': 'upload-doc-type', 'index': MATCH}, 'filename'),
+    State({'type': 'upload-doc-type', 'index': MATCH}, 'id')
+)
+def update_uploaded_doc_name(contents, filename, id_dict):
+    if contents is not None:
+        uploaded_doc_contents[id_dict['index']] = (contents, filename)
+        logging.info(f"{id_dict['index']} file uploaded: {filename}")
+        return filename, contents
+    return "", None
+
+@app.callback(
+    Output('document-preview', 'children', allow_duplicate=True),
+    Output('loading-output', 'children', allow_duplicate=True),
+    Output('status-bar', 'children', allow_duplicate=True),
+    Input({'type': 'btn-generate-doc', 'index': ALL}, 'n_clicks'),
+    State({'type': 'btn-generate-doc', 'index': ALL}, 'id'),
+    State({'type': 'radio-doc-source', 'index': ALL}, 'value'),
+    State({'type': 'upload-doc-type', 'index': ALL}, 'contents'),
+    State({'type': 'upload-doc-type', 'index': ALL}, 'filename'),
     prevent_initial_call=True
 )
-def generate_document_preview(*args):
+def generate_other_doc(n_clicks_list, btn_ids, radio_values, upload_contents, upload_filenames):
     global current_document, document_type, shredded_document, pink_review_document
-    ctx = dash.callback_context
+    ctx = callback_context
     if not ctx.triggered:
         raise dash.exceptions.PreventUpdate
-    button_id = ctx.triggered[0]['prop_id'].split('.')[0]
-    document_type = button_id.replace('btn-', '').replace('-', '_').title()
-    pink_review_file = args[-1]
-
-    if not uploaded_files and document_type != "Shred":
-        return html.Div("Please upload and shred a document first."), "", "Please upload and shred a document first.", {'display': 'none'}
+    idx = [i for i, x in enumerate(n_clicks_list) if x]
+    if not idx:
+        raise dash.exceptions.PreventUpdate
+    idx = idx[-1]
+    doc_type = btn_ids[idx]['index']
+    document_type = doc_type
 
-    if document_type == "Shred":
-        if not uploaded_files:
-            return html.Div("Please upload a document before shredding."), "", "Please upload a document before shredding.", {'display': 'none'}
-        file_contents = list(uploaded_files.values())
-        try:
-            shredded_document = generate_document(document_type, file_contents)
-            return dcc.Markdown(shredded_document), f"{document_type} generated", "Document shredded. You can now proceed with other operations.", {'display': 'none'}
-        except Exception as e:
-            logging.error(f"Error generating document: {str(e)}")
-            return html.Div(f"Error generating document: {str(e)}"), "Error", "An error occurred while shredding the document.", {'display': 'none'}
+    if doc_type == "Shred":
+        raise dash.exceptions.PreventUpdate
 
     if shredded_document is None:
-        return html.Div("Please shred a document first."), "", "Please shred a document first.", {'display': 'none'}
+        return html.Div("Please shred a document first."), "", "Please shred a document first."
 
-    if document_type == "Pink Review":
-        return html.Div("Please upload a Pink Team document or use the generated one."), "", "Please upload a Pink Team document or use the generated one.", {'display': 'block'}
+    source = radio_values[idx] if radio_values and len(radio_values) > idx else 'loaded'
+    doc_content = None
 
-    if document_type in ["Red", "Red Review"] and pink_review_document is None:
-        return html.Div("Please complete Pink Review first."), "", "Please complete Pink Review first.", {'display': 'none'}
-
-    try:
-        if document_type == "Pink Review" and pink_review_file:
-            current_document = generate_document(document_type, [pink_review_file, shredded_document])
-        elif document_type in ["Red", "Red Review"]:
-            current_document = generate_document(document_type, [pink_review_document, shredded_document])
+    if source == 'uploaded':
+        if upload_contents and len(upload_contents) > idx and upload_contents[idx] and upload_filenames and len(upload_filenames) > idx and upload_filenames[idx]:
+            doc_content = process_document(upload_contents[idx], upload_filenames[idx])
+        else:
+            return html.Div("Please upload a document to use as source."), "", "Please upload a document to use as source."
+    else:
+        if doc_type == "Pink":
+            doc_content = shredded_document
+        elif doc_type == "Pink Review":
+            doc_content = pink_review_document if pink_review_document else ""
+        elif doc_type == "Red":
+            doc_content = pink_review_document if pink_review_document else ""
+        elif doc_type == "Red Review":
+            doc_content = pink_review_document if pink_review_document else ""
+        elif doc_type == "Gold":
+            doc_content = shredded_document
+        elif doc_type == "Gold Review":
+            doc_content = shredded_document
+        elif doc_type == "Virtual Board":
+            doc_content = shredded_document
+        elif doc_type == "LOE":
+            doc_content = shredded_document
         else:
-            current_document = generate_document(document_type, [shredded_document])
+            doc_content = shredded_document
 
-        if document_type == "Pink Review":
+    try:
+        if doc_type == "Pink Review":
+            current_document = generate_document(doc_type, [doc_content, shredded_document])
             pink_review_document = current_document
-
-        logging.info(f"{document_type} document generated successfully.")
-        return dcc.Markdown(current_document), f"{document_type} generated", f"{document_type} document generated successfully.", {'display': 'none'}
+        elif doc_type in ["Red", "Red Review"]:
+            current_document = generate_document(doc_type, [doc_content, shredded_document])
+        else:
+            current_document = generate_document(doc_type, [doc_content])
+        logging.info(f"{doc_type} document generated successfully.")
+        return dcc.Markdown(current_document), f"{doc_type} generated", f"{doc_type} document generated successfully."
     except Exception as e:
         logging.error(f"Error generating document: {str(e)}")
-        return html.Div(f"Error generating document: {str(e)}"), "Error", "An error occurred while generating the document.", {'display': 'none'}
-
-@app.callback(
-    Output('pink-review-file-name', 'children'),
-    Input('upload-pink-review', 'contents'),
-    State('upload-pink-review', 'filename')
-)
-def update_pink_review_filename(contents, filename):
-    if contents is not None:
-        logging.info(f"Pink Review file uploaded: {filename}")
-        return filename
-    return ""
+        return html.Div(f"Error generating document: {str(e)}"), "Error", "An error occurred while generating the document."
 
 @app.callback(
     Output('chat-output', 'children'),