Major redesign: AI Investigation agent with conversational multi-turn flow, follow-up questions, and better error handling

ai_investigation.py

@@ -1,956 +1,501 @@
 """
-AI INVESTIGATION ASSISTANT
-
-Analyzes vehicle behavior patterns and generates intelligence summaries.
-Works like a surveillance analyst - not SQL generation, but pattern analysis.
-
-RAG System:
-1. Collect detection data from database
-2. Analyze patterns and behavior
-3. Generate intelligence narratives using Mistral
-4. Provide actionable recommendations
+🤖 CONVERSATIONAL AI INVESTIGATION AGENT
+=====================================
+Professional multi-turn conversational agent for vehicle surveillance analysis.
+Works like ChatGPT - asks clarifying questions, provides insights, suggests related queries.
+
+Architecture:
+1. Query Understanding - Parse user intent
+2. Data Collection - Use advanced NLP-to-SQL from database.py
+3. Data Analysis - Extract insights from retrieved data
+4. Response Generation - Answer user question with data
+5. Clarification - Ask follow-up questions or suggest related queries
+6. Memory - Keep conversation context for multi-turn interactions
 """
 
 import pandas as pd
-from sqlalchemy import text
-from datetime import datetime, timedelta
-from database import engine
 import json
-import os
+import traceback
+from datetime import datetime
+from sqlalchemy import text
 
 # =====================================================
-# MISTRAL AI SETUP (OPTIONAL)
+# IMPORTS
 # =====================================================
 
+from database import engine, ask_llm
 try:
     from mistralai.client import MistralClient
     from mistralai.models.chat_message import ChatMessage
     MISTRAL_AVAILABLE = True
 except ImportError:
     MISTRAL_AVAILABLE = False
-    print("⚠️ Mistral AI not installed - using fallback narrative generation")
-except Exception as e:
-    MISTRAL_AVAILABLE = False
-    print(f"⚠️ Mistral AI unavailable: {e} - using fallback mode")
+    print("⚠️ Mistral not available - fallback mode")
+
+# =====================================================
+# CONFIGURATION
+# =====================================================
+
+MISTRAL_MODEL = "mistral-small"
+MISTRAL_TIMEOUT = 10
+DATA_SAMPLE_SIZE = 5  # Show top 5 records in response
+FOLLOW_UP_QUESTIONS = 3  # Max follow-up questions to suggest
 
 
 # =====================================================
-# DATA COLLECTION & ANALYSIS
+# MISTRAL CLIENT
 # =====================================================
 
-def get_vehicle_detections(plate, days=30):
-    """Get vehicle detection history from database"""
+def get_mistral_client():
+    """Get or create Mistral client"""
     try:
-        if not plate:
-            return pd.DataFrame()
-        
-        query = """
-        SELECT 
-            plate,
-            state,
-            vehicle_type,
-            vehicle_conf,
-            location,
-            date,
-            timestamp,
-            camera_id,
-            hour,
-            day
-        FROM vehicle_logs
-        WHERE REPLACE(REPLACE(UPPER(plate), ' ', ''), '-', '') = :plate
-        AND date >= CURRENT_DATE - INTERVAL '1 day' * :days
-        ORDER BY timestamp DESC
-        """
-        
-        clean_plate = plate.replace(" ", "").replace("-", "").upper()
-        
-        with engine.connect() as conn:
-            result = conn.execute(
-                text(query),
-                {"plate": clean_plate, "days": days}
-            )
-            rows = result.fetchall()
-            df = pd.DataFrame(rows, columns=result.keys())
-        
-        return df
+        api_key = "any"  # Using HF inference
+        client = MistralClient(api_key=api_key)
+        return client
     except Exception as e:
-        print(f"Error fetching detections: {e}")
-        return pd.DataFrame()
+        print(f"⚠️ Mistral client error: {e}")
+        return None
 
 
-def get_area_activity(location, days=7):
-    """Get all vehicle activity in a specific area"""
-    try:
-        query = """
-        SELECT 
-            plate,
-            state,
-            vehicle_type,
-            vehicle_conf,
-            location,
-            date,
-            timestamp,
-            hour
-        FROM vehicle_logs
-        WHERE location ILIKE :location
-        AND date >= CURRENT_DATE - INTERVAL '1 day' * :days
-        ORDER BY timestamp DESC
-        """
-        
-        with engine.connect() as conn:
-            result = conn.execute(
-                text(query),
-                {"location": f"%{location}%", "days": days}
-            )
-            rows = result.fetchall()
-            df = pd.DataFrame(rows, columns=result.keys())
-        
-        return df
-    except Exception as e:
-        print(f"Error fetching area activity: {e}")
-        return pd.DataFrame()
+# =====================================================
+# STEP 1: QUERY UNDERSTANDING
+# =====================================================
+
+def understand_query(user_question):
+    """
+    Understand user intent and extract what they're looking for.
+    Returns classification of query type.
+    """
+    q = user_question.lower()
+    
+    intent_types = {
+        "tracking": any(w in q for w in ["where", "location", "found", "been", "detected", "show"]),
+        "counting": any(w in q for w in ["how many", "count", "total", "number of"]),
+        "pattern": any(w in q for w in ["pattern", "habit", "regular", "often", "frequent"]),
+        "comparison": any(w in q for w in ["compare", "vs", "versus", "different", "same"]),
+        "timeline": any(w in q for w in ["when", "time", "date", "hour", "period", "last", "first"]),
+        "alert": any(w in q for w in ["suspicious", "alert", "warning", "anomaly", "unusual"]),
+        "statistics": any(w in q for w in ["statistics", "stats", "analysis", "breakdown", "distribution"])
+    }
+    
+    return {
+        "question": user_question,
+        "primary_intent": max(intent_types.items(), key=lambda x: x[1])[0] if any(intent_types.values()) else "general",
+        "all_intents": {k: v for k, v in intent_types.items() if v}
+    }
+
 
+# =====================================================
+# STEP 2: DATA COLLECTION (ADVANCED)
+# =====================================================
 
-def get_midnight_activity(hours_start=22, hours_end=4, days=7):
-    """Get suspicious midnight activity"""
+def collect_data(user_question):
+    """
+    Collect data using the advanced NLP-to-SQL engine from database.py
+    Handles complex queries with multiple filters.
+    """
     try:
