Update app.py via AI Editor

app.py

@@ -3,344 +3,339 @@ import io
 import os
 import pandas as pd
 from docx import Document
-from io import BytesIO, StringIO
-import dash # Version 3.0.3
-import dash_bootstrap_components as dbc # Version 2.0.2
-from dash import html, dcc, Input, Output, State, callback_context, ALL, no_update
-from dash.exceptions import PreventUpdate
+from io import BytesIO
+import dash
+import dash_bootstrap_components as dbc
+from dash import html, dcc, Input, Output, State, callback_context
 import google.generativeai as genai
 from docx.shared import Pt
 from docx.enum.style import WD_STYLE_TYPE
 from PyPDF2 import PdfReader
-import logging
-import uuid
-import xlsxwriter # Needed for Excel export engine
-import threading # For multi-threading
-import time # For progress indicator
+from io import StringIO
 
-# --- 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])
 
-# --- Initialize Dash app ---
-# dash==3.0.3
-# dash-bootstrap-components==2.0.2
-app = dash.Dash(__name__,
-                external_stylesheets=[dbc.themes.BOOTSTRAP],
-                suppress_callback_exceptions=True,
-                meta_tags=[{"name": "viewport", "content": "width=device-width, initial-scale=1"}])
-server = app.server # Expose server for deployment
+# Configure Gemini AI
+genai.configure(api_key=os.environ["GEMINI_API_KEY"])
+model = genai.GenerativeModel('gemini-2.5-pro-preview-03-25')
 
-# --- Configure Gemini AI ---
-# IMPORTANT: Set the GEMINI_API_KEY environment variable.
-try:
-    # Prefer direct CUDA GPU configuration in app.py - Note: Not directly applicable for cloud APIs like Gemini.
-    # Configuration happens via environment variable or direct API key setting.
-    api_key = os.environ.get("GEMINI_API_KEY")
-    if not api_key:
-        logging.warning("GEMINI_API_KEY environment variable not found. AI features will be disabled.")
-        model = None
-    else:
-        genai.configure(api_key=api_key)
-        # Using 'gemini-1.5-pro-latest' or similar advanced model is recommended.
-        # Using the user-specified model: gemini-2.5-pro-preview-03-25
-        model = genai.GenerativeModel('gemini-2.5-pro-preview-03-25')
-        logging.info("Gemini AI configured successfully using 'gemini-2.5-pro-preview-03-25'.")
-except Exception as e:
-    logging.error(f"Error configuring Gemini AI: {e}", exc_info=True)
-    model = None
-
-# --- Global Variables ---
-# Using dcc.Store for more robust state management is recommended for production,
-# but for simplicity and current scope, using global variables with locks for threading.
-# A lock for thread-safe access to shared global variables
-data_lock = threading.Lock()
-
-# {session_id: {filename: content_text}} - Store uploaded files per session
+# Global variables
 uploaded_files = {}
+current_document = None
+document_type = None
+shredded_document = None
+pink_review_document = None
 
