Update app.py

app.py

@@ -11,7 +11,7 @@ import gradio as gr
 import pandas as pd
 from datetime import datetime
 import regex as re2
-import re  # Standard library regex module
+import re
 
 # --- BACKEND IMPORTS ---
 from langchain_cohere import ChatCohere
@@ -44,36 +44,26 @@ def _sanitize_text(s: str) -> str:
 def _create_python_script(user_scenario: str, schema_context: str) -> str:
     """Asks the AI to write a Python script that outputs raw, structured JSON."""
     
-    # --- THE FINAL ALIGNMENT AND BUG FIX IS HERE ---
-    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, you cannot assume the zone is in the facility list. You must first identify the high-priority zone from the beds data, and 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.
-"""
-
+    # --- THE FINAL, DEFINITIVE PROMPT FIX ---
     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}
+You are an expert Python data scientist. Your job is to write a script to analyze data that has ALREADY been loaded.
 
---- DATA SCHEMA ---
+--- DATA CONTEXT ---
+The data has been pre-loaded into a Python list of pandas DataFrames called `dfs`.
 {schema_context}
---- END SCHEMA ---
+--- END DATA CONTEXT ---
 
 CRITICAL RULES:
-1.  Your script's ONLY output should be a single JSON object printed to stdout containing the raw data findings.
-2.  Use the exact, case-sensitive column names from the schema.
-3.  Before converting strings to numbers, you MUST robustly clean them of all non-numeric characters (e.g., $, %, ~) using `re.sub()`.
+1.  **DO NOT READ FILES:** You MUST NOT include any code that reads files (e.g., `pd.read_csv`). The data is ALREADY loaded in the `dfs` variable. You MUST use this variable as your starting point. 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. This JSON must contain the raw data findings for each of the user's tasks.
+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 JSON object with the results.
+Now, write the complete Python script that analyzes the `dfs` variable and prints a single, serializable JSON object.
 ```python
 """
     generated_text = cohere_chat(prompt_for_coder)
@@ -101,12 +91,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.
+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.
 
-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'S ORIGINAL SCENARIO ---
 {user_scenario}
 --- END SCENARIO ---
 
@@ -118,7 +105,7 @@ 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.
+4.  Ensure you fully address ALL evaluation questions.
 """
     return _generate_long_report(prompt_for_writer)
 
@@ -234,14 +221,12 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
 
         chat_with_user_msg = _append_msg(chat_history_list, "user", prompt)
         
-        # This is a dummy update function for now, as we're not streaming mid-process
         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()
         
-        # The handle function is now called with the dummy update function
         ai_response_text = handle(prompt, files, dummy_update)
         
         final_chat = _append_msg(chat_with_user_msg, "assistant", ai_response_text)
@@ -256,10 +241,9 @@ 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(" - ") # Safer split
+        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['files'])
@@ -271,7 +255,7 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
         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=[display_history])
     clear_btn.click(lambda: (None, None, [], []), outputs=[prompt_input, files_input, chat_history_output, assessment_history])
     ping_btn.click(ping_cohere, outputs=[ping_out])
     privacy_link.click(lambda: gr.update(visible=True), outputs=[privacy_modal])