rswarm 0.1.8

A Rust implementation of the Swarm framework
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

//! Registry for discovering agents by capability.

use crate::types::{Agent, AgentRef};
use std::collections::HashMap;
use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};

/// Thread-safe registry mapping [`AgentRef`]s to [`Agent`] instances.
///
/// Used by team-formation logic to discover agents that satisfy required
/// capabilities. All operations are `&self` via an internal `RwLock`.
#[derive(Clone, Default)]
pub struct AgentRegistry {
    agents: Arc<RwLock<HashMap<AgentRef, Arc<Agent>>>>,
}

impl AgentRegistry {
    fn read_agents(&self) -> RwLockReadGuard<'_, HashMap<AgentRef, Arc<Agent>>> {
        self.agents.read().unwrap_or_else(|poisoned| {
            tracing::warn!("agent registry lock poisoned; continuing with recovered state");
            poisoned.into_inner()
        })
    }

    fn write_agents(&self) -> RwLockWriteGuard<'_, HashMap<AgentRef, Arc<Agent>>> {
        self.agents.write().unwrap_or_else(|poisoned| {
            tracing::warn!("agent registry lock poisoned; continuing with recovered state");
            poisoned.into_inner()
        })
    }

    pub fn new() -> Self {
        Default::default()
    }

    /// Register an agent, keyed by its [`AgentRef`] (derived from its name).
    ///
    /// Overwrites any previously registered agent with the same ref.
    pub fn register(&self, agent: Arc<Agent>) {
        let key = agent.agent_ref();
        self.write_agents().insert(key, agent);
    }

    /// Look up an agent by its [`AgentRef`].
    pub fn get(&self, r: &AgentRef) -> Option<Arc<Agent>> {
        self.read_agents().get(r).cloned()
    }

    /// Return all [`AgentRef`]s whose agents declare the given capability.
    pub fn find_by_capability(&self, cap: &str) -> Vec<AgentRef> {
        self.read_agents()
            .iter()
            .filter(|(_, agent)| agent.has_capability(cap))
            .map(|(r, _)| r.clone())
            .collect()
    }

    /// Return all registered [`AgentRef`]s.
    pub fn all_refs(&self) -> Vec<AgentRef> {
        self.read_agents().keys().cloned().collect()
    }

    /// Number of registered agents.
    pub fn len(&self) -> usize {
        self.read_agents().len()
    }

    pub fn is_empty(&self) -> bool {
        self.read_agents().is_empty()
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{Agent, Instructions};

    fn make_agent(name: &str, caps: Vec<&str>) -> Arc<Agent> {
        Arc::new(
            Agent::new(
                name,
                "gpt-4o",
                Instructions::Text(format!("{} agent", name)),
            )
            .unwrap()
            .with_capabilities(caps.into_iter().map(String::from).collect()),
        )
    }

    #[test]
    fn test_register_and_get() {
        let reg = AgentRegistry::new();
        let agent = make_agent("planner", vec!["planning"]);
        reg.register(agent.clone());

        let found = reg.get(&AgentRef::new("planner")).unwrap();
        assert_eq!(found.name(), "planner");
    }

    #[test]
    fn test_find_by_capability() {
        let reg = AgentRegistry::new();
        reg.register(make_agent("planner", vec!["planning", "reasoning"]));
        reg.register(make_agent("coder", vec!["coding", "reasoning"]));
        reg.register(make_agent("reviewer", vec!["review"]));

        let mut planners = reg.find_by_capability("planning");
        planners.sort_by(|a, b| a.as_str().cmp(b.as_str()));
        assert_eq!(planners, vec![AgentRef::new("planner")]);

        let mut reasoners = reg.find_by_capability("reasoning");
        reasoners.sort_by(|a, b| a.as_str().cmp(b.as_str()));
        assert_eq!(
            reasoners,
            vec![AgentRef::new("coder"), AgentRef::new("planner")]
        );

        assert!(reg.find_by_capability("nonexistent").is_empty());
    }

    #[test]
    fn test_get_missing() {
        let reg = AgentRegistry::new();
        assert!(reg.get(&AgentRef::new("nobody")).is_none());
    }
}