Update app.py

app.py

@@ -1,79 +1,59 @@
-# app.py
-from __future__ import annotations
+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 as pd
 from datetime import datetime
 import regex as re2
 import re
-
-# --- BACKEND IMPORTS ---
+--- BACKEND IMPORTS ---
 from langchain_cohere import ChatCohere
-
-# --- LOCAL MODULE IMPORTS ---
+--- LOCAL MODULE IMPORTS ---
 from settings import (
-    GENERAL_CONVERSATION_PROMPT,
-    COHERE_MODEL_PRIMARY, COHERE_TIMEOUT_S, USE_OPEN_FALLBACKS
+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
-
-# --- UTILITY FUNCTIONS ---
-
+--- UTILITY FUNCTIONS ---
 def load_markdown_text(filepath: str) -> str:
-    """Safely loads text content from a markdown file."""
-    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."
-
+"""Safely loads text content from a markdown file."""
+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)
-
-# --- THE "ANALYST-WRITER" PIPELINE ---
-
+if not isinstance(s, str): return s
+return re2.sub(r'[\p{C}--[\n\t]]+', '', s)
+--- THE "ANALYST-WRITER" PIPELINE ---
 def _create_python_script(user_scenario: str, schema_context: str) -> str:
-    """Asks the AI to write a Python script that outputs raw, structured JSON."""
-    
-    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 ---
+"""Asks the AI to write a Python script that outputs raw, structured JSON."""
+code
+Code
+# --- THE FINAL, MOST ROBUST PROMPT ---
+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 and print the findings as a single JSON object.
+--- DATA CONTEXT ---
+The data is pre-loaded into a Python list of pandas DataFrames called dfs.
 {schema_context}
---- END DATA SCHEMA ---
-
+--- END DATA CONTEXT ---
 CRITICAL RULES:
-1.  **DO NOT READ FILES:** You MUST NOT include `pd.read_csv`. The data is ALREADY loaded in the `dfs` variable.
-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.
-
+DO NOT READ FILES: You MUST NOT include pd.read_csv. The data is in the dfs variable.
+JSON OUTPUT ONLY: Your script's ONLY output must be a single JSON object printed to stdout.
+JSON SERIALIZATION (VERY IMPORTANT): The json library can only handle standard Python types. Before creating the final dictionary, ensure all values are standard types. If a value is a pandas/numpy number (like int64), convert it to a standard Python int or float using .item(). If a value is a pandas Series, convert it using .tolist().
+DEFENSIVE CODING (CRITICAL): Before passing a variable to a function, be paranoid. For example, if you write a helper function that expects a dictionary, DO NOT pass it a list. If a function expects a single item, DO NOT pass it a whole dataframe. Always check the data type of your variables. This will prevent AttributeError crashes.
+BE PRECISE: Use the exact, case-sensitive column names from the schema and robustly clean strings (re.sub()) before converting them to numbers.
 --- 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
+Now, write the complete, robust, and defensive Python script that analyzes the dfs variable and prints a single, serializable JSON object.
+code
+Python
 """
     generated_text = cohere_chat(prompt_for_coder)
     match = re2.search(r"```python\n(.*?)```", generated_text, re2.DOTALL)
@@ -98,12 +78,9 @@ def _generate_long_report(prompt: str) -> str:
 def _generate_final_report(user_scenario: str, raw_data_json: str) -> str:
     """Asks the AI to act as a consultant and write a polished report from the raw data."""
     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.
+You are an expert management consultant. A data science script has extracted key findings. Your task is to synthesize these findings into a professional report that answers the user's questions.
 
---- USER'S ORIGINAL SCENARIO & DELIVERABLES ---
+--- USER'S ORIGINAL SCENARIO ---
 {user_scenario}
 --- END SCENARIO ---
 
@@ -222,8 +199,6 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
         privacy_link = gr.Button("Privacy Policy", variant="link")
         terms_link = gr.Button("Terms of Service", variant="link")
 
-    # --- DEFINITIVE FIX: ALL UI LOGIC IS NOW CORRECTLY INDENTED ---
-
     def run_analysis_wrapper(prompt, files, chat_history_list, history_state_list):
         if not prompt:
             gr.Warning("Please enter a prompt.")
@@ -244,7 +219,7 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
         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]
+            file_names = [os.path.basename(fn.name if hasattr(fn, 'name') else fn) for fn 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]
@@ -252,15 +227,17 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
         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']}
+
+   def view_history(selection, history_state_list):
+    if not selection or not history_state_list:
+        return ""
+    # THE FIX IS HERE: Correctly extract just the timestamp (the first part)
+    selected_id = selection.split(" - ")[0]
+    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']}
 **Files Used:**
 - {file_list_md}
 ---
@@ -270,21 +247,19 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
 **AI Generated Response:**
 {selected_assessment['response']}
 """
-        return "Could not find the selected assessment."
+    return "Could not find the selected assessment."
 
     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]
-    )
+    history_dropdown.change(view_history, inputs=[history_dropdown, assessment_history], outputs=[history_display])
+   # We remove 'assessment_history' from the list of outputs, so it is no longer cleared.
+    # We remove 'assessment_history' from the list of outputs, so it is no longer cleared.
     clear_btn.click(
-        lambda: (None, None, []), 
-        outputs=[prompt_input, files_input, chat_history_output]
+      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])