Update app.py

app.py

@@ -1,66 +1,60 @@
 # app.py
-
-from __future__ import annotations
-import os
-import io
-import json
-import traceback
-from contextlib import redirect_stdout
-from typing import List, Dict, Any
-
-import gradio as gr
-import pandas a# app.py
 #
-# This file defines a Gradio-based AI data analyst application with
-# support for maintaining a persistent chat and assessment history.
-# The history feature preserves each chat/assessment session (prompt,
-# associated files, generated response, and full conversation) so that
-# users can revisit past analyses without losing any existing
-# functionality. A dropdown selector in the "Assessment History" tab
-# allows users to select and review previous sessions, including the
-# complete chat transcript.
+# Gradio-based AI data analyst app with persistent chat & assessment history.
+# Each session stores: timestamp, prompt, files (if any), final response, and full chat transcript.
 
 from __future__ import annotations
-import os
+
 import io
 import json
+import os
 import traceback
 from contextlib import redirect_stdout
-from typing import List, Dict, Any
+from datetime import datetime
+from typing import Any, Dict, List
 
 import gradio as gr
 import pandas as pd
-from datetime import datetime
 import regex as re2
 import re
 
-from langchain_cohere import ChatCohere
+from langchain_cohere import ChatCohere  # noqa: F401
 
 from settings import (
     GENERAL_CONVERSATION_PROMPT,
-    COHERE_MODEL_PRIMARY, COHERE_TIMEOUT_S, USE_OPEN_FALLBACKS
+    COHERE_MODEL_PRIMARY,
+    COHERE_TIMEOUT_S,  # noqa: F401
+    USE_OPEN_FALLBACKS,  # noqa: F401
 )
 from audit_log import log_event
 from privacy import safety_filter, refusal_reply
 from llm_router import cohere_chat, _co_client, cohere_embed
 
+
 def load_markdown_text(filepath: str) -> str:
     try:
-        with open(filepath, 'r', encoding='utf-8') as f:
+        with open(filepath, "r", encoding="utf-8") as f:
             return f.read()
     except FileNotFoundError:
         return f"**Error:** Document `{os.path.basename(filepath)}` not found."
 
+
 def _sanitize_text(s: str) -> str:
-    if not isinstance(s, str): return s
-    return re2.sub(r'[\p{C}--[\n\t]]+', '', s)
+    if not isinstance(s, str):
+        return s
+    # Remove control characters (except newline and tab)
+    return re2.sub(r"[\p{C}--[\n\t]]+", "", s)
+
 
 def _create_python_script(user_scenario: str, schema_context: str) -> str:
     EXPERT_ANALYTICAL_GUIDELINES = """
 --- EXPERT ANALYTICAL GUIDELINES ---
 When writing your script, you MUST follow these expert business rules:
-1.  **Linking Datasets Rule:** If you need to connect facilities to health zones when the 'zone' column is not in the facility list, you must first identify the high-priority zone from the beds data, then find the major city (by facility count) in the facility list, and *then* assess that city's capacity. Do not try to filter the facility list by a 'zone' column if it does not exist in the schema.
-2.  **Prioritization Rule:** To prioritize locations, you MUST combine the most recent population data with specific high-risk health indicators to create a multi-factor risk score.
+1.  **Linking Datasets Rule:** If you need to connect facilities to health zones when the 'zone' column is not in the facility list,
+    you must first identify the high-priority zone from the beds data, then find the major city (by facility count) in the facility list,
+    and *then* assess that city's capacity. Do not try to filter the facility list by a 'zone' column if it does not exist in the schema.
+2.  **Prioritization Rule:** To prioritize locations, you MUST combine the most recent population data with specific high-risk health indicators
+    to create a multi-factor risk score.
 3.  **Capacity Calculation Rule:** For capacity over a 3-month window, assume **60 working days**.
 4.  **Cost Calculation Rule:** Sum 'Startup cost' and 'Ongoing cost' per person before multiplying.
 """
