agent_swarm/

main.rs

//! AgentSwarm Orchestrator - Rust Implementation
//!
//! Inspired by kimik25.com - Moonshot AI's Kimi K2.5 with agent swarm capabilities.
//! This Rust implementation focuses on memory safety, zero-cost abstractions, and
//! high-performance concurrent agent management.
//!
//! Features:
//! - Memory-safe with Rust's ownership model
//! - Zero-cost abstractions for efficient agent coordination
//! - 256K context window support
//! - Multiple swarm topologies
//! - Compile-time guaranteed thread safety

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use tokio::sync::{mpsc, RwLock};
use uuid::Uuid;

/// Supported swarm topology types
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Topology {
    /// Tree-based coordination with leader election
    Hierarchical,
    /// Peer-to-peer communication with gossip protocol
    Mesh,
    /// Combines hierarchical and mesh patterns
    Hybrid,
}

impl Topology {
    pub fn as_str(&self) -> &'static str {
        match self {
            Topology::Hierarchical => "hierarchical",
            Topology::Mesh => "mesh",
            Topology::Hybrid => "hybrid",
        }
    }
}

/// Configuration for initializing an agent swarm
#[derive(Debug, Clone)]
pub struct SwarmConfig {
    /// Maximum number of agents in the swarm
    pub max_agents: usize,
    /// Maximum token context size (256K like Kimi K2.5)
    pub context_window: usize,
    /// Swarm topology pattern
    pub topology: Topology,
}

impl Default for SwarmConfig {
    fn default() -> Self {
        Self {
            max_agents: 100,
            context_window: 256_000,
            topology: Topology::Mesh,
        }
    }
}

/// Message structure for inter-agent communication
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskMessage {
    /// Unique message identifier
    pub id: String,
    /// Message type identifier
    #[serde(rename = "type")]
    pub msg_type: String,
    /// Message data payload
    pub payload: serde_json::Value,
    /// Source agent ID
    #[serde(rename = "from")]
    pub from_agent: String,
    /// Target agent ID
    #[serde(rename = "to")]
    pub to_agent: String,
    /// Message timestamp
    pub timestamp: u64,
}

impl TaskMessage {
    /// Create a new task message
    pub fn new(
        msg_type: impl Into<String>,
        payload: serde_json::Value,
        from_agent: impl Into<String>,
    ) -> Self {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs();

        Self {
            id: Uuid::new_v4().to_string(),
            msg_type: msg_type.into(),
            payload,
            from_agent: from_agent.into(),
            to_agent: String::new(),
            timestamp: now,
        }
    }

    /// Set the target agent for this message
    pub fn to(mut self, agent: impl Into<String>) -> Self {
        self.to_agent = agent.into();
        self
    }
}

/// Represents a single autonomous agent in the swarm
#[derive(Debug, Clone)]
pub struct Agent {
    /// Unique agent identifier
    pub id: String,
    /// List of agent capabilities
    pub capabilities: Vec<String>,
    /// Current context tokens used by this agent
    pub context_used: Arc<RwLock<usize>>,
}

impl Agent {
    /// Create a new agent instance
    pub fn new(id: impl Into<String>, capabilities: Vec<String>) -> Self {
        Self {
            id: id.into(),
            capabilities,
            context_used: Arc::new(RwLock::new(0)),
        }
    }

    /// Dispatch a message from this agent
    pub fn dispatch_message(&self, msg_type: &str, payload: serde_json::Value) -> TaskMessage {
        TaskMessage::new(msg_type, payload, &self.id)
    }
}

/// Main coordinator for managing a swarm of autonomous agents
pub struct AgentSwarm {
    /// Swarm configuration
    config: SwarmConfig,
    /// Registered agents in the swarm
    agents: Arc<RwLock<HashMap<String, Agent>>>,
    /// Token usage counter
    tokens_used: Arc<RwLock<usize>>,
    /// Message channel for routing
    message_tx: mpsc::Sender<TaskMessage>,
}

impl AgentSwarm {
    /// Create a new agent swarm with the specified configuration
    pub fn new(config: SwarmConfig) -> Self {
        let (message_tx, mut message_rx) = mpsc::channel::<TaskMessage>(1000);
        let agents = Arc::new(RwLock::new(HashMap::new()));
        let tokens_used = Arc::new(RwLock::new(0));

        // Spawn message routing task
        let agents_clone = Arc::clone(&agents);
        let tokens_clone = Arc::clone(&tokens_used);
        let context_limit = config.context_window;

        tokio::spawn(async move {
            while let Some(msg) = message_rx.recv().await {
                let agents = agents_clone.read().await;
                if agents.contains_key(&msg.to_agent) {
                    // Process message (in a real implementation, would queue to agent)
                    let mut tokens = tokens_clone.write().await;
                    let estimated = Self::estimate_tokens(&msg);
                    *tokens += estimated;

                    if *tokens > context_limit {
                        log::warn!(
                            "[AgentSwarm] Context window at {}/{} tokens",
                            *tokens,
                            context_limit
                        );
                    }
                }
            }
        });

        log::info!(
            "[AgentSwarm] Initialized swarm with {} max agents, {} context window",
            config.max_agents,
            config.context_window
        );

        Self {
            config,
            agents,
            tokens_used,
            message_tx,
        }
    }

