mecha10_behavior_patterns/

ensemble.rs

//! Ensemble pattern for multi-model fusion
//!
//! This module implements ensemble learning patterns where multiple AI models
//! can be combined to make more robust decisions.

use async_trait::async_trait;
use mecha10_behavior_runtime::{BehaviorNode, BoxedBehavior, NodeStatus};
use mecha10_core::Context;
use serde::{Deserialize, Serialize};
use tracing::{debug, info};

/// A weighted model in an ensemble.
#[derive(Debug)]
pub struct WeightedModel {
    /// The behavior/model
    pub behavior: BoxedBehavior,
    /// Weight for voting (only used in WeightedVote strategy)
    pub weight: f32,
    /// Optional name for debugging
    pub name: Option<String>,
}

impl WeightedModel {
    mecha10_core::new_with_defaults!(
        /// Create a new weighted model.
        pub fn new(behavior: BoxedBehavior, weight: f32) -> Self {
            name: None
        }
    );

    mecha10_core::builder_setter!(
        /// Set the name of this model.
        into_some, with_name, name, impl Into<String>
    );
}

/// Strategy for combining results from multiple models.
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum EnsembleStrategy {
    /// Weighted voting: success if weighted sum of successes > threshold (default 0.5)
    WeightedVote,

    /// Conservative: all models must succeed
    Conservative,

    /// Optimistic: any model can succeed
    Optimistic,

    /// Majority: more than half must succeed
    Majority,
}

/// Ensemble node for multi-model fusion.
///
/// This node runs multiple behaviors/models in parallel and combines their results
/// using a configurable strategy. This is useful for:
/// - Combining multiple AI models for robust predictions
/// - Redundancy and fault tolerance
/// - Multi-modal sensor fusion
///
/// # Strategies
///
/// - **WeightedVote**: Models vote with weights, success if total > threshold
/// - **Conservative**: All models must succeed (high precision, low recall)
/// - **Optimistic**: Any model can succeed (high recall, low precision)
/// - **Majority**: More than half must succeed (balanced)
///
/// # Example
///
/// ```rust
/// use mecha10_behavior_patterns::prelude::*;
///
/// # #[derive(Debug)]
/// # struct YoloV8;
/// # #[async_trait]
/// # impl BehaviorNode for YoloV8 {
/// #     async fn tick(&mut self, _ctx: &Context) -> anyhow::Result<NodeStatus> {
/// #         Ok(NodeStatus::Success)
/// #     }
/// # }
/// # #[derive(Debug)]
/// # struct YoloV10;
/// # #[async_trait]
/// # impl BehaviorNode for YoloV10 {
/// #     async fn tick(&mut self, _ctx: &Context) -> anyhow::Result<NodeStatus> {
/// #         Ok(NodeStatus::Success)
/// #     }
/// # }
/// # #[derive(Debug)]
/// # struct CustomDetector;
/// # #[async_trait]
/// # impl BehaviorNode for CustomDetector {
/// #     async fn tick(&mut self, _ctx: &Context) -> anyhow::Result<NodeStatus> {
/// #         Ok(NodeStatus::Success)
/// #     }
/// # }
///
/// let ensemble = EnsembleNode::new(EnsembleStrategy::WeightedVote)
///     .add_model(Box::new(YoloV8), 0.4)
///     .add_model(Box::new(YoloV10), 0.3)
///     .add_model(Box::new(CustomDetector), 0.3);
/// ```
#[derive(Debug)]
pub struct EnsembleNode {
    models: Vec<WeightedModel>,
    strategy: EnsembleStrategy,
    threshold: f32,
    model_statuses: Vec<NodeStatus>,
}

impl EnsembleNode {
    /// Create a new ensemble node with the given strategy.
    pub fn new(strategy: EnsembleStrategy) -> Self {
        Self {
            models: Vec::new(),
            strategy,
            threshold: 0.5,
            model_statuses: Vec::new(),
        }
    }

    /// Add a model with a weight.
    pub fn add_model(mut self, behavior: BoxedBehavior, weight: f32) -> Self {
        self.models.push(WeightedModel::new(behavior, weight));
        self.model_statuses.push(NodeStatus::Running);
        self
    }

    /// Set the threshold for weighted voting (default 0.5).
    pub fn with_threshold(mut self, threshold: f32) -> Self {
        self.threshold = threshold;
        self
    }
}

#[async_trait]
impl BehaviorNode for EnsembleNode {
    async fn tick(&mut self, ctx: &Context) -> anyhow::Result<NodeStatus> {
        if self.models.is_empty() {
            return Ok(NodeStatus::Failure);
        }

        for (i, model) in self.models.iter_mut().enumerate() {
            if self.model_statuses[i].is_running() {
                let status = model.behavior.tick(ctx).await?;
                self.model_statuses[i] = status;

                let model_name = model.name.as_deref().unwrap_or("unnamed");
                debug!("Ensemble: model '{}' status: {}", model_name, status);
            }
        }

        let overall_status = match self.strategy {
            EnsembleStrategy::Conservative => self.conservative_combine(),
            EnsembleStrategy::Optimistic => self.optimistic_combine(),
            EnsembleStrategy::Majority => self.majority_combine(),
            EnsembleStrategy::WeightedVote => self.weighted_vote_combine(),
        };

        info!("Ensemble ({:?}): overall status = {}", self.strategy, overall_status);
        Ok(overall_status)
    }

    mecha10_behavior_runtime::reset_child_statuses_field!(model_statuses);

    fn name(&self) -> &str {
        "ensemble"
    }