@@ -90,23 +84,27 @@ Now, write the complete Python script that performs the analysis and prints a si
 """
     generated_text = cohere_chat(prompt_for_coder)
     match = re2.search(r"```python\n(.*?)```", generated_text, re2.DOTALL)
-    if match: return match.group(1).strip()
+    if match:
+        return match.group(1).strip()
     return "print(json.dumps({'error': 'Failed to generate a valid Python script.'}))"
 
+
 def _generate_long_report(prompt: str) -> str:
     try:
         client = _co_client()
-        if not client: return "Error: Cohere client not initialized."
+        if not client:
+            return "Error: Cohere client not initialized."
         response = client.chat(
             model=COHERE_MODEL_PRIMARY,
             message=prompt,
-            max_tokens=4096
+            max_tokens=4096,
         )
         return response.text
     except Exception as e:
         log_event("cohere_chat_error", None, {"err": str(e)})
         return f"Error during final report generation: {e}"
 
+
 def _generate_final_report(user_scenario: str, raw_data_json: str) -> str:
     prompt_for_writer = f"""\
 You are an expert management consultant and data analyst.
@@ -130,57 +128,72 @@ Now, write the final, polished report. The report MUST:
 """
     return _generate_long_report(prompt_for_writer)
 
+
 def _append_msg(h: List[Dict[str, str]], r: str, c: str) -> List[Dict[str, str]]:
     return (h or []) + [{"role": r, "content": c}]
 
+
 def ping_cohere() -> str:
     try:
         cli = _co_client()
-        if not cli: return "Cohere client not initialized."
+        if not cli:
+            return "Cohere client not initialized."
         vecs = cohere_embed(["hello", "world"])
         return f"Cohere OK ✅ (model={COHERE_MODEL_PRIMARY})" if vecs else "Cohere reachable."
-    except Exception as e: return f"Cohere ping failed: {e}"
+    except Exception as e:
+        return f"Cohere ping failed: {e}"
+
 
 def handle(user_msg: str, files: list, yield_update) -> str:
     try:
         safe_in, blocked_in, reason_in = safety_filter(user_msg, mode="input")
-        if blocked_in: return refusal_reply(reason_in)
+        if blocked_in:
+            return refusal_reply(reason_in)
 
         file_paths: List[str] = [getattr(f, "name", None) or f for f in (files or [])]
 
         if file_paths:
             dataframes, schema_parts = [], []
             for i, p in enumerate(file_paths):
-                if p.endswith('.csv'):
-                    try: df = pd.read_csv(p)
-                    except UnicodeDecodeError: df = pd.read_csv(p, encoding='latin1')
+                if p.endswith(".csv"):
+                    try:
+                        df = pd.read_csv(p)
+                    except UnicodeDecodeError:
+                        df = pd.read_csv(p, encoding="latin1")
                     dataframes.append(df)
-                    schema_parts.append(f"DataFrame `dfs[{i}]` (`{os.path.basename(p)}`):\n{df.head().to_markdown()}\n")
+                    schema_parts.append(
+                        f"DataFrame `dfs[{i}]` (`{os.path.basename(p)}`):\n{df.head().to_markdown()}\n"
+                    )
 
-            if not dataframes: return "Please upload at least one CSV file."
+            if not dataframes:
+                return "Please upload at least one CSV file."
 
             schema_context = "\n".join(schema_parts)
 
-            yield_update("```
+            yield_update("""```
 🧠 Generating aligned analysis script...
-```" )
+```""")
             analysis_script = _create_python_script(safe_in, schema_context)
 
-            yield_update("```
+            yield_update("""```
 ⚙️ Executing script to extract raw data...
-```" )
+```""")
             execution_namespace = {"dfs": dataframes, "pd": pd, "re": re, "json": json}
             output_buffer = io.StringIO()
 
             try:
-                with redirect_stdout(output_buffer): exec(analysis_script, execution_namespace)
+                with redirect_stdout(output_buffer):
+                    exec(analysis_script, execution_namespace)
                 raw_data_output = output_buffer.getvalue()
             except Exception as e:
-                return f"An error occurred executing the script: {e}\n\nGenerated Script:\n```python\n{analysis_script}\n```" 
+                return (
+                    f"An error occurred executing the script: {e}\n\nGenerated Script:\n"
+                    f"```python\n{analysis_script}\n```"
+                )
 