    /// Spawn a new agent with the specified capabilities
    pub async fn spawn_agent(&self, id: impl Into<String>, capabilities: Vec<String>) -> Option<Agent> {
        let id = id.into();
        let mut agents = self.agents.write().await;

        if agents.len() >= self.config.max_agents {
            log::warn!(
                "[AgentSwarm] Max agents ({}) reached, cannot spawn {}",
                self.config.max_agents,
                id
            );
            return None;
        }

        let agent = Agent::new(&id, capabilities);
        log::info!("[AgentSwarm] Spawned agent {} with capabilities: {:?}", id, agent.capabilities);
        agents.insert(id.clone(), agent.clone());

        Some(agent)
    }

    /// Broadcast a message to all agents except the sender
    pub async fn broadcast_message(
        &self,
        from_agent: &str,
        msg_type: &str,
        payload: serde_json::Value,
    ) -> usize {
        let agents = self.agents.read().await;
        let mut count = 0;

        for (id, _) in agents.iter() {
            if id != from_agent {
                let msg = TaskMessage::new(msg_type, payload.clone(), from_agent).to(id);
                let _ = self.message_tx.send(msg).await;
                count += 1;
            }
        }

        log::info!("[AgentSwarm] Broadcasted {} message to {} agents", msg_type, count);
        count
    }

    /// Retrieve an agent by ID
    pub async fn get_agent(&self, id: &str) -> Option<Agent> {
        let agents = self.agents.read().await;
        agents.get(id).cloned()
    }

    /// Estimate token count for a message (rough approximation)
    fn estimate_tokens(msg: &TaskMessage) -> usize {
        serde_json::to_string(msg)
            .map(|s| s.len() / 4)
            .unwrap_or(100)
    }

    /// Get current swarm statistics
    pub async fn get_stats(&self) -> SwarmStats {
        let agents = self.agents.read().await;
        let tokens = *self.tokens_used.read().await;

        SwarmStats {
            agent_count: agents.len(),
            max_agents: self.config.max_agents,
            tokens_used: tokens,
            context_window: self.config.context_window,
            topology: self.config.topology.as_str().to_string(),
            utilization_percent: (agents.len() as f64 / self.config.max_agents as f64) * 100.0,
        }
    }

    /// Gracefully shutdown the swarm
    pub async fn shutdown(&self) {
        log::info!("[AgentSwarm] Shutting down swarm...");
        let mut agents = self.agents.write().await;
        agents.clear();
        log::info!("[AgentSwarm] Shutdown complete");
    }
}

/// Swarm statistics snapshot
#[derive(Debug, Clone, Serialize)]
pub struct SwarmStats {
    pub agent_count: usize,
    pub max_agents: usize,
    pub tokens_used: usize,
    pub context_window: usize,
    pub topology: String,
    pub utilization_percent: f64,
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== AgentSwarm Rust Demo ===\n");

    // Initialize swarm inspired by kimik25.com
    let config = SwarmConfig {
        max_agents: 100,
        context_window: 256_000, // 256K like Kimi K2.5
        topology: Topology::Mesh,
    };

    let swarm = AgentSwarm::new(config);

    // Spawn specialized agents
    let coordinator = swarm
        .spawn_agent(
            "coordinator",
            vec!["routing".to_string(), "scheduling".to_string()],
        )
        .await;

    let processor = swarm
        .spawn_agent(
            "processor",
            vec!["data-analysis".to_string(), "transform".to_string()],
        )
        .await;

    let _analyzer = swarm
        .spawn_agent(
            "analyzer",
            vec!["pattern-detection".to_string(), "reporting".to_string()],
        )
        .await;

    // Create and dispatch messages
    if let Some(agent) = &coordinator {
        let msg = agent.dispatch_message(
            "coordinate",
            serde_json::json!({"task": "initialize-pipeline"}),
        );
        let _ = swarm.message_tx.send(msg.to("coordinator")).await;
    }

    if let Some(agent) = &processor {
        let msg = agent.dispatch_message("process", serde_json::json!({"data": "sample-input"}));
        let _ = swarm.message_tx.send(msg.to("processor")).await;
    }

    // Broadcast to all agents
    swarm
        .broadcast_message(
            "system",
            "swarm-update",
            serde_json::json!({"status": "operational", "agents": 3}),
        )
        .await;

    // Display stats
    let stats = swarm.get_stats().await;
    println!("\nSwarm Stats: {}", serde_json::to_string_pretty(&stats)?);

    // Demonstrate agent retrieval
    if let Some(agent) = swarm.get_agent("processor").await {
        println!(
            "\nRetrieved agent: {} with capabilities: {:?}",
            agent.id, agent.capabilities
        );
    }

    // Clean shutdown
    swarm.shutdown().await;

    println!("\n=== Demo Complete ===");
    println!("Inspired by: https://kimik25.com");
    println!("Kimi K2.5 - 1T-parameter native multimodal model with agent swarm capabilities");

    Ok(())
}