main.py

#!/usr/bin/env python3
"""
Agent Monitoring System - Main Entry Point
This script serves as the main entry point for the agent monitoring system.
Updated: Graph visualization sizes optimized for better UI layout.
"""

# Import the complete LiteLLM fix FIRST, before any other imports that might use LiteLLM
from utils.fix_litellm_stop_param import *  # This applies all the patches

# Continue with regular imports
import argparse
import sys
import os
import logging
import subprocess
import signal
import time
import shutil
from pathlib import Path

# Add the current directory to the Python path to ensure imports work correctly
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S"
)
logger = logging.getLogger("agent_monitoring")

# Helper utilities -----------------------------------------------------------

def is_command_available(cmd: str) -> bool:
    """Return True if *cmd* is found in PATH."""
    return shutil.which(cmd) is not None


def in_virtualenv() -> bool:
    """Detect if running inside a virtual environment."""
    return getattr(sys, 'base_prefix', sys.prefix) != sys.prefix


def run_subprocess(cmd: list[str], cwd: str | None = None, env: dict | None = None) -> bool:
    """Run subprocess, stream output, return True on success."""
    try:
        proc = subprocess.run(cmd, cwd=cwd, env=env, capture_output=True, text=True)
        if proc.returncode != 0:
            logger.error(f"Command failed: {' '.join(cmd)}\n{proc.stderr}")
            return False
        return True
    except FileNotFoundError:
        logger.error(f"Executable not found: {cmd[0]}")
        return False

def parse_arguments():
    """Parse command line arguments"""
    parser = argparse.ArgumentParser(description="Agent Monitoring System")
    parser.add_argument("--setup", action="store_true", help="Set up the environment")
    parser.add_argument("--server", action="store_true", help="Start the visualization server")
    parser.add_argument("--dev", action="store_true", help="Start both frontend and backend in development mode")
    parser.add_argument("--frontend", action="store_true", help="Start only the frontend development server")
    parser.add_argument("--init-db", action="store_true", help="Initialize the database")
    parser.add_argument("--port", type=int, default=5280, help="Port for the server")
    parser.add_argument("--frontend-port", type=int, default=3001, help="Port for the frontend dev server")
    parser.add_argument("--host", default="127.0.0.1", help="Host for the server")
    parser.add_argument("--log-level", default="INFO", 
                        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
                        help="Set logging level")
    parser.add_argument("--no-open-browser", action="store_true", help="Do not open browser automatically")
    parser.add_argument("--first-run", action="store_true", help="Run setup + init DB, then start dev server")
    
    return parser, parser.parse_args()

def start_frontend_dev_server(frontend_port=3001):
    """Start the React frontend development server"""
    try:
        logger.info(f"Starting frontend development server on port {frontend_port}...")
        
        # Change to the React app directory
        frontend_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "frontend")
        
        if not os.path.exists(frontend_dir):
            logger.error(f"Frontend directory not found: {frontend_dir}")
            return None
            
        # Check if node_modules exists
        node_modules_dir = os.path.join(frontend_dir, "node_modules")
        if not os.path.exists(node_modules_dir):
            logger.info("Installing frontend dependencies...")
            install_process = subprocess.run(
                ["npm", "install"], 
                cwd=frontend_dir,
                capture_output=True,
                text=True
            )
            if install_process.returncode != 0:
                logger.error(f"Failed to install frontend dependencies: {install_process.stderr}")
                return None
        
        # Start the development server
        env = os.environ.copy()
        env['PORT'] = str(frontend_port)
        
        process = subprocess.Popen(
            ["npm", "run", "dev"],
            cwd=frontend_dir,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            env=env,
            bufsize=1,
            universal_newlines=True
        )
        
        logger.info(f"✅ Frontend development server started (PID: {process.pid})")
        logger.info(f"🌐 Frontend available at: http://localhost:{frontend_port}")
        
        return process
        
    except Exception as e:
        logger.error(f"Failed to start frontend development server: {str(e)}")
        return None

def start_backend_server(args):
    """Start the backend server"""
    try:
        logger.info(f"Starting backend server on port {args.port}...")
        
        from backend.main import main as server_main
        
        # Create a new argv to pass all relevant arguments to the server
        server_args = [sys.argv[0]]
        
        # Add port if specified
        if args.port:
            server_args.extend(["--port", str(args.port)])
            
        # Preserve any other relevant flags
        if args.no_open_browser:
            server_args.append("--no-open-browser")
            
        if args.log_level:
            server_args.extend(["--log-level", args.log_level])
                
        if args.host:
            server_args.extend(["--host", args.host])
        
        # Replace sys.argv with our server-specific arguments
        original_argv = sys.argv.copy()
        sys.argv = server_args
        
        # Start server
        logger.info(f"✅ Backend server starting...")
        logger.info(f"🚀 Backend API available at: http://{args.host}:{args.port}")
        server_main()
        
        # Restore original argv
        sys.argv = original_argv
        
    except Exception as e:
        logger.error(f"Failed to start backend server: {str(e)}")