-            yield_update("```
+            yield_update("""```
 ✍️ Synthesizing final comprehensive report...
-```" )
+```""")
             final_report = _generate_final_report(safe_in, raw_data_output)
             return _sanitize_text(final_report)
         else:
@@ -192,19 +205,21 @@ def handle(user_msg: str, files: list, yield_update) -> str:
         log_event("app_error", None, {"err": str(e), "tb": tb})
         return f"A critical error occurred: {e}"
 
+
 PRIVACY_POLICY_TEXT = load_markdown_text("privacy_policy.md")
 TERMS_OF_SERVICE_TEXT = load_markdown_text("terms_of_service.md")
 
 with gr.Blocks(theme="soft", css="style.css") as demo:
-    # Maintain a persistent history of past assessments or chat sessions.
-    # Each entry in `assessment_history` is a dictionary containing:
-    # - id: timestamp of the session (string)
-    # - prompt: the initial user prompt (string)
-    # - files: list of filenames uploaded by the user (list of str)
-    # - response: the assistant's final response text (string)
-    # - chat_history: the full chat transcript as a list of message dictionaries
+    # Persistent history of past assessments / chat sessions
+    # Each entry:
+    # - id: timestamp
+    # - prompt: original prompt
+    # - files: list of uploaded filenames
+    # - response: final response text
+    # - chat_history: full transcript (list of {role, content})
     assessment_history = gr.State([])
 
+    # Modals
     with gr.Group(visible=False) as privacy_modal:
         with gr.Blocks():
             gr.Markdown(PRIVACY_POLICY_TEXT)
@@ -215,13 +230,25 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
             gr.Markdown(TERMS_OF_SERVICE_TEXT)
             close_terms_btn = gr.Button("Close")
 
+    # UI
     gr.Markdown("# Universal AI Data Analyst")
     with gr.Row(variant="panel"):
         with gr.Column(scale=1):
             gr.Markdown("## New Assessment")
-            gr.Markdown("<p style='font-size:0.9rem; color: #6C757D;'>Upload CSVs for data analysis, or just enter a prompt to chat.</p>")
-            files_input = gr.Files(label="Upload Data Files (.csv)", file_count="multiple", type="filepath", file_types=[".csv"])
-            prompt_input = gr.Textbox(label="Prompt", placeholder="Paste your scenario or question here.", lines=15)
+            gr.Markdown(
+                "<p style='font-size:0.9rem; color: #6C757D;'>Upload CSVs for data analysis, or just enter a prompt to chat.</p>"
+            )
+            files_input = gr.Files(
+                label="Upload Data Files (.csv)",
+                file_count="multiple",
+                type="filepath",
+                file_types=[".csv"],
+            )
+            prompt_input = gr.Textbox(
+                label="Prompt",
+                placeholder="Paste your scenario or question here.",
+                lines=15,
+            )
             with gr.Row():
                 send_btn = gr.Button("▶️ Send / Run Analysis", variant="primary", scale=2)
                 clear_btn = gr.Button("🗑️ Clear")
@@ -230,67 +257,66 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
         with gr.Column(scale=2):
             with gr.Tabs():
                 with gr.TabItem("Current Assessment", id=0):
-                    chat_history_output = gr.Chatbot(label="Analysis Output", type="messages", height=600)
+                    chat_history_output = gr.Chatbot(
+                        label="Analysis Output", type="messages", height=600
+                    )
                 with gr.TabItem("Assessment History", id=1):
                     gr.Markdown("## Review Past Assessments")
-                    history_dropdown = gr.Dropdown(label="Select an assessment to review", choices=[])
-                    # Use Markdown to display details of the selected assessment, including chat transcript.
+                    history_dropdown = gr.Dropdown(
+                        label="Select an assessment to review", choices=[]
+                    )
                     history_display = gr.Markdown(label="Selected Assessment Details")
-    with gr.Row(): gr.Markdown("---")
+
+    with gr.Row():
+        gr.Markdown("---")
+
     with gr.Row():
         privacy_link = gr.Button("Privacy Policy", variant="link")
         terms_link = gr.Button("Terms of Service", variant="link")
 
+    # Logic
+
     def run_analysis_wrapper(prompt, files, chat_history_list, history_state_list):
-        """Handle a new user prompt and update chat and assessment history.
-
-        This wrapper manages the entire lifecycle of a chat or data analysis:
-        1. Append the user's message to the ongoing conversation.
-        2. Dispatch the request to the AI handler and receive a response.
-        3. Construct a new session entry (with timestamp, prompt, files, response and full chat).
-        4. Update the persistent history and dropdown choices.
-
-        Args:
-            prompt (str): The current user prompt.
-            files (list): List of file paths selected by the user.
-            chat_history_list (list): Current chat conversation as a list of message dicts.
-            history_state_list (list): List of past assessment/chat sessions.
-
-        Returns:
-            tuple: Updated chat history list, updated history list, and updated dropdown choices.
+        """
+        Handle a new user prompt and update chat & assessment history.
         """
         if not prompt:
             gr.Warning("Please enter a prompt.")
             yield chat_history_list, history_state_list, gr.update()
             return
 
-        # Append the user's message to the existing chat history
+        # Append user's message
         chat_with_user_msg = _append_msg(chat_history_list, "user", prompt)
 
-        # Provide immediate feedback to the user that analysis is in progress
-        def dummy_update(message):
-            # This callback is intentionally left blank; progress messages are not streamed here
+        # Optional streaming update callback (unused here)
+        def dummy_update(message: str):
             pass
 
-        thinking_message = _append_msg(chat_with_user_msg, "assistant", "```
+        # Show thinking message
+        thinking_message = _append_msg(
+            chat_with_user_msg,
+            "assistant",
+            """```
 🧠 Generating and executing analysis... Please wait.
-```" )
-        # Yield intermediate state showing a thinking message
+```""",
+        )
         yield thinking_message, history_state_list, gr.update()
 
-        # Run the AI handler (analysis or chat) to get the assistant's response
+        # Run analysis/chat
         ai_response_text = handle(prompt, files, dummy_update)
 
-        # Append the assistant's final response to the chat conversation
+        # Append final assistant response
         final_chat = _append_msg(chat_with_user_msg, "assistant", ai_response_text)
         timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
 
-        # Capture uploaded filenames (if any)
+        # Capture filenames (if any)
         file_names: List[str] = []
         if files:
-            file_names = [os.path.basename(f.name if hasattr(f, 'name') else f) for f in files]
+            file_names = [
+                os.path.basename(f.name if hasattr(f, "name") else f) for f in files
+            ]
 
-        # Build a new entry for the assessment/chat history
+        # Create a new history record (always, even for chat-only)
         new_entry = {
             "id": timestamp,
             "prompt": prompt,
@@ -299,324 +325,45 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
             "chat_history": final_chat,
         }
 
-        # Update the history state (initialize if necessary)
         updated_history: List[Dict[str, Any]] = (history_state_list or []) + [new_entry]
+        history_labels = [
+            f"{item['id']} - {item['prompt'][:40]}..." for item in updated_history
+        ]
 
-        # Build dropdown labels showing timestamp and a snippet of the prompt
-        history_labels = [f"{item['id']} - {item['prompt'][:40]}..." for item in updated_history]
-
-        # Return the final chat, updated history, and updated dropdown choices
         yield final_chat, updated_history, gr.update(choices=history_labels)
 
     def view_history(selection: str, history_state_list: List[Dict[str, Any]]) -> str:
-        """Render the details of a selected past assessment or chat session.
-
-        The selection string contains the timestamp and prompt snippet separated by ' - '.
-        This function locates the corresponding history entry and returns a formatted
-        Markdown string with all relevant details, including the full chat transcript.
-
-        Args:
-            selection (str): The selected dropdown label of the form 'timestamp - prompt...'.
-            history_state_list (list): The list of stored history entries.
-
-        Returns:
-            str: Markdown-formatted details of the selected session.
+        """
+        Render details for a selected past assessment/chat session.
         """
         if not selection or not history_state_list:
             return ""
-        # Extract the unique ID (timestamp) from the dropdown label
-        # The dropdown label is of the form "timestamp - snippet..."
+        # Selection label format: "timestamp - prompt..."
         try:
             selected_id = selection.split(" - ", 1)[0]
         except Exception:
             selected_id = selection
-        # Find the matching session in the history
-        selected_assessment = next((item for item in history_state_list if item.get("id") == selected_id), None)
-
-        if selected_assessment:
-            # Prepare file list display
-            file_list = selected_assessment.get('files', [])
-            file_list_md = "\n- ".join(file_list) if file_list else "*(no files uploaded)*"
-
-            # Prepare chat history display: show each role/message pair on its own line
-            chat_entries = selected_assessment.get("chat_history", [])
-            chat_md_lines = []
-            for msg in chat_entries:
-                role = msg.get("role", "").capitalize()
-                content = msg.get("content", "")
-                chat_md_lines.append(f"**{role}:** {content}")
-            chat_md = "\n\n".join(chat_md_lines)
-
-            return f"""### Assessment from: {selected_assessment['id']}
-**Files Used:**\n- {file_list_md}
----
-**Original Prompt:**\n> {selected_assessment['prompt']}
----
-**AI Generated Response:**\n{selected_assessment['response']}
----
-**Chat Transcript:**\n{chat_md}
-"""
-        return "Could not find the selected assessment."
-
-    # Register interaction handlers
-    send_btn.click(
-        run_analysis_wrapper,
-        inputs=[prompt_input, files_input, chat_history_output, assessment_history],
-        outputs=[chat_history_output, assessment_history, history_dropdown]
-    )
-    history_dropdown.change(
-        view_history,
-        inputs=[history_dropdown, assessment_history],
-        outputs=[history_display]
-    )
-    clear_btn.click(
-        lambda: (None, None, []),
-        outputs=[prompt_input, files_input, chat_history_output]
-    )
-    ping_btn.click(ping_cohere, outputs=[ping_out])
-    privacy_link.click(lambda: gr.update(visible=True), outputs=[privacy_modal])
-    close_privacy_btn.click(lambda: gr.update(visible=False), outputs=[privacy_modal])
-    terms_link.click(lambda: gr.update(visible=True), outputs=[terms_modal])
-    close_terms_btn.click(lambda: gr.update(visible=False), outputs=[terms_modal])
-
-if __name__ == "__main__":
-    if not os.getenv("COHERE_API_KEY"):
-        print("🔴 COHERE_API_KEY environment variable not set. Application may not function correctly.")
-    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))s pd
-from datetime import datetime
-import regex as re2
-import re
-
-from langchain_cohere import ChatCohere
-
-from settings import (
-    GENERAL_CONVERSATION_PROMPT,
-    COHERE_MODEL_PRIMARY, COHERE_TIMEOUT_S, USE_OPEN_FALLBACKS
-)
-from audit_log import log_event
-from privacy import safety_filter, refusal_reply
-from llm_router import cohere_chat, _co_client, cohere_embed
-
-def load_markdown_text(filepath: str) -> str:
-    try:
-        with open(filepath, 'r', encoding='utf-8') as f:
-            return f.read()
-    except FileNotFoundError:
-        return f"**Error:** Document `{os.path.basename(filepath)}` not found."
-
-def _sanitize_text(s: str) -> str:
-    if not isinstance(s, str): return s
-    return re2.sub(r'[\p{C}--[\n\t]]+', '', s)
-
-def _create_python_script(user_scenario: str, schema_context: str) -> str:
-    EXPERT_ANALYTICAL_GUIDELINES = """
---- EXPERT ANALYTICAL GUIDELINES ---
-When writing your script, you MUST follow these expert business rules:
-1.  **Linking Datasets Rule:** If you need to connect facilities to health zones when the 'zone' column is not in the facility list, you must first identify the high-priority zone from the beds data, then find the major city (by facility count) in the facility list, and *then* assess that city's capacity. Do not try to filter the facility list by a 'zone' column if it does not exist in the schema.
-2.  **Prioritization Rule:** To prioritize locations, you MUST combine the most recent population data with specific high-risk health indicators to create a multi-factor risk score.
-3.  **Capacity Calculation Rule:** For capacity over a 3-month window, assume **60 working days**.
-4.  **Cost Calculation Rule:** Sum 'Startup cost' and 'Ongoing cost' per person before multiplying.
-"""
-
-    prompt_for_coder = f"""
-You are an expert Python data scientist. Your job is to write a script to extract the data needed to answer the user's request.
-You have dataframes in a list `dfs`.
-
-{EXPERT_ANALYTICAL_GUIDELINES}
 
---- DATA SCHEMA ---
-{schema_context}
---- END DATA SCHEMA ---
-
-CRITICAL RULES:
-1.  **DO NOT READ FILES:** You MUST NOT include `pd.read_csv`. The data is ALREADY loaded in the `dfs` variable. You MUST use this variable. Failure to do so will cause a fatal error.
-2.  **JSON OUTPUT ONLY:** Your script's ONLY output must be a single JSON object printed to stdout containing the raw data findings.
-3.  **BE PRECISE:** Use the exact, case-sensitive column names from the schema and robustly clean strings (`re.sub()`) before converting to numbers.
-4.  **JSON SERIALIZATION:** Before adding data to your final dictionary for JSON conversion, you MUST convert any pandas-specific types (like `int64`) to standard Python types using `.item()` for single values or `.tolist()` for lists.
-
---- USER'S SCENARIO ---
-{user_scenario}
-
---- PYTHON SCRIPT ---
-Now, write the complete Python script that performs the analysis and prints a single, serializable JSON object.
-```python
-"""
-    generated_text = cohere_chat(prompt_for_coder)
-    match = re2.search(r"```python\n(.*?)```", generated_text, re2.DOTALL)
-    if match: return match.group(1).strip()
-    return "print(json.dumps({'error': 'Failed to generate a valid Python script.'}))"
-
-def _generate_long_report(prompt: str) -> str:
-    try:
-        client = _co_client()
-        if not client: return "Error: Cohere client not initialized."
-        response = client.chat(
-            model=COHERE_MODEL_PRIMARY,
-            message=prompt,
-            max_tokens=4096
+        selected_assessment = next(
+            (item for item in history_state_list if item.get("id") == selected_id),
+            None,
         )
-        return response.text
-    except Exception as e:
-        log_event("cohere_chat_error", None, {"err": str(e)})
-        return f"Error during final report generation: {e}"
-
-def _generate_final_report(user_scenario: str, raw_data_json: str) -> str:
-    prompt_for_writer = f"""
-You are an expert management consultant and data analyst.
-A data science script has run to extract key findings. You have the user's original request and the raw JSON data.
-
-Your task is to synthesize these raw findings into a single, comprehensive, and professional report that directly answers all of the user's questions with detailed justifications.
-
---- USER'S ORIGINAL SCENARIO & DELIVERABLES ---
-{user_scenario}
---- END SCENARIO ---
-
---- RAW DATA FINDINGS (JSON) ---
-{raw_data_json}
---- END RAW DATA ---
-
-Now, write the final, polished report. The report MUST:
-1.  Follow the "Expected Output Format" requested by the user.
-2.  Use tables, bullet points, and DETAILED narrative justifications for each recommendation.
-3.  Synthesize the raw data into actionable insights. Do not just copy the raw numbers; interpret them.
-4.  Ensure you fully address ALL evaluation questions, especially the final recommendations.
-"""
-    return _generate_long_report(prompt_for_writer)
-
-def _append_msg(h: List[Dict[str, str]], r: str, c: str) -> List[Dict[str, str]]:
-    return (h or []) + [{"role": r, "content": c}]
-
-def ping_cohere() -> str:
-    try:
-        cli = _co_client()
-        if not cli: return "Cohere client not initialized."
-        vecs = cohere_embed(["hello", "world"])
-        return f"Cohere OK ✅ (model={COHERE_MODEL_PRIMARY})" if vecs else "Cohere reachable."
-    except Exception as e: return f"Cohere ping failed: {e}"
-
-def handle(user_msg: str, files: list, yield_update) -> str:
-    try:
-        safe_in, blocked_in, reason_in = safety_filter(user_msg, mode="input")
-        if blocked_in: return refusal_reply(reason_in)
-
-        file_paths: List[str] = [getattr(f, "name", None) or f for f in (files or [])]
-
-        if file_paths:
-            dataframes, schema_parts = [], []
-            for i, p in enumerate(file_paths):
-                if p.endswith('.csv'):
-                    try: df = pd.read_csv(p)
-                    except UnicodeDecodeError: df = pd.read_csv(p, encoding='latin1')
-                    dataframes.append(df)
-                    schema_parts.append(f"DataFrame `dfs[{i}]` (`{os.path.basename(p)}`):\n{df.head().to_markdown()}\n")
-            
-            if not dataframes: return "Please upload at least one CSV file."
-
-            schema_context = "\n".join(schema_parts)
-            
-            yield_update("```\n🧠 Generating aligned analysis script...\n```")
-            analysis_script = _create_python_script(safe_in, schema_context)
-
-            yield_update("```\n⚙️ Executing script to extract raw data...\n```")
-            execution_namespace = {"dfs": dataframes, "pd": pd, "re": re, "json": json}
-            output_buffer = io.StringIO()
-            
-            try:
-                with redirect_stdout(output_buffer): exec(analysis_script, execution_namespace)
-                raw_data_output = output_buffer.getvalue()
-            except Exception as e:
-                return f"An error occurred executing the script: {e}\n\nGenerated Script:\n```python\n{analysis_script}\n```"
-
-            yield_update("```\n✍️ Synthesizing final comprehensive report...\n```")
-            final_report = _generate_final_report(safe_in, raw_data_output)
-            return _sanitize_text(final_report)
-        else:
-            prompt = f"{GENERAL_CONVERSATION_PROMPT}\n\nUser: {safe_in}\nAssistant:"
-            return _sanitize_text(cohere_chat(prompt) or "How can I help further?")
-
-    except Exception as e:
-        tb = traceback.format_exc()
-        log_event("app_error", None, {"err": str(e), "tb": tb})
-        return f"A critical error occurred: {e}"
-
-PRIVACY_POLICY_TEXT = load_markdown_text("privacy_policy.md")
-TERMS_OF_SERVICE_TEXT = load_markdown_text("terms_of_service.md")
-
-with gr.Blocks(theme="soft", css="style.css") as demo:
-    assessment_history = gr.State([])
-    
-    with gr.Group(visible=False) as privacy_modal:
-        with gr.Blocks():
-            gr.Markdown(PRIVACY_POLICY_TEXT)
-            close_privacy_btn = gr.Button("Close")
-
-    with gr.Group(visible=False) as terms_modal:
-        with gr.Blocks():
-            gr.Markdown(TERMS_OF_SERVICE_TEXT)
-            close_terms_btn = gr.Button("Close")
-
-    gr.Markdown("# Universal AI Data Analyst")
-    with gr.Row(variant="panel"):
-        with gr.Column(scale=1):
-            gr.Markdown("## New Assessment")
-            gr.Markdown("<p style='font-size:0.9rem; color: #6C757D;'>Upload CSVs for data analysis, or just enter a prompt to chat.</p>")
-            files_input = gr.Files(label="Upload Data Files (.csv)", file_count="multiple", type="filepath", file_types=[".csv"])
-            prompt_input = gr.Textbox(label="Prompt", placeholder="Paste your scenario or question here.", lines=15)
-            with gr.Row():
-                send_btn = gr.Button("▶️ Send / Run Analysis", variant="primary", scale=2)
-                clear_btn = gr.Button("🗑️ Clear")
-            ping_btn = gr.Button("Ping Cohere")
-            ping_out = gr.Markdown()
-        with gr.Column(scale=2):
-            with gr.Tabs():
-                with gr.TabItem("Current Assessment", id=0):
-                    chat_history_output = gr.Chatbot(label="Analysis Output", type="messages", height=600)
-                with gr.TabItem("Assessment History", id=1):
-                    gr.Markdown("## Review Past Assessments")
-                    history_dropdown = gr.Dropdown(label="Select an assessment to review", choices=[])
-                    history_display = gr.Markdown(label="Selected Assessment Details")
-    with gr.Row(): gr.Markdown("---")
-    with gr.Row():
-        privacy_link = gr.Button("Privacy Policy", variant="link")
-        terms_link = gr.Button("Terms of Service", variant="link")
-
-    def run_analysis_wrapper(prompt, files, chat_history_list, history_state_list):
-        if not prompt:
-            gr.Warning("Please enter a prompt.")
-            yield chat_history_list, history_state_list, gr.update()
-            return
-
-        chat_with_user_msg = _append_msg(chat_history_list, "user", prompt)
-        
-        def dummy_update(message):
-            pass
-
-        thinking_message = _append_msg(chat_with_user_msg, "assistant", "```\n🧠 Generating and executing analysis... Please wait.\n```")
-        yield thinking_message, history_state_list, gr.update()
-        
-        ai_response_text = handle(prompt, files, dummy_update)
-        
-        final_chat = _append_msg(chat_with_user_msg, "assistant", ai_response_text)
-        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-        
-        if files:
-            file_names = [os.path.basename(f.name if hasattr(f, 'name') else f) for f in files]
-            new_assessment = {"id": timestamp, "prompt": prompt, "files": file_names, "response": ai_response_text}
-            updated_history = (history_state_list or []) + [new_assessment]
-            history_labels = [f"{item['id']} - {item['prompt'][:40]}..." for item in updated_history]
-            yield final_chat, updated_history, gr.update(choices=history_labels)
-        else:
-            yield final_chat, history_state_list, gr.update()
-
-    def view_history(selection, history_state_list):
-        if not selection or not history_state_list:
-            return ""
-        selected_id = selection.split(" - ")
-        selected_assessment = next((item for item in history_state_list if item["id"] == selected_id), None)
-        
-        if selected_assessment:
-            file_list_md = "\n- ".join(selected_assessment.get('files', []))
-            return f"""### Assessment from: {selected_assessment['id']}
+        if not selected_assessment:
+            return "Could not find the selected assessment."
+
+        file_list = selected_assessment.get("files", [])
+        file_list_md = "\n- ".join(file_list) if file_list else "*(no files uploaded)*"
+
+        # Chat transcript (role + content)
+        chat_entries = selected_assessment.get("chat_history", [])
+        chat_md_lines = []
+        for msg in chat_entries:
+            role = msg.get("role", "").capitalize()
+            content = msg.get("content", "")
+            chat_md_lines.append(f"**{role}:** {content}")
+        chat_md = "\n\n".join(chat_md_lines)
+
+        return f"""### Assessment from: {selected_assessment['id']}
 **Files Used:**
 - {file_list_md}
 ---
@@ -625,22 +372,23 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
 ---
 **AI Generated Response:**
 {selected_assessment['response']}
+---
+**Chat Transcript:**
+{chat_md}
 """
-        return "Could not find the selected assessment."
 
+    # Wire up UI events
     send_btn.click(
         run_analysis_wrapper,
         inputs=[prompt_input, files_input, chat_history_output, assessment_history],
-        outputs=[chat_history_output, assessment_history, history_dropdown]
+        outputs=[chat_history_output, assessment_history, history_dropdown],
     )
     history_dropdown.change(
-        view_history,
-        inputs=[history_dropdown, assessment_history],
-        outputs=[history_display]
+        view_history, inputs=[history_dropdown, assessment_history], outputs=[history_display]
     )
     clear_btn.click(
-        lambda: (None, None, []), 
-        outputs=[prompt_input, files_input, chat_history_output]
+        lambda: (None, None, []),  # clear prompt, files, and chat
+        outputs=[prompt_input, files_input, chat_history_output],
     )
     ping_btn.click(ping_cohere, outputs=[ping_out])
     privacy_link.click(lambda: gr.update(visible=True), outputs=[privacy_modal])
@@ -651,4 +399,4 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
 if __name__ == "__main__":
     if not os.getenv("COHERE_API_KEY"):
         print("🔴 COHERE_API_KEY environment variable not set. Application may not function correctly.")
-    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))
+    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))