Update app.py

app.py

@@ -1,3 +1,4 @@
+```python
 import base64
 import io
 import os
@@ -6,7 +7,7 @@ from docx import Document
 from io import BytesIO, StringIO
 import dash # Version 3.0.3 (or compatible)
 import dash_bootstrap_components as dbc # Version 2.0.2 (or compatible)
-from dash import html, dcc, Input, Output, State, callback_context, ALL
+from dash import html, dcc, Input, Output, State, callback_context, ALL, no_update # Import no_update
 from dash.exceptions import PreventUpdate # <-- Import PreventUpdate from here
 import google.generativeai as genai
 from docx.shared import Pt
@@ -14,6 +15,7 @@ from docx.enum.style import WD_STYLE_TYPE
 from PyPDF2 import PdfReader
 import logging
 import uuid # For unique IDs if needed with pattern matching
+import xlsxwriter # Needed for Excel export engine
 
 # --- Logging Configuration ---
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -39,9 +41,9 @@ try:
     else:
         genai.configure(api_key=api_key)
         # Specify a model compatible with function calling or more advanced generation if needed.
-        # Using 'gemini-pro' as a generally available and capable model.
-        model = genai.GenerativeModel('gemini-2.5-pro-preview-03-25')
-        logging.info("Gemini AI configured successfully using 'gemini-pro'.")
+        # Using 'gemini-1.5-pro-latest' for potential better performance, check compatibility if issues arise.
+        model = genai.GenerativeModel('gemini-2.5-pro-preview-03-25'') # Updated model
+        logging.info("Gemini AI configured successfully using 'gemini-1.5-pro-latest'.")
 except Exception as e:
     logging.error(f"Error configuring Gemini AI: {e}", exc_info=True)
     model = None
@@ -182,7 +184,7 @@ app.layout = dbc.Container(fluid=True, className="dbc", children=[
                             )
                         ]), className="mb-3"
                     ),
-                    dbc.Button("Download Output", id="btn-download", color="success", className="mt-3", style={'display': 'none'}), # Hidden initially
+                    dbc.Button("Download Output", id="btn-download", color="success", className="mt-3 me-2", style={'display': 'none'}), # Hidden initially, add margin
                     dcc.Download(id="download-document"),
 
                     html.Hr(),
@@ -197,7 +199,11 @@ app.layout = dbc.Container(fluid=True, className="dbc", children=[
                                 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.Button("Send Chat", id="btn-send-chat", color="secondary", className="mb-3"), # Use secondary style
+                                    # Button Group for Send and Clear Chat
+                                    dbc.ButtonGroup([
+                                        dbc.Button("Send Chat", id="btn-send-chat", color="secondary"), # Use secondary style
+                                        dbc.Button("Clear Chat", id="btn-clear-chat", color="tertiary") # Use tertiary style
+                                    ], className="mb-3"),
                                     html.Div(id="chat-output", style={'whiteSpace': 'pre-wrap', 'marginTop': '10px', 'border': '1px solid #eee', 'padding': '10px', 'borderRadius': '5px', 'minHeight': '50px'}) # Add border/padding
                                 ]
                             )
@@ -271,8 +277,8 @@ def generate_ai_document(doc_type, input_docs, context_docs=None):
     if not model:
         logging.error("Gemini AI model not initialized.")
         return "Error: AI Model not configured. Please check API Key."
-    if not input_docs or not any(input_docs): # Check if list exists and has content
-        logging.warning(f"generate_ai_document called for {doc_type} with no input documents.")
+    if not input_docs or not any(doc.strip() for doc in input_docs if doc): # Check if list exists and has non-empty content
+        logging.warning(f"generate_ai_document called for {doc_type} with no valid input documents.")
         return f"Error: Missing required input document(s) for {doc_type} generation."
 
     # Combine input documents into a single string
@@ -287,7 +293,7 @@ def generate_ai_document(doc_type, input_docs, context_docs=None):
 **Core Instructions:**
 1.  **Adhere Strictly to the Task:** Generate *only* the content for the '{doc_type}'. Do not add introductions, summaries, or conversational filler unless it's part of the requested document format itself.
 2.  **Follow Format Guidelines:**
-    *   **Spreadsheet Types (Shred, Reviews, LOE, Board):** Structure output clearly. Use Markdown tables or a delimited format (like CSV) suitable for parsing. Define clear columns (e.g., `PWS_Section | Requirement | Finding | Recommendation` for reviews; `Section | Task | Estimated_Hours | Resource_Type` for LOE).
+    *   **Spreadsheet Types (Shred, Reviews, LOE, Board):** Structure output clearly. Use Markdown tables or a delimited format (like CSV) suitable for parsing. Define clear columns (e.g., `PWS_Section | Requirement | Finding | Recommendation` for reviews; `Section | Task | Estimated_Hours | Resource_Type` for LOE). Use '|' as the primary delimiter for tables.
     *   **Proposal Sections (Pink, Red, Gold):** Write professional, compelling prose. Use active voice ("MicroHealth will..."). Directly address requirements from context (Shredded PWS). Detail the 'how' (technical approach, methodology, workflow, tools). Incorporate innovation and benefits (efficiency, quality, outcomes). Substantiate claims (e.g., cite Gartner, Forrester if applicable). Clearly state roles/responsibilities (labor categories). Ensure compliance with Section L/M (Evaluation Criteria) from context. Avoid vague terms ('might', 'could', 'potentially'); be assertive and confident. Use paragraphs primarily; limit bullet points to lists where essential.
 3.  **Utilize Provided Documents:**
     *   **Context Document(s):** Use these as the primary reference or baseline (e.g., Shredded Requirements are the basis for compliance).
@@ -308,7 +314,7 @@ def generate_ai_document(doc_type, input_docs, context_docs=None):
 **Detailed Instructions for '{doc_type}':**
 {document_types.get(doc_type, "Generate the requested document based on the inputs and context.")}
 
-**Begin '{doc_type}' Output:**
+**Begin '{doc_type}' Output (Use Markdown table format for spreadsheet types):**
 """
 
     logging.info(f"Generating AI document for: {doc_type}")
@@ -319,11 +325,29 @@ def generate_ai_document(doc_type, input_docs, context_docs=None):
         response = model.generate_content(prompt) # Consider adding request_options={'timeout': 300} if needed
 
         # Handle potential safety blocks or empty responses
-        if not response.parts:
-             logging.warning(f"Gemini AI returned no parts for {doc_type}. Potential safety block or empty response.")
-             generated_text = f"Error: AI returned no content for {doc_type}. This might be due to safety filters or an issue with the prompt/input."
-        else:
-            generated_text = response.text
+        # Accessing response text might differ slightly based on API version/model behavior
+        generated_text = ""
+        try:
+             if hasattr(response, 'text'):
+                 generated_text = response.text
+             elif hasattr(response, 'parts') and response.parts:
+                 # Concatenate text from parts if necessary
+                 generated_text = "".join(part.text for part in response.parts if hasattr(part, 'text'))
+             else:
+                 # Check for finish_reason if available
+                 finish_reason = getattr(response, 'prompt_feedback', {}).get('block_reason') or getattr(response, 'candidates', [{}])[0].get('finish_reason')
+                 logging.warning(f"Gemini AI response for {doc_type} has no text/parts. Finish Reason: {finish_reason}. Response: {response}")
+                 generated_text = f"Error: AI returned no content for {doc_type}. Possible reason: {finish_reason}. Check Gemini safety settings or prompt complexity."
+
+        except Exception as resp_err:
+             logging.error(f"Error extracting text from Gemini response for {doc_type}: {resp_err}", exc_info=True)
+             generated_text = f"Error: Could not parse AI response for {doc_type}."
+
+
+        if not generated_text.strip() and not generated_text.startswith("Error:"):
+             logging.warning(f"Gemini AI returned empty text for {doc_type}.")
+             generated_text = f"Error: AI returned empty content for {doc_type}. Please try again or adjust the input documents."
+
 
         logging.info(f"Successfully generated document for: {doc_type}")
         # Update global state for download/chat *only if successful*
@@ -331,7 +355,7 @@ def generate_ai_document(doc_type, input_docs, context_docs=None):
              current_display_document = generated_text
              current_display_type = doc_type
         else:
-             # Ensure error message is displayed if AI returns an error internally
+             # Ensure error message is displayed if AI returns an error internally or extraction failed
              current_display_document = generated_text
              current_display_type = doc_type # Still set type so user knows what failed
 
@@ -458,7 +482,8 @@ def handle_file_remove(n_clicks, current_file_list_display):
         raise PreventUpdate
 
     # Check if any click count is > 0 (or just check the specific one that triggered)
-    if not any(nc > 0 for nc in n_clicks if nc is not None): # Ensure click happened
+    # Ensure n_clicks is a list before using any()
+    if not n_clicks or not any(nc for nc in n_clicks if nc is not None): # Check if any click occurred
          raise PreventUpdate
 
 
@@ -507,7 +532,8 @@ def handle_file_remove(n_clicks, current_file_list_display):
     else:
         status_message += "Ready for 'Shred' or other actions."
 
-    return updated_file_list_display, status_message
+    # If the list is now empty, return an empty list instead of None
+    return updated_file_list_display if updated_file_list_display else [], status_message
 
 
 # 3. Handle Action Button Clicks (Show Controls or Trigger Generation)
@@ -526,13 +552,14 @@ def handle_action_button(n_clicks):
     global shredded_document, pink_document, red_document, gold_document, pink_review_document, red_review_document, gold_review_document, loe_document, virtual_board_document
     # Reset potentially uploaded review files when a *new* main action is selected from the left nav
     global uploaded_pink_content, uploaded_red_content, uploaded_gold_content
+    global current_display_document, current_display_type # Need to update preview/download state
 
     triggered_id_dict = callback_context.triggered_id
     if not triggered_id_dict or not isinstance(triggered_id_dict, dict) or 'index' not in triggered_id_dict:
         raise PreventUpdate
 
     # Check if any click count is > 0
-    if not any(nc > 0 for nc in n_clicks if nc is not None):
+    if not n_clicks or not any(nc for nc in n_clicks if nc is not None):
          raise PreventUpdate
 
     action_type = triggered_id_dict['index']
@@ -542,8 +569,8 @@ def handle_action_button(n_clicks):
     review_controls_style = {'display': 'none'} # Hide review controls by default
     review_controls_children = []
     status_message = f"Selected action: {action_type}"
-    doc_preview_children = dash.no_update # Avoid clearing preview unless needed
-    loading_style = {'visibility':'hidden'} # Hide loading dots by default
+    doc_preview_children = no_update # Avoid clearing preview unless needed
+    loading_style = {'visibility':'hidden', 'height': '30px'} # Hide loading dots by default
     download_style = {'display': 'none'} # Hide download button by default
 
 
@@ -568,6 +595,8 @@ def handle_action_button(n_clicks):
         if not shredded_document:
             status_message = "Error: Please 'Shred' the source documents first."
             doc_preview_children = html.Div(status_message, className="text-danger")
+            current_display_document = None # Clear potentially stale preview
+            current_display_type = None
             return review_controls_style, [], status_message, doc_preview_children, loading_style, download_style
 
         # Check specific prerequisites and get the document to review
@@ -588,6 +617,8 @@ def handle_action_button(n_clicks):
         if prereq_doc is None and action_type != "Pink Review": # Shred is checked above
              status_message = f"Error: Please complete '{prereq_doc_name.replace(' Document','')}' first."
              doc_preview_children = html.Div(status_message, className="text-danger")
+             current_display_document = None
+             current_display_type = None
              return review_controls_style, [], status_message, doc_preview_children, loading_style, download_style
 
 
@@ -632,15 +663,19 @@ def handle_action_button(n_clicks):
         # Clear preview when showing controls, provide instruction
         doc_preview_children = html.Div(f"Configure input source for {base_doc_type} document and click 'Generate {action_type}'.", style={'padding':'10px'})
         status_message = f"Ready to configure input for {action_type}."
+        current_display_document = None # Clear internal state for preview/download too
+        current_display_type = None
+        download_style = {'display': 'none'} # Ensure download button is hidden
 
 
     # --- Actions Triggering Direct Generation (Shred, Pink, Red, Gold, LOE, Virtual Board) ---
     else:
         review_controls_style = {'display': 'none'} # Hide review controls
         review_controls_children = []
-        loading_style = {'visibility':'visible'} # Show loading dots
+        loading_style = {'visibility':'visible', 'height': '30px'} # Show loading dots
         doc_preview_children = "" # Clear preview while loading/generating
         status_message = f"Generating {action_type}..."
+        download_style = {'display': 'none'} # Hide download during generation
 
         # Determine inputs based on action type
         input_docs = []
@@ -693,8 +728,11 @@ def handle_action_button(n_clicks):
         if generation_possible:
             result_doc = generate_ai_document(action_type, input_docs, context_docs)
 
-            # Store result in the correct global variable
+            # Store result in the correct global variable AND current_display vars
             if result_doc and not result_doc.startswith("Error:"):
+                current_display_document = result_doc # Set for preview/download/chat
+                current_display_type = action_type
+
                 if action_type == "Shred": shredded_document = result_doc
                 elif action_type == "Pink": pink_document = result_doc
                 elif action_type == "Red": red_document = result_doc
@@ -703,20 +741,25 @@ def handle_action_button(n_clicks):
                 elif action_type == "Virtual Board": virtual_board_document = result_doc
                 # Reviews are handled separately
 
-                doc_preview_children = dcc.Markdown(result_doc, style={'wordWrap': 'break-word'})
+                doc_preview_children = dcc.Markdown(result_doc, style={'wordWrap': 'break-word', 'overflowX': 'auto'}) # Allow horizontal scroll for wide tables
                 status_message = f"{action_type} generated successfully."
-                download_style = {'display': 'inline-block'} # Show download button on success
+                download_style = {'display': 'inline-block', 'marginRight': '10px'} # Show download button on success
             else:
                 # If generation failed, result_doc contains the error message from generate_ai_document
                 doc_preview_children = html.Div(result_doc, className="text-danger") # Display error in preview
                 status_message = f"Error generating {action_type}. See preview for details."
+                current_display_document = result_doc # Show error in preview but prevent download/chat
+                current_display_type = action_type
                 download_style = {'display': 'none'} # Hide download button on error
 
         else:
             # Generation not possible due to prerequisites
             doc_preview_children = html.Div(status_message, className="text-danger")
+            current_display_document = None # No doc generated
+            current_display_type = None
+            download_style = {'display': 'none'}
 
-        loading_style = {'visibility':'hidden'} # Hide loading dots when finished/failed
+        loading_style = {'visibility':'hidden', 'height': '30px'} # Hide loading dots when finished/failed
 
 
     return review_controls_style, review_controls_children, status_message, doc_preview_children, loading_style, download_style
@@ -733,17 +776,21 @@ def toggle_review_upload_visibility(radio_value, current_style):
     # Preserves existing style attributes while toggling 'display'
     new_style = current_style.copy() if current_style else {}
     should_display = (radio_value == 'upload')
-
-    if should_display:
-        new_style['display'] = 'block'
-    else:
-        new_style['display'] = 'none'
+    new_display_value = 'block' if should_display else 'none'
 
     # Prevent update if display style is already correct
-    if ('display' in current_style and current_style['display'] == new_style['display']):
+    if current_style and ('display' in current_style and current_style['display'] == new_display_value):
         raise PreventUpdate
     else:
-        logging.debug(f"Toggling review upload visibility. Radio: {radio_value}, New Style: {new_style}")
+        logging.debug(f"Toggling review upload visibility. Radio: {radio_value}, New display: {new_display_value}")
+        new_style['display'] = new_display_value
+        # Make sure other style defaults are present if creating new style dict
+        if not current_style:
+            new_style.update({
+                'width': '100%', 'height': '60px', 'lineHeight': '60px', 'borderWidth': '1px',
+                'borderStyle': 'dashed', 'borderRadius': '5px', 'textAlign': 'center', 'margin': '10px 0',
+                'backgroundColor': '#f8f9fa'
+            })
         return new_style
 
 
@@ -754,30 +801,37 @@ def toggle_review_upload_visibility(radio_value, current_style):
     Input('upload-review-doc', 'contents'),
     State('upload-review-doc', 'filename'),
     # Get the current review type from the button ID that generated the controls
-    State({'type': 'generate-review-button', 'index': ALL}, 'id'),
+    # We need the ID that triggered the controls to be shown, not necessarily the state of the generate button itself
+    # Let's get the action_type from the state of the review controls (which are generated by action buttons)
+    State('review-controls', 'children'),
     prevent_initial_call=True
 )
-def handle_review_upload(contents, filename, button_ids):
+def handle_review_upload(contents, filename, review_controls_children):
     global uploaded_pink_content, uploaded_red_content, uploaded_gold_content
 
-    if contents is None or filename is None or not button_ids:
-        # No file uploaded or controls not fully rendered yet
+    if contents is None or filename is None or not review_controls_children:
+        # No file uploaded or controls not populated yet (e.g., user selected action but hasn't uploaded)
         raise PreventUpdate
 
-    # Determine which review type this upload is for based on the visible button ID
-    # Assumes only one set of review controls is visible, thus only one button ID in the list
-    if not isinstance(button_ids, list) or len(button_ids) == 0:
-         logging.warning("handle_review_upload: Could not determine review type from button ID state.")
-         raise PreventUpdate
+    # Determine which review type this upload is for. Find the H5 title.
+    review_type = None
+    base_type = None
+    try:
+        # Assuming the first child is the H5 title like "Configure Input for Pink Review"
+        if isinstance(review_controls_children, list) and len(review_controls_children) > 0 and isinstance(review_controls_children[0], html.H5):
+            title_text = review_controls_children[0].children
+            # Extract "Pink Review", "Red Review", etc.
+            parts = title_text.split(" for ")
+            if len(parts) > 1:
+                review_type = parts[1].strip()
+                base_type = review_type.split(" ")[0] # Pink, Red, Gold
+    except Exception as e:
+        logging.warning(f"Could not reliably determine review type from review controls children: {e}")
 
-    # Extract the first ID dictionary from the list
-    button_id_dict = button_ids[0]
-    if not isinstance(button_id_dict, dict) or 'index' not in button_id_dict:
-        logging.warning("handle_review_upload: Invalid button ID format.")
-        raise PreventUpdate
+    if not review_type or not base_type:
+         logging.warning("handle_review_upload: Could not determine review type from controls.")
+         return html.Div("Error determining review type.", className="text-danger small"), "Internal error: Could not determine review type."
 
-    review_type = button_id_dict['index'] # e.g., "Pink Review"
-    base_type = review_type.split(" ")[0] # e.g., "Pink"
 
     logging.info(f"Handling upload of file '{filename}' for {review_type} input.")
 
@@ -794,6 +848,9 @@ def handle_review_upload(contents, filename, button_ids):
     if error:
         status_bar_message = f"Error processing uploaded {base_type} file: {error}"
         upload_status_display = html.Div(f"Failed to load {filename}: {error}", className="text-danger small")
+    elif file_content_text is None: # Handle case where process_document returns (None, None)
+        status_bar_message = f"Error processing uploaded {base_type} file: No content could be extracted."
+        upload_status_display = html.Div(f"Failed to load {filename}: No text extracted.", className="text-danger small")
     else:
         status_bar_message = f"Uploaded '{filename}' successfully for {review_type} input."
         upload_status_display = html.Div(f"Using uploaded file: {filename}", className="text-success small")
@@ -823,13 +880,14 @@ def generate_review_document(n_clicks, source_option, button_ids):
     global shredded_document, pink_document, red_document, gold_document
     global pink_review_document, red_review_document, gold_review_document
     global uploaded_pink_content, uploaded_red_content, uploaded_gold_content
+    global current_display_document, current_display_type # Update preview state
 
     triggered_id_dict = callback_context.triggered_id
     if not triggered_id_dict or not isinstance(triggered_id_dict, dict) or 'index' not in triggered_id_dict:
         raise PreventUpdate
 
      # Check if any click count is > 0
-    if not any(nc > 0 for nc in n_clicks if nc is not None):
+    if not n_clicks or not any(nc for nc in n_clicks if nc is not None):
          raise PreventUpdate
 
     review_type = triggered_id_dict['index'] # e.g., "Pink Review"
@@ -838,14 +896,16 @@ def generate_review_document(n_clicks, source_option, button_ids):
 
     doc_preview_children = "" # Clear preview
     status_message = f"Generating {review_type}..."
-    loading_style = {'visibility':'visible'} # Show loading dots
+    loading_style = {'visibility':'visible', 'height': '30px'} # Show loading dots
     download_style = {'display': 'none'} # Hide download initially
+    current_display_document = None # Clear display state
+    current_display_type = None
 
 
     # --- Prerequisite Check ---
     if not shredded_document:
         status_message = "Error: 'Shredded' document is missing. Please perform 'Shred' first."
-        loading_style = {'visibility':'hidden'}
+        loading_style = {'visibility':'hidden', 'height': '30px'}
         doc_preview_children = html.Div(status_message, className="text-danger")
         return doc_preview_children, status_message, loading_style, download_style
 
@@ -860,8 +920,8 @@ def generate_review_document(n_clicks, source_option, button_ids):
         elif base_type == "Gold": input_document_content = gold_document
 
         if not input_document_content:
-            status_message = f"Error: Cannot use 'generated' option. The {input_doc_source_name} was not found (was it generated successfully?)."
-            loading_style = {'visibility':'hidden'}
+            status_message = f"Error: Cannot use 'generated' option. The {input_doc_source_name} was not found (was it generated successfully?). Please generate or upload it."
+            loading_style = {'visibility':'hidden', 'height': '30px'}
             doc_preview_children = html.Div(status_message, className="text-danger")
             return doc_preview_children, status_message, loading_style, download_style
 
@@ -872,14 +932,13 @@ def generate_review_document(n_clicks, source_option, button_ids):
         elif base_type == "Gold": input_document_content = uploaded_gold_content
 
         if not input_document_content:
-            # Check if the upload component has children (file name displayed)
-            status_message = f"Error: Cannot use 'upload' option. No {base_type} document was successfully uploaded and processed for this review step."
-            loading_style = {'visibility':'hidden'}
+            status_message = f"Error: Cannot use 'upload' option. No {base_type} document was successfully uploaded and processed for this review step. Please upload a valid file."
+            loading_style = {'visibility':'hidden', 'height': '30px'}
             doc_preview_children = html.Div(status_message, className="text-danger")
             return doc_preview_children, status_message, loading_style, download_style
     else:
         status_message = f"Error: Invalid source option '{source_option}' selected."
-        loading_style = {'visibility':'hidden'}
+        loading_style = {'visibility':'hidden', 'height': '30px'}
         doc_preview_children = html.Div(status_message, className="text-danger")
         return doc_preview_children, status_message, loading_style, download_style
 
@@ -892,56 +951,76 @@ def generate_review_document(n_clicks, source_option, button_ids):
     review_result = generate_ai_document(review_type, [input_document_content], context_docs=[shredded_document])
 
     if review_result and not review_result.startswith("Error:"):
-        doc_preview_children = dcc.Markdown(review_result, style={'wordWrap': 'break-word'})
+        doc_preview_children = dcc.Markdown(review_result, style={'wordWrap': 'break-word', 'overflowX': 'auto'}) # Allow horizontal scroll
         status_message = f"{review_type} generated successfully using {input_doc_source_name}."
         # Store the result in the correct global variable
         if review_type == "Pink Review": pink_review_document = review_result
         elif review_type == "Red Review": red_review_document = review_result
         elif review_type == "Gold Review": gold_review_document = review_result
-        download_style = {'display': 'inline-block'} # Show download button
+        # Update display state
+        current_display_document = review_result
+        current_display_type = review_type
+        download_style = {'display': 'inline-block', 'marginRight': '10px'} # Show download button
     else:
         # review_result contains the error message
         doc_preview_children = html.Div(f"Error generating {review_type}: {review_result}", className="text-danger")
         status_message = f"Failed to generate {review_type}. See preview for details."
+        current_display_document = review_result # Show error, but don't allow download/chat
+        current_display_type = review_type
         download_style = {'display': 'none'}
 
-    loading_style = {'visibility':'hidden'} # Hide loading dots
+    loading_style = {'visibility':'hidden', 'height': '30px'} # Hide loading dots
     return doc_preview_children, status_message, loading_style, download_style
 
 
-# 7. Handle Chat Interaction to Refine Displayed Document
+# 7. Handle Chat Interaction (Send and Clear)
 @app.callback(
-    Output('chat-output', 'children'),                   # Display chat confirmation/error
-    Output('document-preview', 'children', allow_duplicate=True), # Update the preview
+    Output('chat-output', 'children', allow_duplicate=True),                   # Display chat confirmation/error or clear
+    Output('document-preview', 'children', allow_duplicate=True), # Update the preview on successful refinement
     Output('status-bar', 'children', allow_duplicate=True),     # Update main status
+    Output('chat-input', 'value'), # Clear input field after send/clear
     Input('btn-send-chat', 'n_clicks'),
+    Input('btn-clear-chat', 'n_clicks'),
     State('chat-input', 'value'),                       # Get the chat instruction
     prevent_initial_call=True
 )
-def handle_chat(n_clicks, chat_input):
+def handle_chat(send_clicks, clear_clicks, chat_input):
     global current_display_document, current_display_type
     # Also need to update the specific underlying document variable (e.g., pink_document)
     # so the chat changes persist if that document is used later.
     global shredded_document, pink_document, red_document, gold_document, pink_review_document, red_review_document, gold_review_document, loe_document, virtual_board_document
 
-    if not n_clicks or not chat_input or not chat_input.strip():
-        # No click or empty input
+    # Determine which button was pressed
+    ctx = callback_context
+    if not ctx.triggered:
         raise PreventUpdate
-    if not current_display_document or not current_display_type:
-         # No document currently loaded in the preview to refine
-         return html.Div("Error: No document is currently displayed to refine.", className="text-warning"), dash.no_update, "Cannot refine: No document loaded in preview."
+    button_id = ctx.triggered[0]['prop_id'].split('.')[0]
+
+    # --- Handle Clear Chat ---
+    if button_id == 'btn-clear-chat' and clear_clicks > 0:
+        logging.info("Clear chat button clicked.")
+        return html.Div("Chat cleared.", className="text-muted fst-italic"), no_update, "Chat cleared.", "" # Clear output, no preview change, update status, clear input
 
-    logging.info(f"Chat refinement requested for displayed document type: {current_display_type}. Instruction: '{chat_input[:100]}...'")
+    # --- Handle Send Chat ---
+    if button_id == 'btn-send-chat' and send_clicks > 0:
+        if not chat_input or not chat_input.strip():
+            # No input to send
+            return html.Div("Please enter an instruction.", className="text-warning small"), no_update, "Chat instruction empty.", chat_input # Keep input
+        if not current_display_document or not current_display_type:
+             # No document currently loaded in the preview to refine
+             return html.Div("Error: No document is currently displayed to refine.", className="text-warning"), no_update, "Cannot refine: No document loaded in preview.", "" # Clear input
 
-    # Construct prompt for refinement
-    prompt = f"""**Objective:** Refine the following '{current_display_type}' document based *only* on the user's instruction below.
+        logging.info(f"Chat refinement requested for displayed document type: {current_display_type}. Instruction: '{chat_input[:100]}...'")
+
+        # Construct prompt for refinement
+        prompt = f"""**Objective:** Refine the following '{current_display_type}' document based *only* on the user's instruction below.
 
 **Your Role:** Act as an editor making precise changes.
 
 **Core Instructions:**
 1.  **Apply Instruction:** Modify the 'Original Document' solely based on the 'User Instruction'.
 2.  **Maintain Context:** Preserve the overall structure, tone, and format of the original document unless the instruction explicitly directs otherwise.
-3.  **Output Only Updated Document:** Provide *only* the complete, updated '{current_display_type}' document. Do not add any conversational text, preamble, or explanation of changes.
+3.  **Output Only Updated Document:** Provide *only* the complete, updated '{current_display_type}' document. Do not add any conversational text, preamble, or explanation of changes. Ensure the output format matches the original (e.g., Markdown table if original was a table).
 
 **Original Document:**
 ```text