-        query = """
-        SELECT 
-            plate,
-            state,
-            vehicle_type,
-            vehicle_conf,
-            location,
-            date,
-            timestamp,
-            hour
-        FROM vehicle_logs
-        WHERE (hour >= :hour_start OR hour <= :hour_end)
-        AND date >= CURRENT_DATE - INTERVAL '1 day' * :days
-        ORDER BY timestamp DESC
-        """
+        print(f"\n🔍 COLLECTING DATA FOR: {user_question}")
+        
+        # Use the advanced ask_llm function from database.py
+        sql_query = ask_llm(user_question)
         
+        if not sql_query:
+            return None, None, "Could not parse your question. Try being more specific (e.g., 'show bikes in adyar')"
+        
+        print(f"📊 Generated SQL:\n{sql_query}")
+        
+        # Execute query
         with engine.connect() as conn:
-            result = conn.execute(
-                text(query),
-                {"hour_start": hours_start, "hour_end": hours_end, "days": days}
-            )
+            conn.execute(text("SET statement_timeout = 30000"))  # 30 second timeout
+            result = conn.execute(text(sql_query))
             rows = result.fetchall()
-            df = pd.DataFrame(rows, columns=result.keys())
         
-        return df
+        if not rows:
+            return None, sql_query, "No data found matching your criteria"
+        
+        df = pd.DataFrame(rows, columns=result.keys())
+        return df, sql_query, None
+        
     except Exception as e:
-        print(f"Error fetching midnight activity: {e}")
-        return pd.DataFrame()
+        print(f"❌ Data collection error: {e}")
+        traceback.print_exc()
+        error_msg = str(e)
+        if "timeout" in error_msg.lower():
+            return None, None, "Query took too long. Try narrowing down your search (add location, date, or vehicle type)"
+        elif "syntax" in error_msg.lower():
+            return None, None, "Query syntax error. Please rephrase your question"
+        else:
+            return None, None, f"Data collection failed: {error_msg[:80]}"
 
 
 # =====================================================
-# PATTERN ANALYSIS
+# STEP 3: DATA ANALYSIS
 # =====================================================
 
-def analyze_vehicle_behavior(df):
-    """Analyze vehicle behavior patterns"""
-    if df.empty:
-        return {}
-    
-    df['timestamp'] = pd.to_datetime(df['timestamp'])
-    df['date'] = pd.to_datetime(df['date'])
-    
-    analysis = {
-        "total_detections": len(df),
-        "unique_locations": df["location"].nunique(),
-        "unique_dates": df["date"].nunique(),
-        "date_range": f"{df['date'].min().date()} to {df['date'].max().date()}",
-        "vehicle_type": df["vehicle_type"].iloc[0] if len(df) > 0 else "Unknown",
-        "state": df["state"].iloc[0] if len(df) > 0 else "Unknown",
-        "avg_confidence": float(df["vehicle_conf"].mean()),
-        "detection_frequency": len(df) / max(df["date"].nunique(), 1),
-    }
-    
-    # Time pattern analysis
-    df['hour'] = pd.to_datetime(df['timestamp']).dt.hour
-    midnight_activity = len(df[(df['hour'] >= 22) | (df['hour'] <= 4)])
-    analysis["midnight_detections"] = midnight_activity
-    analysis["midnight_percentage"] = round((midnight_activity / len(df) * 100), 2)
-    
-    # Location clustering
-    location_counts = df["location"].value_counts()
-    analysis["top_locations"] = location_counts.head(5).to_dict()
-    analysis["location_clustering"] = len(df[df["location"] == location_counts.index[0]]) / len(df) * 100
-    
-    # Time gaps between detections
-    time_diffs = df['timestamp'].diff().dt.total_seconds() / 3600
-    analysis["avg_time_gap_hours"] = float(time_diffs.mean())
-    analysis["min_time_gap_hours"] = float(time_diffs.min())
-    
-    return analysis
-
-
-def analyze_area_patterns(df):
-    """Analyze patterns in a specific area"""
-    if df.empty:
+def analyze_data(df, user_question):
+    """
+    Analyze collected data to extract insights.
+    Returns structured analysis with key metrics and findings.
+    """
+    if df is None or df.empty:
         return {}
     
-    df['timestamp'] = pd.to_datetime(df['timestamp'])
-    df['hour'] = df['timestamp'].dt.hour
-    
     analysis = {
-        "total_vehicles": df["plate"].nunique(),
-        "total_detections": len(df),
-        "unique_states": df["state"].nunique(),
-        "peak_hour": int(df["hour"].mode()[0]) if len(df["hour"].mode()) > 0 else 0,
+        "total_records": len(df),
+        "columns": list(df.columns),
+        "unique_vehicles": df["plate"].nunique() if "plate" in df.columns else 0,
+        "unique_locations": df["location"].nunique() if "location" in df.columns else 0,
+        "vehicle_types": [],
+        "states": [],
+        "date_range": None,
+        "time_range": None,
+        "key_findings": []
     }
     
-    # Repeated vehicles
-    plate_counts = df["plate"].value_counts()
-    analysis["most_repeated_vehicle"] = plate_counts.index[0] if len(plate_counts) > 0 else None
-    analysis["most_repeated_count"] = int(plate_counts.iloc[0]) if len(plate_counts) > 0 else 0
-    analysis["repeated_vehicles"] = int((plate_counts >= 5).sum())
-    
-    # Time analysis
-    night_detections = len(df[(df["hour"] >= 22) | (df["hour"] <= 4)])
-    analysis["night_activity_percentage"] = round((night_detections / len(df) * 100), 2)
-    
-    return analysis
-
-
-def analyze_midnight_patterns(df):
-    """Analyze midnight/suspicious timing patterns"""
-    if df.empty:
-        return {}
-    
-    df['timestamp'] = pd.to_datetime(df['timestamp'])
+    # Vehicle types
+    if "vehicle_type" in df.columns:
+        type_counts = df["vehicle_type"].value_counts()
+        analysis["vehicle_types"] = [(v, int(c)) for v, c in type_counts.head(5).items()]
     
-    analysis = {
-        "total_midnight_detections": len(df),
-        "unique_vehicles": df["plate"].nunique(),
-        "unique_locations": df["location"].nunique(),
-        "vehicles_with_multiple_visits": int((df["plate"].value_counts() >= 3).sum()),
-    }
+    # States
+    if "state" in df.columns:
+        state_counts = df["state"].value_counts()
+        analysis["states"] = [(s, int(c)) for s, c in state_counts.head(5).items()]
     
