llama-rs 0.16.1

//! Council configuration — TOML/JSON deserializable, validated.

use std::path::PathBuf;

use serde::Deserialize;

#[derive(Debug, Clone, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct CouncilConfig {
    #[serde(default = "default_min_rounds")]
    pub min_rounds: u32,

    #[serde(default = "default_max_rounds")]
    pub max_rounds: u32,

    #[serde(default = "default_threshold")]
    pub convergence_threshold: f32,

    pub embedder: EmbedderConfig,

    #[serde(rename = "agent")]
    pub agents: Vec<AgentConfig>,

    #[serde(default)]
    pub failure: FailureConfig,

    #[serde(default)]
    pub sampling: SamplingConfig,

    /// Optional shared system prompt sent to every expert and the synthesizer.
    /// When `None`, no system message is sent.
    #[serde(default)]
    pub system_prompt: Option<String>,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
pub enum EmbedderConfig {
    /// Local GGUF embedding model loaded by the orchestrator.
    LocalGguf {
        path: PathBuf,
    },
}

#[derive(Debug, Clone, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct AgentConfig {
    pub role: AgentRole,
    pub endpoint: String,
    pub model: String,
    pub timeout_ms: u64,

    /// Environment variable name to read the bearer token from
    /// (only honored for HTTP endpoints).
    #[serde(default)]
    pub api_key_env: Option<String>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AgentRole {
    Expert,
    Synthesizer,
}

#[derive(Debug, Clone, Deserialize, Default)]
#[serde(deny_unknown_fields)]
pub struct FailureConfig {
    /// Absolute count of experts required per round. When `None`,
    /// `floor(N_experts / 2) + 1` is used (simple majority).
    #[serde(default)]
    pub min_quorum: Option<u32>,
}

#[derive(Debug, Clone, Default, PartialEq, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct SamplingConfig {
    #[serde(default)]
    pub temperature: Option<f32>,
    #[serde(default)]
    pub top_p: Option<f32>,
    #[serde(default)]
    pub max_tokens: Option<u32>,
    #[serde(default)]
    pub seed: Option<u64>,
}

#[derive(Debug, thiserror::Error)]
pub enum ConfigError {
    #[error("failed to parse council config: {0}")]
    Parse(String),
    #[error("at least 2 experts are required; got {0}")]
    TooFewExperts(usize),
    #[error("exactly one synthesizer is required; got {0}")]
    SynthesizerCount(usize),
    #[error("min_rounds ({min}) must be >= 1 and <= max_rounds ({max})")]
    InvalidRounds { min: u32, max: u32 },
    #[error("convergence_threshold ({0}) must be in [0.0, 1.0]")]
    InvalidThreshold(f32),
    #[error("min_quorum ({quorum}) must be in [1, {experts}] (count of experts)")]
    InvalidQuorum { quorum: u32, experts: u32 },
    #[error("unsupported endpoint scheme in `{0}`; expected grpc://, http://, or https://")]
    UnsupportedScheme(String),
    #[error("agent timeout_ms must be > 0; got 0 for endpoint {0}")]
    ZeroTimeout(String),
}

fn default_min_rounds() -> u32 {
    2
}
fn default_max_rounds() -> u32 {
    4
}
fn default_threshold() -> f32 {
    0.92
}

impl CouncilConfig {
    /// Parse a TOML string and validate it.
    pub fn from_toml_str(s: &str) -> Result<Self, ConfigError> {
        let cfg: Self = toml::from_str(s).map_err(|e| ConfigError::Parse(e.to_string()))?;
        cfg.validate()?;
        Ok(cfg)
    }

    /// Number of expert-role agents.
    pub fn expert_count(&self) -> u32 {
        self.agents
            .iter()
            .filter(|a| a.role == AgentRole::Expert)
            .count() as u32
    }

    /// Effective `min_quorum` — either the explicit config value or the
    /// simple-majority default.
    pub fn effective_min_quorum(&self) -> u32 {
        match self.failure.min_quorum {
            Some(q) => q,
            None => self.expert_count() / 2 + 1,
        }
    }

