mecha10_behavior_patterns/
ensemble.rs

1//! Ensemble pattern for multi-model fusion
2//!
3//! This module implements ensemble learning patterns where multiple AI models
4//! can be combined to make more robust decisions.
5
6use 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)]
13pub struct WeightedModel {
14    pub behavior: BoxedBehavior,
15    pub weight: f32,
16    pub name: Option<String>,
17}
18
19impl WeightedModel {
20    mecha10_core::new_with_defaults!(
21        /// Create a new weighted model.
22        pub fn new(behavior: BoxedBehavior, weight: f32) -> Self {
23            name: None
24        }
25    );
26
27    mecha10_core::builder_setter!(into_some, with_name, name, impl Into<String>);
28}
29
30#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
31#[serde(rename_all = "snake_case")]
32pub enum EnsembleStrategy {
33    WeightedVote,
34
35    Conservative,
36
37    Optimistic,
38
39    Majority,
40}
41
42#[derive(Debug)]
43pub struct EnsembleNode {
44    models: Vec<WeightedModel>,
45    strategy: EnsembleStrategy,
46    threshold: f32,
47    model_statuses: Vec<NodeStatus>,
48}
49
50impl EnsembleNode {
51    /// Create a new ensemble node with the given strategy.
52    pub fn new(strategy: EnsembleStrategy) -> Self {
53        Self {
54            models: Vec::new(),
55            strategy,
56            threshold: 0.5,
57            model_statuses: Vec::new(),
58        }
59    }
60
61    /// Add a model with a weight.
62    pub fn add_model(mut self, behavior: BoxedBehavior, weight: f32) -> Self {
63        self.models.push(WeightedModel::new(behavior, weight));
64        self.model_statuses.push(NodeStatus::Running);
65        self
66    }
67
68    /// Set the threshold for weighted voting (default 0.5).
69    pub fn with_threshold(mut self, threshold: f32) -> Self {
70        self.threshold = threshold;
71        self
72    }
73}
74
75#[async_trait]
76impl BehaviorNode for EnsembleNode {
77    async fn tick(&mut self, ctx: &Context) -> anyhow::Result<NodeStatus> {
78        if self.models.is_empty() {
79            return Ok(NodeStatus::Failure);
80        }
81
82        for (i, model) in self.models.iter_mut().enumerate() {
83            if self.model_statuses[i].is_running() {
84                let status = model.behavior.tick(ctx).await?;
85                self.model_statuses[i] = status;
86
87                let model_name = model.name.as_deref().unwrap_or("unnamed");
88                debug!("Ensemble: model '{}' status: {}", model_name, status);
89            }
90        }
91
92        let overall_status = match self.strategy {
93            EnsembleStrategy::Conservative => self.conservative_combine(),
94            EnsembleStrategy::Optimistic => self.optimistic_combine(),
95            EnsembleStrategy::Majority => self.majority_combine(),
96            EnsembleStrategy::WeightedVote => self.weighted_vote_combine(),
97        };
98
99        info!("Ensemble ({:?}): overall status = {}", self.strategy, overall_status);
100        Ok(overall_status)
101    }
102
103    mecha10_behavior_runtime::reset_child_statuses_field!(model_statuses);
104
105    fn name(&self) -> &str {
106        "ensemble"
107    }
108
109    mecha10_behavior_runtime::children_mut_via_wrapped_behavior!(models, model);
110}
111
112impl EnsembleNode {
113    /// Shared shape behind `conservative_combine`/`optimistic_combine`: both pick an overall
114    /// status if *any* model already reports `veto`, otherwise fall back to `Running` while any
115    /// model is still in flight, otherwise settle on `otherwise` - they only differ in which
116    /// status acts as the veto vs. the settled fallback (cargo-dupes flagged the two near-
117    /// identical bodies as duplicates of each other).
118    fn combine_if_any(&self, veto: NodeStatus, otherwise: NodeStatus) -> NodeStatus {
119        let vetoed = self.model_statuses.contains(&veto);
120        let running = self.model_statuses.iter().any(|s| s.is_running());
121
122        if vetoed {
123            veto
124        } else if running {
125            NodeStatus::Running
126        } else {
127            otherwise
128        }
129    }
130
131    fn conservative_combine(&self) -> NodeStatus {
132        self.combine_if_any(NodeStatus::Failure, NodeStatus::Success)
133    }
134
135    fn optimistic_combine(&self) -> NodeStatus {
136        self.combine_if_any(NodeStatus::Success, NodeStatus::Failure)
137    }
138
139    fn majority_combine(&self) -> NodeStatus {
140        let successes = self.model_statuses.iter().filter(|s| s.is_success()).count();
141        let failures = self.model_statuses.iter().filter(|s| s.is_failure()).count();
142        let running = self.model_statuses.iter().filter(|s| s.is_running()).count();
143
144        let _total_decided = successes + failures;
145        let majority_threshold = self.models.len().div_ceil(2);
146
147        if successes >= majority_threshold {
148            NodeStatus::Success
149        } else if failures >= majority_threshold {
150            NodeStatus::Failure
151        } else if running > 0 {
152            NodeStatus::Running
153        } else {
154            NodeStatus::Failure
155        }
156    }
157
158    fn weighted_vote_combine(&self) -> NodeStatus {
159        let mut success_weight = 0.0;
160        let mut _failure_weight = 0.0;
161        let mut total_weight = 0.0;
162        let mut any_running = false;
163
164        for (model, status) in self.models.iter().zip(&self.model_statuses) {
165            total_weight += model.weight;
166            match status {
167                NodeStatus::Success => success_weight += model.weight,
168                NodeStatus::Failure => _failure_weight += model.weight,
169                NodeStatus::Running => any_running = true,
170            }
171        }
172
173        if any_running {
174            return NodeStatus::Running;
175        }
176
177        let success_ratio = if total_weight > 0.0 {
178            success_weight / total_weight
179        } else {
180            0.0
181        };
182
183        if success_ratio >= self.threshold {
184            NodeStatus::Success
185        } else {
186            NodeStatus::Failure
187        }
188    }
189}
190
191#[allow(dead_code)]
192#[derive(Debug, Clone, Serialize, Deserialize)]
193pub struct EnsembleModelConfig {
194    pub node: String,
195    pub weight: f32,
196    #[serde(skip_serializing_if = "Option::is_none")]
197    pub name: Option<String>,
198    #[serde(default)]
199    pub config: serde_json::Value,
200}
201
202#[cfg(test)]
203mod tests {
204    use super::*;
205
206    fn node_with_statuses(strategy: EnsembleStrategy, statuses: Vec<NodeStatus>) -> EnsembleNode {
207        EnsembleNode {
208            models: Vec::new(),
209            strategy,
210            threshold: 0.5,
211            model_statuses: statuses,
212        }
213    }
214
215    #[test]
216    fn conservative_combine_fails_if_any_model_failed() {
217        let node = node_with_statuses(
218            EnsembleStrategy::Conservative,
219            vec![NodeStatus::Success, NodeStatus::Failure, NodeStatus::Success],
220        );
221        assert_eq!(node.conservative_combine(), NodeStatus::Failure);
222    }
223
224    #[test]
225    fn conservative_combine_runs_while_any_model_is_running() {
226        let node = node_with_statuses(
227            EnsembleStrategy::Conservative,
228            vec![NodeStatus::Success, NodeStatus::Running],
229        );
230        assert_eq!(node.conservative_combine(), NodeStatus::Running);
231    }
232
233    #[test]
234    fn conservative_combine_succeeds_if_all_models_succeeded() {
235        let node = node_with_statuses(
236            EnsembleStrategy::Conservative,
237            vec![NodeStatus::Success, NodeStatus::Success],
238        );
239        assert_eq!(node.conservative_combine(), NodeStatus::Success);
240    }
241
242    #[test]
243    fn optimistic_combine_succeeds_if_any_model_succeeded() {
244        let node = node_with_statuses(
245            EnsembleStrategy::Optimistic,
246            vec![NodeStatus::Failure, NodeStatus::Success, NodeStatus::Failure],
247        );
248        assert_eq!(node.optimistic_combine(), NodeStatus::Success);
249    }
250
251    #[test]
252    fn optimistic_combine_runs_while_any_model_is_running() {
253        let node = node_with_statuses(
254            EnsembleStrategy::Optimistic,
255            vec![NodeStatus::Failure, NodeStatus::Running],
256        );
257        assert_eq!(node.optimistic_combine(), NodeStatus::Running);
258    }
259
260    #[test]
261    fn optimistic_combine_fails_if_all_models_failed() {
262        let node = node_with_statuses(
263            EnsembleStrategy::Optimistic,
264            vec![NodeStatus::Failure, NodeStatus::Failure],
265        );
266        assert_eq!(node.optimistic_combine(), NodeStatus::Failure);
267    }
268}
mecha10_behavior_patterns/ensemble.rs

mecha10_behavior_patterns/
ensemble.rs