def run_fullstack_dev(args):
    """Run both frontend and backend in development mode"""
    logger.info("🚀 Starting full-stack development environment...")
    logger.info("=" * 60)
    
    # Start frontend in a separate process
    frontend_process = start_frontend_dev_server(args.frontend_port)
    
    if frontend_process is None:
        logger.error("Failed to start frontend. Exiting.")
        return
    
    # Wait a moment for frontend to start
    time.sleep(2)
    
    # Setup signal handlers for graceful shutdown
    def signal_handler(signum, frame):
        logger.info("\n🛑 Shutting down full-stack development environment...")
        if frontend_process:
            logger.info("Stopping frontend development server...")
            frontend_process.terminate()
            try:
                frontend_process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                frontend_process.kill()
        logger.info("✅ Shutdown complete")
        sys.exit(0)
    
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    
    # Print startup summary
    logger.info("🎉 Full-stack development environment ready!")
    logger.info("=" * 60)
    logger.info(f"🌐 Frontend (React):     http://localhost:{args.frontend_port}")
    logger.info(f"🚀 Backend (FastAPI):   http://{args.host}:{args.port}")
    logger.info(f"📚 API Documentation:   http://{args.host}:{args.port}/docs")
    logger.info("=" * 60)
    logger.info("Press Ctrl+C to stop both servers")
    logger.info("=" * 60)
    
    try:
        # Start backend server (this will block)
        start_backend_server(args)
    except KeyboardInterrupt:
        signal_handler(signal.SIGINT, None)

def main():
    """Main entry point"""
    parser, args = parse_arguments()
    
    # Update logging level if specified
    if args.log_level:
        logging.getLogger().setLevel(getattr(logging, args.log_level))
        logger.setLevel(getattr(logging, args.log_level))
    
    # --------------------------------------------------
    # First-run combo flag: environment setup + init DB + dev server
    # --------------------------------------------------
    if args.first_run:
        logger.info("🔰 First-run initialization requested")

        # 0. Ensure we're inside a virtual environment; if not, create one with uv
        if not in_virtualenv():
            if not is_command_available("uv"):
                logger.error("❌ 'uv' CLI not found. Please install it first (e.g., 'pipx install uv' or 'brew install uv').")
                return

            logger.info("📦 Creating virtual environment with uv ...")
            if not run_subprocess(["uv", "venv", ".venv"]):
                logger.error("Failed to create virtualenv via uv")
                return

            # Re-execute this script inside the new virtualenv
            new_python = os.path.join(".venv", "Scripts" if os.name == "nt" else "bin", "python")
            logger.info("🔄 Re-execing inside virtualenv ...")
            os.execv(new_python, [new_python] + sys.argv)

        # 0.5 Ensure dependencies installed (run once)
        deps_marker = Path(".venv/.deps_installed")
        if not deps_marker.exists():
            logger.info("🔧 Installing project dependencies with uv ...")
            if not run_subprocess(["uv", "pip", "install", "-e", "."]):
                return
            deps_marker.touch()

        # 1. Environment setup
        try:
            from setup_env import setup_environment
            logger.info("🔧 Running environment setup...")
            setup_environment()
            logger.info("✅ Environment setup complete")
        except Exception as e:
            logger.error(f"Environment setup failed: {e}")
            return

        # 2. Database initialization
        try:
            from backend.database.init_db import main as init_db_main
            logger.info("🗄️  Initializing the database...")
            _orig_argv = sys.argv.copy()
            sys.argv = [sys.argv[0]]  # clear extra flags for init_db parser
            try:
                init_db_main()
            finally:
                sys.argv = _orig_argv
            logger.info("✅ Database initialization complete")
        except Exception as e:
            logger.error(f"Database initialization failed: {e}")
            return

        # 3. Decide runtime mode
        if not any([args.server, args.frontend]):
            args.dev = True  # default to dev if nothing else specified

        # 4. Verify npm if dev (frontend) requested
        if args.dev and not is_command_available("npm"):
            logger.warning("⚠️  'npm' not found. Frontend will be skipped. Install Node.js & npm to enable full-stack mode.")
            args.dev = False
            args.server = True

        # Prevent redundant checks later in the script
        args.setup = False
        args.init_db = False
        # Proceed to regular argument handling below
    
    # Environment setup
    if args.setup:
        from setup_env import setup_environment
        setup_environment()
        return
    
    # Initialize database
    if args.init_db:
        from backend.database.init_db import main as init_db_main
        init_db_main()
        return
    
    # Start full-stack development environment
    if args.dev:
        run_fullstack_dev(args)
        return
    
    # Start only frontend
    if args.frontend:
        logger.info("🌐 Starting frontend development server only...")
        frontend_process = start_frontend_dev_server(args.frontend_port)
        if frontend_process:
            try:
                # Wait for the process and handle Ctrl+C
                frontend_process.wait()
            except KeyboardInterrupt:
                logger.info("\n🛑 Stopping frontend development server...")
                frontend_process.terminate()
                try:
                    frontend_process.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    frontend_process.kill()
                logger.info("✅ Frontend stopped")
        return
    
    # Start server
    if args.server:
        start_backend_server(args)
        return
    
    # If no arguments are provided, show help
    parser.print_help()

if __name__ == "__main__":
    main()