    fn validate(&self) -> Result<(), ConfigError> {
        let expert_count = self.expert_count();
        if expert_count < 2 {
            return Err(ConfigError::TooFewExperts(expert_count as usize));
        }

        let synth_count = self
            .agents
            .iter()
            .filter(|a| a.role == AgentRole::Synthesizer)
            .count();
        if synth_count != 1 {
            return Err(ConfigError::SynthesizerCount(synth_count));
        }

        if self.min_rounds < 1 || self.max_rounds < self.min_rounds {
            return Err(ConfigError::InvalidRounds {
                min: self.min_rounds,
                max: self.max_rounds,
            });
        }

        if !(0.0..=1.0).contains(&self.convergence_threshold) {
            return Err(ConfigError::InvalidThreshold(self.convergence_threshold));
        }

        if let Some(q) = self.failure.min_quorum
            && (q < 1 || q > expert_count)
        {
            return Err(ConfigError::InvalidQuorum {
                quorum: q,
                experts: expert_count,
            });
        }

        for agent in &self.agents {
            if !is_supported_scheme(&agent.endpoint) {
                return Err(ConfigError::UnsupportedScheme(agent.endpoint.clone()));
            }
            if agent.timeout_ms == 0 {
                return Err(ConfigError::ZeroTimeout(agent.endpoint.clone()));
            }
        }

        Ok(())
    }
}

