mecha10-behavior-patterns 0.8.2

//! 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};

#[derive(Debug)]
pub struct WeightedModel {
    pub behavior: BoxedBehavior,
    pub weight: f32,
    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!(into_some, with_name, name, impl Into<String>);
}

#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum EnsembleStrategy {
    WeightedVote,

    Conservative,

    Optimistic,

    Majority,
}

#[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
        }
    }
}

#[allow(dead_code)]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnsembleModelConfig {
    pub node: String,
    pub weight: f32,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
    #[serde(default)]
    pub config: serde_json::Value,
}

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

    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);
    }
}