qae_kernel/

declarative.rs

// SPDX-License-Identifier: BUSL-1.1
//! Declarative constraint channels — JSON/TOML-defined constraints evaluated at runtime.
//!
//! A `DeclarativeChannel` wraps a `ConstraintDefinition` (which can be loaded from
//! JSON or TOML) and evaluates it against a state vector using a generic interpreter.
//! This enables constraint management without recompilation.

use crate::constraint::ConstraintChannel;
use crate::KernelResult;
use serde::{Deserialize, Serialize};

/// A constraint defined declaratively (JSON/TOML), not compiled Rust.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConstraintDefinition {
    /// Unique identifier for this constraint.
    pub id: String,
    /// Human-readable name.
    pub name: String,
    /// Description shown in UI.
    pub description: String,
    /// Domain tag (e.g., "finance", "agentic", "custom").
    pub domain_tag: String,
    /// Names of state dimensions this constraint operates on.
    pub dimensions: Vec<String>,
    /// The margin computation rule.
    pub rule: MarginRule,
    /// Optional per-channel threshold overrides.
    #[serde(default)]
    pub thresholds: ThresholdOverrides,
}

/// How to compute the margin from a state vector.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum MarginRule {
    /// Margin = value / budget. Produces 1.0 when value == 0, 0.0 when value >= budget.
    /// `dimension_index` selects which state dimension to read.
    #[serde(rename = "budget_ratio")]
    BudgetRatio {
        dimension_index: usize,
        budget: f64,
    },

    /// Margin = how far value is from the [min, max] boundary, normalized to [0, 1].
    /// Inside range → margin based on distance to nearest edge.
    #[serde(rename = "range_bound")]
    RangeBound {
        dimension_index: usize,
        min: f64,
        max: f64,
    },

    /// Margin = 1.0 if value matches pattern, else fallback.
    /// Useful for categorical checks (e.g., "is this in an approved list?").
    #[serde(rename = "pattern_match")]
    PatternMatch {
        dimension_index: usize,
        /// Values that produce margin = 1.0
        approved_values: Vec<f64>,
        /// Tolerance for floating-point comparison
        #[serde(default = "default_tolerance")]
        tolerance: f64,
        /// Margin when value is not approved
        #[serde(default = "default_fallback_margin")]
        fallback_margin: f64,
    },

    /// Margin computed from a simple arithmetic expression over state dimensions.
    /// Supports: weighted sum clamped to [0, 1].
    #[serde(rename = "weighted_sum")]
    WeightedSum {
        /// (dimension_index, weight) pairs
        weights: Vec<(usize, f64)>,
        /// Offset added after weighted sum
        #[serde(default)]
        offset: f64,
    },

    /// Constant margin — useful for testing or placeholder channels.
    #[serde(rename = "constant")]
    Constant { margin: f64 },
}

fn default_tolerance() -> f64 {
    1e-10
}

fn default_fallback_margin() -> f64 {
    0.0
}

/// Optional per-channel threshold overrides.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ThresholdOverrides {
    pub safe_threshold: Option<f64>,
    pub caution_threshold: Option<f64>,
    pub block_threshold: Option<f64>,
}

/// A constraint channel backed by a declarative definition.
pub struct DeclarativeChannel {
    definition: ConstraintDefinition,
}

impl DeclarativeChannel {
    /// Create a new declarative channel from a definition.
    pub fn new(definition: ConstraintDefinition) -> Self {
        Self { definition }
    }

    /// Get the underlying definition.
    pub fn definition(&self) -> &ConstraintDefinition {
        &self.definition
    }

    /// Parse a constraint definition from JSON.
    pub fn from_json(json: &str) -> KernelResult<Self> {
        let definition: ConstraintDefinition =
            serde_json::from_str(json).map_err(|e| crate::KernelError::DeclarativeError(e.to_string()))?;
        Ok(Self::new(definition))
    }
}

impl ConstraintChannel for DeclarativeChannel {
    fn name(&self) -> &str {
        &self.definition.name
    }

    fn evaluate(&self, state: &[f64]) -> KernelResult<f64> {
        let margin = evaluate_rule(&self.definition.rule, state)?;
        Ok(margin.clamp(0.0, 1.0))
    }

    fn dimension_names(&self) -> Vec<String> {
        self.definition.dimensions.clone()
    }
}