-    # Most suspicious vehicles at midnight
-    plate_counts = df["plate"].value_counts()
-    analysis["top_suspicious_vehicles"] = plate_counts.head(5).to_dict()
+    # Date range
+    if "timestamp" in df.columns:
+        try:
+            df["timestamp_dt"] = pd.to_datetime(df["timestamp"], errors='coerce')
+            valid_times = df["timestamp_dt"].dropna()
+            if len(valid_times) > 0:
+                analysis["date_range"] = {
+                    "start": str(valid_times.min()),
+                    "end": str(valid_times.max())
+                }
+        except:
+            pass
+    
+    # Key findings
+    try:
+        if "plate" in df.columns:
+            plate_counts = df["plate"].value_counts()
+            if len(plate_counts) > 0:
+                analysis["key_findings"].append(
+                    f"Most detected: {plate_counts.index[0]} ({plate_counts.iloc[0]} times)"
+                )
+    except:
+        pass
     
-    # Industrial zone activity
-    analysis["industrial_zones_active"] = df["location"].nunique()
+    try:
+        if "location" in df.columns:
+            loc_counts = df["location"].value_counts()
+            if len(loc_counts) > 0:
+                analysis["key_findings"].append(
+                    f"Hottest location: {loc_counts.index[0]} ({loc_counts.iloc[0]} detections)"
+                )
+    except:
+        pass
     
     return analysis
 
 
 # =====================================================
-# MISTRAL AI INVESTIGATION NARRATIVES
+# STEP 4: RESPONSE GENERATION
 # =====================================================
 
-def get_mistral_client():
-    """Initialize Mistral client"""
-    if not MISTRAL_AVAILABLE:
+def generate_answer(user_question, df, analysis, client):
+    """
+    Generate intelligent answer using Mistral with data context.
+    Answers the user's original question based on retrieved data.
+    """
+    if df is None or df.empty:
         return None
     
-    try:
-        api_key = os.environ.get("MISTRAL_API_KEY", "")
-        if not api_key:
-            return None
-        client = MistralClient(api_key=api_key)
-        return client
-    except Exception as e:
-        print(f"Error initializing Mistral: {e}")
-        return None
+    # Build data context for Mistral
+    context = f"""
+You are a professional vehicle surveillance analyst. A user asked you a question and you retrieved vehicle detection data from a database.
 
+USER QUESTION: {user_question}
 
-def generate_investigation_narrative(analysis, investigation_type="vehicle"):
-    """Generate AI investigation narrative using Mistral"""
+DATA RETRIEVED:
+- Total Records: {analysis.get('total_records', 0)}
+- Unique Vehicles: {analysis.get('unique_vehicles', 0)}
+- Unique Locations: {analysis.get('unique_locations', 0)}
+- Unique States: {len(analysis.get('states', []))}
+- Data Time Range: {analysis.get('date_range', 'Unknown')}
+
+KEY FINDINGS:
+"""
     
-    client = get_mistral_client()
-    if not client:
-        return _generate_fallback_narrative(analysis, investigation_type)
+    for finding in analysis.get("key_findings", []):
+        context += f"- {finding}\n"
     
-    # Build context from analysis
-    context = json.dumps(analysis, indent=2, default=str)
+    # Add sample records
+    context += "\nSAMPLE RECORDS:\n"
+    for idx, row in df.head(DATA_SAMPLE_SIZE).iterrows():
+        row_dict = dict(row)
+        # Show only key columns
+        key_cols = ["plate", "location", "vehicle_type", "timestamp", "state"]
+        filtered = {k: v for k, v in row_dict.items() if k in key_cols and k in row_dict}
+        context += f"- {filtered}\n"
     