-# {session_id: document_content} - Stores the *results* of generation/review steps per session
-shredded_document = {}
-pink_review_document = {}
-red_review_document = {}
-gold_review_document = {}
-loe_document = {}
-virtual_board_document = {}
-
-# {session_id: document_content} - Stores the *generated* proposal drafts per session
-pink_document = {}
-red_document = {}
-gold_document = {}
-
-# {session_id: content_text} - Store uploaded content specifically for review inputs per session
-uploaded_pink_content = {}
-uploaded_red_content = {}
-uploaded_gold_content = {}
-
-# {session_id: {'doc': content, 'type': doc_type, 'format': format}} - Store the currently displayed document, its type, and format for download/chat per session
-current_display_document = {}
-
-# --- Document Types ---
+# Document types and their descriptions
 document_types = {
-    "Shred": "Generate a requirements spreadsheet from the PWS/Source Docs, identifying action words (shall, will, perform, etc.) by section.",
-    "Pink": "Create a compliant and compelling Pink Team proposal draft based on the Shredded requirements.",
-    "Pink Review": "Evaluate a Pink Team draft against Shredded requirements. Output findings (compliance, gaps, recommendations) in a spreadsheet.",
-    "Red": "Create a Red Team proposal draft, addressing feedback from the Pink Review and enhancing compliance/compellingness.",
-    "Red Review": "Evaluate a Red Team draft against Shredded requirements and Pink Review findings. Output findings in a spreadsheet.",
-    "Gold": "Create a Gold Team proposal draft, addressing feedback from the Red Review for final compliance and polish.",
-    "Gold Review": "Perform a final compliance review of the Gold Team draft against Shredded requirements and Red Review findings. Output findings.",
-    "Virtual Board": "Simulate a source selection board evaluation of the final proposal against PWS/Shred requirements and evaluation criteria (Sec L/M). Output evaluation.",
-    "LOE": "Generate a Level of Effort (LOE) estimate spreadsheet based on the Shredded requirements."
+    "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.",
+    "Pink Review": "Evaluate compliance of the Pink Team document against the requirements and output a spreadsheet of non compliant findings by pws number, the goal of that pws section, what made it non compliant and your recommendations for recovery",
+    "Red": "Produce a Red Team document based on the Pink Review by pws sections. Your goal is to be compliant and compelling by recovering all the findings in Pink Review",
+    "Red Review": "Evaluate compliance of the Red Team document against the requirements and output a spreadsheet of non compliant findings by pws number, the goal of that pws section, what made it non compliant and your recommendations for recovery",
+    "Gold": "Create a Pink Team document based on the PWS response by pws sections. Your goal is to be compliant and compelling by recovering all the findings in Red Review",
+    "Gold Review": "Perform a final compliance review against the requirements and output a spreadsheet of non compliant findings by pws number, the goal of that pws section, what made it non compliant and your recommendations for recovery",
+    "Virtual Board": "Based on the requirements and in particular the evaulation criteria, you will evaluate the proposal as if you were a contracting office and provide section by section evaluation as unsatisfactory, satisfactory, good, very good, excellent and why in a spreadsheet",
+    "LOE": "Generate a Level of Effort (LOE) breakdown as a spreadsheet"
 }
 