/// Evaluate a margin rule against a state vector.
fn evaluate_rule(rule: &MarginRule, state: &[f64]) -> KernelResult<f64> {
    match rule {
        MarginRule::BudgetRatio {
            dimension_index,
            budget,
        } => {
            let value = get_state_value(state, *dimension_index)?;
            if *budget <= 0.0 {
                return Ok(0.0);
            }
            // margin = 1 - (value / budget), clamped
            Ok(1.0 - (value / budget))
        }

        MarginRule::RangeBound {
            dimension_index,
            min,
            max,
        } => {
            let value = get_state_value(state, *dimension_index)?;
            if max <= min {
                return Ok(0.0);
            }
            let range = max - min;
            let midpoint = (min + max) / 2.0;
            let distance_to_edge = (range / 2.0) - (value - midpoint).abs();
            Ok(distance_to_edge / (range / 2.0))
        }

        MarginRule::PatternMatch {
            dimension_index,
            approved_values,
            tolerance,
            fallback_margin,
        } => {
            let value = get_state_value(state, *dimension_index)?;
            for approved in approved_values {
                if (value - approved).abs() <= *tolerance {
                    return Ok(1.0);
                }
            }
            Ok(*fallback_margin)
        }

        MarginRule::WeightedSum { weights, offset } => {
            let mut sum = *offset;
            for (dim_index, weight) in weights {
                let value = get_state_value(state, *dim_index)?;
                sum += value * weight;
            }
            Ok(sum)
        }

        MarginRule::Constant { margin } => Ok(*margin),
    }
}