-    prompts = {
-        "vehicle": f"""You are a professional surveillance analyst. Analyze this vehicle behavior data and provide an intelligence summary:
-
-DATA:
-{context}
-
-Generate a professional investigation narrative that:
-1. Summarizes the vehicle's movement patterns
-2. Identifies suspicious behavior (if any)
-3. Notes timing patterns and location clustering
-4. Provides recommendations for further investigation
-
-Keep it concise, professional, and actionable. Focus on behavioral patterns, not raw numbers.""",
-
-        "area": f"""You are a professional surveillance analyst. Analyze this area activity data:
-
-DATA:
-{context}
-
-Generate an area intelligence report that:
-1. Summarizes activity patterns in this zone
-2. Identifies unusual vehicle concentrations
-3. Notes peak activity times
-4. Highlights most suspicious vehicles
-5. Provides zone security recommendations
-
-Keep it concise and actionable.""",
+    prompt = f"""{context}
 
-        "midnight": f"""You are a surveillance intelligence specialist. Analyze this midnight activity data:
+Now analyze this data and answer the user's question: "{user_question}"
 
-DATA:
-{context}
-
-Generate a night activity intelligence report that:
-1. Analyzes suspicious midnight vehicle movements
-2. Identifies patterns suggesting reconnaissance or illegal activity
-3. Notes industrial zone activity
-4. Highlights repeat offenders
-5. Provides tactical recommendations
+Your response should:
+1. Directly answer what the user asked
+2. Cite specific data points from the results
+3. Provide relevant insights or patterns
+4. Be concise but comprehensive (2-3 paragraphs)
+5. Use professional but conversational language
 
-Be professional and factual."""
-    }
-    
-    prompt = prompts.get(investigation_type, prompts["vehicle"])
+Provide your detailed analysis:"""
     
     try:
-        message = ChatMessage(role="user", content=prompt)
-        response = client.chat(
-            model="mistral-small",
-            messages=[message],
-            temperature=0.3,
-            max_tokens=500
-        )
-        
-        return response.choices[0].message.content
+        if client:
+            response = client.chat(
+                model=MISTRAL_MODEL,
+                messages=[{"role": "user", "content": prompt}],
+                temperature=0.3,
+                max_tokens=1000
+            )
+            return response.choices[0].message.content.strip()
     except Exception as e:
-        print(f"Error generating narrative: {e}")
-        return _generate_fallback_narrative(analysis, investigation_type)
-
-
-def _generate_fallback_narrative(analysis, investigation_type="vehicle"):
-    """Fallback narrative generation (no Mistral)"""
+        print(f"⚠️ Mistral error: {e}")
     
-    if investigation_type == "vehicle":
-        narrative = f"""
-VEHICLE INVESTIGATION SUMMARY
+    # Fallback response
+    return generate_fallback_answer(user_question, analysis, df)
 
-Detection Activity:
-• Total Detections: {analysis.get('total_detections', 0)}
-• Unique Locations: {analysis.get('unique_locations', 0)}
-• Detection Span: {analysis.get('date_range', 'Unknown')}
-• Average Detection Confidence: {analysis.get('avg_confidence', 0):.2f}
 
-Behavior Patterns:
-• Detection Frequency: {analysis.get('detection_frequency', 0):.2f} per day
-• Midnight Activity: {analysis.get('midnight_percentage', 0)}% of detections
-• Average Time Gap: {analysis.get('avg_time_gap_hours', 0):.1f} hours
-• Location Clustering: {analysis.get('location_clustering', 0):.1f}% concentrated in top location
-
-Top Locations:
-"""
-        for loc, count in list(analysis.get('top_locations', {}).items())[:5]:
-            narrative += f"• {loc}: {count} detections\n"
-        
-        narrative += "\nInvestigation Recommendations:\n"
-        if analysis.get('midnight_percentage', 0) > 30:
-            narrative += "🚨 High midnight activity - Consider night surveillance verification\n"
-        if analysis.get('unique_locations', 0) > 8:
-            narrative += "🚨 Multi-area movement - Could indicate surveillance-style patterns\n"
-        if analysis.get('detection_frequency', 0) > 3:
-            narrative += "🚨 High detection frequency - Recommend blacklist verification\n"
-        if analysis.get('avg_confidence', 0) < 0.7:
-            narrative += "⚠️ Low confidence readings - Verify plate clarity\n"
-        
-        return narrative
-    
-    elif investigation_type == "area":
-        narrative = f"""
-AREA INTELLIGENCE REPORT
-
-Activity Overview:
-• Total Vehicles: {analysis.get('total_vehicles', 0)} unique plates
-• Total Detections: {analysis.get('total_detections', 0)}
-• Origin States: {analysis.get('unique_states', 0)}
-• Peak Activity Hour: {analysis.get('peak_hour', 0)}:00
-
-Suspicious Vehicles:
-• Most Repeated: {analysis.get('most_repeated_vehicle', 'N/A')} ({analysis.get('most_repeated_count', 0)} times)
-• Vehicles with 5+ visits: {analysis.get('repeated_vehicles', 0)}
-• Night Activity: {analysis.get('night_activity_percentage', 0)}% of detections
-
-Recommendations:
-• Increase monitoring during peak hours
-• Flag vehicles with 5+ repeated visits
-• Cross-reference interstate plates with blacklist
-• Verify CCTV coverage during night hours
-"""
-        return narrative
+def generate_fallback_answer(user_question, analysis, df):
+    """Fallback answer generation without Mistral"""
+    answer = f"Based on {analysis['total_records']} records:\n\n"
     
-    elif investigation_type == "midnight":
-        narrative = f"""
-MIDNIGHT ACTIVITY INTELLIGENCE REPORT
-
-Night Activity Summary:
-• Total Midnight Detections: {analysis.get('total_midnight_detections', 0)}
-• Unique Vehicles: {analysis.get('unique_vehicles', 0)}
-• Unique Locations: {analysis.get('unique_locations', 0)}
-• Vehicles with Multiple Visits: {analysis.get('vehicles_with_multiple_visits', 0)}
-
-Suspicious Vehicles (Multiple Visits):
-"""
-        for plate, count in list(analysis.get('top_suspicious_vehicles', {}).items())[:5]:
-            narrative += f"• {plate}: {count} midnight detections\n"
-        
-        narrative += f"""
-Activity Pattern:
-• Industrial Zone Activity: {analysis.get('industrial_zones_active', 0)} locations
-• Concentration Level: High (multiple zones active)
-
-Tactical Recommendations:
-• Deploy night patrol units to monitor top 3 vehicles
-• Coordinate with neighboring jurisdictions for interstate tracking
-• Verify industrial zone security protocols
-• Consider checkpoints during peak midnight hours (10PM-3AM)
-"""
-        return narrative
+    if analysis["key_findings"]:
+        answer += "Key Findings:\n"
+        for finding in analysis["key_findings"][:3]:
+            answer += f"• {finding}\n"
     
-    return "Investigation data incomplete"
-
-
-# =====================================================
-# MAIN INVESTIGATION FUNCTIONS
-# =====================================================
-
-def investigate_vehicle(plate):
-    """Investigate single vehicle"""
-    df = get_vehicle_detections(plate, days=30)
-    if df.empty:
-        return {
-            "status": "error",
-            "message": f"No detections found for {plate}"
-        }
+    if analysis["vehicle_types"]:
+        answer += "\nVehicle Breakdown:\n"
+        for vtype, count in analysis["vehicle_types"][:3]:
+            answer += f"• {vtype}: {count} detections\n"
     
-    analysis = analyze_vehicle_behavior(df)
-    narrative = generate_investigation_narrative(analysis, "vehicle")
+    if analysis["states"]:
+        answer += "\nState Distribution:\n"
+        for state, count in analysis["states"][:3]:
+            answer += f"• {state}: {count} detections\n"
     
-    return {
-        "status": "success",
-        "plate": plate,
-        "analysis": analysis,
-        "investigation_narrative": narrative,
-        "data": df.to_dict(orient="records")[:20]  # Top 20 records
-    }
+    return answer
 
 
-def investigate_area(location):
-    """Investigate area activity"""
-    df = get_area_activity(location, days=7)
-    if df.empty:
-        return {
-            "status": "error",
-            "message": f"No activity found in {location}"
-        }
-    
-    analysis = analyze_area_patterns(df)
-    narrative = generate_investigation_narrative(analysis, "area")
-    
-    return {
-        "status": "success",
-        "location": location,
-        "analysis": analysis,
-        "investigation_narrative": narrative,
-        "top_vehicles": analysis.get("repeated_vehicles", [])
-    }
-
+# =====================================================
+# STEP 5: FOLLOW-UP QUESTIONS & CLARIFICATIONS
+# =====================================================
 
-def investigate_midnight_activity():
-    """Investigate suspicious midnight activity"""
-    df = get_midnight_activity(hours_start=22, hours_end=4, days=7)
-    if df.empty:
-        return {
-            "status": "error",
-            "message": "No midnight activity found"
-        }
-    
-    analysis = analyze_midnight_patterns(df)
-    narrative = generate_investigation_narrative(analysis, "midnight")
+def generate_follow_ups(user_question, analysis, intent):
+    """
+    Generate intelligent follow-up questions or clarifications.
+    Suggests related queries the user might be interested in.
+    """
+    follow_ups = []
     
-    return {
-        "status": "success",
-        "analysis": analysis,
-        "investigation_narrative": narrative,
-        "suspicious_vehicles": analysis.get("top_suspicious_vehicles", {})
-    }
-
-
-def find_most_suspicious_today():
-    """Find most suspicious vehicle today"""
     try:
-        query = """
-        SELECT 
-            plate,
-            state,
-            vehicle_type,
-            vehicle_conf,
-            location,
-            date,
-            timestamp,
-            hour
-        FROM vehicle_logs
-        WHERE date = CURRENT_DATE
-        ORDER BY timestamp DESC
-        """
+        primary_intent = intent.get("primary_intent", "general")
         
-        with engine.connect() as conn:
-            result = conn.execute(text(query))
-            rows = result.fetchall()
-            df = pd.DataFrame(rows, columns=result.keys())
+        # Extract key entities from question
+        q = user_question.lower()
         
-        if df.empty:
-            return {
-                "status": "error",
-                "message": "No detections today"
-            }
+        # Suggest follow-ups based on data and intent
+        if analysis.get("unique_vehicles", 0) > 1 and "plate" not in q:
+            follow_ups.append("Would you like to focus on a specific vehicle plate?")
         
-        # Find vehicle with most detections
-        plate_counts = df["plate"].value_counts()
-        if plate_counts.empty:
-            return {
-                "status": "error",
-                "message": "No data available"
-            }
+        if analysis.get("unique_locations", 0) > 1 and "location" not in q:
+            follow_ups.append("Would you like details about a specific location?")
         
-        top_plate = plate_counts.index[0]
-        return investigate_vehicle(top_plate)
-    
-    except Exception as e:
-        print(f"Error finding suspicious vehicle: {e}")
-        return {
-            "status": "error",
-            "message": str(e)
-        }
-
-
-# =====================================================
-# FORMAT OUTPUT
-# =====================================================
-
-# =====================================================
-# NLP TO SQL CONVERSION WITH MISTRAL
-# =====================================================
-
-def convert_nlp_to_sql(question):
-    """Convert natural language question to SQL query using Mistral"""
-    
-    client = get_mistral_client()
-    if not client:
-        return None
-    
-    sql_prompt = f"""You are a PostgreSQL expert for vehicle surveillance database.
-
-Database schema:
-- vehicle_logs table has: plate, state, vehicle_type, location, timestamp, date, hour, camera_id, vehicle_conf, day
-- All text fields are case-insensitive, use LOWER() for matching
-- timestamp and date store when vehicle was detected
-
-User Question: {question}
-
-Generate ONLY a valid PostgreSQL SELECT query with these rules:
-1. MUST use LOWER() for case-insensitive matching on text fields
-2. MUST use LIKE '%pattern%' for partial/flexible matching
-3. MUST handle multiple conditions with AND/OR appropriately
-4. Multiple items (vehicles/locations) → use OR inside parentheses
-5. Single item with multiple types (e.g., "bike and mini truck") → use OR for vehicle_type
-6. Always ORDER BY timestamp DESC for latest records
-7. Always LIMIT 100 for safety
-8. SELECT plate, location, vehicle_type, timestamp, state
-9. No explanations, only the SQL query
-10. Handle these patterns:
-    - "show bike in location" → WHERE LOWER(vehicle_type) LIKE '%bike%' AND LOWER(location) LIKE '%location%'
-    - "show bike and car in location" → WHERE (LOWER(vehicle_type) LIKE '%bike%' OR LOWER(vehicle_type) LIKE '%car%') AND LOWER(location) LIKE '%location%'
-    - "show bikes in location1 and location2" → WHERE LOWER(vehicle_type) LIKE '%bike%' AND (LOWER(location) LIKE '%location1%' OR LOWER(location) LIKE '%location2%')
-    - "vehicle_type in multiple areas" → build OR condition for locations
-
-Output ONLY the SQL query, no markdown or explanation:"""
-
-    try:
-        response = client.chat(
-            model="mistral-small",
-            messages=[{"role": "user", "content": sql_prompt}],
-            temperature=0.1,
-            max_tokens=400
-        )
+        if primary_intent == "tracking":
+            follow_ups.append("Do you want to see the complete route history of this vehicle?")
         
-        sql_query = response.choices[0].message.content.strip()
-        # Clean up the query
-        sql_query = sql_query.replace("```sql", "").replace("```", "").replace("```postgresql", "").strip()
+        if primary_intent == "counting":
+            follow_ups.append("Would you like to see the breakdown by location or vehicle type?")
+        
+        if primary_intent == "pattern":
+            follow_ups.append("Should I analyze time patterns (peak hours, days)?")
+        
+        if analysis.get("total_records", 0) > 20:
+            follow_ups.append("The results have many records. Would you like me to focus on specific criteria?")
+        
+        if analysis.get("states") and len(analysis["states"]) > 1:
+            follow_ups.append("I found vehicles from multiple states. Want to focus on one state?")
         
-        print(f"Generated SQL: {sql_query}")
-        return sql_query
-    
     except Exception as e:
-        print(f"Error converting NLP to SQL: {e}")
-        return None
-
-
-def execute_custom_sql(sql_query):
-    """Execute custom SQL query with error handling"""
-    
-    try:
-        with engine.connect() as conn:
-            # Add safety checks
-            query_upper = sql_query.upper().strip()
-            
-            # Only allow SELECT queries
-            if not query_upper.startswith("SELECT"):
-                return None, "Query must be a SELECT statement"
-            
-            # Prevent dangerous operations
-            dangerous = ["DROP", "DELETE", "UPDATE", "INSERT", "TRUNCATE", "ALTER", "CREATE"]
-            if any(d in query_upper for d in dangerous):
-                return None, "Query contains restricted operations"
-            
-            # Check for SQL injection patterns
-            if ";" in sql_query.rstrip(";")[len("SELECT"):]:  # semicolons after SELECT are dangerous
-                return None, "Invalid query format"
-            
-            result = conn.execute(text(sql_query))
-            rows = result.fetchall()
-            
-            if not rows:
-                return None, "No records found"
-            
-            df = pd.DataFrame(rows, columns=result.keys())
-            return df, None
+        print(f"Follow-up generation error: {e}")
     
-    except Exception as e:
-        error_msg = str(e)
-        print(f"SQL Execution Error: {error_msg}")
-        # Return user-friendly error message
-        if "syntax" in error_msg.lower():
-            return None, "SQL syntax error in generated query"
-        elif "column" in error_msg.lower():
-            return None, "Invalid column reference in query"
-        else:
-            return None, f"Query error: {error_msg[:100]}"
+    return follow_ups[:FOLLOW_UP_QUESTIONS]
 
 
 # =====================================================
-# UNIFIED AI AGENT - NATURAL LANGUAGE QUESTIONS
+# STEP 6: MULTI-TURN CONVERSATION
 # =====================================================
 
-def ask_investigation_question(question):
+def ask_investigation_question(question, conversation_history=None):
     """
-    Advanced AI Investigation Agent with NLP-to-SQL
+    Main conversational agent function.
+    Handles single or multi-turn conversations.
     
-    Flow:
-    1. Uses Mistral to convert NLP to SQL (handles complex queries)
-    2. Executes query safely with validation
-    3. Analyzes retrieved data
-    4. Uses Mistral for deep analysis with RAG to answer user question
+    Returns structured response with answer, data, and follow-ups.
     """
     
     if not question or len(question.strip()) < 2:
         return {
             "status": "error",
-            "message": "Please enter a valid question"
+            "message": "Please enter a valid question",
+            "answer": None,
+            "data": None,
+            "follow_ups": []
         }
     
-    client = get_mistral_client()
+    print(f"\n{'='*60}")
+    print(f"🚀 PROCESSING: {question}")
+    print(f"{'='*60}")
     
     try:
-        print(f"\n=== Processing Question ===")
-        print(f"Question: {question}")
+        # Step 1: Understand query
+        print("📖 Step 1: Understanding query...")
+        intent = understand_query(question)
+        print(f"   Intent: {intent['primary_intent']}")
         
-        # Step 1: Convert NLP to SQL
-        print("Step 1: Converting natural language to SQL...")
-        sql_query = convert_nlp_to_sql(question)
+        # Step 2: Collect data
+        print("📊 Step 2: Collecting data...")
+        df, sql, error = collect_data(question)
         
-        if not sql_query:
+        if error:
+            print(f"   Error: {error}")
             return {
                 "status": "error",
-                "message": "Could not generate SQL query from your question. Please try rephrasing."
+                "message": error,
+                "sql_generated": sql,
+                "answer": None,
+                "data": None,
+                "follow_ups": [
+                    "Try being more specific (e.g., 'show bikes in adyar')",
+                    "Or specify a date range (e.g., 'bikes from 10-05-2026 to 18-05-2026')",
+                    "Or add a vehicle type (e.g., 'show vehicles in location X')"
+                ]
             }
         
-        print(f"Generated SQL: {sql_query}")
+        print(f"   ✅ Retrieved {len(df)} records")
         
-        # Step 2: Execute SQL query
-        print("Step 2: Executing query...")
-        df, error = execute_custom_sql(sql_query)
-        
-        data_context = {"summary": "", "analysis": {}, "findings": [], "raw_data": None}
-        
-        if error:
-            print(f"Execution error: {error}")
-            return {
-                "status": "error",
-                "message": f"Query execution failed: {error}"
-            }
+        # Step 3: Analyze data
+        print("🔍 Step 3: Analyzing data...")
+        analysis = analyze_data(df, question)
+        print(f"   Unique vehicles: {analysis['unique_vehicles']}")
+        print(f"   Unique locations: {analysis['unique_locations']}")
         
-        if df is None or df.empty:
-            return {
-                "status": "error",
-                "message": "No records found matching your query criteria"
-            }
+        # Step 4: Generate answer
+        print("💬 Step 4: Generating answer...")
+        client = get_mistral_client()
+        answer = generate_answer(question, df, analysis, client)
         
-        print(f"Retrieved {len(df)} records")
+        if not answer:
+            answer = generate_fallback_answer(question, analysis, df)
         
-        # Step 3: Analyze the data
-        print("Step 3: Analyzing data...")
-        data_context = _analyze_custom_query_results(df, question)
-        print(f"Analysis: {data_context['summary']}")
+        print(f"   ✅ Answer generated")
         
-        # Step 4: Generate intelligent answer with RAG
-        print("Step 4: Generating analysis with AI...")
-        answer = _generate_advanced_rag_answer(question, data_context, client, df)
+        # Step 5: Generate follow-ups
+        print("❓ Step 5: Generating follow-up questions...")
+        follow_ups = generate_follow_ups(question, analysis, intent)
+        print(f"   {len(follow_ups)} follow-ups suggested")
         
-        if not answer:
-            answer = _generate_fallback_analysis(question, data_context, df)
+        # Prepare response data
+        data_preview = df.head(10).to_dict(orient="records") if df is not None else None
         
-        print(f"=== Analysis Complete ===\n")
+        print(f"{'='*60}✅ COMPLETE\n")
         
         return {
             "status": "success",
             "question": question,
+            "intent": intent,
             "answer": answer,
-            "data_summary": data_context.get("summary", ""),
-            "analysis": data_context.get("analysis", {}),
-            "findings": data_context.get("findings", []),
-            "raw_data": df.to_dict(orient="records")[:20] if df is not None else None
+            "analysis": analysis,
+            "data_preview": data_preview,
+            "total_records": len(df),
+            "sql_generated": sql,
+            "follow_ups": follow_ups,
+            "raw_data": df.to_dict(orient="records") if df is not None else None
         }
         
     except Exception as e:
-        print(f"Error in investigation: {e}")
-        import traceback
+        print(f"❌ Error: {e}")
         traceback.print_exc()
         return {
             "status": "error",
-            "message": f"Investigation failed: {str(e)[:100]}"
+            "message": f"Investigation error: {str(e)[:100]}",
+            "answer": None,
+            "data": None,
+            "follow_ups": ["Try rephrasing your question with more specific details"]
         }
 
 
-def _analyze_custom_query_results(df, question):
-    """Analyze results from custom SQL queries with robust handling"""
-    
-    if df.empty:
-        return {"summary": "No data found matching your query", "analysis": {}, "findings": [], "raw_data": None}
-    
-    # Build analysis with safe column access
-    analysis = {
-        "total_records": len(df),
-        "unique_vehicles": df["plate"].nunique() if "plate" in df.columns else 0,
-        "unique_locations": df["location"].nunique() if "location" in df.columns else 0,
-        "vehicle_types": [],
-        "date_range": "Unknown"
-    }
-    
-    # Safe extraction of vehicle types
-    if "vehicle_type" in df.columns:
-        types = df["vehicle_type"].unique()
-        analysis["vehicle_types"] = [str(t) for t in types if pd.notna(t)]
-    
-    # Safe date range calculation
-    if "timestamp" in df.columns:
-        valid_times = pd.to_datetime(df["timestamp"], errors='coerce')
-        valid_times = valid_times[valid_times.notna()]
-        if len(valid_times) > 0:
-            analysis["date_range"] = f"{valid_times.min()} to {valid_times.max()}"
-    
-    # Extract findings with robust error handling
-    findings = []
-    
-    try:
-        if "plate" in df.columns:
-            plate_counts = df["plate"].value_counts()
-            if len(plate_counts) > 0:
-                findings.append(f"Most frequently detected: {plate_counts.index[0]} ({plate_counts.iloc[0]} detections)")
-                if len(plate_counts) > 1:
-                    findings.append(f"Total unique vehicles: {len(plate_counts)}")
-    except:
-        pass
-    
-    try:
-        if "location" in df.columns:
-            loc_counts = df["location"].value_counts()
-            if len(loc_counts) > 0:
-                findings.append(f"Most active location: {loc_counts.index[0]} ({loc_counts.iloc[0]} detections)")
-                if len(loc_counts) > 1:
-                    findings.append(f"Covered {len(loc_counts)} different locations")
-    except:
-        pass
-    
-    try:
-        if "vehicle_type" in df.columns:
-            vehicles = df["vehicle_type"].value_counts()
-            if len(vehicles) > 0:
-                types_str = ", ".join([f"{v} ({c})" for v, c in vehicles.head(3).items()])
-                findings.append(f"Vehicle breakdown: {types_str}")
-    except:
-        pass
-    
-    try:
-        if "timestamp" in df.columns and "hour" in df.columns:
-            df["hour_numeric"] = pd.to_datetime(df["timestamp"], errors='coerce').dt.hour
-            peak_hour = df["hour_numeric"].mode()
-            if len(peak_hour) > 0:
-                findings.append(f"Peak activity hour: {int(peak_hour[0])}:00")
-    except:
-        pass
-    
-    # Build summary
-    summary_parts = [f"Retrieved {len(df)} records"]
-    if analysis["unique_vehicles"] > 0:
-        summary_parts.append(f"from {analysis['unique_vehicles']} unique vehicles")
-    if analysis["unique_locations"] > 0:
-        summary_parts.append(f"across {analysis['unique_locations']} locations")
-    if analysis["vehicle_types"]:
-        summary_parts.append(f"(Types: {', '.join(analysis['vehicle_types'][:3])})")
-    
-    summary = " ".join(summary_parts)
-    
-    return {
-        "summary": summary,
-        "analysis": analysis,
-        "findings": findings[:5],  # Top 5 findings
-        "raw_data": df.to_dict(orient="records")
-    }
-
+# =====================================================
+# UTILITY FUNCTIONS
+# =====================================================
 
-def _generate_advanced_rag_answer(question, data_context, client, df):
-    """Generate advanced RAG answer using Mistral with full data context"""
+def format_investigation_output(result):
+    """Format investigation result for display (Gradio compatible)"""
     
-    analysis = data_context.get("analysis", {})
-    summary = data_context.get("summary", "")
-    findings = data_context.get("findings", [])
+    if result["status"] == "error":
+        return result["message"], "", json.dumps(result.get("follow_ups", []), indent=2), "{}"
     
-    # Build detailed context for Mistral
-    context_details = f"""Question: {question}
-
-Retrieved Data Summary:
-- {summary}
+    # Main answer
+    answer_text = result.get("answer", "No answer generated")
+    
+    # Data structure summary
+    analysis = result.get("analysis", {})
+    data_structure = f"""
+RETRIEVED DATA SUMMARY:
+- Total Records: {result.get('total_records', 0)}
+- Unique Vehicles: {analysis.get('unique_vehicles', 0)}
+- Unique Locations: {analysis.get('unique_locations', 0)}
+- Date Range: {analysis.get('date_range', 'Unknown')}
 
-Data Statistics:
+KEY METRICS:
 """
     
-    for key, value in analysis.items():
-        if isinstance(value, list):
-            context_details += f"  - {key}: {', '.join(str(v) for v in value[:5])}\n"
-        else:
-            context_details += f"  - {key}: {value}\n"
-    
-    # Add sample records for context
-    if df is not None and not df.empty:
-        context_details += f"\nSample Data (showing {min(3, len(df))} records):\n"
-        for idx, row in df.head(3).iterrows():
-            row_dict = dict(row)
-            record_str = ", ".join([f"{k}: {v}" for k, v in list(row_dict.items())[:6]])
-            context_details += f"  - {record_str}\n"
+    if analysis.get("vehicle_types"):
+        data_structure += "\nVehicle Types:\n"
+        for vtype, count in analysis["vehicle_types"]:
+            data_structure += f"  • {vtype}: {count}\n"
     
-    if findings:
-        context_details += f"\nInitial Findings:\n"
-        for finding in findings[:3]:
-            context_details += f"  - {finding}\n"
+    if analysis.get("states"):
+        data_structure += "\nStates:\n"
+        for state, count in analysis["states"]:
+            data_structure += f"  • {state}: {count}\n"
     
-    # Better prompt for detailed analysis
-    prompt = f"""You are a professional traffic surveillance analyst. 
-The user asked a specific question and we retrieved vehicle detection data from the database.
-
-{context_details}
-
-Now analyze this data and answer the user's original question: "{question}"
-
-Your response should:
-1. Directly answer what the user asked
-2. Provide specific data-backed insights
-3. Highlight patterns, anomalies, or important findings
-4. Be concise but comprehensive
-5. Use natural conversational language (not bullet points)
-
-Provide a detailed professional analysis:"""
-
-    try:
-        if client:
-            response = client.chat(
-                model="mistral-small",
-                messages=[{"role": "user", "content": prompt}],
-                temperature=0.3,
-                max_tokens=1200
-            )
-            answer = response.choices[0].message.content.strip()
-            
-            # Ensure answer is not empty
-            if not answer or answer.lower() == "none":
-                return _generate_fallback_analysis(question, data_context, df)
-            return answer
-        else:
-            return _generate_fallback_analysis(question, data_context, df)
+    # Follow-up questions
+    follow_ups = result.get("follow_ups", [])
+    follow_up_text = json.dumps(follow_ups, indent=2)
     
-    except Exception as e:
-        print(f"Error generating answer: {e}")
-        return _generate_fallback_analysis(question, data_context, df)
-
-
-def _generate_fallback_analysis(question, data_context, df):
-    """Fallback analysis without Mistral - provides structured response"""
-    
-    summary = data_context.get('summary', 'No data available')
-    analysis = data_context.get('analysis', {})
-    findings = data_context.get("findings", [])
-    
-    answer = f"Analysis for: {question}\n\n"
-    answer += f"Summary: {summary}\n\n"
-    
-    if df is not None and not df.empty:
-        answer += f"Retrieved Data Details:\n"
-        answer += f"- Total Records: {len(df)}\n"
-    
-    if analysis:
-        answer += f"\nKey Metrics:\n"
-        for key, value in list(analysis.items())[:8]:
-            if isinstance(value, list):
-                if value:
-                    answer += f"- {key}: {', '.join(str(v) for v in value[:3])}\n"
-            elif isinstance(value, dict):
-                if value:
-                    answer += f"- {key}: {len(value)} items\n"
-            else:
-                answer += f"- {key}: {value}\n"
-    
-    if findings:
-        answer += f"\nKey Findings:\n"
-        for finding in findings:
-            answer += f"- {finding}\n"
-    
-    if df is not None and not df.empty and len(df) > 0:
-        answer += f"\nSample Records:\n"
-        for idx, row in df.head(3).iterrows():
-            record_dict = dict(row)
-            record_str = ", ".join([f"{k}: {v}" for k, v in list(record_dict.items())[:5]])
-            answer += f"- {record_str}\n"
+    # Data preview
+    data_preview = result.get("data_preview", [])
+    data_json = json.dumps(data_preview[:5], indent=2, default=str)
     
-    return answer
+    return answer_text, data_structure, follow_up_text, data_json
 
 
-def format_investigation_output(result):
-    """Format investigation result for display"""
-    if result.get("status") == "error":
-        return {
-            "error": result.get("message", "Investigation failed")
-        }
-    
-    return {
-        "query": result.get("plate") or result.get("location") or "Investigation",
-        "answer": result.get("answer", ""),
-        "analysis": result.get("analysis", {}),
-        "key_findings": result.get("findings", [])
-    }
+# =====================================================
+# LEGACY FUNCTIONS (kept for compatibility)
+# =====================================================
 
+def format_investigation_output_tuple(result):
+    """Returns tuple of (answer, data_structure, follow_ups_json, data_json)"""
+    return format_investigation_output(result)
 
-if __name__ == "__main__":
-    # Test the investigation assistant
-    print("AI Investigation Assistant Ready\n")
-    
-    # Example: Ask a question
-    result = ask_investigation_question("show bikes in Adyar")
-    print(json.dumps(result, indent=2, default=str))
+
+# For backwards compatibility with app.py
+def investigate_with_npl(question):
+    """Wrapper function that matches app.py expectation"""
+    result = ask_investigation_question(question)
+    return format_investigation_output(result)

app.py

@@ -648,60 +648,75 @@ with gr.Blocks(
                     label="Records Retrieved",
                     interactive=False
                 )
+            
+            # Follow-up Questions
+            with gr.Tab("❓ Next Questions"):
+                follow_up_questions = gr.Textbox(
+                    label="Suggested Follow-ups & Clarifications",
+                    lines=6,
+                    interactive=False
+                )
         
         def investigate_with_npl(question):
             """Investigate with NLP-to-SQL and advanced analysis"""
             if not question or len(question.strip()) < 2:
-                return ("Please ask a question", "", "", None)
+                return ("Please ask a question", "", "", None, "")
             
             result = ask_investigation_question(question)
             
             if result.get("status") == "error":
                 error_msg = result.get("message", "Investigation failed")
-                return (f"Error: {error_msg}", "", f"Error: {error_msg}", None)
+                follow_ups = result.get("follow_ups", [])
+                suggestions = "\n".join([f"• {f}" for f in follow_ups]) if follow_ups else "Try with more specific details"
+                return (f"❌ {error_msg}", "No data collected", "", None, f"💡 Try this:\n{suggestions}")
             
-            # Format data structure info
-            data_info = result.get("data_summary", "")
-            analysis = result.get("analysis", {})
+            # Main AI answer
+            answer_text = result.get("answer", "")
             
-            # Build data structure with execution flow
-            structure = f"""Data Execution Flow
-━━━━━━━━━━━━━━━━━━━━━━
-Summary: {data_info}
-
-Data Structure Retrieved:
-"""
+            # Data structure info (enhanced)
+            analysis = result.get("analysis", {})
+            data_info = f"""📊 DATA COLLECTED:
+────────────────────────
+• Total Records: {result.get('total_records', 0)}
+• Unique Vehicles: {analysis.get('unique_vehicles', 0)}
+• Unique Locations: {analysis.get('unique_locations', 0)}
+• Time Range: {analysis.get('date_range', 'Unknown')}
+
+🔍 KEY FINDINGS:"""
             
-            if analysis:
-                for key, value in list(analysis.items())[:10]:
-                    if isinstance(value, list):
-                        structure += f"→ {key}: {', '.join(str(v) for v in value[:5])}\n"
-                    elif isinstance(value, dict):
-                        structure += f"→ {key}: {len(value)} items\n"
-                    else:
-                        structure += f"→ {key}: {value}\n"
+            for finding in analysis.get("key_findings", []):
+                data_info += f"\n  • {finding}"
             
-            # Format response (remove markdown)
-            response = result.get("answer", "").strip()
-            if not response:
-                response = f"Analysis: {data_info}"
+            if analysis.get("vehicle_types"):
+                data_info += "\n\n🚗 VEHICLE BREAKDOWN:"
+                for vtype, count in analysis.get("vehicle_types", []):
+                    data_info += f"\n  • {vtype}: {count} detections"
             
-            # Format findings
-            findings_list = result.get("findings", [])
-            findings_text = "\n".join([f"{f}" for f in findings_list]) if findings_list else "No specific findings"
+            # Insights (findings reformatted)
+            insights = "\n".join(analysis.get("key_findings", []))
+            if not insights:
+                insights = "Data analysis complete - see data structure for details"
             
             # Get raw data for table
             df_data = None
-            if result.get("raw_data"):
+            if result.get("data_preview"):
                 import pandas as pd
-                df_data = pd.DataFrame(result.get("raw_data", []))
+                df_data = pd.DataFrame(result.get("data_preview", []))
+            
+            # Follow-up questions (new feature!)
+            follow_ups = result.get("follow_ups", [])
+            follow_up_text = ""
+            if follow_ups:
+                follow_up_text = "💬 FOLLOW-UP QUESTIONS:\n────────────────────\n"
+                for idx, q in enumerate(follow_ups, 1):
+                    follow_up_text += f"{idx}. {q}\n"
             
-            return (response, findings_text, structure, df_data)
+            return (answer_text, data_info, insights, df_data, follow_up_text)
         
         investigate_btn.click(
             fn=investigate_with_npl,
             inputs=[question_input],
-            outputs=[ai_response, findings, data_structure, data_table]
+            outputs=[ai_response, data_structure, findings, data_table, follow_up_questions]
         )
 
     # =====================================================