main.py

import asyncio
import argparse
from rich.console import Console
from rich.panel import Panel
from rich.progress import Progress, SpinnerColumn, TextColumn
from orchestrator.langgraph_orchestrator import LangGraphOrchestrator

console = Console()

async def main():
    # Parse command line arguments
    parser = argparse.ArgumentParser(description='Stateful Multi-Agent Orchestration')
    parser.add_argument('task', nargs='?', default='persistent stateful multi-agent orchestration',
                       help='The task to process')
    parser.add_argument('--redis-host', default='localhost', help='Redis server host')
    parser.add_argument('--redis-port', type=int, default=6379, help='Redis server port')
    parser.add_argument('--session-id', help='Optional session ID for state persistence')
    
    args = parser.parse_args()
    
    # Initialize the orchestrator with Redis configuration
    redis_config = {
        'host': args.redis_host,
        'port': args.redis_port
    }
    
    try:
        # Show initialization status
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            transient=True,
        ) as progress:
            task = progress.add_task("Initializing orchestrator...", total=None)
            orchestrator = LangGraphOrchestrator(redis_config=redis_config)
            progress.update(task, completed=1, description="Orchestrator ready!")
        
        # Run the orchestration
        console.rule("🚀 Starting Multi-Agent Orchestration")
        console.print(f"📝 Task: [bold cyan]{args.task}[/]")
        if args.session_id:
            console.print(f"🔗 Session ID: [bold yellow]{args.session_id}[/]")
        
        # Execute the workflow
        result = await orchestrator.run(args.task, session_id=args.session_id)
        
        # Display final status
        if result.get('status') == 'completed':
            console.print("\n✅ [bold green]Orchestration completed successfully![/]")
        else:
            console.print("\n❌ [bold red]Orchestration completed with issues[/]")
        
        return result
    
    except Exception as e:
        console.print(f"\n[bold red]Error:[/] {str(e)}")
        if hasattr(e, '__traceback__'):
            import traceback
            console.print(traceback.format_exc())
        return {'status': 'error', 'message': str(e)}

if __name__ == "__main__":
    asyncio.run(main())