void / organon / setup_group.py

#!/usr/bin/env python3
"""
Simple tool to set up the Organon ecosystem group without running the firehose listener.
"""

import os
import sys
from dotenv import load_dotenv
from letta_client import Letta, SupervisorManager
from rich.console import Console

# Add parent directory to path for imports
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config_loader import get_config

load_dotenv()

def main():
    console = Console()
    
    try:
        # Initialize configuration and client
        config = get_config()
        
        client = Letta(
            base_url=config.get('letta.base_url', os.environ.get('LETTA_BASE_URL')),
            token=config.get('letta.api_key', os.environ.get('LETTA_API_KEY')),
            timeout=config.get('letta.timeout', 30)
        )
        
        project_id = config.get('letta.project_id', os.environ.get('LETTA_PROJECT_ID'))
        if not project_id:
            raise ValueError("Project ID must be set in config.yaml under letta.project_id or as LETTA_PROJECT_ID environment variable")
        
        # Get Organon central agent
        console.print("[blue]Finding Organon Central agent...[/blue]")
        try:
            organon_agents = client.agents.list(project_id=project_id, name="organon-central")
        except:
            # Fallback for self-hosted without project support
            organon_agents = client.agents.list(name="organon-central")
        if not organon_agents:
            console.print("[red]❌ Organon Central agent not found. Run create_organon.py first.[/red]")
            return
        
        organon_central_id = organon_agents[0].id
        console.print(f"[green]✅ Found Organon Central: {organon_central_id[:8]}[/green]")
        
        # Get Organon shards
        console.print("[blue]Finding Organon shards...[/blue]")
        try:
            shard_agents = client.agents.list(project_id=project_id, tags=["organon-shard"])
        except:
            # Fallback for self-hosted without project support
            shard_agents = client.agents.list(tags=["organon-shard"])
        if not shard_agents:
            console.print("[red]❌ No Organon shards found. Run create_organon.py to create shards.[/red]")
            return
        
        console.print(f"[green]✅ Found {len(shard_agents)} shards:[/green]")
        for shard in shard_agents:
            console.print(f"   • {shard.name} ({shard.id[:8]})")
        
        # Check if group already exists
        console.print("[blue]Checking for existing groups...[/blue]")
        try:
            groups = client.groups.list(project_id=project_id)
        except:
            # Fallback for self-hosted without project support
            groups = client.groups.list()
        
        existing_group = None
        for group in groups:
            if (group.description and 'organon ecosystem' in group.description.lower()) or \
               (hasattr(group, 'manager_config') and group.manager_config and 
                hasattr(group.manager_config, 'manager_agent_id') and 
                group.manager_config.manager_agent_id == organon_central_id):
                existing_group = group
                break
        
        if existing_group:
            console.print(f"[yellow]Group already exists: {existing_group.id[:12]}[/yellow]")
            return
        
        # Create the supervisor group
        console.print("[blue]Creating Organon ecosystem group...[/blue]")
        worker_agent_ids = [shard.id for shard in shard_agents]
        
        group = client.groups.create(
            agent_ids=worker_agent_ids,
            description="Supervisor group for the Organon ecosystem with organon-central managing all shards",
            manager_config=SupervisorManager(
                manager_agent_id=organon_central_id
            )
        )
        
        console.print(f"[green]✅ Created Organon ecosystem group: {group.id[:12]}[/green]")
        console.print(f"   Supervisor: organon-central ({organon_central_id[:8]})")
        console.print(f"   Workers: {len(worker_agent_ids)} shards")
        
        # Verify the group was actually created
        console.print("[blue]Verifying group creation...[/blue]")
        try:
            retrieved_group = client.groups.retrieve(group_id=group.id)
            console.print(f"[green]✅ Group verified: {retrieved_group.id[:12]}[/green]")
            
            # Also check if it shows up in the list
            try:
                all_groups = client.groups.list(project_id=project_id)
            except:
                all_groups = client.groups.list()
            found_in_list = any(g.id == group.id for g in all_groups)
            console.print(f"[{'green' if found_in_list else 'red'}]{'✅' if found_in_list else '❌'} Group appears in list: {found_in_list}[/{'green' if found_in_list else 'red'}]")
            
        except Exception as e:
            console.print(f"[red]❌ Error verifying group: {e}[/red]")
        
        console.print("\n[bold green]Setup complete! You can now use:[/bold green]")
        console.print("• [cyan]python organon/chat_with_organon.py[/cyan] - Chat with the ecosystem")
        console.print("• [cyan]python organon/list_groups.py[/cyan] - View group status")
        console.print("• [cyan]python organon/firehose_listener.py[/cyan] - Start the firehose listener")
    
    except Exception as e:
        console.print(f"[red]Error: {e}[/red]")
        sys.exit(1)

if __name__ == "__main__":
    main()