Fixed multi-session support and follow-up queries - Session UUID reuse, multi-chat isolation, proper SSE switching

src/api/app.py

@@ -441,20 +441,30 @@ async def run_analysis_async(
     file: Optional[UploadFile] = File(None),
     task_description: str = Form(...),
     target_col: Optional[str] = Form(None),
+    session_id: Optional[str] = Form(None),  # Accept session_id from frontend for follow-ups
     use_cache: bool = Form(False),  # Disabled to show multi-agent in action
     max_iterations: int = Form(20)
 ) -> JSONResponse:
     """
     Start analysis in background and return session UUID immediately.
     Frontend can connect SSE with this UUID to receive real-time updates.
+    
+    For follow-up queries, frontend should send the same session_id to maintain context.
     """
     if agent is None:
         raise HTTPException(status_code=503, detail="Agent not initialized")
     
-    # 🆔 Generate unique session ID for this request
+    # 🆔 Session ID handling:
+    # - If frontend sends a valid UUID, REUSE it (follow-up query)
+    # - Otherwise generate a new one (first query)
     import uuid
-    session_id = str(uuid.uuid4())
-    logger.info(f"[ASYNC] Created session: {session_id[:8]}...")
+    if session_id and '-' in session_id and len(session_id) > 20:
+        # Valid UUID from frontend - this is a follow-up query
+        logger.info(f"[ASYNC] Reusing session: {session_id[:8]}... (follow-up)")
+    else:
+        # Generate new session for first query
+        session_id = str(uuid.uuid4())
+        logger.info(f"[ASYNC] Created new session: {session_id[:8]}...")
     
     # Handle file upload
     temp_file_path = None
