mecha10_behavior_patterns/
ensemble.rs1use async_trait::async_trait;
7use mecha10_behavior_runtime::{BehaviorNode, BoxedBehavior, NodeStatus};
8use mecha10_core::Context;
9use serde::{Deserialize, Serialize};
10use tracing::{debug, info};
11
12#[derive(Debug)]
14pub struct WeightedModel {
15 pub behavior: BoxedBehavior,
17 pub weight: f32,
19 pub name: Option<String>,
21}
22
23impl WeightedModel {
24 pub fn new(behavior: BoxedBehavior, weight: f32) -> Self {
26 Self {
27 behavior,
28 weight,
29 name: None,
30 }
31 }
32
33 mecha10_core::builder_setter!(
34 into_some, with_name, name, impl Into<String>
36 );
37}
38
39#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
41#[serde(rename_all = "snake_case")]
42pub enum EnsembleStrategy {
43 WeightedVote,
45
46 Conservative,
48
49 Optimistic,
51
52 Majority,
54}
55
56#[derive(Debug)]
107pub struct EnsembleNode {
108 models: Vec<WeightedModel>,
109 strategy: EnsembleStrategy,
110 threshold: f32,
111 model_statuses: Vec<NodeStatus>,
112}
113
114impl EnsembleNode {
115 pub fn new(strategy: EnsembleStrategy) -> Self {
117 Self {
118 models: Vec::new(),
119 strategy,
120 threshold: 0.5,
121 model_statuses: Vec::new(),
122 }
123 }
124
125 pub fn add_model(mut self, behavior: BoxedBehavior, weight: f32) -> Self {
127 self.models.push(WeightedModel::new(behavior, weight));
128 self.model_statuses.push(NodeStatus::Running);
129 self
130 }
131
132 pub fn with_threshold(mut self, threshold: f32) -> Self {
134 self.threshold = threshold;
135 self
136 }
137}
138
139#[async_trait]
140impl BehaviorNode for EnsembleNode {
141 async fn tick(&mut self, ctx: &Context) -> anyhow::Result<NodeStatus> {
142 if self.models.is_empty() {
143 return Ok(NodeStatus::Failure);
144 }
145
146 for (i, model) in self.models.iter_mut().enumerate() {
148 if self.model_statuses[i].is_running() {
149 let status = model.behavior.tick(ctx).await?;
150 self.model_statuses[i] = status;
151
152 let model_name = model.name.as_deref().unwrap_or("unnamed");
153 debug!("Ensemble: model '{}' status: {}", model_name, status);
154 }
155 }
156
157 let overall_status = match self.strategy {
159 EnsembleStrategy::Conservative => self.conservative_combine(),
160 EnsembleStrategy::Optimistic => self.optimistic_combine(),
161 EnsembleStrategy::Majority => self.majority_combine(),
162 EnsembleStrategy::WeightedVote => self.weighted_vote_combine(),
163 };
164
165 info!("Ensemble ({:?}): overall status = {}", self.strategy, overall_status);
166 Ok(overall_status)
167 }
168
169 async fn reset(&mut self) -> anyhow::Result<()> {
170 for (i, model) in self.models.iter_mut().enumerate() {
171 model.behavior.reset().await?;
172 self.model_statuses[i] = NodeStatus::Running;
173 }
174 Ok(())
175 }
176
177 fn name(&self) -> &str {
178 "ensemble"
179 }
180
181 fn children_mut(&mut self) -> Vec<&mut dyn BehaviorNode> {
182 self.models
183 .iter_mut()
184 .map(|model| model.behavior.as_mut() as &mut dyn BehaviorNode)
185 .collect()
186 }
187}
188
189impl EnsembleNode {
190 fn conservative_combine(&self) -> NodeStatus {
191 let failures = self.model_statuses.iter().filter(|s| s.is_failure()).count();
193 let running = self.model_statuses.iter().filter(|s| s.is_running()).count();
194
195 if failures > 0 {
196 NodeStatus::Failure
197 } else if running > 0 {
198 NodeStatus::Running
199 } else {
200 NodeStatus::Success
201 }
202 }
203
204 fn optimistic_combine(&self) -> NodeStatus {
205 let successes = self.model_statuses.iter().filter(|s| s.is_success()).count();
207 let running = self.model_statuses.iter().filter(|s| s.is_running()).count();
208
209 if successes > 0 {
210 NodeStatus::Success
211 } else if running > 0 {
212 NodeStatus::Running
213 } else {
214 NodeStatus::Failure
215 }
216 }
217
218 fn majority_combine(&self) -> NodeStatus {
219 let successes = self.model_statuses.iter().filter(|s| s.is_success()).count();
221 let failures = self.model_statuses.iter().filter(|s| s.is_failure()).count();
222 let running = self.model_statuses.iter().filter(|s| s.is_running()).count();
223
224 let _total_decided = successes + failures;
225 let majority_threshold = self.models.len().div_ceil(2);
226
227 if successes >= majority_threshold {
228 NodeStatus::Success
229 } else if failures >= majority_threshold {
230 NodeStatus::Failure
231 } else if running > 0 {
232 NodeStatus::Running
233 } else {
234 NodeStatus::Failure
236 }
237 }
238
239 fn weighted_vote_combine(&self) -> NodeStatus {
240 let mut success_weight = 0.0;
242 let mut _failure_weight = 0.0;
243 let mut total_weight = 0.0;
244 let mut any_running = false;
245
246 for (model, status) in self.models.iter().zip(&self.model_statuses) {
247 total_weight += model.weight;
248 match status {
249 NodeStatus::Success => success_weight += model.weight,
250 NodeStatus::Failure => _failure_weight += model.weight,
251 NodeStatus::Running => any_running = true,
252 }
253 }
254
255 if any_running {
256 return NodeStatus::Running;
257 }
258
259 let success_ratio = if total_weight > 0.0 {
260 success_weight / total_weight
261 } else {
262 0.0
263 };
264
265 if success_ratio >= self.threshold {
266 NodeStatus::Success
267 } else {
268 NodeStatus::Failure
269 }
270 }
271}
272
273#[allow(dead_code)] #[derive(Debug, Clone, Serialize, Deserialize)]
276pub struct EnsembleModelConfig {
277 pub node: String,
279 pub weight: f32,
281 #[serde(skip_serializing_if = "Option::is_none")]
283 pub name: Option<String>,
284 #[serde(default)]
286 pub config: serde_json::Value,
287}