    mecha10_behavior_runtime::children_mut_via_wrapped_behavior!(models, model);
}

impl EnsembleNode {
    /// Shared shape behind `conservative_combine`/`optimistic_combine`: both pick an overall
    /// status if *any* model already reports `veto`, otherwise fall back to `Running` while any
    /// model is still in flight, otherwise settle on `otherwise` - they only differ in which
    /// status acts as the veto vs. the settled fallback (cargo-dupes flagged the two near-
    /// identical bodies as duplicates of each other).
    fn combine_if_any(&self, veto: NodeStatus, otherwise: NodeStatus) -> NodeStatus {
        let vetoed = self.model_statuses.contains(&veto);
        let running = self.model_statuses.iter().any(|s| s.is_running());

        if vetoed {
            veto
        } else if running {
            NodeStatus::Running
        } else {
            otherwise
        }
    }

    fn conservative_combine(&self) -> NodeStatus {
        self.combine_if_any(NodeStatus::Failure, NodeStatus::Success)
    }

    fn optimistic_combine(&self) -> NodeStatus {
        self.combine_if_any(NodeStatus::Success, NodeStatus::Failure)
    }

    fn majority_combine(&self) -> NodeStatus {
        let successes = self.model_statuses.iter().filter(|s| s.is_success()).count();
        let failures = self.model_statuses.iter().filter(|s| s.is_failure()).count();
        let running = self.model_statuses.iter().filter(|s| s.is_running()).count();

        let _total_decided = successes + failures;
        let majority_threshold = self.models.len().div_ceil(2);

        if successes >= majority_threshold {
            NodeStatus::Success
        } else if failures >= majority_threshold {
            NodeStatus::Failure
        } else if running > 0 {
            NodeStatus::Running
        } else {
            NodeStatus::Failure
        }
    }

    fn weighted_vote_combine(&self) -> NodeStatus {
        let mut success_weight = 0.0;
        let mut _failure_weight = 0.0;
        let mut total_weight = 0.0;
        let mut any_running = false;

        for (model, status) in self.models.iter().zip(&self.model_statuses) {
            total_weight += model.weight;
            match status {
                NodeStatus::Success => success_weight += model.weight,
                NodeStatus::Failure => _failure_weight += model.weight,
                NodeStatus::Running => any_running = true,
            }
        }

        if any_running {
            return NodeStatus::Running;
        }

        let success_ratio = if total_weight > 0.0 {
            success_weight / total_weight
        } else {
            0.0
        };

        if success_ratio >= self.threshold {
            NodeStatus::Success
        } else {
            NodeStatus::Failure
        }
    }
}

/// Configuration for an ensemble model (for JSON deserialization).
#[allow(dead_code)]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnsembleModelConfig {
    /// Node type to instantiate
    pub node: String,
    /// Weight for this model
    pub weight: f32,
    /// Optional name for this model
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
    /// Configuration for the node
    #[serde(default)]
    pub config: serde_json::Value,
}

#[cfg(test)]
mod tests {
    use super::*;

    // These exercise `combine_if_any` (the shared helper behind `conservative_combine` and
    // `optimistic_combine`) directly against `model_statuses`, without needing a real `Context`
    // or Redis - the `tests/ensemble_tests.rs` integration tests cover the full tick loop but
    // are ignored in environments without Redis, so unit coverage here is what actually runs.
    fn node_with_statuses(strategy: EnsembleStrategy, statuses: Vec<NodeStatus>) -> EnsembleNode {
        EnsembleNode {
            models: Vec::new(),
            strategy,
            threshold: 0.5,
            model_statuses: statuses,
        }
    }

    #[test]
    fn conservative_combine_fails_if_any_model_failed() {
        let node = node_with_statuses(
            EnsembleStrategy::Conservative,
            vec![NodeStatus::Success, NodeStatus::Failure, NodeStatus::Success],
        );
        assert_eq!(node.conservative_combine(), NodeStatus::Failure);
    }

    #[test]
    fn conservative_combine_runs_while_any_model_is_running() {
        let node = node_with_statuses(
            EnsembleStrategy::Conservative,
            vec![NodeStatus::Success, NodeStatus::Running],
        );
        assert_eq!(node.conservative_combine(), NodeStatus::Running);
    }

    #[test]
    fn conservative_combine_succeeds_if_all_models_succeeded() {
        let node = node_with_statuses(
            EnsembleStrategy::Conservative,
            vec![NodeStatus::Success, NodeStatus::Success],
        );
        assert_eq!(node.conservative_combine(), NodeStatus::Success);
    }

    #[test]
    fn optimistic_combine_succeeds_if_any_model_succeeded() {
        let node = node_with_statuses(
            EnsembleStrategy::Optimistic,
            vec![NodeStatus::Failure, NodeStatus::Success, NodeStatus::Failure],
        );
        assert_eq!(node.optimistic_combine(), NodeStatus::Success);
    }

    #[test]
    fn optimistic_combine_runs_while_any_model_is_running() {
        let node = node_with_statuses(
            EnsembleStrategy::Optimistic,
            vec![NodeStatus::Failure, NodeStatus::Running],
        );
        assert_eq!(node.optimistic_combine(), NodeStatus::Running);
    }

    #[test]
    fn optimistic_combine_fails_if_all_models_failed() {
        let node = node_with_statuses(
            EnsembleStrategy::Optimistic,
            vec![NodeStatus::Failure, NodeStatus::Failure],
        );
        assert_eq!(node.optimistic_combine(), NodeStatus::Failure);
    }
}