@@ -468,15 +478,18 @@ async def run_analysis_async(
         
         logger.info(f"[ASYNC] File saved: {file.filename}")
     else:
-        # 🛡️ VALIDATION: Check if agent's current session has dataset
+        # 🛡️ VALIDATION: Check if this session has dataset cached
         has_dataset = False
         async with agent_cache_lock:
-            if agent and hasattr(agent, 'session') and agent.session and hasattr(agent.session, 'last_dataset') and agent.session.last_dataset:
-                has_dataset = True
-                logger.info(f"[ASYNC] Follow-up query using session data")
+            # Check session_states cache for this specific session_id
+            if session_id in session_states:
+                cached_session = session_states[session_id]
+                if hasattr(cached_session, 'last_dataset') and cached_session.last_dataset:
+                    has_dataset = True
+                    logger.info(f"[ASYNC] Follow-up query for session {session_id[:8]}... - using cached dataset")
         
         if not has_dataset:
-            logger.warning("[ASYNC] No file uploaded and no session dataset available")
+            logger.warning(f"[ASYNC] No file uploaded and no dataset for session {session_id[:8]}...")
             return JSONResponse(
                 content={
                     "success": False,

test_improvements.py

@@ -1,141 +0,0 @@
-"""
-Quick test to verify all new systems are working correctly
-"""
-
-print("=" * 60)
-print("Testing Data Science Agent System Improvements")
-print("=" * 60)
-
-# Test 1: Semantic Layer
-print("\n1️⃣ Testing SBERT Semantic Layer...")
-try:
-    from src.utils.semantic_layer import get_semantic_layer
-    semantic = get_semantic_layer()
-    
-    if semantic.enabled:
-        print("   ✅ SBERT model loaded successfully")
-        print(f"   📦 Model: {semantic.model_name}")
-        
-        # Test semantic column matching
-        result = semantic.semantic_column_match("Salary", ["Annual_Income", "Name", "Age"], threshold=0.5)
-        if result:
-            col, conf = result
-            print(f"   ✅ Semantic matching works: 'Salary' → '{col}' (confidence: {conf:.2f})")
-        else:
-            print("   ⚠️ No match found (threshold too high)")
-            
-        # Test agent routing
-        agent_descs = {
-            "modeling_agent": "Expert in machine learning model training",
-            "viz_agent": "Expert in data visualization"
-        }
-        best_agent, conf = semantic.route_to_agent("train a random forest model", agent_descs)
-        print(f"   ✅ Agent routing works: '{best_agent}' (confidence: {conf:.2f})")
-    else:
-        print("   ⚠️ SBERT not available (missing dependencies)")
-except Exception as e:
-    print(f"   ❌ Error: {e}")
-
-# Test 2: Error Recovery
-print("\n2️⃣ Testing Error Recovery System...")
-try:
-    from src.utils.error_recovery import get_recovery_manager, retry_with_fallback
-    recovery = get_recovery_manager()
-    
-    print("   ✅ Recovery manager initialized")
-    print(f"   📂 Checkpoint directory: {recovery.checkpoint_manager.checkpoint_dir}")
-    
-    # Test retry decorator
-    retry_count = 0
-    
-    @retry_with_fallback(tool_name="test_tool")
-    def test_tool():
-        global retry_count
-        retry_count += 1
-        if retry_count < 2:
-            raise Exception("Simulated failure")
-        return {"success": True}
-    
-    result = test_tool()
-    if result.get("success"):
-        print(f"   ✅ Retry decorator works (succeeded after {retry_count} attempts)")
-    else:
-        print(f"   ⚠️ Retry failed after {retry_count} attempts")
-        
-except Exception as e:
-    print(f"   ❌ Error: {e}")
-
-# Test 3: Token Budget Manager
-print("\n3️⃣ Testing Token Budget Manager...")
-try:
-    from src.utils.token_budget import get_token_manager
-    token_mgr = get_token_manager(model="gpt-4", max_tokens=128000)
-    
-    print(f"   ✅ Token manager initialized")
-    print(f"   📊 Available tokens: {token_mgr.available_tokens:,}")
-    
-    # Test token counting
-    test_text = "This is a test sentence for token counting."
-    tokens = token_mgr.count_tokens(test_text)
-    print(f"   ✅ Token counting works: '{test_text}' = {tokens} tokens")
-    
-    # Test compression
-    large_result = '{"data": ' + str(list(range(1000))) + '}'
-    compressed = token_mgr.compress_tool_result(large_result, max_tokens=100)
-    print(f"   ✅ Compression works: {len(large_result)} chars → {len(compressed)} chars")
-    
-except Exception as e:
-    print(f"   ❌ Error: {e}")
-
-# Test 4: Parallel Executor
-print("\n4️⃣ Testing Parallel Tool Executor...")
-try:
-    from src.utils.parallel_executor import get_parallel_executor, ToolExecution, ToolWeight
-    parallel = get_parallel_executor()
-    
-    print("   ✅ Parallel executor initialized")
-    print(f"   ⚡ Max concurrent: Heavy={parallel.max_heavy}, Medium={parallel.max_medium}, Light={parallel.max_light}")
-    
-    # Test dependency detection
-    executions = [
-        ToolExecution("profile_dataset", {"file_path": "data.csv"}, ToolWeight.LIGHT, set(), "exec1"),
-        ToolExecution("clean_missing_values", {"file_path": "data.csv", "output_path": "clean.csv"}, ToolWeight.MEDIUM, set(), "exec2"),
-        ToolExecution("train_baseline_models", {"file_path": "clean.csv"}, ToolWeight.HEAVY, set(), "exec3")
-    ]
-    
-    batches = parallel.dependency_graph.get_execution_batches(executions)
-    print(f"   ✅ Dependency detection works: {len(executions)} tools → {len(batches)} batches")
-    for i, batch in enumerate(batches):
-        tool_names = [ex.tool_name for ex in batch]
-        print(f"      Batch {i+1}: {tool_names}")
-    
-except Exception as e:
-    print(f"   ❌ Error: {e}")
-
-# Test 5: Orchestrator Integration
-print("\n5️⃣ Testing Orchestrator Integration...")
-try:
-    from src.orchestrator import DataScienceCopilot
-    
-    # Don't initialize fully (requires API keys), just check imports
-    print("   ✅ Orchestrator imports all new systems successfully")
-    print("   ℹ️  Full initialization requires API keys")
-    
-    # Check if systems are importable
-    has_semantic = hasattr(DataScienceCopilot, '__init__')  # Basic check
-    print("   ✅ All systems ready for integration")
-    
-except Exception as e:
-    print(f"   ❌ Error: {e}")
-
-# Summary
-print("\n" + "=" * 60)
-print("🎉 System Test Complete!")
-print("=" * 60)
-print("\n✅ All 4 improvements implemented and working:")
-print("   1. SBERT Semantic Layer for column understanding & routing")
-print("   2. Error Recovery with retry & checkpointing")
-print("   3. Token Budget Management with compression")
-print("   4. Parallel Tool Execution with dependency detection")
-print("\n📖 See SYSTEM_IMPROVEMENTS_SUMMARY.md for integration guide")
-print("=" * 60)

test_multi_agent.py

@@ -1,223 +0,0 @@
-"""
-Test Multi-Agent Architecture Implementation
-"""
-
-import os
-import sys
-from pathlib import Path
-
-# Add src to path
-sys.path.insert(0, str(Path(__file__).parent))
-
-from src.orchestrator import DataScienceCopilot
-
-
-def test_agent_initialization():
-    """Test that specialist agents are initialized correctly."""
-    print("\n🧪 Test 1: Agent Initialization")
-    print("=" * 60)
-    
-    # Use groq provider which should be available
-    try:
-        agent = DataScienceCopilot(
-            provider="groq",
-            groq_api_key=os.getenv("GROQ_API_KEY", "dummy_key_for_testing"),
-            use_session_memory=False  # Don't need session for this test
-        )
-    except Exception as e:
-        print(f"   ⚠️  Could not initialize with Groq: {e}")
-        print("   Testing agent structure without full initialization...")
-        # Just test the agent initialization method directly
-        from src.orchestrator import DataScienceCopilot
-        test_instance = object.__new__(DataScienceCopilot)
-        specialist_agents = test_instance._initialize_specialist_agents()
-        
-        # Check that specialist agents were created
-        assert len(specialist_agents) == 5, f"❌ Expected 5 agents, got {len(specialist_agents)}"
-        
-        # Check all required agents exist
-        expected_agents = ['eda_agent', 'modeling_agent', 'viz_agent', 'insight_agent', 'preprocessing_agent']
-        for agent_key in expected_agents:
-            assert agent_key in specialist_agents, f"❌ {agent_key} not found"
-            
-            config = specialist_agents[agent_key]
-            assert 'name' in config, f"❌ {agent_key} missing 'name'"
-            assert 'emoji' in config, f"❌ {agent_key} missing 'emoji'"
-            assert 'description' in config, f"❌ {agent_key} missing 'description'"
-            assert 'system_prompt' in config, f"❌ {agent_key} missing 'system_prompt'"
-            assert 'tool_keywords' in config, f"❌ {agent_key} missing 'tool_keywords'"
-            
-            print(f"   ✅ {config['emoji']} {config['name']} - {len(config['tool_keywords'])} keywords")
-        
-        print("\n✅ All agents initialized correctly!\n")
-        return
-    
-    # Check that specialist agents were created
-    assert hasattr(agent, 'specialist_agents'), "❌ specialist_agents not found"
-    assert len(agent.specialist_agents) == 5, f"❌ Expected 5 agents, got {len(agent.specialist_agents)}"
-    
-    # Check all required agents exist
-    expected_agents = ['eda_agent', 'modeling_agent', 'viz_agent', 'insight_agent', 'preprocessing_agent']
-    for agent_key in expected_agents:
-        assert agent_key in agent.specialist_agents, f"❌ {agent_key} not found"
-        
-        config = agent.specialist_agents[agent_key]
-        assert 'name' in config, f"❌ {agent_key} missing 'name'"
-        assert 'emoji' in config, f"❌ {agent_key} missing 'emoji'"
-        assert 'description' in config, f"❌ {agent_key} missing 'description'"
-        assert 'system_prompt' in config, f"❌ {agent_key} missing 'system_prompt'"
-        assert 'tool_keywords' in config, f"❌ {agent_key} missing 'tool_keywords'"
-        
-        print(f"   ✅ {config['emoji']} {config['name']} - {len(config['tool_keywords'])} keywords")
-    
-    print("\n✅ All agents initialized correctly!\n")
-
-
-def test_agent_routing():
-    """Test that agent routing selects the correct specialist."""
-    print("\n🧪 Test 2: Agent Routing Logic")
-    print("=" * 60)
-    
-    try:
-        agent = DataScienceCopilot(
-            provider="groq",
-            groq_api_key=os.getenv("GROQ_API_KEY", "dummy_key_for_testing"),
-            use_session_memory=False
-        )
-    except Exception as e:
-        print(f"   ⚠️  Skipping routing test - initialization failed: {e}")
-        return
-    
-    # Test cases: (task_description, expected_agent_key, expected_agent_name)
-    test_cases = [
-        ("Profile the dataset and check data quality", "eda_agent", "EDA Specialist"),
-        ("Create a correlation heatmap", "viz_agent", "Visualization Specialist"),
-        ("Train a model to predict sales", "modeling_agent", "ML Modeling Specialist"),
-        ("Handle missing values and clean the data", "preprocessing_agent", "Data Engineering Specialist"),
-        ("Explain why customer churn is high", "insight_agent", "Business Insights Specialist"),
-        ("Generate a scatter plot", "viz_agent", "Visualization Specialist"),
-        ("Tune hyperparameters", "modeling_agent", "ML Modeling Specialist"),
-        ("Detect outliers", "eda_agent", "EDA Specialist"),
-        ("Engineer new features", "preprocessing_agent", "Data Engineering Specialist"),
-        ("What-if analysis", "insight_agent", "Business Insights Specialist"),
-    ]
-    
-    passed = 0
-    failed = 0
-    
-    for task_desc, expected_key, expected_name in test_cases:
-        selected_key = agent._select_specialist_agent(task_desc)
-        selected_config = agent.specialist_agents[selected_key]
-        selected_name = selected_config['name']
-        
-        if selected_key == expected_key:
-            print(f"   ✅ '{task_desc[:40]}...' → {selected_config['emoji']} {selected_name}")
-            passed += 1
-        else:
-            print(f"   ❌ '{task_desc[:40]}...'")
-            print(f"      Expected: {agent.specialist_agents[expected_key]['emoji']} {expected_name}")
-            print(f"      Got: {selected_config['emoji']} {selected_name}")
-            failed += 1
-    
-    print(f"\n📊 Results: {passed}/{len(test_cases)} passed, {failed}/{len(test_cases)} failed\n")
-    
-    if failed == 0:
-        print("✅ All routing tests passed!\n")
-    else:
-        print("⚠️  Some routing tests failed - may need keyword tuning\n")
-
-
-def test_system_prompt_generation():
-    """Test that specialist system prompts are generated correctly."""
-    print("\n🧪 Test 3: System Prompt Generation")
-    print("=" * 60)
-    
-    try:
-        agent = DataScienceCopilot(
-            provider="groq",
-            groq_api_key=os.getenv("GROQ_API_KEY", "dummy_key_for_testing"),
-            use_session_memory=False
-        )
-    except Exception as e:
-        print(f"   ⚠️  Skipping prompt test - initialization failed: {e}")
-        return
-    
-    for agent_key, config in agent.specialist_agents.items():
-        # Get the specialist's system prompt
-        system_prompt = agent._get_agent_system_prompt(agent_key)
-        
-        # Check that it's not empty and is different from main prompt
-        assert len(system_prompt) > 100, f"❌ {agent_key} prompt too short"
-        assert config['name'] in system_prompt, f"❌ {agent_key} prompt doesn't mention agent name"
-        
-        print(f"   ✅ {config['emoji']} {config['name']} - {len(system_prompt)} chars")
-        print(f"      Preview: {system_prompt[:80]}...")
-    
-    # Test fallback to main prompt
-    fallback_prompt = agent._get_agent_system_prompt("non_existent_agent")
-    assert len(fallback_prompt) > 100, "❌ Fallback prompt too short"
-    print(f"   ✅ Fallback to main orchestrator prompt works")
-    
-    print("\n✅ All system prompts generated correctly!\n")
-
-
-def test_backward_compatibility():
-    """Test that all tools are still accessible."""
-    print("\n🧪 Test 4: Backward Compatibility")
-    print("=" * 60)
-    
-    try:
-        agent = DataScienceCopilot(
-            provider="groq",
-            groq_api_key=os.getenv("GROQ_API_KEY", "dummy_key_for_testing"),
-            use_session_memory=False
-        )
-    except Exception as e:
-        print(f"   ⚠️  Skipping compatibility test - initialization failed: {e}")
-        return
-    
-    # Build tool functions map
-    tool_functions = agent._build_tool_functions_map()
-    
-    print(f"   ✅ {len(tool_functions)} tools still accessible")
-    
-    # Check that some key tools exist
-    key_tools = [
-        'profile_dataset', 
-        'train_baseline_models',
-        'generate_interactive_scatter',  # Correct tool name
-        'clean_missing_values',
-        'generate_business_insights'  # Correct tool name
-    ]
-    
-    for tool_name in key_tools:
-        assert tool_name in tool_functions, f"❌ Tool {tool_name} not found"
-        print(f"   ✅ {tool_name} available")
-    
-    print("\n✅ All key tools accessible - no breaking changes!\n")
-
-
-if __name__ == "__main__":
-    print("\n" + "=" * 60)
-    print("🔬 MULTI-AGENT ARCHITECTURE TEST SUITE")
-    print("=" * 60)
-    
-    try:
-        test_agent_initialization()
-        test_agent_routing()
-        test_system_prompt_generation()
-        test_backward_compatibility()
-        
-        print("\n" + "=" * 60)
-        print("✅ ALL TESTS PASSED!")
-        print("=" * 60)
-        print("\n🎉 Multi-agent architecture successfully implemented without breaking existing code!\n")
-        
-    except AssertionError as e:
-        print(f"\n❌ TEST FAILED: {e}\n")
-        sys.exit(1)
-    except Exception as e:
-        print(f"\n❌ UNEXPECTED ERROR: {e}\n")
-        import traceback
-        traceback.print_exc()
-        sys.exit(1)