multi_agents.py

import os
import boto3
import concurrent.futures
import uuid
from typing import List, Dict, Tuple

class AgentConfig:
    def __init__(self, name: str, agent_id: str, alias_id: str):
        self.name = name
        self.agent_id = agent_id
        self.alias_id = alias_id
    
    def __repr__(self):
        return f"Agent({self.name}: {self.agent_id}/{self.alias_id})"

def load_agents() -> List[AgentConfig]:
    """Load all configured agents from environment variables."""
    agents = []
    
    # Check for agents numbered 1-10
    for i in range(1, 11):
        agent_id = os.getenv(f"BEDROCK_AGENT_ID_{i}")
        alias_id = os.getenv(f"BEDROCK_AGENT_ALIAS_ID_{i}")
        name = os.getenv(f"AGENT_NAME_{i}", f"Agent{i}")
        
        if agent_id and alias_id:
            agents.append(AgentConfig(name, agent_id, alias_id))
    
    return agents

def call_agent(agent: AgentConfig, prompt: str, region: str, session_id: str) -> Tuple[str, str, str]:
    """Call a single agent and return (agent_name, response_text, error)."""
    try:
        print(f"🤖 Calling {agent.name}...")
        
        # Create client for this call
        client = boto3.client("bedrock-agent-runtime", region_name=region)
        
        # Call the agent
        response = client.invoke_agent(
            agentId=agent.agent_id,
            agentAliasId=agent.alias_id,
            inputText=prompt,
            sessionId=f"{session_id}-{agent.name.lower()}"
        )
        
        # Extract text from streaming response
        completion = response.get("completion")
        if not completion:
            return agent.name, "", "No completion stream found"
        
        text_parts = []
        for event in completion:
            if "chunk" in event:
                chunk = event["chunk"]
                if "bytes" in chunk:
                    text = chunk["bytes"].decode("utf-8")
                    text_parts.append(text)
        
        response_text = "".join(text_parts)
        return agent.name, response_text, ""
        
    except Exception as e:
        return agent.name, "", str(e)

def call_agents_synchronously(agents: List[AgentConfig], prompt: str, region: str) -> Dict[str, Dict]:
    """Call all agents one after another (synchronously)."""
    session_id = str(uuid.uuid4())[:8]
    results = {}
    
    print(f"📨 Calling {len(agents)} agents synchronously...")
    print(f"📝 Prompt: {prompt}")
    print(f"🌍 Region: {region}")
    print(f"🎯 Session base: {session_id}")
    print("=" * 80)
    
    for agent in agents:
        agent_name, response_text, error = call_agent(agent, prompt, region, session_id)
        
        results[agent_name] = {
            "agent_id": agent.agent_id,
            "alias_id": agent.alias_id,
            "response": response_text,
            "error": error,
            "success": not bool(error)
        }
        
        # Print result immediately
        if error:
            print(f"❌ {agent_name}: ERROR - {error}")
        else:
            print(f"✅ {agent_name}: {response_text}")
        print("-" * 40)
    
    return results

def call_agents_in_parallel(agents: List[AgentConfig], prompt: str, region: str) -> Dict[str, Dict]:
    """Call all agents in parallel using ThreadPoolExecutor."""
    session_id = str(uuid.uuid4())[:8]
    results = {}
    
    print(f"🚀 Calling {len(agents)} agents in parallel...")
    print(f"📝 Prompt: {prompt}")
    print(f"🌍 Region: {region}")
    print(f"🎯 Session base: {session_id}")
    print("=" * 80)
    
    # Use ThreadPoolExecutor for parallel execution
    with concurrent.futures.ThreadPoolExecutor(max_workers=len(agents)) as executor:
        # Submit all agent calls
        future_to_agent = {
            executor.submit(call_agent, agent, prompt, region, session_id): agent 
            for agent in agents
        }
        
        # Collect results as they complete
        for future in concurrent.futures.as_completed(future_to_agent):
            agent = future_to_agent[future]
            try:
                agent_name, response_text, error = future.result()
                
                results[agent_name] = {
                    "agent_id": agent.agent_id,
                    "alias_id": agent.alias_id,
                    "response": response_text,
                    "error": error,
                    "success": not bool(error)
                }
                
                # Print result as it completes
                if error:
                    print(f"❌ {agent_name}: ERROR - {error}")
                else:
                    print(f"✅ {agent_name}: {response_text}")
                print("-" * 40)
                
            except Exception as exc:
                print(f"❌ {agent.name}: EXCEPTION - {exc}")
                results[agent.name] = {
                    "agent_id": agent.agent_id,
                    "alias_id": agent.alias_id,
                    "response": "",
                    "error": str(exc),
                    "success": False
                }
    
    return results

def execute_rollcall(prompt: str = None, parallel: bool = False) -> Dict[str, Dict]:
    """
    Execute rollcall across all configured agents.
    
    Args:
        prompt: Custom prompt to send to agents. If None, uses default from env or fallback.
        parallel: Whether to call agents in parallel (True) or synchronously (False).
    
    Returns:
        Dictionary with agent results containing response, error, and success status.
    """
    # Load configuration
    region = os.getenv("AWS_REGION", "us-west-2")
    if prompt is None:
        prompt = os.getenv("PROMPT", "Hello from Large Marge! Please report in.")
    
    # Load agents from environment
    agents = load_agents()
    
    if not agents:
        return {"error": "No agents configured. Add BEDROCK_AGENT_ID_X and BEDROCK_AGENT_ALIAS_ID_X to .env file."}
    
    # Execute based on mode
    if parallel:
        results = call_agents_in_parallel(agents, prompt, region)
    else:
        results = call_agents_synchronously(agents, prompt, region)
    
    return results

def get_rollcall_summary(results: Dict[str, Dict]) -> Dict[str, any]:
    """
    Generate a summary of rollcall results.
    
    Args:
        results: Results dictionary from execute_rollcall()
    
    Returns:
        Summary dictionary with counts and status information.
    """
    if "error" in results:
        return {"error": results["error"], "total": 0, "successful": 0, "failed": 0}
    
    total = len(results)
    successful = sum(1 for r in results.values() if r.get("success", False))
    failed = total - successful
    
    return {
        "total": total,
        "successful": successful, 
        "failed": failed,
        "success_rate": successful / total if total > 0 else 0,
        "agents": list(results.keys())
    }

# Keep main function for standalone execution
def main():
    """Main function for standalone execution (when run directly)."""
    results = execute_rollcall()
    summary = get_rollcall_summary(results)
    
    if "error" in results:
        print(f"❌ Error: {results['error']}")
        return
    
    print("\n" + "=" * 80)
    print("📊 ROLLCALL SUMMARY")
    print("=" * 80)
    print(f"✅ Successful: {summary['successful']}/{summary['total']}")
    print(f"❌ Failed: {summary['failed']}/{summary['total']}")
    print(f"📈 Success Rate: {summary['success_rate']:.1%}")
    
    for name, result in results.items():
        status = "✅" if result["success"] else "❌"
        response_length = len(result.get("response", ""))
        print(f"{status} {name}: {response_length} chars")

if __name__ == "__main__":
    main()