Simplify sample data loading with direct SQL to fix HF Spaces connection issues

backend/database/__init__.py

@@ -80,278 +80,139 @@ def test_database_connection():
         return False
 
 def add_sample_data_for_hf():
-    """Add sample data for HF Spaces using SQLAlchemy."""
+    """Add simple sample data for HF Spaces using direct SQL."""
     if not os.getenv("SPACE_ID"):
         return  # Only run on HF Spaces
         
-    # First test database connection
-    if not test_database_connection():
-        print("❌ Database connection failed, skipping sample data")
-        return
+    print(f"🔍 HF Spaces environment check:")
+    print(f"   • SPACE_ID: {os.getenv('SPACE_ID')}")
+    print(f"   • Database URL: {DATABASE_URL}")
     
-    import json
-    import uuid
-    import hashlib
-    from backend.database.models import Trace, KnowledgeGraph, Entity, Relation
-    
-    # Use SessionLocal instead of scoped Session for better control
-    session = SessionLocal()
     try:
-        # Ensure tables exist first
-        Base.metadata.create_all(bind=engine)
-        
-        # Check if data already exists
-        existing_traces = session.query(Trace).count()
-        existing_kgs = session.query(KnowledgeGraph).count()
-        
-        print(f"🔍 HF Spaces environment check:")
-        print(f"   • SPACE_ID: {os.getenv('SPACE_ID')}")
-        print(f"   • Database URL: {DATABASE_URL}")
-        print(f"   • Existing traces: {existing_traces}")
-        print(f"   • Existing KGs: {existing_kgs}")
-        
-        if existing_traces > 0 or existing_kgs > 0:
-            print("📊 Sample data already exists, skipping...")
-            return  # Data already exists
-        
-        print("🎯 HF Spaces: Adding sample data to in-memory database...")
-        
-        # Sample trace content - Multi-agent customer service scenario
-        sample_trace_content = '''[
-    {
-        "role": "user",
-        "content": "I need help with my order #12345. It was supposed to arrive yesterday but I still haven't received it. This is very frustrating as I need it for an important meeting tomorrow.",
-        "timestamp": "2024-08-31T10:00:00Z",
-        "metadata": {"priority": "high", "emotion": "frustrated"}
-    },
-    {
-        "role": "assistant", 
-        "name": "RouterAgent",
-        "content": "I understand your frustration about the delayed order. Let me route this to our order tracking specialist who can provide detailed information about your shipment.",
-        "timestamp": "2024-08-31T10:00:15Z",
-        "actions": [{"type": "route", "target": "OrderTrackingAgent", "priority": "high"}]
-    },
-    {
-        "role": "assistant",
-        "name": "OrderTrackingAgent", 
-        "content": "I'll look up order #12345 immediately. Checking our tracking system... I found the issue - there was a delay at the distribution center due to weather conditions. Your package is currently in transit and should arrive by 2 PM today.",
-        "timestamp": "2024-08-31T10:01:00Z",
-        "actions": [
-            {"type": "database_query", "query": "SELECT * FROM orders WHERE order_id = '12345'"},
-            {"type": "api_call", "endpoint": "shipping_tracker", "response": "in_transit"}
-        ],
-        "findings": {
-            "issue": "Weather delay at distribution center", 
-            "current_status": "In transit",
-            "estimated_delivery": "2024-08-31T14:00:00Z"
-        }
-    },
-    {
-        "role": "assistant",
-        "name": "CompensationAgent",
-        "content": "Given this delay caused inconvenience, I'm authorizing a $10 credit to your account and providing free expedited shipping on your next order. I'm also escalating this to ensure better communication about delays in the future.",
-        "timestamp": "2024-08-31T10:02:00Z", 
-        "actions": [
-            {"type": "credit_account", "amount": 10, "reason": "delivery_delay"},
-            {"type": "add_benefit", "type": "free_expedited_shipping"},
-            {"type": "escalate", "department": "operations", "issue": "delayed_delivery_communication"}
-        ],
-        "error": "Failed to process account credit - payment system temporarily unavailable"
-    },
-    {
-        "role": "assistant",
-        "name": "SupervisorAgent",
-        "content": "I see there was an error processing the account credit. I'm manually flagging your account for the $10 credit and expedited shipping benefit. Our technical team will resolve the payment system issue. You should see the credit within 24 hours.",
-        "timestamp": "2024-08-31T10:03:00Z",
-        "actions": [
-            {"type": "manual_flag", "account_credit": 10},
-            {"type": "technical_escalation", "system": "payment_processor"},
-            {"type": "follow_up_schedule", "timeframe": "24_hours"}
-        ],
-        "optimization_opportunity": "Implement payment system redundancy to prevent credit processing failures"
-    },
-    {
-        "role": "user",
-        "content": "Thank you for the quick response and compensation. I appreciate you taking care of this issue promptly.",
-        "timestamp": "2024-08-31T10:04:00Z",
-        "sentiment": "satisfied"
-    },
-    {
-        "role": "assistant",
-        "name": "QualityAgent", 
-        "content": "Session completed successfully. Customer satisfaction restored. Identified system improvement: need better proactive communication about shipping delays and payment system reliability backup.",
-        "timestamp": "2024-08-31T10:05:00Z",
-        "session_analysis": {
-            "resolution_time": "5 minutes",
-            "customer_satisfaction": "high",
-            "agents_involved": 4,
-            "system_errors": 1,
-            "optimization_recommendations": [
-                "Implement proactive delay notifications",
-                "Add payment system redundancy", 
-                "Improve agent handoff protocols"
-            ]
-        }
-    }
-]'''
-
-        # Generate trace metadata
-        trace_id = str(uuid.uuid4())
-        content_hash = hashlib.sha256(sample_trace_content.encode()).hexdigest()
-        
-        # Create and save sample trace
-        sample_trace = Trace(
-            trace_id=trace_id,
-            filename="sample_customer_service.json", 
-            title="Multi-Agent Customer Service Resolution",
-            description="Demonstration of multi-agent system handling customer complaint with error handling and optimization opportunities",
-            content=sample_trace_content,
-            content_hash=content_hash,
-            uploader="AgentGraph Demo",
-            trace_type="multi_agent",
-            trace_source="sample",
-            character_count=len(sample_trace_content),
-            turn_count=6,
-            status="processed",
-            tags=json.dumps(["sample", "customer_service", "multi_agent", "error_handling", "optimization"]),
-            trace_metadata=json.dumps({
-                "scenario": "customer_service",
-                "agents": ["RouterAgent", "OrderTrackingAgent", "CompensationAgent", "SupervisorAgent", "QualityAgent"],
-                "domain": "e_commerce",
-                "complexity": "high"
-            })
-        )
-        
-        session.add(sample_trace)
-        session.flush()  # Get the trace ID
-        
-        # Create sample knowledge graph
-        sample_kg = KnowledgeGraph(
-            filename="sample_customer_service_kg.json",
-            creator="AgentGraph Demo",
-            entity_count=15,
-            relation_count=18,
-            namespace="customer_service_demo",
-            system_name="Multi-Agent Customer Service System",
-            system_summary="An intelligent customer service system featuring multiple specialized agents working together to resolve customer issues, handle errors, and identify optimization opportunities. The system demonstrates sophisticated agent coordination, error recovery mechanisms, and continuous improvement processes.",
-            status="completed",
-            trace_id=trace_id,
-            window_index=0,
-            window_total=1,
-            processing_run_id="sample_demo_run"
-        )
-        
-        session.add(sample_kg)
-        session.flush()  # Get the KG ID
-        
-        # Sample entities with rich properties
-        entities_data = [
-            # Agents
-            ("agent_1", "agent", "RouterAgent", {"role": "traffic_routing", "specialization": "request_classification", "priority_handling": True}),
-            ("agent_2", "agent", "OrderTrackingAgent", {"role": "order_management", "specialization": "shipping_tracking", "data_sources": ["internal_db", "shipping_apis"]}),
-            ("agent_3", "agent", "CompensationAgent", {"role": "customer_retention", "specialization": "compensation_authorization", "max_credit_limit": 50}),
-            ("agent_4", "agent", "SupervisorAgent", {"role": "escalation_handling", "specialization": "system_error_recovery", "override_authority": True}),
-            ("agent_5", "agent", "QualityAgent", {"role": "quality_assurance", "specialization": "session_analysis", "improvement_tracking": True}),
-            
-            # Systems and Tools
-            ("system_1", "system", "OrderDatabase", {"type": "database", "function": "order_storage", "performance": "high"}),
-            ("system_2", "system", "ShippingTracker", {"type": "external_api", "function": "package_tracking", "reliability": "99.5%"}),
-            ("system_3", "system", "PaymentProcessor", {"type": "financial_system", "function": "account_credits", "status": "temporarily_unavailable"}),
-            
-            # Issues and Problems
-            ("issue_1", "issue", "DeliveryDelay", {"severity": "medium", "cause": "weather_conditions", "impact": "customer_satisfaction"}),
-            ("issue_2", "issue", "PaymentSystemFailure", {"severity": "high", "cause": "system_unavailability", "impact": "compensation_processing"}),
+        # Use direct connection to avoid session issues
+        with engine.connect() as conn:
+            # Begin transaction for atomic operations
+            trans = conn.begin()
             
-            # Actions and Processes
-            ("action_1", "action", "RouteRequest", {"type": "traffic_management", "success_rate": "98%"}),
-            ("action_2", "action", "TrackPackage", {"type": "information_retrieval", "data_sources": 2}),
-            ("action_3", "action", "AuthorizeCredit", {"type": "financial_transaction", "approval_required": True}),
-            ("action_4", "action", "EscalateIssue", {"type": "process_escalation", "department": "operations"}),
-            
-            # Improvements and Optimizations
-            ("improvement_1", "improvement", "ProactiveNotifications", {"priority": "high", "implementation_effort": "medium", "expected_impact": "reduce_complaints_by_30%"})
-        ]
-        
-        # Add entities
-        for entity_id, entity_type, name, properties in entities_data:
-            entity = Entity(
-                graph_id=sample_kg.id,
-                entity_id=entity_id,
-                type=entity_type,
-                name=name,
-                properties=json.dumps(properties),
-                knowledge_graph_namespace="customer_service_demo"
-            )
-            session.add(entity)
-        
-        # Sample relations showing complex interactions
-        relations_data = [
-            # Agent interactions
-            ("rel_1", "agent_1", "routes_to", "agent_2", {"context": "order_inquiry", "priority": "high", "success": True}),
-            ("rel_2", "agent_2", "collaborates_with", "agent_3", {"context": "customer_compensation", "coordination": "automated"}),
-            ("rel_3", "agent_3", "escalates_to", "agent_4", {"context": "system_error", "escalation_reason": "payment_failure"}),
-            ("rel_4", "agent_4", "coordinates_with", "agent_5", {"context": "quality_improvement", "outcome": "optimization_identified"}),
-            
-            # System interactions
-            ("rel_5", "agent_2", "queries", "system_1", {"query_type": "order_lookup", "response_time": "0.5s", "success": True}),
-            ("rel_6", "agent_2", "calls", "system_2", {"api_endpoint": "track_package", "response_time": "1.2s", "success": True}),
-            ("rel_7", "agent_3", "attempts_transaction", "system_3", {"transaction_type": "credit", "amount": 10, "success": False}),
-            
-            # Problem identification and resolution
-            ("rel_8", "agent_2", "identifies", "issue_1", {"detection_method": "system_query", "severity_assessed": "medium"}),
-            ("rel_9", "agent_3", "encounters", "issue_2", {"error_handling": "automatic_escalation", "recovery_action": "manual_override"}),
-            ("rel_10", "agent_4", "resolves", "issue_2", {"resolution_method": "manual_flag", "permanent_fix": False}),
-            
-            # Action execution
-            ("rel_11", "agent_1", "executes", "action_1", {"execution_time": "15s", "outcome": "successful_routing"}),
-            ("rel_12", "agent_2", "performs", "action_2", {"data_retrieved": True, "accuracy": "100%"}),
-            ("rel_13", "agent_3", "initiates", "action_3", {"authorization_level": "standard", "blocked_by": "system_error"}),
-            ("rel_14", "agent_4", "triggers", "action_4", {"escalation_department": "operations", "follow_up_required": True}),
-            
-            # Improvement opportunities
-            ("rel_15", "agent_5", "identifies", "improvement_1", {"analysis_method": "session_review", "confidence": "high"}),
-            ("rel_16", "issue_1", "leads_to", "improvement_1", {"causal_relationship": "direct", "prevention_potential": "high"}),
-            ("rel_17", "issue_2", "exposes", "system_3", {"vulnerability_type": "single_point_of_failure", "risk_level": "high"}),
-            ("rel_18", "improvement_1", "would_prevent", "issue_1", {"prevention_mechanism": "early_warning", "effectiveness": "85%"})
-        ]
-        
-        # Add relations
-        for relation_id, from_entity, relation_type, to_entity, properties in relations_data:
-            relation = Relation(
-                graph_id=sample_kg.id,
-                relation_id=relation_id,
-                from_entity_id=from_entity,
-                relation_type=relation_type,
-                to_entity_id=to_entity,
-                properties=json.dumps(properties),
-                knowledge_graph_namespace="customer_service_demo"
-            )
-            session.add(relation)
-        
-        # Commit all data
-        session.commit()
-        
-        # Verify data was added
-        final_traces = session.query(Trace).count()
-        final_kgs = session.query(KnowledgeGraph).count()
-        final_entities = session.query(Entity).count()
-        final_relations = session.query(Relation).count()
-        
-        print("✅ Sample data added successfully to HF Spaces in-memory database!")
-        print(f"   • Traces added: {final_traces}")
-        print(f"   • Knowledge graphs added: {final_kgs}")
-        print(f"   • Entities added: {final_entities}")
-        print(f"   • Relations added: {final_relations}")
-        print(f"   • Scenario: Multi-Agent Customer Service with error handling and optimization")
-        
+            try:
+                # Check if data already exists
+                result = conn.execute("SELECT COUNT(*) FROM traces").fetchone()
+                existing_traces = result[0] if result else 0
+                
+                result = conn.execute("SELECT COUNT(*) FROM knowledge_graphs").fetchone()  
+                existing_kgs = result[0] if result else 0
+                
+                print(f"   • Existing traces: {existing_traces}")
+                print(f"   • Existing KGs: {existing_kgs}")
+                
+                if existing_traces > 0 or existing_kgs > 0:
+                    print("📊 Sample data already exists, skipping...")
+                    return
+                
+                print("🎯 Adding sample data using direct SQL...")
+                
+                import json
+                import uuid
+                import hashlib
+                from datetime import datetime
+                
+                # Simple trace content
+                sample_trace_content = '''[
+    {"role": "user", "content": "I need help with my delayed order #12345. This is frustrating!", "timestamp": "2024-08-31T10:00:00Z"},
+    {"role": "assistant", "name": "RouterAgent", "content": "Let me route you to our order specialist.", "timestamp": "2024-08-31T10:00:15Z"},
+    {"role": "assistant", "name": "OrderAgent", "content": "I found the issue - weather delay. Your package arrives today at 2 PM.", "timestamp": "2024-08-31T10:01:00Z"},
+    {"role": "assistant", "name": "CompensationAgent", "content": "I'll authorize a $10 credit for the inconvenience.", "timestamp": "2024-08-31T10:02:00Z", "error": "Payment system unavailable"},
+    {"role": "assistant", "name": "SupervisorAgent", "content": "I'll manually flag your account for the credit. Technical team notified.", "timestamp": "2024-08-31T10:03:00Z"},
+    {"role": "user", "content": "Thank you for the quick resolution!", "timestamp": "2024-08-31T10:04:00Z", "sentiment": "satisfied"}
+]'''
+                
+                # Generate IDs
+                trace_id = str(uuid.uuid4())
+                content_hash = hashlib.sha256(sample_trace_content.encode()).hexdigest()
+                now = datetime.utcnow()
+                
+                # Insert trace
+                conn.execute(
+                    """INSERT INTO traces (trace_id, filename, title, description, content, content_hash, 
+                       upload_timestamp, update_timestamp, uploader, trace_type, trace_source, 
+                       character_count, turn_count, status, tags, trace_metadata) 
+                       VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
+                    (trace_id, "sample_demo.json", "Multi-Agent Customer Service Demo", 
+                     "Demo showing agent coordination and error handling", sample_trace_content, content_hash,
+                     now, now, "AgentGraph Demo", "multi_agent", "sample", len(sample_trace_content), 6, 
+                     "processed", '["demo", "customer_service", "multi_agent"]', 
+                     '{"scenario": "customer_service", "agents": ["RouterAgent", "OrderAgent", "CompensationAgent", "SupervisorAgent"]}')
+                )
+                
+                # Insert knowledge graph
+                conn.execute(
+                    """INSERT INTO knowledge_graphs (filename, creator, entity_count, relation_count, namespace,
+                       system_name, system_summary, status, trace_id, window_index, window_total, processing_run_id)
+                       VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
+                    ("demo_kg.json", "AgentGraph Demo", 5, 4, "demo", "Customer Service System",
+                     "Multi-agent customer service system with error handling and coordination", 
+                     "completed", trace_id, 0, 1, "demo_run")
+                )
+                
+                # Get KG ID
+                kg_result = conn.execute("SELECT id FROM knowledge_graphs WHERE trace_id = ?", (trace_id,))
+                kg_id = kg_result.fetchone()[0]
+                
+                # Insert sample entities
+                entities = [
+                    ("agent_1", "agent", "RouterAgent", '{"role": "routing", "priority_handling": true}'),
+                    ("agent_2", "agent", "OrderAgent", '{"role": "order_tracking", "data_sources": ["db", "api"]}'),
+                    ("agent_3", "agent", "CompensationAgent", '{"role": "compensation", "max_credit": 50}'),
+                    ("agent_4", "agent", "SupervisorAgent", '{"role": "escalation", "override_authority": true}'),
+                    ("issue_1", "issue", "PaymentSystemFailure", '{"severity": "high", "impact": "service_disruption"}')
+                ]
+                
+                for entity_id, entity_type, name, properties in entities:
+                    conn.execute(
+                        """INSERT INTO entities (graph_id, entity_id, type, name, properties, knowledge_graph_namespace)
+                           VALUES (?, ?, ?, ?, ?, ?)""",
+                        (kg_id, entity_id, entity_type, name, properties, "demo")
+                    )
+                
+                # Insert sample relations
+                relations = [
+                    ("rel_1", "agent_1", "routes_to", "agent_2", '{"priority": "high", "success": true}'),
+                    ("rel_2", "agent_2", "escalates_to", "agent_3", '{"reason": "compensation_needed"}'),
+                    ("rel_3", "agent_3", "escalates_to", "agent_4", '{"reason": "system_error"}'),
+                    ("rel_4", "agent_4", "resolves", "issue_1", '{"method": "manual_override"}')
+                ]
+                
+                for relation_id, from_entity, relation_type, to_entity, properties in relations:
+                    conn.execute(
+                        """INSERT INTO relations (graph_id, relation_id, from_entity_id, relation_type, to_entity_id, 
+                           properties, knowledge_graph_namespace) VALUES (?, ?, ?, ?, ?, ?, ?)""",
+                        (kg_id, relation_id, from_entity, relation_type, to_entity, properties, "demo")
+                    )
+                
+                # Commit transaction
+                trans.commit()
+                
+                # Verify data
+                final_traces = conn.execute("SELECT COUNT(*) FROM traces").fetchone()[0]
+                final_kgs = conn.execute("SELECT COUNT(*) FROM knowledge_graphs").fetchone()[0]
+                final_entities = conn.execute("SELECT COUNT(*) FROM entities").fetchone()[0]
+                final_relations = conn.execute("SELECT COUNT(*) FROM relations").fetchone()[0]
+                
+                print("✅ Sample data added successfully!")
+                print(f"   • Traces: {final_traces}")
+                print(f"   • Knowledge graphs: {final_kgs}")
+                print(f"   • Entities: {final_entities}")
+                print(f"   • Relations: {final_relations}")
+                
+            except Exception as e:
+                trans.rollback()
+                raise e
+                
     except Exception as e:
         print(f"❌ Failed to add sample data: {e}")
         import traceback
         print("Full error traceback:")
         traceback.print_exc()
-        session.rollback()
-    finally:
-        session.close()
 
 __all__ = [
     'get_db',