vibesurf 0.1.0__py3-none-any.whl

vibe_surf/backend/migrations/init_db.py

@@ -0,0 +1,303 @@
+"""
+Database Initialization and Migration Scripts - Simplified
+
+Handles database schema creation for Task, LLMProfile, and UploadedFile tables.
+"""
+
+import asyncio
+from typing import Optional
+from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
+from sqlalchemy import text
+from sqlalchemy.exc import SQLAlchemyError
+
+from ..database.models import Base
+
+async def init_database(database_url: str = "sqlite+aiosqlite:///./vibe_surf.db") -> bool:
+    """
+    Initialize the database with Task, LLMProfile, and UploadedFile tables and indexes.
+    
+    Args:
+        database_url: Database connection URL
+        
+    Returns:
+        bool: True if initialization was successful
+    """
+    try:
+        # Create async engine
+        engine = create_async_engine(
+            database_url,
+            echo=False,
+            future=True,
+            pool_pre_ping=True
+        )
+        
+        # Create all tables
+        success = await create_tables(engine)
+        
+        if success:
+            print("✅ Database initialized successfully")
+            print(f"📍 Database URL: {database_url}")
+            
+            # Print table information
+            await print_table_info(engine)
+        else:
+            print("❌ Database initialization failed")
+        
+        await engine.dispose()
+        return success
+        
+    except Exception as e:
+        print(f"❌ Database initialization error: {e}")
+        return False
+
+async def create_tables(engine: AsyncEngine) -> bool:
+    """
+    Create all tables (Task, LLMProfile, UploadedFile) and indexes.
+    
+    Args:
+        engine: SQLAlchemy async engine
+        
+    Returns:
+        bool: True if creation was successful
+    """
+    try:
+        async with engine.begin() as conn:
+            # Create all tables (Task, LLMProfile, UploadedFile)
+            await conn.run_sync(Base.metadata.create_all)
+            
+            # Create additional indexes for performance
+            await create_performance_indexes(conn)
+            
+        print("✅ Database tables created successfully")
+        return True
+        
+    except SQLAlchemyError as e:
+        print(f"❌ Error creating tables: {e}")
+        return False
+
+async def drop_tables(engine: AsyncEngine) -> bool:
+    """
+    Drop all database tables (use with caution).
+    
+    Args:
+        engine: SQLAlchemy async engine
+        
+    Returns:
+        bool: True if drop was successful
+    """
+    try:
+        async with engine.begin() as conn:
+            await conn.run_sync(Base.metadata.drop_all)
+            
+        print("✅ Database tables dropped successfully")
+        return True
+        
+    except SQLAlchemyError as e:
+        print(f"❌ Error dropping tables: {e}")
+        return False
+
+async def create_performance_indexes(conn):
+    """
+    Create additional performance indexes for all tables.
+    
+    Args:
+        conn: Database connection
+    """
+    # Performance indexes for all tables
+    indexes = [
+        # Tasks - for session task tracking
+        "CREATE INDEX IF NOT EXISTS idx_tasks_session_status ON tasks(session_id, status)",
+        "CREATE INDEX IF NOT EXISTS idx_tasks_session_created ON tasks(session_id, created_at DESC)",
+        "CREATE INDEX IF NOT EXISTS idx_tasks_status_created ON tasks(status, created_at DESC)",
+        "CREATE INDEX IF NOT EXISTS idx_tasks_completed ON tasks(completed_at DESC)",
+        
+        # LLMProfile - for profile lookups
+        "CREATE INDEX IF NOT EXISTS idx_llm_profiles_name ON llm_profiles(profile_name)",
+        "CREATE INDEX IF NOT EXISTS idx_llm_profiles_provider ON llm_profiles(provider)",
+        "CREATE INDEX IF NOT EXISTS idx_llm_profiles_created ON llm_profiles(created_at DESC)",
+        
+        # UploadedFile - for file management
+        "CREATE INDEX IF NOT EXISTS idx_uploaded_files_session ON uploaded_files(session_id)",
+        "CREATE INDEX IF NOT EXISTS idx_uploaded_files_session_active ON uploaded_files(session_id, is_deleted)",
+        "CREATE INDEX IF NOT EXISTS idx_uploaded_files_upload_time ON uploaded_files(upload_time DESC)",
+        "CREATE INDEX IF NOT EXISTS idx_uploaded_files_deleted ON uploaded_files(is_deleted, deleted_at)",
+        "CREATE INDEX IF NOT EXISTS idx_uploaded_files_filename ON uploaded_files(original_filename)",
+    ]
+    
+    for index_sql in indexes:
+        try:
+            await conn.execute(text(index_sql))
+            print(f"✅ Created index: {index_sql.split('idx_')[1].split(' ')[0]}")
+        except Exception as e:
+            print(f"⚠️  Index creation warning: {e}")
+
+async def print_table_info(engine: AsyncEngine):
+    """
+    Print information about created tables.
+    
+    Args:
+        engine: SQLAlchemy async engine
+    """
+    try:
+        async with engine.connect() as conn:
+            # Get table list
+            if "sqlite" in str(engine.url):
+                result = await conn.execute(text("SELECT name FROM sqlite_master WHERE type='table'"))
+            else:
+                result = await conn.execute(text("SELECT tablename FROM pg_tables WHERE schemaname='public'"))
+            
+            tables = [row[0] for row in result]
+            
+            print(f"\n📊 Created {len(tables)} tables:")
+            for table in sorted(tables):
+                if not table.startswith('sqlite_'):  # Skip SQLite system tables
+                    print(f"   • {table}")
+            
+            # Get approximate row counts (for SQLite)
+            if "sqlite" in str(engine.url):
+                print(f"\n📈 Table sizes:")
+                for table in sorted(tables):
+                    if not table.startswith('sqlite_'):
+                        try:
+                            count_result = await conn.execute(text(f"SELECT COUNT(*) FROM {table}"))
+                            count = count_result.scalar()
+                            print(f"   • {table}: {count} rows")
+                        except Exception:
+                            print(f"   • {table}: - rows")
+    
+    except Exception as e:
+        print(f"⚠️  Could not retrieve table information: {e}")
+
+async def verify_database_schema(database_url: str = "sqlite+aiosqlite:///./vibe_surf.db") -> bool:
+    """
+    Verify that the database schema is correctly set up.
+    
+    Args:
+        database_url: Database connection URL
+        
+    Returns:
+        bool: True if schema is valid
+    """
+    try:
+        engine = create_async_engine(database_url, echo=False)
+        
+        # Expected tables (Task, LLMProfile, UploadedFile)
+        expected_tables = ['tasks', 'llm_profiles', 'uploaded_files']
+        
+        async with engine.connect() as conn:
+            # Check if all expected tables exist
+            if "sqlite" in str(engine.url):
+                result = await conn.execute(text("SELECT name FROM sqlite_master WHERE type='table'"))
+            else:
+                result = await conn.execute(text("SELECT tablename FROM pg_tables WHERE schemaname='public'"))
+            
+            existing_tables = {row[0] for row in result}
+            missing_tables = set(expected_tables) - existing_tables
+            
+            if missing_tables:
+                print(f"❌ Missing tables: {missing_tables}")
+                await engine.dispose()
+                return False
+            
+            # Test basic operations
+            await conn.execute(text("SELECT 1"))
+            
+        await engine.dispose()
+        print("✅ Database schema verification passed")
+        return True
+        
+    except Exception as e:
+        print(f"❌ Database schema verification failed: {e}")
+        return False
+
+async def reset_database(database_url: str = "sqlite+aiosqlite:///./vibe_surf.db") -> bool:
+    """
+    Reset the database by dropping and recreating all tables.
+    
+    Args:
+        database_url: Database connection URL
+        
+    Returns:
+        bool: True if reset was successful
+    """
+    try:
+        engine = create_async_engine(database_url, echo=False)
+        
+        print("🔄 Resetting database...")
+        
+        # Drop all tables
+        if not await drop_tables(engine):
+            await engine.dispose()
+            return False
+        
+        # Recreate all tables
+        if not await create_tables(engine):
+            await engine.dispose()
+            return False
+        
+        await engine.dispose()
+        print("✅ Database reset completed successfully")
+        return True
+        
+    except Exception as e:
+        print(f"❌ Database reset failed: {e}")
+        return False
+
+async def migrate_database(database_url: str = "sqlite+aiosqlite:///./vibe_surf.db") -> bool:
+    """
+    Apply database migrations and updates.
+    
+    Args:
+        database_url: Database connection URL
+        
+    Returns:
+        bool: True if migration was successful
+    """
+    try:
+        engine = create_async_engine(database_url, echo=False)
+        
+        print("🔄 Running database migrations...")
+        
+        # For now, this is equivalent to create_tables
+        # In the future, this could handle schema updates
+        success = await create_tables(engine)
+        
+        await engine.dispose()
+        
+        if success:
+            print("✅ Database migrations completed successfully")
+        else:
+            print("❌ Database migrations failed")
+            
+        return success
+        
+    except Exception as e:
+        print(f"❌ Database migration failed: {e}")
+        return False
+
+# CLI functions for direct usage
+async def main():
+    """Main function for running migrations from command line."""
+    import sys
+    
+    if len(sys.argv) < 2:
+        print("Usage: python -m backend.migrations.init_db [init|reset|verify|migrate]")
+        return
+    
+    command = sys.argv[1]
+    database_url = sys.argv[2] if len(sys.argv) > 2 else "sqlite+aiosqlite:///./vibe_surf.db"
+    
+    if command == "init":
+        await init_database(database_url)
+    elif command == "reset":
+        await reset_database(database_url)
+    elif command == "verify":
+        await verify_database_schema(database_url)
+    elif command == "migrate":
+        await migrate_database(database_url)
+    else:
+        print(f"Unknown command: {command}")
+        print("Available commands: init, reset, verify, migrate")
+
+if __name__ == "__main__":
+    asyncio.run(main())