-# --- Layout Definition ---
-app.layout = dbc.Container(fluid=True, className="dbc", children=[
-    dcc.Store(id='session-id', storage_type='session'), # Store for unique session ID
-    # Title Row
-    dbc.Row(
-        dbc.Col(html.H1("Proposal AI Assistant", className="text-center my-4"), width=12)
-    ),
-
-    # Progress Indicator Row
-    dbc.Row(
-        dbc.Col(
-            dcc.Loading(
-                id="loading-indicator",
-                type="dots", # Changed to dots
-                children=[html.Div(id="loading-output", style={'height': '10px'})], # Simplified children
-                overlay_style={"visibility":"hidden", "opacity": 0}, # Hide default overlay
-                style={'visibility':'hidden'}, # Initially hidden
-                parent_style={'minHeight': '30px'}, # Ensure space is allocated
-                fullscreen=False,
-            ),
-            width=12,
-            className="text-center mb-3"
-        )
-    ),
-
-    # Main Content Row
+app.layout = dbc.Container([
     dbc.Row([
-        # Left Column (Nav/Upload) - lg=4 (approx 33%)
-        dbc.Col(
-            dbc.Card(
-                dbc.CardBody([
-                    html.H4("1. Upload Source Documents", className="card-title"),
-                    dcc.Upload(
-                        id='upload-document',
-                        children=html.Div(['Drag and Drop or ', html.A('Select PWS/Source Files')]),
-                        style={ # Basic styling, colors/backgrounds handled by CSS
-                            'width': '100%', 'height': '60px', 'lineHeight': '60px',
-                            'borderWidth': '1px', 'borderStyle': 'dashed', 'borderRadius': '5px',
-                            'textAlign': 'center', 'margin': '10px 0'
-                        },
-                        multiple=True
-                    ),
-                    dbc.Card( # Inner card for file list
-                        dbc.CardBody(
-                           html.Div(id='file-list', style={'maxHeight': '150px', 'overflowY': 'auto', 'fontSize': '0.9em'})
-                        ), className="mb-3" # Removed inline style
-                    ),
-                    html.Hr(),
-                    html.H4("2. Select Action", className="card-title mt-3"),
-                    dbc.Card( # Inner card for buttons
-                        dbc.CardBody([
-                            *[dbc.Button(
-                                doc_type,
-                                id={'type': 'action-button', 'index': doc_type},
-                                color="primary", # Use bootstrap classes
-                                className="mb-2 w-100 d-block",
-                                style={'textAlign': 'left', 'whiteSpace': 'normal', 'height': 'auto', 'wordWrap': 'break-word'} # Style for word wrap
-                              ) for doc_type in document_types.keys()]
-                         ])
-                    )
+        dbc.Col([
+            html.H4("Proposal Documents", className="mt-3 mb-4"),
+            dcc.Upload(
+                id='upload-document',
+                children=html.Div([
+                    'Drag and Drop or ',
+                    html.A('Select Files')
                 ]),
-                # color="light", # Let CSS handle background
-                className="h-100 left-nav-card", # Add custom class for CSS targeting
+                style={
+                    'width': '100%',
+                    'height': '60px',
+                    'lineHeight': '60px',
+                    'borderWidth': '1px',
+                    'borderStyle': 'dashed',
+                    'borderRadius': '5px',
+                    'textAlign': 'center',
+                    'margin': '10px 0'
+                },
+                multiple=True
             ),
-            width=12, lg=4, # Full width on small, 4/12 on large
-            className="mb-3 mb-lg-0",
-            style={'paddingRight': '15px'} # Add padding between columns
-        ),
-
-        # Right Column (Status/Preview/Controls/Chat) - lg=8 (approx 67%)
-        dbc.Col(
-            dbc.Card(
-                dbc.CardBody([
-                    dbc.Alert(id='status-bar', children="Upload source documents and select an action.", color="info"),
-                    dbc.Card(id='review-controls-card', children=[dbc.CardBody(id='review-controls')], className="mb-3", style={'display': 'none'}), # Initially hidden review controls
-                    dbc.Card( # Card for preview
-                        dbc.CardBody([
-                            html.H5("Document Preview / Output", className="card-title"),
-                             dcc.Loading(
-                                 id="loading-preview",
-                                 type="circle",
-                                 children=[html.Div(id='document-preview', style={'whiteSpace': 'pre-wrap', 'wordWrap': 'break-word', 'maxHeight': '400px', 'overflowY': 'auto', 'border': '1px solid #ccc', 'padding': '10px', 'borderRadius': '5px'})] # Added wordWrap
-                            )
-                        ]), className="mb-3"
-                    ),
-                    dbc.Button("Download Output", id="btn-download", color="success", className="mt-3 me-2", style={'display': 'none'}), # Initially hidden download
-                    dcc.Download(id="download-document"),
-                    html.Hr(),
-                    dbc.Card( # Card for chat
-                         dbc.CardBody([
-                            html.H5("Refine Output (Chat)", className="card-title"),
-                            dcc.Loading(
-                                id="chat-loading",
-                                type="circle",
-                                children=[
-                                    dbc.Textarea(id="chat-input", placeholder="Enter instructions to refine the document shown above...", className="mb-2", style={'whiteSpace': 'normal', 'wordWrap': 'break-word'}), # Ensure word wrap
-                                    dbc.ButtonGroup([
-                                        dbc.Button("Send Chat", id="btn-send-chat", color="secondary"),
-                                        dbc.Button("Clear Chat", id="btn-clear-chat", color="tertiary")
-                                    ], className="mb-3"),
-                                    html.Div(id="chat-output", style={'whiteSpace': 'pre-wrap', 'wordWrap': 'break-word', 'marginTop': '10px', 'border': '1px solid #eee', 'padding': '10px', 'borderRadius': '5px', 'minHeight': '50px'}) # Added wordWrap
-                                ]
-                            )
-                         ]), className="mb-3"
-                     )
-                ]),
-                # color="white", # Let CSS handle background
-                className="h-100 right-nav-card", # Add custom class for CSS targeting
+            html.Div(id='file-list'),
+            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([
+            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")]
             ),
-            width=12, lg=8, # Full width on small, 8/12 on large
-            style={'paddingLeft': '15px'} # Add padding between columns
-        )
+            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(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",
+                children=[
+                    dbc.Input(id="chat-input", type="text", placeholder="Chat with AI to update document...", className="mb-2"),
+                    dbc.Button("Send", id="btn-send-chat", color="primary", className="mb-3"),
+                    html.Div(id="chat-output")
+                ]
+            )
+        ], width=9)
     ])
-], style={'padding': '0 15px'}) # Add padding around the container
+], fluid=True)
 
+def process_document(contents, filename):
+    content_type, content_string = contents.split(',')
+    decoded = base64.b64decode(content_string)
+    try:
+        if filename.lower().endswith('.docx'):
+            doc = Document(BytesIO(decoded))
+            text = "\n".join([para.text for para in doc.paragraphs])
+            return text
+        elif filename.lower().endswith('.pdf'):
+            pdf = PdfReader(BytesIO(decoded))
+            text = ""
+            for page in pdf.pages:
+                text += page.extract_text()
+            return text
+        else:
+            return f"Unsupported file format: {filename}. Please upload a PDF or DOCX file."
+    except Exception as e:
+        return f"Error processing document: {str(e)}"
 
-# --- Helper Functions ---
+@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, shredded_document
+    if list_of_contents is not None:
+        new_files = []
+        for i, (content, name) in enumerate(zip(list_of_contents, list_of_names)):
+            file_content = process_document(content, name)
+            uploaded_files[name] = file_content
+            new_files.append(html.Div([
+                html.Button('×', id={'type': 'remove-file', 'index': name}, style={'marginRight': '5px', 'fontSize': '10px'}),
+                html.Span(name)
+            ]))
+        if existing_files is None:
+            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."
 
-def get_session_id(session_id_value=None):
-    """Gets the current session ID or generates a new one."""
-    if session_id_value:
-        return session_id_value
-    # Fallback for initial load or if session ID is missing
-    new_id = str(uuid.uuid4())
-    logging.info(f"Generated new session ID: {new_id}")
-    return new_id
+@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, 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)
+    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 parse_generated_content(content_text):
-    """Attempts to parse AI-generated content into a DataFrame if it looks like a table."""
-    try:
-        # Simple check: does it contain multiple lines and pipe characters?
-        if content_text and '\n' in content_text and '|' in content_text:
-            # Try parsing as Markdown-like table (skip lines that don't fit)
-            lines = [line.strip() for line in content_text.strip().split('\n')]
-            # Remove separator lines like |---|---|
-            lines = [line for line in lines if not all(c in '-| ' for c in line)]
-            if len(lines) > 1:
-                # Use the first line as header, split by '|'
-                header = [h.strip() for h in lines[0].strip('|').split('|')]
-                data_rows = []
-                for line in lines[1:]:
-                    values = [v.strip() for v in line.strip('|').split('|')]
-                    if len(values) == len(header): # Ensure matching column count
-                        data_rows.append(values)
-                    else:
-                        logging.warning(f"Skipping row due to mismatched columns: {line}")
+def generate_document(document_type, file_contents):
+    prompt = f"""Generate a {document_type} based on the following project artifacts:
+{' '.join(file_contents)}
+Instructions:
+1. Create the {document_type} as a detailed document.
+2. Use proper formatting and structure.
+3. Include all necessary sections and details.
+4. Start the output immediately with the document content.
+5. IMPORTANT: If the document type is Pink, Red, Gold and not review type, loe or board
+then your goal is to be compliant and compelling based on the
+requrements, write in paragraph in active voice as
+MicroHealth, limit bullets, answer the
+requirement with what MicroHealth will do 
+to satisfy the requirement, the technical 
+approach with innovation for efficiency, 
+productivity, quality and measurable 
+outcomes, the industry standard that 
+methodology is based on if applicable, 
+detail the workflow or steps to accomplish 
+the requirement with labor categories that 
+will do those tasks in that workflow, 
+reference reputable research like gartner, 
+forrester, IDC, Deloitte, Accenture etc 
+with measures of success and substantiation 
+of MicroHealth's approach. Never use soft words
+like maybe, could be, should, possible be definitive in your language and confident.
+6. you must also take into account section L&M of the document which is the evaluation criteria
+to be sure we address them.
+Now, generate the {document_type}:
+"""
 
-                if data_rows:
-                    df = pd.DataFrame(data_rows, columns=header)
-                    logging.info("Successfully parsed generated content as DataFrame.")
-                    return df
-    except Exception as e:
-        logging.warning(f"Could not parse content into DataFrame: {e}. Treating as plain text.")
-    # If parsing fails or it doesn't look like a table, return None
-    logging.info("Content does not appear to be a table or parsing failed. Treating as plain text.")
-    return None
+    response = model.generate_content(prompt)
+    return response.text
 
-def process_document(contents, filename):
-    """Processes uploaded file content (PDF or DOCX) and returns text, or None and error message."""
-    if contents is None:
-        logging.warning(f"process_document called with None contents for {filename}")
-        return None, f"Error: No content provided for {filename}."
+@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, 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()
+    pink_review_file = args[-1]
 
-    try:
-        content_type, content_string = contents.split(',')
-        decoded = base64.b64decode(content_string)
-        logging.info(f"Processing file: {filename}")
-        text = None
-        error_message = None
+    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 filename.lower().endswith('.docx'):
-            doc = Document(io.BytesIO(decoded))
-            text = "\n".join([para.text for para in doc.paragraphs if para.text.strip()])
-            logging.info(f"Successfully processed DOCX: {filename}")
-        elif filename.lower().endswith('.pdf'):
-            pdf = PdfReader(io.BytesIO(decoded))
-            extracted_pages = []
-            for i, page in enumerate(pdf.pages):
-                try:
-                    page_text = page.extract_text()
-                    if page_text:
-                        extracted_pages.append(page_text)
-                except Exception as page_e:
-                    logging.warning(f"Could not extract text from page {i+1} of {filename}: {page_e}")
-            text = "\n\n".join(extracted_pages)
-            if not text:
-                 logging.warning(f"No text extracted from PDF: {filename}. It might be image-based or corrupted.")
-                 error_message = f"Error: No text could be extracted from PDF {filename}. It might be image-based or require OCR."
-            else:
-                logging.info(f"Successfully processed PDF: {filename}")
-        else:
-            logging.warning(f"Unsupported file format: {filename}")
-            error_message = f"Unsupported file format: {filename}. Please upload PDF or DOCX."
+    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'}
 
-        return text, error_message
-    except Exception as e:
-        logging.error(f"Error processing document {filename}: {e}", exc_info=True)
-        return None, f"Error processing file {filename}: {str(e)}"
+    if shredded_document is None:
+        return html.Div("Please shred a document first."), "", "Please shred a document first.", {'display': 'none'}
 
-def get_combined_uploaded_text(session_id, file_dict):
-    """Combines text content of files in the provided dictionary for a session."""
-    with data_lock:
-        session_files = file_dict.get(session_id, {})
-        if not session_files:
-            return ""
-        # Combine content, adding filenames for context if multiple files
-        if len(session_files) > 1:
-            return "\n\n--- FILE BREAK ---\n\n".join(
-                f"**File: {fname}**\n\n{content}" for fname, content in session_files.items()
-            )
+    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'}
+
+    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])
         else:
-            return next(iter(session_files.values()), "")
+            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", "An error occurred while generating the document.", {'display': 'none'}
 
-def generate_ai_document(session_id, doc_type, input_docs, context_docs=None):
-    """Generates document using Gemini AI. Returns generated content and format ('text' or 'dataframe')."""
-    if not model:
-        logging.error(f"[{session_id}] Gemini AI model not initialized.")
-        return "Error: AI Model not configured. Please check API Key.", 'text'
-    if not input_docs or not any(doc.strip() for doc in input_docs if doc):
-        logging.warning(f"[{session_id}] generate_ai_document called for {doc_type} with no valid input documents.")
-        return f"Error: Missing required input document(s) for {doc_type} generation.", 'text'
+@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 ""
 
-    combined_input = "\n\n---\n\n".join(filter(None, input_docs))
-    combined_context = "\n\n---\n\n".join(filter(None, context_docs)) if context_docs else ""
+@app.callback(
+    Output('chat-output', 'children'),
+    Output('document-preview', 'children', allow_duplicate=True),
+    Input('btn-send-chat', 'n_clicks'),
+    State('chat-input', 'value'),
+    prevent_initial_call=True
+)
+def update_document_via_chat(n_clicks, chat_input):
+    global current_document, document_type
+    if not chat_input or current_document is None:
+        raise dash.exceptions.PreventUpdate
+    
+    prompt = f"""Update the following {document_type} based on this instruction: {chat_input}
+Current document:
+{current_document}
+Instructions:
+1. Provide the updated document content.
+2. Maintain proper formatting and structure.
+3. Incorporate the requested changes seamlessly.
+Now, provide the updated {document_type}:
+"""
+    
+    response = model.generate_content(prompt)
+    current_document = response.text
+    
+    return f"Document updated based on: {chat_input}", dcc.Markdown(current_document)
 
-    # Define expected output format based on doc_type
-    is_spreadsheet_type = doc_type in ["Shred", "Pink Review", "Red Review", "Gold Review", "LOE", "Virtual Board"]
-    output_format_instruction = """**Output Format:** Structure the output as a clear, parseable Markdown table. Use '|' as the column delimiter. Define meaningful column headers relevant to the task (e.g., PWS_Section, Requirement, Action_Verb for Shred; Section, Requirement, Compliance_Status, Finding, Recommendation for Reviews; Section, Task, Estimated_Hours, Resource_Type for LOE). Ensure each row corresponds to a distinct item (e.g., requirement, finding, task).""" if is_spreadsheet_type else """**Output Format:** Write professional, compelling proposal prose. Use clear paragraphs and standard formatting. Address all requirements logically. Avoid tables unless explicitly part of the proposal structure."""
+@app.callback(
+    Output("download-document", "data"),
+    Input("btn-download", "n_clicks"),
+    prevent_initial_call=True
+)
+def download_document(n_clicks):
+    global current_document, document_type
+    if current_document is None:
+        raise dash.exceptions.PreventUpdate
+    
+    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")
 
-    prompt = f"""**Objective:** Generate the '{doc_type}' document.
-**Your Role:** Act as an expert proposal writer/analyst specialized in government contracting.
-**Core Instructions:**
-1.  **Adhere Strictly to the Task:** Generate *only* the content for the '{doc_type}'. Do not add introductions, summaries, explanations, or conversational filler unless it's part of the requested document format itself (e.g., an executive summary within a proposal draft).
-2.  **Follow Format Guidelines:** {output_format_instruction}
-3.  **Content Requirements:**
-    *   **Shred:** Identify requirements (explicit and implied), action verbs (shall, will, must, provide, perform, etc.), and PWS section references.
-    *   **Proposal Sections (Pink, Red, Gold):** Write compliant and compelling content. Directly address requirements from the Context Document(s). Detail the 'how' (approach, methodology, tools). Incorporate win themes, strengths, and discriminators. Substantiate claims. Use active voice ("Our team will..."). Ensure compliance with evaluation criteria (e.g., Section L/M). Clearly map responses back to PWS requirements.
-    *   **Reviews (Pink, Red, Gold):** Evaluate the submitted draft against the requirements (Shred/PWS) and previous review findings (if applicable). Identify compliance issues, gaps, weaknesses, and areas for improvement. Provide actionable recommendations. Be specific and reference relevant sections.
-    *   **LOE:** Estimate the Level of Effort (hours, resource types) required to fulfill each major task or requirement identified in the Shred/PWS. Justify estimates briefly if necessary.
-    *   **Virtual Board:** Simulate a source selection evaluation. Assess the final proposal against the PWS/Shred and evaluation criteria (Sec L/M). Assign strengths, weaknesses, deficiencies, risks. Provide a summary evaluation.
-4.  **Utilize Provided Documents:**
-    *   **Context Document(s):** These provide the baseline or reference material (e.g., Shredded Requirements, PWS Section L/M, Previous Review Findings). Refer to them diligently.
-    *   **Primary Input Document(s):** This is the main subject of the task (e.g., the PWS text to be Shredded, the Pink draft to be Reviewed, the Red Review findings to incorporate into the Gold draft). Analyze and process this document according to the task.
-**Provided Documents:**
-**Context Document(s):**
-```text
-{combined_context if combined_context else "N/A"}
-```
-**Primary Input Document(s):**
-```text
-{combined_input}
+if __name__ == '__main__':
+    print("Starting the Dash application...")
+    app.run(debug=False, host='0.0.0.0', port=7860)
+    print("Dash application has finished running.")