Vapora/crates/vapora-swarm/examples/01-agent-registration.rs

//! # Swarm Agent Registration Example
//!
//! Demonstrates how to register agents in a swarm and query their capabilities
//! and load.
//!
//! ## What This Example Shows
//! - Creating a `SwarmCoordinator`
//! - Registering multiple agents with different roles and capabilities
//! - Querying swarm statistics
//! - Understanding agent load balancing basics
//!
//! ## Run
//! ```bash
//! cargo run --example 01-agent-registration -p vapora-swarm
//! ```
//!
//! ## Expected Output
//! ```text
//! === Swarm Agent Registration Example ===
//!
//! Created SwarmCoordinator
//!
//! Registering agents:
//! - developer_1 (role: developer, load: 0.30)
//! - developer_2 (role: developer, load: 0.10)
//! - reviewer_1 (role: reviewer, load: 0.00)
//!
//! === Swarm Statistics ===
//! Total agents: 3
//! Developers: 2, Reviewers: 1
//! Average load: 0.13
//! Busiest agent: developer_1 (load: 0.30)
//! ```

use vapora_swarm::{AgentProfile, SwarmCoordinator};

fn main() {
    println!("=== Swarm Agent Registration Example ===\n");

    // Step 1: Create SwarmCoordinator
    let coordinator = SwarmCoordinator::new();
    println!("Created SwarmCoordinator\n");

    // Step 2: Define agent profiles
    let agents = vec![
        AgentProfile {
            id: "developer_1".to_string(),
            roles: vec!["developer".to_string()],
            capabilities: vec!["coding".to_string(), "testing".to_string()],
            current_load: 0.30,
            success_rate: 0.92,
            availability: true,
        },
        AgentProfile {
            id: "developer_2".to_string(),
            roles: vec!["developer".to_string()],
            capabilities: vec!["coding".to_string(), "documentation".to_string()],
            current_load: 0.10,
            success_rate: 0.85,
            availability: true,
        },
        AgentProfile {
            id: "reviewer_1".to_string(),
            roles: vec!["reviewer".to_string()],
            capabilities: vec!["code_review".to_string(), "quality_check".to_string()],
            current_load: 0.00,
            success_rate: 0.95,
            availability: true,
        },
    ];

    // Step 3: Register agents
    println!("Registering agents:");
    for agent in &agents {
        match coordinator.register_agent(agent.clone()) {
            Ok(_) => {
                println!(
                    "  ✓ {} (roles: {}, load: {:.2})",
                    agent.id,
                    agent.roles.join("+"),
                    agent.current_load
                );
            }
            Err(e) => {
                eprintln!("  ✗ Failed to register {}: {}", agent.id, e);
            }
        }
    }
    println!();

    // Step 4: Query swarm statistics
    let stats = coordinator.get_swarm_stats();
    println!("=== Swarm Statistics ===");
    println!("Total agents: {}", stats.total_agents);
    println!("Available agents: {}", stats.available_agents);

    // Count agents by role
    let mut role_counts: std::collections::HashMap<String, usize> =
        std::collections::HashMap::new();
    for agent in &agents {
        for role in &agent.roles {
            *role_counts.entry(role.clone()).or_insert(0) += 1;
        }
    }
    let roles_str = role_counts
        .iter()
        .map(|(role, count)| format!("{}={}", role, count))
        .collect::<Vec<_>>()
        .join(", ");
    println!("Agents by role: {}", roles_str);

    println!("Average load: {:.2}", stats.avg_load);
    println!("Active tasks: {}", stats.active_tasks);
    println!("Active coalitions: {}", stats.active_coalitions);

    // Step 5: Query agents by capability
    println!("\n=== Agents by Capability ===");
    for capability in &["coding", "testing", "code_review", "documentation"] {
        let agents_with_cap = agents
            .iter()
            .filter(|a| a.capabilities.contains(&capability.to_string()))
            .count();
        println!("  {}: {} agents", capability, agents_with_cap);
    }

    // Step 6: Demonstrate load-aware selection
    println!("\n=== Load-Aware Selection ===");
    let score_formula = "success_rate / (1 + current_load)";
    println!("Scoring formula: {}\n", score_formula);

    for agent in &agents {
        let score = agent.success_rate / (1.0 + agent.current_load);
        println!(
            "{}: score = {:.2} / (1 + {:.2}) = {:.3}",
            agent.id, agent.success_rate, agent.current_load, score
        );
    }

    println!("\n  → developer_2 has best score (lowest load + high success rate)");
}