Update app.py

app.py

@@ -24,6 +24,8 @@ model = genai.GenerativeModel('gemini-2.5-pro-preview-03-25')
 uploaded_files = {}
 current_document = None
 document_type = None
+shredded_document = None
+pink_review_document = None
 
 # Document types and their descriptions
 document_types = {
@@ -73,7 +75,7 @@ app.layout = dbc.Container([
             ])
         ], width=3),
         dbc.Col([
-            html.Div(style={"height": "20px"}),  # Added small gap
+            html.Div(id='status-bar', className="alert alert-info", style={'marginBottom': '20px'}),
             dcc.Loading(
                 id="loading-indicator",
                 type="dot",
@@ -83,7 +85,24 @@ 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(style={"height": "20px"}),  # Added small gap
+            html.Div(id='pink-review-upload', style={'display': 'none'}, children=[
+                dcc.Upload(
+                    id='upload-pink-review',
+                    children=html.Div(['Drag and Drop or ', html.A('Select Pink Review File')]),
+                    style={
+                        'width': '100%',
+                        'height': '60px',
+                        'lineHeight': '60px',
+                        'borderWidth': '1px',
+                        'borderStyle': 'dashed',
+                        'borderRadius': '5px',
+                        'textAlign': 'center',
+                        'margin': '10px 0'
+                    },
+                    multiple=False
+                ),
+                html.Div(id='pink-review-file-name')
+            ]),
             dcc.Loading(
                 id="chat-loading",
                 type="dot",
@@ -118,12 +137,13 @@ def process_document(contents, filename):
 
 @app.callback(
     Output('file-list', 'children'),
+    Output('status-bar', 'children'),
     Input('upload-document', 'contents'),
     State('upload-document', 'filename'),
     State('file-list', 'children')
 )
 def update_output(list_of_contents, list_of_names, existing_files):
-    global uploaded_files
+    global uploaded_files, shredded_document
     if list_of_contents is not None:
         new_files = []
         for i, (content, name) in enumerate(zip(list_of_contents, list_of_names)):
@@ -135,23 +155,26 @@ def update_output(list_of_contents, list_of_names, existing_files):
             ]))
         if existing_files is None:
             existing_files = []
-        return existing_files + new_files
-    return existing_files
+        shredded_document = None  # Reset shredded document when new files are uploaded
+        return existing_files + new_files, "Document uploaded. Please click 'Shred' to proceed."
+    return existing_files, "Please upload a document and click 'Shred' to begin."
 
 @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'),
     State('file-list', 'children'),
     prevent_initial_call=True
 )
 def remove_file(n_clicks, existing_files):
-    global uploaded_files
+    global uploaded_files, shredded_document
     ctx = dash.callback_context
     if not ctx.triggered:
         raise dash.exceptions.PreventUpdate
     removed_file = ctx.triggered[0]['prop_id'].split(',')[0].split(':')[-1].strip('}')
     uploaded_files.pop(removed_file, None)
-    return [file for file in existing_files if file['props']['children'][1]['props']['children'] != removed_file]
+    shredded_document = None  # Reset shredded document when a file is removed
+    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."
 
 def generate_document(document_type, file_contents):
     prompt = f"""Generate a {document_type} based on the following project artifacts:
@@ -170,28 +193,69 @@ Now, generate the {document_type}:
 @app.callback(
     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'),
     prevent_initial_call=True
 )
 def generate_document_preview(*args):
-    global current_document, document_type
+    global current_document, document_type, shredded_document, pink_review_document
     ctx = dash.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()
-    
-    if not uploaded_files:
-        return html.Div("Please upload project artifacts before generating a document."), ""
+    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'}
+
+    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:
+            print(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 shredded_document is None:
+        return html.Div("Please shred a document first."), "", "Please shred a document first.", {'display': 'none'}
+
+    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'}
+
+    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'}
 
-    file_contents = list(uploaded_files.values())
-    
     try:
-        current_document = generate_document(document_type, file_contents)
-        return dcc.Markdown(current_document), f"{document_type} generated"
+        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])
+        else:
+            current_document = generate_document(document_type, [shredded_document])
+
+        if document_type == "Pink Review":
+            pink_review_document = current_document
+
+        return dcc.Markdown(current_document), f"{document_type} generated", f"{document_type} document generated successfully.", {'display': 'none'}
     except Exception as e:
         print(f"Error generating document: {str(e)}")
-        return html.Div(f"Error generating document: {str(e)}"), "Error"
+        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:
+        return filename
+    return ""
 
 @app.callback(
     Output('chat-output', 'children'),
@@ -230,15 +294,26 @@ def download_document(n_clicks):
     if current_document is None:
         raise dash.exceptions.PreventUpdate
     
-    # Create an in-memory Word document
-    doc = Document()
-    doc.add_paragraph(current_document)
-    
-    # Save the document to a BytesIO object
-    output = BytesIO()
-    doc.save(output)
-    
-    return dcc.send_bytes(output.getvalue(), f"{document_type}.docx")
+    if document_type == "LOE":
+        # Create a pandas DataFrame for LOE
+        df = pd.read_csv(StringIO(current_document))
+        
+        # Save the DataFrame to an Excel file
+        output = BytesIO()
+        with pd.ExcelWriter(output, engine='xlsxwriter') as writer:
+            df.to_excel(writer, sheet_name='LOE', index=False)
+        
+        return dcc.send_bytes(output.getvalue(), f"{document_type}.xlsx")
+    else:
+        # Create an in-memory Word document
+        doc = Document()
+        doc.add_paragraph(current_document)
+        
+        # Save the document to a BytesIO object
+        output = BytesIO()
+        doc.save(output)
+        
+        return dcc.send_bytes(output.getvalue(), f"{document_type}.docx")
 
 if __name__ == '__main__':
     print("Starting the Dash application...")