backend/database/__init__.py

"""
Database functionality for the agent monitoring system.
This package provides database access and utilities for agent monitoring.
"""

import os
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, scoped_session

# Get the absolute path to the project root directory
ROOT_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# Database URL - Use in-memory SQLite for HF Spaces to avoid data sharing between users
# This ensures each container restart gets a fresh database
import uuid
session_id = str(uuid.uuid4())[:8]

# For HF Spaces: Use in-memory database to prevent user data sharing
# For local development: Use file database
if os.getenv("SPACE_ID"):  # HF Spaces environment
    DATABASE_URL = "sqlite:///:memory:"
    print(f"🔒 HF Spaces: Using in-memory database (session: {session_id})")
else:
    DATABASE_URL = f"sqlite:///{os.path.join(ROOT_DIR, 'datasets/db/agent_monitoring.db')}"
    print(f"💾 Local: Using persistent database")

# Create engine
engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})

# Create session factory
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

# Create a scoped session for thread safety
Session = scoped_session(SessionLocal)

# Base class for models
Base = declarative_base()

# Function to get database session
def get_db():
    db = Session()
    try:
        yield db
    finally:
        db.close()

# Import models and utility functions
from backend.database.models import KnowledgeGraph, Entity, Relation, Base
from backend.database.utils import (
    save_knowledge_graph, 
    update_knowledge_graph_status,
    get_knowledge_graph,
    get_all_knowledge_graphs,
    delete_knowledge_graph
)

# Function to initialize database
def init_db():
    """Initialize the database by creating all tables."""
    Base.metadata.create_all(bind=engine)

def test_database_connection():
    """Test if database connection is working."""
    try:
        session = SessionLocal()
        # Try to create tables
        Base.metadata.create_all(bind=engine)
        
        # Test a simple query
        result = session.execute("SELECT 1").fetchone()
        print(f"✅ Database connection test successful: {result}")
        
        session.close()
        return True
    except Exception as e:
        print(f"❌ Database connection test failed: {e}")
        import traceback
        traceback.print_exc()
        return False

def add_sample_data_for_hf():
    """Add simple sample data for HF Spaces using direct SQL."""
    if not os.getenv("SPACE_ID"):
        return  # Only run on HF Spaces
        
    print(f"🔍 HF Spaces environment check:")
    print(f"   • SPACE_ID: {os.getenv('SPACE_ID')}")
    print(f"   • Database URL: {DATABASE_URL}")
    
    try:
        # Use direct connection to avoid session issues
        with engine.connect() as conn:
            # Begin transaction for atomic operations
            trans = conn.begin()
            
            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()

__all__ = [
    'get_db',
    'models',
    'init_db',
    'test_database_connection',
    'add_sample_data_for_hf',
    'save_knowledge_graph',
    'update_knowledge_graph_status',
    'get_knowledge_graph',
    'get_all_knowledge_graphs',
    'delete_knowledge_graph',
    'KnowledgeGraph',
    'Entity',
    'Relation'
]