@@ -956,26 +1035,46 @@ def handle_chat(n_clicks, chat_input):
 **Begin Updated '{current_display_type}' Output:**
 """
 
-    try:
-        # Show loading indicator for chat refinement? (Optional, maybe use inner loading)
-        status_message = f"Refining {current_display_type} based on chat instruction..."
-        # Note: Consider adding a loading state specifically for the chat output area.
+        try:
+            # Show loading indicator for chat refinement? (Optional, maybe use inner loading)
+            status_message = f"Refining {current_display_type} based on chat instruction..."
+            # Note: Using a separate loading indicator for chat output is possible but adds complexity.
+            # For now, rely on the main status bar.
 
-        response = model.generate_content(prompt)
+            response = model.generate_content(prompt)
 
-        # Handle potential safety blocks or empty responses
-        if not response.parts:
-             logging.warning(f"Gemini AI returned no parts for chat refinement of {current_display_type}.")
-             updated_document = f"Error: AI returned no content during refinement. This might be due to safety filters or an issue with the instruction."
-             chat_response_display = html.Div(updated_document, className="text-danger")
-             status_message = f"Error refining {current_display_type} via chat."
-             return chat_response_display, dash.no_update, status_message # Don't update preview if refinement failed
-        else:
-            updated_document = response.text
+            # Handle potential safety blocks or empty responses (refined logic)
+            updated_document = ""
+            try:
+                if hasattr(response, 'text'):
+                    updated_document = response.text
+                elif hasattr(response, 'parts') and response.parts:
+                    updated_document = "".join(part.text for part in response.parts if hasattr(part, 'text'))
+                else:
+                    finish_reason = getattr(response, 'prompt_feedback', {}).get('block_reason') or getattr(response, 'candidates', [{}])[0].get('finish_reason')
+                    logging.warning(f"Gemini AI response for chat refinement of {current_display_type} has no text/parts. Finish Reason: {finish_reason}. Response: {response}")
+                    updated_document = f"Error: AI returned no content during refinement. Possible reason: {finish_reason}. Check Gemini safety settings or instruction complexity."
+
+            except Exception as resp_err:
+                logging.error(f"Error extracting text from Gemini refinement response for {current_display_type}: {resp_err}", exc_info=True)
+                updated_document = f"Error: Could not parse AI response during refinement for {current_display_type}."
+
+            if not updated_document.strip() and not updated_document.startswith("Error:"):
+                logging.warning(f"Gemini AI returned empty text for chat refinement of {current_display_type}.")
+                updated_document = f"Error: AI returned empty content during refinement for {current_display_type}. Please try a different instruction."
+
+
+            # If refinement failed, show error in chat output, don't update preview
+            if updated_document.startswith("Error:"):
+                chat_response_display = html.Div(updated_document, className="text-danger")
+                status_message = f"Error refining {current_display_type} via chat."
+                return chat_response_display, no_update, status_message, "" # Clear input
+
+
+            # --- Successful Refinement ---
             logging.info(f"Successfully refined {current_display_type} via chat.")
 
-            # --- CRITICAL: Update the correct underlying global variable ---
-            # This ensures the refined document is used in subsequent steps if needed.
+            # CRITICAL: Update the correct underlying global variable
             original_doc_updated = False
             if current_display_type == "Shred": shredded_document = updated_document; original_doc_updated = True
             elif current_display_type == "Pink": pink_document = updated_document; original_doc_updated = True
@@ -999,22 +1098,24 @@ def handle_chat(n_clicks, chat_input):
             # Display confirmation in chat output area
             chat_response_display = html.Div([
                 html.Strong("Refinement applied successfully."),
-                #html.P(f'Instruction: "{chat_input}"'), # Optional: Echo instruction
                 html.Hr(),
                 html.Em("Preview above has been updated. The changes will be used if this document is input for subsequent steps.")
             ])
             status_message = f"{current_display_type} updated via chat instruction."
             # Update the document preview itself
-            doc_preview_update = dcc.Markdown(updated_document, style={'wordWrap': 'break-word'})
+            doc_preview_update = dcc.Markdown(updated_document, style={'wordWrap': 'break-word', 'overflowX': 'auto'})
 
-            return chat_response_display, doc_preview_update, status_message
+            return chat_response_display, doc_preview_update, status_message, "" # Clear input field
 
-    except Exception as e:
-        logging.error(f"Error during chat refinement call for {current_display_type}: {e}", exc_info=True)
-        chat_response_display = html.Div(f"Error refining document via chat: {str(e)}", className="text-danger")
-        status_message = f"Error refining {current_display_type} via chat."
-        # Do not update the main document preview if chat refinement fails
-        return chat_response_display, dash.no_update, status_message
+        except Exception as e:
+            logging.error(f"Error during chat refinement call for {current_display_type}: {e}", exc_info=True)
+            chat_response_display = html.Div(f"Error refining document via chat: {str(e)}", className="text-danger")
+            status_message = f"Error refining {current_display_type} via chat."
+            # Do not update the main document preview if chat refinement fails
+            return chat_response_display, no_update, status_message, "" # Clear input
+
+    # If no button was triggered (shouldn't normally happen with prevent_initial_call)
+    raise PreventUpdate
 
 
 # 8. Handle Download Button Click
@@ -1027,8 +1128,8 @@ def download_generated_document(n_clicks):
     """Prepares the currently displayed document for download."""
     global current_display_document, current_display_type
 
-    if not n_clicks or current_display_document is None or current_display_type is None:
-        # No clicks or nothing to download
+    if not n_clicks or current_display_document is None or current_display_type is None or current_display_document.startswith("Error:"):
+        # No clicks, nothing to download, or current display is an error message
         raise PreventUpdate
 
     logging.info(f"Download requested for displayed document: {current_display_type}")
@@ -1047,57 +1148,85 @@ def download_generated_document(n_clicks):
             data_io = StringIO(current_display_document)
             df = None
 
-            # Attempt 1: Try parsing as CSV (or detect delimiter)
-            try:
-                 # Read a sample to sniff delimiter
-                 sniffer_sample = data_io.read(2048) # Read more data for better sniffing
-                 data_io.seek(0) # Reset pointer after reading sample
-                 dialect = pd.io.parsers.readers.csv.Sniffer().sniff(sniffer_sample, delimiters=',|\t') # Sniff common delimiters
-                 df = pd.read_csv(data_io, sep=dialect.delimiter)
-                 logging.info(f"Successfully parsed {current_display_type} using detected delimiter '{dialect.delimiter}'.")
-            except Exception as e_csv:
-                 logging.warning(f"Could not parse {current_display_type} as standard CSV/TSV ({e_csv}). Trying Markdown Table parsing.")
+            # Refined parsing logic: prioritize Markdown tables, then CSV/TSV
+            lines = [line.strip() for line in data_io.readlines() if line.strip()]
+            data_io.seek(0) # Reset pointer
+
+            # Attempt 1: Parse as Markdown table
+            header = []
+            data = []
+            header_found = False
+            separator_found = False
+            potential_header_line = -1
+            potential_sep_line = -1
+
+            for i, line in enumerate(lines):
+                 if line.startswith('|') and line.endswith('|'):
+                     parts = [p.strip() for p in line.strip('|').split('|')]
+                     if not header_found and '---' not in line: # Potential header
+                         header = parts
+                         header_found = True
+                         potential_header_line = i
+                     elif header_found and '---' in line and potential_header_line == i - 1: # Separator line immediately follows header
+                         separator_found = True
+                         potential_sep_line = i
+                         # Optional: Check column count match
+                         if len(line.strip('|').split('|')) != len(header):
+                              logging.warning("Markdown table header/separator column count mismatch detected.")
+                              # Reset, might not be a valid table
+                              header_found = False
+                              separator_found = False
+                              header = []
+                              continue
+                     elif header_found and separator_found and potential_sep_line == i - 1: # Data line follows separator
+                         if len(parts) == len(header):
+                             data.append(parts)
+                         else:
+                             logging.warning(f"Markdown table row data mismatch (expected {len(header)}, got {len(parts)}): {line}")
+                             # Decide whether to continue or break parsing for this row
+                     else: # Doesn't fit the pattern, reset if we were in the middle of parsing a potential table
+                         if header_found or separator_found:
+                            logging.debug(f"Resetting Markdown parse state at line: {line}")
+                            header_found = False
+                            separator_found = False
+                            header = []
+                            data = []
+
+
+            if header and data:
+                df = pd.DataFrame(data, columns=header)
+                logging.info(f"Successfully parsed {current_display_type} as Markdown Table.")
+            else:
+                 # Attempt 2: Try parsing as CSV/TSV
+                 logging.warning(f"Could not parse {current_display_type} as Markdown Table. Attempting CSV/TSV parsing.")
                  data_io.seek(0) # Reset pointer
-
-                 # Attempt 2: Try parsing as a Markdown table (simple version)
-                 lines = [line.strip() for line in data_io.readlines() if line.strip()]
-                 header = []
-                 data = []
-                 header_found = False
-                 separator_found = False
-
-                 for line in lines:
-                      if line.startswith('|') and line.endswith('|'):
-                          parts = [p.strip() for p in line.strip('|').split('|')]
-                          if not header_found:
-                              header = parts
-                              header_found = True
-                          elif '---' in line: # Detect separator line
-                              separator_found = True
-                              # Optional: Check if separator alignment matches header count
-                              if len(line.strip('|').split('|')) != len(header):
-                                   logging.warning("Markdown table header/separator mismatch detected.")
-                                   # Decide whether to proceed or fail parsing
-                          elif header_found and separator_found: # Only add data lines after header and separator
-                               if len(parts) == len(header):
-                                  data.append(parts)
-                               else:
-                                   logging.warning(f"Markdown table row data mismatch (expected {len(header)}, got {len(parts)}): {line}")
-
-
-                 if header and data:
-                     df = pd.DataFrame(data, columns=header)
-                     logging.info(f"Successfully parsed {current_display_type} as Markdown Table.")
-                 else:
-                     logging.warning(f"Could not parse {current_display_type} as Markdown Table after CSV attempt failed.")
-                     # Fallback: If no DataFrame could be created, send as text
-                     return dict(content=current_display_document, filename=f"{safe_filename_base}.txt")
+                 try:
+                      # Read a sample to sniff delimiter
+                      sniffer_sample = data_io.read(4096) # Read more data for better sniffing
+                      data_io.seek(0) # Reset pointer after reading sample
+                      dialect = pd.io.parsers.readers.csv.Sniffer().sniff(sniffer_sample, delimiters=',|\t') # Sniff common delimiters
+                      df = pd.read_csv(data_io, sep=dialect.delimiter)
+                      logging.info(f"Successfully parsed {current_display_type} using detected delimiter '{dialect.delimiter}'.")
+                 except Exception as e_csv:
+                      logging.warning(f"Could not parse {current_display_type} as standard CSV/TSV after Markdown attempt failed ({e_csv}). Sending as text.")
+                      # Fallback: If no DataFrame could be created, send as text
+                      return dict(content=current_display_document, filename=f"{safe_filename_base}.txt")
 
             # If DataFrame was created, save to Excel
             output = BytesIO()
-            # Use xlsxwriter engine for better compatibility/features if needed
+            # Use xlsxwriter engine for better compatibility/features
             with pd.ExcelWriter(output, engine='xlsxwriter') as writer:
-                df.to_excel(writer, sheet_name=current_display_type[:31], index=False) # Sheet name limit is 31 chars
+                sheet_name = safe_filename_base[:31] # Use sanitized name, limit 31 chars
+                df.to_excel(writer, sheet_name=sheet_name, index=False)
+                 # Auto-adjust column widths (optional, can be slow for large files)
+                worksheet = writer.sheets[sheet_name]
+                for idx, col in enumerate(df):  # loop through columns
+                    series = df[col]
+                    max_len = max((
+                        series.astype(str).map(len).max(),  # len of largest item
+                        len(str(series.name))  # len of column name/header
+                    )) + 1  # adding a little extra space
+                    worksheet.set_column(idx, idx, max_len)  # set column width
             logging.info(f"Sending {filename} (Excel format)")
             return dcc.send_bytes(output.getvalue(), filename)
 
@@ -1117,8 +1246,8 @@ def download_generated_document(n_clicks):
                 if paragraph_text.strip():
                     doc.add_paragraph(paragraph_text)
                 else:
-                     # Optionally add an empty paragraph to represent blank lines, or skip
-                     doc.add_paragraph() # Add blank lines if they were intentional
+                     # Add an empty paragraph to represent blank lines
+                     doc.add_paragraph()
 
             # Save the document to an in-memory BytesIO object
             output = BytesIO()
@@ -1138,7 +1267,8 @@ if __name__ == '__main__':
     # Set debug=False for production/deployment environments like Hugging Face Spaces
     # Set host='0.0.0.0' to make the app accessible on the network (required for Docker/Spaces)
     # Default port 8050, using 7860 as often used for ML demos/Spaces
-    # Use server=app.server for Gunicorn compatibility (multi-threading via workers)
     # Multi-threading for multiple simultaneous user support is handled by the deployment server (e.g., Gunicorn with workers), not directly in app.run for production.
+    # Set debug=True locally for development, False for HF Spaces
     app.run(debug=False, host='0.0.0.0', port=7860)
-    print("Dash application has finished running.")
+    print("Dash application has finished running.")
+```