Update app.py

app.py

@@ -124,27 +124,41 @@ First, analyze the user's scenario and the provided data schemas. Identify the k
 Based on your mapping, formulate a step-by-step plan. Describe the data cleaning, merging, grouping, and aggregation steps needed to answer the user's request using the columns you identified.
 
 **Step 3: Write the Python Script.**
-Based on your plan, write a complete Python script that performs the analysis.
+Based on your plan, write a complete Python script.
 
 CRITICAL SCRIPTING RULES:
-1.  **DYNAMIC DATAFRAME IDENTIFICATION:** The order of DataFrames in the `dfs` list is NOT guaranteed. Your script MUST identify the correct DataFrame to use for each part of the analysis by checking for the presence of the columns you mapped in Step 1. Do NOT use hardcoded indices like `dfs[0]`.
-2.  **VERIFY COLUMN EXISTENCE:** Only use columns that you have explicitly identified and mapped in your plan. This will prevent `KeyError`.
-3.  **NO FILE READING:** The data is already in the `dfs` list.
-4.  **STRICTLY JSON OUTPUT:** The script's ONLY output must be a single JSON object.
-5.  **ROBUST & GENERIC:** Write robust code that can handle potential missing data (`errors='coerce'`, checking for `None`) and is not hardcoded to specific values from this single request.
+1.  **DYNAMIC DATAFRAME IDENTIFICATION:** Your script MUST identify the correct DataFrame by checking for the presence of the columns you mapped in Step 1. Do NOT use hardcoded indices like `dfs[0]`.
+2.  **ROBUST SUCCESS CHECK (MOST IMPORTANT TO PREVENT AMBIGUITY ERROR):** After attempting to find a DataFrame, you MUST check for success by comparing the result to `None`. Do NOT use `if not my_dataframe:` as this is ambiguous.
+    ```python
+    # Good, robust code
+    def find_df_by_cols(dfs, required_cols):
+        for df in dfs:
+            if all(col in df.columns for col in required_cols):
+                return df
+        return None
+
+    primary_df = find_df_by_cols(dfs, ['user_id', 'transaction_amount'])
+    
+    # This is the correct way to check for failure
+    if primary_df is None:
+        raise ValueError("Could not find the primary dataframe based on its columns.")
+    ```
+3.  **VERIFY COLUMN EXISTENCE:** Only use columns that you have explicitly identified and mapped.
+4.  **NO FILE READING:** The data is already in the `dfs` list.
+5.  **STRICTLY JSON OUTPUT:** The script's ONLY output must be a single JSON object.
+6.  **ROBUST & GENERIC:** Write robust code that can handle potential missing data (`errors='coerce'`, checking for `None`).
 
 Now, provide your response in the following format:
 
 **ANALYSIS PLAN:**
 ```text
 **1. Concept-to-Column Mapping:**
-- Concept: [e.g., 'Hospitals' or 'Facilities'] -> Mapped Column: [e.g., `Facility` from the wait times dataframe]
-- Concept: [e.g., 'Surgical Wait Time' Metric] -> Mapped Column: [e.g., `Surgery_Median` from the wait times dataframe]
-- Concept: [e.g., 'Geographic Locations'] -> Mapped Columns: [e.g., `latitude`, `longitude` from the facilities dataframe]
+- Concept: [e.g., 'Hospitals'] -> Mapped Column: [e.g., `Facility`]
+- Concept: [e.g., 'Surgical Wait Time'] -> Mapped Column: [e.g., `Surgery_Median`]
 
 **2. Step-by-Step Analysis:**
-1.  **Data Identification:** Identify the necessary dataframes by checking for the mapped columns (e.g., find the DF with 'Surgery_Median', find the DF with 'facility_name').
-2.  **Data Cleaning:** [Describe steps, e.g., "Convert metric columns to numeric using `pd.to_numeric`..."]
+1.  **Data Identification:** [e.g., "Define a helper function to find dataframes by checking for key columns..."]
+2.  **Data Cleaning:** [e.g., "Convert metric columns to numeric..."]
 3.  **Analysis Step A:** [e.g., "Group the primary dataframe by the 'Facility' column and calculate the mean of the 'Surgery_Median' column..."]
 4.  ... 
 
@@ -158,7 +172,7 @@ import re
 # Main analysis logic...
 # ...
 # Final print statement
-print(json.dumps(final_data_structure, indent=4))
+print(json.dumps(final_data_structure, indent=4))```
 """
     generated_text = cohere_chat(prompt_for_coder)
     # This regex is more robust for extracting the final code block