fn is_supported_scheme(endpoint: &str) -> bool {
    endpoint.starts_with("grpc://")
        || endpoint.starts_with("http://")
        || endpoint.starts_with("https://")
}

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

    fn good_config() -> &'static str {
        r#"
            min_rounds = 2
            max_rounds = 4
            convergence_threshold = 0.92

            [embedder]
            kind = "local_gguf"
            path = "/tmp/embedder.gguf"

            [[agent]]
            role = "expert"
            endpoint = "grpc://node1:50051"
            model = "qwen3-32b"
            timeout_ms = 30000

            [[agent]]
            role = "expert"
            endpoint = "http://node2:8080"
            model = "llama-3.3-70b"
            api_key_env = "OPENAI_API_KEY"
            timeout_ms = 30000

            [[agent]]
            role = "synthesizer"
            endpoint = "grpc://node3:50051"
            model = "deepseek-v3"
            timeout_ms = 60000

            [failure]
            min_quorum = 2
        "#
    }

    #[test]
    fn parses_valid_config() {
        let cfg = CouncilConfig::from_toml_str(good_config()).expect("config should parse");
        assert_eq!(cfg.min_rounds, 2);
        assert_eq!(cfg.max_rounds, 4);
        assert!((cfg.convergence_threshold - 0.92).abs() < 1e-6);
        assert_eq!(cfg.agents.len(), 3);
        assert_eq!(cfg.expert_count(), 2);
        assert_eq!(cfg.effective_min_quorum(), 2);
        assert!(matches!(cfg.embedder, EmbedderConfig::LocalGguf { .. }));
    }

    #[test]
    fn defaults_apply_when_omitted() {
        let toml = r#"
            [embedder]
            kind = "local_gguf"
            path = "/x.gguf"

            [[agent]]
            role = "expert"
            endpoint = "grpc://a:1"
            model = "m1"
            timeout_ms = 1000

            [[agent]]
            role = "expert"
            endpoint = "grpc://b:1"
            model = "m2"
            timeout_ms = 1000

            [[agent]]
            role = "expert"
            endpoint = "grpc://c:1"
            model = "m3"
            timeout_ms = 1000

            [[agent]]
            role = "synthesizer"
            endpoint = "grpc://s:1"
            model = "ms"
            timeout_ms = 1000
        "#;
        let cfg = CouncilConfig::from_toml_str(toml).expect("defaults config should parse");
        assert_eq!(cfg.min_rounds, 2, "min_rounds default");
        assert_eq!(cfg.max_rounds, 4, "max_rounds default");
        assert_eq!(cfg.convergence_threshold, 0.92);
        // 3 experts → floor(3/2)+1 = 2
        assert_eq!(cfg.effective_min_quorum(), 2);
        assert!(cfg.system_prompt.is_none());
    }

    #[test]
    fn rejects_too_few_experts() {
        let toml = r#"
            [embedder]
            kind = "local_gguf"
            path = "/x.gguf"

            [[agent]]
            role = "expert"
            endpoint = "grpc://a:1"
            model = "m1"
            timeout_ms = 1000

            [[agent]]
            role = "synthesizer"
            endpoint = "grpc://s:1"
            model = "ms"
            timeout_ms = 1000
        "#;
        let err = CouncilConfig::from_toml_str(toml).unwrap_err();
        assert!(matches!(err, ConfigError::TooFewExperts(1)), "got {err:?}");
    }

    #[test]
    fn rejects_no_synthesizer() {
        let toml = r#"
            [embedder]
            kind = "local_gguf"
            path = "/x.gguf"

            [[agent]]
            role = "expert"
            endpoint = "grpc://a:1"
            model = "m1"
            timeout_ms = 1000

            [[agent]]
            role = "expert"
            endpoint = "grpc://b:1"
            model = "m2"
            timeout_ms = 1000
        "#;
        let err = CouncilConfig::from_toml_str(toml).unwrap_err();
        assert!(matches!(err, ConfigError::SynthesizerCount(0)), "got {err:?}");
    }

    #[test]
    fn rejects_two_synthesizers() {
        let toml = r#"
            [embedder]
            kind = "local_gguf"
            path = "/x.gguf"

            [[agent]]
            role = "expert"
            endpoint = "grpc://a:1"
            model = "m1"
            timeout_ms = 1000

            [[agent]]
            role = "expert"
            endpoint = "grpc://b:1"
            model = "m2"
            timeout_ms = 1000

            [[agent]]
            role = "synthesizer"
            endpoint = "grpc://s1:1"
            model = "ms"
            timeout_ms = 1000

            [[agent]]
            role = "synthesizer"
            endpoint = "grpc://s2:1"
            model = "ms2"
            timeout_ms = 1000
        "#;
        let err = CouncilConfig::from_toml_str(toml).unwrap_err();
        assert!(matches!(err, ConfigError::SynthesizerCount(2)), "got {err:?}");
    }

    #[test]
    fn rejects_invalid_threshold() {
        let bad = good_config().replace("convergence_threshold = 0.92", "convergence_threshold = 1.5");
        let err = CouncilConfig::from_toml_str(&bad).unwrap_err();
        assert!(matches!(err, ConfigError::InvalidThreshold(_)));
    }

    #[test]
    fn rejects_invalid_rounds() {
        let bad = good_config()
            .replace("min_rounds = 2", "min_rounds = 5")
            .replace("max_rounds = 4", "max_rounds = 4");
        let err = CouncilConfig::from_toml_str(&bad).unwrap_err();
        assert!(
            matches!(err, ConfigError::InvalidRounds { min: 5, max: 4 }),
            "got {err:?}"
        );
    }

    #[test]
    fn rejects_zero_timeout() {
        let bad = good_config().replacen("timeout_ms = 30000", "timeout_ms = 0", 1);
        let err = CouncilConfig::from_toml_str(&bad).unwrap_err();
        assert!(matches!(err, ConfigError::ZeroTimeout(_)));
    }

    #[test]
    fn rejects_unsupported_scheme() {
        let bad = good_config().replace("grpc://node1:50051", "ftp://node1:50051");
        let err = CouncilConfig::from_toml_str(&bad).unwrap_err();
        assert!(matches!(err, ConfigError::UnsupportedScheme(_)));
    }

    #[test]
    fn rejects_quorum_above_expert_count() {
        let bad = good_config().replace("min_quorum = 2", "min_quorum = 5");
        let err = CouncilConfig::from_toml_str(&bad).unwrap_err();
        assert!(
            matches!(err, ConfigError::InvalidQuorum { quorum: 5, experts: 2 }),
            "got {err:?}"
        );
    }

    #[test]
    fn min_quorum_default_is_simple_majority() {
        // 4 experts should default to floor(4/2)+1 = 3
        let toml = r#"
            [embedder]
            kind = "local_gguf"
            path = "/x.gguf"

            [[agent]]
            role = "expert"
            endpoint = "grpc://a:1"
            model = "m1"
            timeout_ms = 1000

            [[agent]]
            role = "expert"
            endpoint = "grpc://b:1"
            model = "m2"
            timeout_ms = 1000

            [[agent]]
            role = "expert"
            endpoint = "grpc://c:1"
            model = "m3"
            timeout_ms = 1000

            [[agent]]
            role = "expert"
            endpoint = "grpc://d:1"
            model = "m4"
            timeout_ms = 1000

            [[agent]]
            role = "synthesizer"
            endpoint = "grpc://s:1"
            model = "ms"
            timeout_ms = 1000
        "#;
        let cfg = CouncilConfig::from_toml_str(toml).expect("valid");
        assert_eq!(cfg.expert_count(), 4);
        assert_eq!(cfg.effective_min_quorum(), 3);
    }

    #[test]
    fn rejects_unknown_top_level_field() {
        let bad = format!("{}\nrandom_field = 42\n", good_config());
        let err = CouncilConfig::from_toml_str(&bad).unwrap_err();
        assert!(matches!(err, ConfigError::Parse(_)));
    }
}