vibe_surf/backend/migrations/seed_data.py

@@ -0,0 +1,236 @@
+"""
+Database Seed Data Scripts - Simplified Single Task Model
+
+Creates sample data for development and testing purposes using the simplified Task table.
+"""
+
+import asyncio
+from datetime import datetime, timedelta
+from typing import List
+
+from ..database.models import Task, TaskStatus
+from .. import shared_state
+
+async def seed_sample_tasks(database_url: str = "sqlite+aiosqlite:///./vibe_surf.db") -> bool:
+    """
+    Seed the database with sample tasks for development.
+    
+    Args:
+        database_url: Database connection URL
+        
+    Returns:
+        bool: True if seeding was successful
+    """
+    try:
+        # Use shared_state db_manager if available, otherwise create temporary one
+        if shared_state.db_manager:
+            db_manager = shared_state.db_manager
+        else:
+            from ..database.manager import DatabaseManager
+            db_manager = DatabaseManager(database_url)
+        
+        async for db in db_manager.get_session():
+            # Check if data already exists
+            result = await db.execute("SELECT COUNT(*) FROM tasks")
+            count = result.scalar()
+            
+            if count > 0:
+                print(f"⚠️  Database already contains {count} tasks. Skipping seed data.")
+                return True
+            
+            print("🌱 Seeding sample tasks...")
+            
+            # Create sample tasks
+            tasks = await create_sample_tasks(db)
+            
+            await db.commit()
+            
+        print("✅ Sample tasks seeded successfully")
+        return True
+        
+    except Exception as e:
+        print(f"❌ Seeding failed: {e}")
+        return False
+
+async def create_sample_tasks(db) -> List[Task]:
+    """Create sample tasks for testing."""
+    
+    sample_tasks = [
+        Task(
+            session_id="session_001",
+            task_description="Create a simple web scraper to extract product information from an e-commerce website",
+            status=TaskStatus.COMPLETED,
+            upload_files_path="./uploads/session_001/requirements.pdf",
+            mcp_server_config={
+                "exclude_actions": [],
+                "max_actions_per_task": 100,
+                "display_files_in_done_text": True
+            },
+            llm_config={
+                "model": "gpt-4o-mini",
+                "provider": "openai",
+                "temperature": 0.1,
+                "max_tokens": 4000
+            },
+            task_result="Successfully created web scraper that extracts product names, prices, and descriptions. Generated 500 product records.",
+            report_path="./reports/session_001/scraper_report.html",
+            started_at=datetime.now() - timedelta(hours=2),
+            completed_at=datetime.now() - timedelta(hours=1),
+            task_metadata={
+                "execution_duration_seconds": 3600.0,
+                "total_actions": 45,
+                "created_via": "api"
+            }
+        ),
+        Task(
+            session_id="session_002",
+            task_description="Automate login process for a social media platform and post a scheduled message",
+            status=TaskStatus.RUNNING,
+            upload_files_path="./uploads/session_002/login_credentials.txt",
+            mcp_server_config={
+                "exclude_actions": ["dangerous_action"],
+                "max_actions_per_task": 50,
+                "display_files_in_done_text": True
+            },
+            llm_config={
+                "model": "gpt-4o",
+                "provider": "openai",
+                "temperature": 0.2,
+                "max_tokens": 2000
+            },
+            started_at=datetime.now() - timedelta(minutes=30),
+            task_metadata={
+                "created_via": "api",
+                "estimated_duration": 1800
+            }
+        ),
+        Task(
+            session_id="session_003",
+            task_description="Research and compile information about top AI companies and their latest products",
+            status=TaskStatus.PENDING,
+            mcp_server_config={
+                "exclude_actions": [],
+                "max_actions_per_task": 200,
+                "display_files_in_done_text": True
+            },
+            llm_config={
+                "model": "claude-3-sonnet-20240229",
+                "provider": "anthropic",
+                "temperature": 0.3,
+                "max_tokens": 8000
+            },
+            task_metadata={
+                "created_via": "api",
+                "priority": "high"
+            }
+        ),
+        Task(
+            session_id="session_004",
+            task_description="Fill out and submit an online form with provided customer data",
+            status=TaskStatus.FAILED,
+            upload_files_path="./uploads/session_004/customer_data.csv",
+            mcp_server_config={
+                "exclude_actions": [],
+                "max_actions_per_task": 30,
+                "display_files_in_done_text": True
+            },
+            llm_config={
+                "model": "gpt-3.5-turbo",
+                "provider": "openai",
+                "temperature": 0.0,
+                "max_tokens": 1000
+            },
+            error_message="Form submission failed due to CAPTCHA protection",
+            started_at=datetime.now() - timedelta(hours=6),
+            completed_at=datetime.now() - timedelta(hours=5, minutes=45),
+            task_metadata={
+                "execution_duration_seconds": 900.0,
+                "total_actions": 15,
+                "created_via": "api",
+                "error_recovery_attempts": 3
+            }
+        ),
+        Task(
+            session_id="session_005",
+            task_description="Monitor a website for price changes and send notifications when target price is reached",
+            status=TaskStatus.PAUSED,
+            mcp_server_config={
+                "exclude_actions": [],
+                "max_actions_per_task": 1000,
+                "display_files_in_done_text": False
+            },
+            llm_config={
+                "model": "gpt-4o-mini",
+                "provider": "openai",
+                "temperature": 0.1,
+                "max_tokens": 2000
+            },
+            started_at=datetime.now() - timedelta(hours=12),
+            task_metadata={
+                "created_via": "api",
+                "monitoring_interval": 3600,
+                "target_price": 299.99
+            }
+        )
+    ]
+    
+    for task in sample_tasks:
+        db.add(task)
+    
+    await db.flush()
+    return sample_tasks
+
+async def clear_sample_data(database_url: str = "sqlite+aiosqlite:///./vibe_surf.db") -> bool:
+    """
+    Clear all sample data from the database.
+    
+    Args:
+        database_url: Database connection URL
+        
+    Returns:
+        bool: True if clearing was successful
+    """
+    try:
+        # Use shared_state db_manager if available, otherwise create temporary one
+        if shared_state.db_manager:
+            db_manager = shared_state.db_manager
+        else:
+            from ..database.manager import DatabaseManager
+            db_manager = DatabaseManager(database_url)
+        
+        async for db in db_manager.get_session():
+            print("🧹 Clearing sample tasks...")
+            
+            # Delete all tasks
+            await db.execute("DELETE FROM tasks")
+            await db.commit()
+            
+        print("✅ Sample data cleared successfully")
+        return True
+        
+    except Exception as e:
+        print(f"❌ Clearing sample data failed: {e}")
+        return False
+
+# CLI functions
+async def main():
+    """Main function for running seed operations from command line."""
+    import sys
+    
+    if len(sys.argv) < 2:
+        print("Usage: python -m backend.migrations.seed_data [seed|clear]")
+        return
+    
+    command = sys.argv[1]
+    database_url = sys.argv[2] if len(sys.argv) > 2 else "sqlite+aiosqlite:///./vibe_surf.db"
+    
+    if command == "seed":
+        await seed_sample_tasks(database_url)
+    elif command == "clear":
+        await clear_sample_data(database_url)
+    else:
+        print(f"Unknown command: {command}")
+        print("Available commands: seed, clear")
+
+if __name__ == "__main__":
+    asyncio.run(main())