tirea_agentos/composition/registry/
behavior.rs

1use super::sorted_registry_ids;
2use super::traits::{BehaviorRegistry, BehaviorRegistryError};
3use crate::contracts::runtime::AgentBehavior;
4use std::collections::HashMap;
5use std::sync::Arc;
6
7#[derive(Clone, Default)]
8pub struct InMemoryBehaviorRegistry {
9    behaviors: HashMap<String, Arc<dyn AgentBehavior>>,
10}
11
12impl std::fmt::Debug for InMemoryBehaviorRegistry {
13    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
14        f.debug_struct("InMemoryBehaviorRegistry")
15            .field("len", &self.behaviors.len())
16            .finish()
17    }
18}
19
20impl InMemoryBehaviorRegistry {
21    pub fn new() -> Self {
22        Self::default()
23    }
24
25    pub fn register(
26        &mut self,
27        behavior: Arc<dyn AgentBehavior>,
28    ) -> Result<(), BehaviorRegistryError> {
29        let id = behavior.id().to_string();
30        if self.behaviors.contains_key(&id) {
31            return Err(BehaviorRegistryError::BehaviorIdConflict(id));
32        }
33        self.behaviors.insert(id, behavior);
34        Ok(())
35    }
36
37    pub fn register_named(
38        &mut self,
39        id: impl Into<String>,
40        behavior: Arc<dyn AgentBehavior>,
41    ) -> Result<(), BehaviorRegistryError> {
42        let key = id.into();
43        let behavior_id = behavior.id().to_string();
44        if key != behavior_id {
45            return Err(BehaviorRegistryError::BehaviorIdMismatch { key, behavior_id });
46        }
47        if self.behaviors.contains_key(&key) {
48            return Err(BehaviorRegistryError::BehaviorIdConflict(key));
49        }
50        self.behaviors.insert(key, behavior);
51        Ok(())
52    }
53
54    pub fn extend_named(
55        &mut self,
56        behaviors: HashMap<String, Arc<dyn AgentBehavior>>,
57    ) -> Result<(), BehaviorRegistryError> {
58        for (key, behavior) in behaviors {
59            self.register_named(key, behavior)?;
60        }
61        Ok(())
62    }
63
64    pub fn extend_registry(
65        &mut self,
66        other: &dyn BehaviorRegistry,
67    ) -> Result<(), BehaviorRegistryError> {
68        self.extend_named(other.snapshot())
69    }
70}
71
72impl BehaviorRegistry for InMemoryBehaviorRegistry {
73    fn len(&self) -> usize {
74        self.behaviors.len()
75    }
76
77    fn get(&self, id: &str) -> Option<Arc<dyn AgentBehavior>> {
78        self.behaviors.get(id).cloned()
79    }
80
81    fn ids(&self) -> Vec<String> {
82        sorted_registry_ids(&self.behaviors)
83    }
84
85    fn snapshot(&self) -> HashMap<String, Arc<dyn AgentBehavior>> {
86        self.behaviors.clone()
87    }
88}
89
90#[derive(Clone, Default)]
91pub struct CompositeBehaviorRegistry {
92    merged: InMemoryBehaviorRegistry,
93}
94
95impl std::fmt::Debug for CompositeBehaviorRegistry {
96    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
97        f.debug_struct("CompositeBehaviorRegistry")
98            .field("len", &self.merged.len())
99            .finish()
100    }
101}
102
103impl CompositeBehaviorRegistry {
104    pub fn try_new(
105        regs: impl IntoIterator<Item = Arc<dyn BehaviorRegistry>>,
106    ) -> Result<Self, BehaviorRegistryError> {
107        let mut merged = InMemoryBehaviorRegistry::new();
108        for r in regs {
109            merged.extend_registry(r.as_ref())?;
110        }
111        Ok(Self { merged })
112    }
113}
114
115impl BehaviorRegistry for CompositeBehaviorRegistry {
116    fn len(&self) -> usize {
117        self.merged.len()
118    }
119
120    fn get(&self, id: &str) -> Option<Arc<dyn AgentBehavior>> {
121        self.merged.get(id)
122    }
123
124    fn ids(&self) -> Vec<String> {
125        sorted_registry_ids(&self.merged.behaviors)
126    }
127
128    fn snapshot(&self) -> HashMap<String, Arc<dyn AgentBehavior>> {
129        self.merged.snapshot()
130    }
131}