/// Safely get a state value by dimension index.
fn get_state_value(state: &[f64], index: usize) -> KernelResult<f64> {
    state
        .get(index)
        .copied()
        .ok_or_else(|| crate::KernelError::DeclarativeError(
            format!("Dimension index {} out of bounds (state has {} dimensions)", index, state.len()),
        ))
}

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

    fn budget_channel() -> DeclarativeChannel {
        DeclarativeChannel::new(ConstraintDefinition {
            id: "budget_test".into(),
            name: "Budget Test".into(),
            description: "Test budget ratio rule".into(),
            domain_tag: "test".into(),
            dimensions: vec!["spend".into()],
            rule: MarginRule::BudgetRatio {
                dimension_index: 0,
                budget: 1000.0,
            },
            thresholds: ThresholdOverrides::default(),
        })
    }

    #[test]
    fn budget_ratio_zero_usage() {
        let ch = budget_channel();
        let margin = ch.evaluate(&[0.0]).unwrap();
        assert!((margin - 1.0).abs() < f64::EPSILON);
    }

    #[test]
    fn budget_ratio_half_usage() {
        let ch = budget_channel();
        let margin = ch.evaluate(&[500.0]).unwrap();
        assert!((margin - 0.5).abs() < f64::EPSILON);
    }

    #[test]
    fn budget_ratio_full_usage() {
        let ch = budget_channel();
        let margin = ch.evaluate(&[1000.0]).unwrap();
        assert!((margin - 0.0).abs() < f64::EPSILON);
    }

    #[test]
    fn budget_ratio_over_usage_clamped() {
        let ch = budget_channel();
        let margin = ch.evaluate(&[1500.0]).unwrap();
        assert!((margin - 0.0).abs() < f64::EPSILON); // clamped to 0
    }

    #[test]
    fn range_bound_at_center() {
        let ch = DeclarativeChannel::new(ConstraintDefinition {
            id: "range_test".into(),
            name: "Range Test".into(),
            description: "Test range bound rule".into(),
            domain_tag: "test".into(),
            dimensions: vec!["temperature".into()],
            rule: MarginRule::RangeBound {
                dimension_index: 0,
                min: 0.0,
                max: 100.0,
            },
            thresholds: ThresholdOverrides::default(),
        });
        let margin = ch.evaluate(&[50.0]).unwrap();
        assert!((margin - 1.0).abs() < f64::EPSILON);
    }

    #[test]
    fn range_bound_at_edge() {
        let ch = DeclarativeChannel::new(ConstraintDefinition {
            id: "range_test".into(),
            name: "Range Test".into(),
            description: "".into(),
            domain_tag: "test".into(),
            dimensions: vec!["x".into()],
            rule: MarginRule::RangeBound {
                dimension_index: 0,
                min: 0.0,
                max: 100.0,
            },
            thresholds: ThresholdOverrides::default(),
        });
        let margin = ch.evaluate(&[0.0]).unwrap();
        assert!((margin - 0.0).abs() < f64::EPSILON);
    }

    #[test]
    fn pattern_match_approved() {
        let ch = DeclarativeChannel::new(ConstraintDefinition {
            id: "pattern_test".into(),
            name: "Pattern Test".into(),
            description: "".into(),
            domain_tag: "test".into(),
            dimensions: vec!["code".into()],
            rule: MarginRule::PatternMatch {
                dimension_index: 0,
                approved_values: vec![1.0, 2.0, 3.0],
                tolerance: 1e-10,
                fallback_margin: 0.0,
            },
            thresholds: ThresholdOverrides::default(),
        });
        assert!((ch.evaluate(&[2.0]).unwrap() - 1.0).abs() < f64::EPSILON);
        assert!((ch.evaluate(&[5.0]).unwrap() - 0.0).abs() < f64::EPSILON);
    }

    #[test]
    fn weighted_sum_basic() {
        let ch = DeclarativeChannel::new(ConstraintDefinition {
            id: "ws_test".into(),
            name: "Weighted Sum Test".into(),
            description: "".into(),
            domain_tag: "test".into(),
            dimensions: vec!["a".into(), "b".into()],
            rule: MarginRule::WeightedSum {
                weights: vec![(0, 0.3), (1, 0.7)],
                offset: 0.0,
            },
            thresholds: ThresholdOverrides::default(),
        });
        let margin = ch.evaluate(&[1.0, 1.0]).unwrap();
        assert!((margin - 1.0).abs() < f64::EPSILON);
    }

    #[test]
    fn constant_margin() {
        let ch = DeclarativeChannel::new(ConstraintDefinition {
            id: "const_test".into(),
            name: "Constant".into(),
            description: "".into(),
            domain_tag: "test".into(),
            dimensions: vec![],
            rule: MarginRule::Constant { margin: 0.42 },
            thresholds: ThresholdOverrides::default(),
        });
        let margin = ch.evaluate(&[]).unwrap();
        assert!((margin - 0.42).abs() < f64::EPSILON);
    }

    #[test]
    fn from_json_roundtrip() {
        let json = r#"{
            "id": "json_test",
            "name": "JSON Channel",
            "description": "Loaded from JSON",
            "domain_tag": "agentic",
            "dimensions": ["budget_used"],
            "rule": {
                "type": "budget_ratio",
                "dimension_index": 0,
                "budget": 500.0
            }
        }"#;

        let ch = DeclarativeChannel::from_json(json).unwrap();
        assert_eq!(ch.name(), "JSON Channel");
        assert_eq!(ch.definition().domain_tag, "agentic");

        let margin = ch.evaluate(&[250.0]).unwrap();
        assert!((margin - 0.5).abs() < f64::EPSILON);
    }

    #[test]
    fn dimension_out_of_bounds() {
        let ch = DeclarativeChannel::new(ConstraintDefinition {
            id: "oob_test".into(),
            name: "OOB".into(),
            description: "".into(),
            domain_tag: "test".into(),
            dimensions: vec!["x".into()],
            rule: MarginRule::BudgetRatio {
                dimension_index: 5,
                budget: 100.0,
            },
            thresholds: ThresholdOverrides::default(),
        });
        assert!(ch.evaluate(&[1.0]).is_err());
    }

    #[test]
    fn serialization_roundtrip() {
        let def = ConstraintDefinition {
            id: "ser_test".into(),
            name: "Serializable".into(),
            description: "Test serialization".into(),
            domain_tag: "test".into(),
            dimensions: vec!["a".into(), "b".into()],
            rule: MarginRule::WeightedSum {
                weights: vec![(0, 0.5), (1, 0.5)],
                offset: 0.0,
            },
            thresholds: ThresholdOverrides {
                safe_threshold: Some(0.7),
                caution_threshold: None,
                block_threshold: Some(0.05),
            },
        };

        let json = serde_json::to_string(&def).unwrap();
        let deserialized: ConstraintDefinition = serde_json::from_str(&json).unwrap();
        assert_eq!(deserialized.id, "ser_test");
        assert_eq!(deserialized.thresholds.safe_threshold, Some(0.7));
    }
}