Update app.py

app.py

@@ -1,9 +1,5 @@
-# -------------------------------
-# 📌 FastAPI AI Data Insights App
-# -------------------------------
-
 from fastapi import FastAPI, Request, File, UploadFile, Form
-from fastapi.responses import HTMLResponse, JSONResponse
+from fastapi.responses import StreamingResponse, JSONResponse, HTMLResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
 import pandas as pd
@@ -11,13 +7,13 @@ from google import genai
 from google.genai import types
 import os
 import json
+import asyncio
 
 # -------------------------------
 # 🔑 Configuration
 # -------------------------------
-API_KEY = os.getenv("GEMINI_API_KEY", "YOUR_GEMINI_KEY")
+API_KEY = os.getenv("GEMINI_API_KEY", "AIzaSyB1jgGCuzg7ELPwNEEwaluQZoZhxhgLmAs")
 MODEL = "gemini-2.5-flash-lite"
-
 client = genai.Client(api_key=API_KEY)
 
 # -------------------------------
@@ -27,12 +23,10 @@ app = FastAPI()
 app.mount("/static", StaticFiles(directory="static"), name="static")
 templates = Jinja2Templates(directory="templates")
 
-
 # -------------------------------
 # 🛠️ Helper Functions
 # -------------------------------
 def get_metadata(df: pd.DataFrame):
-    """Extract lightweight metadata for prompting."""
     return {
         "columns": list(df.columns),
         "dtypes": df.dtypes.apply(lambda x: str(x)).to_dict(),
@@ -43,81 +37,39 @@ def get_metadata(df: pd.DataFrame):
         "sample_rows": df.head(3).to_dict(orient="records"),
     }
 
+async def stream_insights(user_query, metadata):
+    """Generator that yields insights step by step as text/json strings."""
 
-def generate_json_insight(user_query, metadata, df):
-    """Ask Gemini for structured JSON insights across fixed sections."""
+    # Step 1: Start message
+    yield json.dumps({"status": "started", "message": "File received. Extracting metadata..."}) + "\n"
+    await asyncio.sleep(0.5)
+
+    # Step 2: Metadata
+    yield json.dumps({"status": "metadata", "metadata": metadata}) + "\n"
+    await asyncio.sleep(0.5)
+
+    # Step 3: Call Gemini for structured insights
     system_prompt = """
-You are a data analysis assistant.
-You will be given dataset metadata and a user request.
-
-Always return JSON with exactly these 5 sections:
-1. Efficiency Analysis (bar chart, actual vs target if available)
-2. Cumulative Performance (line chart over time if possible)
-3. Process Issues (pie chart breakdown if available)
-4. Planning vs Projection (comparison planned vs projected values)
-5. Loss Analysis (summary with stats: total, avg, min, max)
-
-Schema:
-{
-  "insights": [
-    {
-      "title": "Efficiency Analysis",
-      "type": "bar",
-      "description": "Actual vs Target Efficiency",
-      "chartData": [{"x": "...", "y": ..., "target": ...}],
-      "stats": {}
-    },
-    {
-      "title": "Cumulative Performance",
-      "type": "line",
-      "description": "Cumulative trend over time",
-      "chartData": [],
-      "stats": {}
-    },
-    {
-      "title": "Process Issues",
-      "type": "pie",
-      "description": "Breakdown of process issues",
-      "chartData": [],
-      "stats": {}
-    },
+    You are a data analysis assistant.
+    Always return JSON with this schema:
     {
-      "title": "Planning vs Projection",
-      "type": "comparison",
-      "description": "Planned vs Projected output",
-      "chartData": [],
-      "stats": {}
-    },
-    {
-      "title": "Loss Analysis",
-      "type": "summary",
-      "description": "Key loss statistics",
-      "chartData": [],
-      "stats": {}
+      "excel_info": {...},
+      "data_type_context": "...",
+      "auto_insights": {
+        "insights": [
+          {... Efficiency Analysis ...},
+          {... Cumulative Performance ...},
+          {... Process Issues ...},
+          {... Planning vs Projection ...},
+          {... Loss Analysis ...}
+        ]
+      },
+      "query_insights": {...}
     }
-  ]
-}
-
-Rules:
-- All 5 sections must be present in the JSON.
-- If no data available, return empty arrays/objects.
-- Do NOT output Python code or text explanations, JSON only.
-"""
-
-    user_prompt = f"""
-Dataset metadata:
-Columns: {metadata['columns']}
-Data types: {metadata['dtypes']}
-Null counts: {metadata['null_counts']}
-Unique counts: {metadata['unique_counts']}
-Sample rows: {metadata['sample_rows']}
-
-User request: {user_query}
-"""
-
-    contents = [
-        types.Content(role="user", parts=[types.Part.from_text(text=user_prompt)])
-    ]
+    """
+    user_prompt = f"Dataset metadata: {metadata}\nUser request: {user_query}"
+
+    contents = [types.Content(role="user", parts=[types.Part.from_text(text=user_prompt)])]
     config = types.GenerateContentConfig(
         temperature=0,
         max_output_tokens=2000,
@@ -125,16 +77,34 @@ User request: {user_query}
     )
 
     result = ""
-    for chunk in client.models.generate_content_stream(
-        model=MODEL, contents=contents, config=config
-    ):
+    for chunk in client.models.generate_content_stream(model=MODEL, contents=contents, config=config):
         if chunk.text:
             result += chunk.text
 
     try:
-        return json.loads(result)
+        parsed = json.loads(result)
     except Exception:
-        return {"raw_output": result}
+        yield json.dumps({"status": "error", "raw_output": result}) + "\n"
+        return
+
+    # Step 4: Excel info
+    yield json.dumps({"status": "excel_info", "excel_info": parsed.get("excel_info", {})}) + "\n"
+    await asyncio.sleep(0.5)
+
+    # Step 5: Data type context
+    yield json.dumps({"status": "context", "data_type_context": parsed.get("data_type_context", "")}) + "\n"
+    await asyncio.sleep(0.5)
+
+    # Step 6: Stream each insight one by one
+    for insight in parsed.get("auto_insights", {}).get("insights", []):
+        yield json.dumps({"status": "insight", "insight": insight}) + "\n"
+        await asyncio.sleep(0.5)
+
+    # Step 7: Query insights
+    yield json.dumps({"status": "query", "query_insights": parsed.get("query_insights", {})}) + "\n"
+
+    # Step 8: Completed
+    yield json.dumps({"status": "completed", "message": "All insights generated"}) + "\n"
 
 
 # -------------------------------
@@ -144,18 +114,13 @@ User request: {user_query}
 async def home(request: Request):
     return templates.TemplateResponse("index.html", {"request": request})
 
-
-@app.post("/generate_insight_file")
-async def generate_insight_file(
-    file: UploadFile = File(...), query: str = Form("Analyze the dataset")
-):
-    """Upload Excel, generate structured JSON insights."""
+@app.post("/stream_insights")
+async def stream_insight_file(file: UploadFile = File(...), query: str = Form("Analyze the dataset")):
+    """Stream structured JSON insights step by step."""
     try:
         df = pd.read_excel(file.file)
     except Exception as e:
         return JSONResponse({"success": False, "error": f"Failed to read file: {str(e)}"})
 
     metadata = get_metadata(df)
-    insights = generate_json_insight(query, metadata, df)
-
-    return JSONResponse({"success": True, "insights": insights})
+    return StreamingResponse(stream_insights(query, metadata), media_type="application/json")