Update app.py

app.py

@@ -1,4 +1,4 @@
-
+# app.py
 from fastapi import FastAPI, UploadFile, File, HTTPException, Query
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
@@ -16,6 +16,7 @@ from google.genai import types
 import logging
 import asyncio
 from concurrent.futures import ThreadPoolExecutor
+from functools import partial
 import re
 import traceback
 import motor.motor_asyncio
@@ -116,7 +117,6 @@ def _find_balanced_json(s: str):
             stack.append('{')
         elif ch == '}':
             if not stack:
-                # malformed, but continue
                 return None
             stack.pop()
             if not stack:
@@ -125,7 +125,6 @@ def _find_balanced_json(s: str):
 
 
 def _escape_problematic_backslashes(s: str) -> str:
-    # Escape backslashes that are not followed by valid JSON escape char
     return re.sub(r'\\(?!["\\/bfnrtu])', r'\\\\', s)
 
 
@@ -189,8 +188,8 @@ def stream_save_and_hash(upload_file: UploadFile, tmp_path: str, size_limit: Opt
 async def save_preprocessed_df(df: pd.DataFrame, snapshot_id: str) -> str:
     path = os.path.join(SNAPSHOT_BUCKET, f"{snapshot_id}.csv")
     loop = asyncio.get_running_loop()
-    # df.to_csv(path, index=False) is blocking -> run in executor
-    await loop.run_in_executor(EXECUTOR, df.to_csv, path, False, None)
+    # Use functools.partial to pass keyword args to to_csv (handles pandas 3.0+ keyword-only changes)
+    await loop.run_in_executor(EXECUTOR, partial(df.to_csv, path, index=False))
     return path
 
 
@@ -240,7 +239,8 @@ def get_metadata(df: pd.DataFrame) -> dict:
 
 # ---------- AI generation (blocking) ----------
 def generate_summary_blocking(meta, fiverow, system_prompt_override: Optional[str] = None):
-    api_key = "AIzaSyB1jgGCuzg7ELPwNEEwaluQZoZhxhgLmAs"
+    # using provided API key (kept as-is per deployment)
+    api_key = os.getenv("GEMINI_API_KEY") or "AIzaSyB1jgGCuzg7ELPwNEEwaluQZoZhxhgLmAs"
     if not api_key:
         raise RuntimeError("GEMINI_API_KEY environment variable not set")
     client = genai.Client(api_key=api_key)
@@ -254,9 +254,37 @@ Input contains:
 You must output JSON with the following structure:
 {
   "summary": "<short natural language overview of dataset>",
-  "recommended_charts": [ ... ]
+  "recommended_charts": [
+    {
+      "type": "<one of: bar, pie, timeseries, histogram, scatter, multiple_columns, stacked_bar, heatmap>",
+      "title": "<short title for chart>",
+      "columns": ["<col1>", "<col2>", "..."],
+      "python_code": "<full runnable Python code using seaborn/matplotlib that produces the chart>"
+    },
+    ...
+  ]
 }
-Always produce syntactically valid JSON ONLY.
+Mandatory rules:
+- Always produce syntactically valid JSON ONLY. No text outside the JSON object.
+- Provide at least these chart types somewhere in recommended_charts: bar, pie, timeseries, histogram, scatter, multiple_columns, stacked_bar, heatmap.
+- Use only column names that appear in meta['column_names'].
+- The python_code string must be self-contained and runnable assuming a variable `df` exists containing the full cleaned DataFrame. Start the code with imports:
+    import pandas as pd
+    import seaborn as sns
+    import matplotlib.pyplot as plt
+  and include any necessary preprocessing steps (e.g., parsing dates).
+- For timeseries charts ensure the datetime column is parsed (`pd.to_datetime`) before plotting.
+- For multiple_columns provide a pairplot or facetgrid example that uses up to 4 numeric columns or sensible categorical splits.
+- For stacked_bar, show aggregation code (groupby + unstack) and plotting with df.plot(kind='bar', stacked=True).
+- For heatmap, compute correlation matrix and plot sns.heatmap with annotations.
+- For pie charts, ensure grouping/aggregation when there are >20 unique categories (group small categories into 'Other').
+- For histogram and scatter include axis labels and tight_layout; include plt.show() at the end.
+- Keep code minimal but complete so a user can copy-paste and run (assume seaborn, matplotlib, pandas installed).
+- For each chart add a sensible "columns" list showing which columns the code uses.
+- Do not include examples using columns not present in meta.
+- Do not include more than 10 recommended_charts.
+- Ensure strings inside the JSON are escaped properly so the JSON parses.
+Produce concise natural-language one-line summary in "summary". Ensure JSON is parseable by json.loads in Python.
 """
 
     user_prompt = {"meta": meta, "fiverow": fiverow}
@@ -344,7 +372,6 @@ def extract_chart_data_json_by_type(parsed_summary: dict, df: pd.DataFrame):
                 elif chart_type == "timeseries":
                     df_copy = df[columns].copy()
                     for c in columns:
-                        # attempt parse only if not already datetime dtype
                         if not pd.api.types.is_datetime64_any_dtype(df_copy[c]):
                             df_copy[c] = pd.to_datetime(df_copy[c], errors='coerce')
                     chart_data = df_copy.astype(str).to_dict(orient="records")
@@ -460,7 +487,6 @@ async def analyze_data(file: UploadFile = File(...)):
 
         if existing:
             snapshot_id_return = existing.get("snapshot_id") or str(existing.get("_id"))
-            # ensure structure matches response_model
             summary = existing.get("summary") or {}
             chart_data = existing.get("chart_data") or {}
             metadata = existing.get("metadata") or meta