converge_core/

model_selection.rs

// Copyright 2024-2025 Aprio One AB, Sweden
// Author: Kenneth Pernyer, kenneth@aprio.one
// SPDX-License-Identifier: LicenseRef-Proprietary
// All rights reserved. This source code is proprietary and confidential.
// Unauthorized copying, modification, or distribution is strictly prohibited.

//! Model selection based on agent requirements.
//!
//! This module provides the **abstract interface** for model selection.
//! Concrete implementations (with provider-specific metadata) are in `converge-provider`.
//!
//! # Architecture
//!
//! - **Core (this module)**: Abstract requirements and selection trait
//! - **Provider crate**: Concrete selector with all provider metadata
//!
//! This separation ensures core remains provider-agnostic while allowing
//! injection of provider-specific selection logic.

use crate::llm::LlmError;

/// Cost classification for model selection.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum CostClass {
    /// Very low cost (e.g., Haiku, GPT-3.5 Turbo, Gemini Flash)
    VeryLow,
    /// Low cost (e.g., Sonnet, GPT-4 Turbo)
    Low,
    /// Medium cost (e.g., Opus, GPT-4)
    Medium,
    /// High cost (e.g., Opus-4, GPT-4o)
    High,
    /// Very high cost (e.g., specialized models)
    VeryHigh,
}

impl CostClass {
    /// Returns the maximum cost class that satisfies this requirement.
    ///
    /// For example, `CostClass::Low` allows VeryLow and Low.
    #[must_use]
    pub fn allowed_classes(self) -> Vec<CostClass> {
        match self {
            Self::VeryLow => vec![Self::VeryLow],
            Self::Low => vec![Self::VeryLow, Self::Low],
            Self::Medium => vec![Self::VeryLow, Self::Low, Self::Medium],
            Self::High => vec![Self::VeryLow, Self::Low, Self::Medium, Self::High],
            Self::VeryHigh => vec![
                Self::VeryLow,
                Self::Low,
                Self::Medium,
                Self::High,
                Self::VeryHigh,
            ],
        }
    }
}

/// Data sovereignty requirements.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DataSovereignty {
    /// No specific requirements (default).
    Any,
    /// Data must remain in EU/EEA.
    EU,
    /// Data must remain in Switzerland.
    Switzerland,
    /// Data must remain in China.
    China,
    /// Data must remain in US.
    US,
    /// Self-hosted or on-premises.
    OnPremises,
}

/// Compliance and explainability requirements.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ComplianceLevel {
    /// No specific compliance requirements (default).
    None,
    /// GDPR compliance required.
    GDPR,
    /// SOC 2 compliance required.
    SOC2,
    /// HIPAA compliance required.
    HIPAA,
    /// High explainability (audit trails, provenance).
    HighExplainability,
}

/// Requirements for an agent's LLM usage.
///
/// Agents specify their requirements, and the model selector
/// finds the best matching model.
#[derive(Debug, Clone, PartialEq)]
pub struct AgentRequirements {
    /// Maximum acceptable cost class.
    pub max_cost_class: CostClass,
    /// Maximum acceptable latency in milliseconds.
    pub max_latency_ms: u32,
    /// Whether the agent requires advanced reasoning capabilities.
    pub requires_reasoning: bool,
    /// Whether the agent requires web search capabilities.
    pub requires_web_search: bool,
    /// Minimum quality threshold (0.0-1.0).
    /// Higher values prefer more capable models.
    pub min_quality: f64,
    /// Data sovereignty requirements.
    pub data_sovereignty: DataSovereignty,
    /// Compliance and explainability requirements.
    pub compliance: ComplianceLevel,
    /// Whether the agent requires multi-language support.
    pub requires_multilingual: bool,
}

impl AgentRequirements {
    /// Creates requirements for a fast, cheap agent (many instances).
    ///
    /// Use case: High-volume agents that need quick, cost-effective responses.
    #[must_use]
    pub fn fast_cheap() -> Self {
        Self {
            max_cost_class: CostClass::VeryLow,
            max_latency_ms: 2000,
            requires_reasoning: false,
            requires_web_search: false,
            min_quality: 0.6,
            data_sovereignty: DataSovereignty::Any,
            compliance: ComplianceLevel::None,
            requires_multilingual: false,
        }
    }

    /// Creates requirements for a deep research agent.
    ///
    /// Use case: Agents that need thorough analysis and reasoning.
    #[must_use]
    pub fn deep_research() -> Self {
        Self {
            max_cost_class: CostClass::High,
            max_latency_ms: 30000, // 30 seconds
            requires_reasoning: true,
            requires_web_search: true,
            min_quality: 0.9,
            data_sovereignty: DataSovereignty::Any,
            compliance: ComplianceLevel::None,
            requires_multilingual: false,
        }
    }

    /// Creates requirements for a balanced agent.
    ///
    /// Use case: General-purpose agents with moderate requirements.
    #[must_use]
    pub fn balanced() -> Self {
        Self {
            max_cost_class: CostClass::Medium,
            max_latency_ms: 5000,
            requires_reasoning: false,
            requires_web_search: false,
            min_quality: 0.7,
            data_sovereignty: DataSovereignty::Any,
            compliance: ComplianceLevel::None,
            requires_multilingual: false,
        }
    }

    /// Creates custom requirements.
    #[must_use]
    pub fn new(
        max_cost_class: CostClass,
        max_latency_ms: u32,
        requires_reasoning: bool,
    ) -> Self {
        Self {
            max_cost_class,
            max_latency_ms,
            requires_reasoning,
            requires_web_search: false,
            min_quality: 0.7,
            data_sovereignty: DataSovereignty::Any,
            compliance: ComplianceLevel::None,
            requires_multilingual: false,
        }
    }

    /// Sets web search requirement.
    #[must_use]
    pub fn with_web_search(mut self, requires: bool) -> Self {
        self.requires_web_search = requires;
        self
    }

    /// Sets minimum quality threshold.
    #[must_use]
    pub fn with_min_quality(mut self, quality: f64) -> Self {
        self.min_quality = quality.clamp(0.0, 1.0);
        self
    }

    /// Sets data sovereignty requirement.
    #[must_use]
    pub fn with_data_sovereignty(mut self, sovereignty: DataSovereignty) -> Self {
        self.data_sovereignty = sovereignty;
        self
    }

    /// Sets compliance requirement.
    #[must_use]
    pub fn with_compliance(mut self, compliance: ComplianceLevel) -> Self {
        self.compliance = compliance;
        self
    }

    /// Sets multilingual requirement.
    #[must_use]
    pub fn with_multilingual(mut self, requires: bool) -> Self {
        self.requires_multilingual = requires;
        self
    }
}

// ModelMetadata and ModelSelector implementations are in converge-provider.
// Core only provides the abstract interface (ModelSelectorTrait) and requirements.

/// Trait for model selection based on agent requirements.
///
/// This trait allows injecting provider-specific model selection logic
/// without coupling core to concrete providers.
///
/// Concrete implementations (with provider metadata) are in `converge-provider`.
pub trait ModelSelectorTrait: Send + Sync {
    /// Selects the best model for the given requirements.
    ///
    /// Returns `(provider, model)` if a suitable model is found.
    ///
    /// # Errors
    ///
    /// Returns error if no model satisfies the requirements.
    fn select(
        &self,
        requirements: &AgentRequirements,
    ) -> Result<(String, String), LlmError>;
}

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

    #[test]
    fn test_fast_cheap_requirements() {
        let reqs = AgentRequirements::fast_cheap();
        assert_eq!(reqs.max_cost_class, CostClass::VeryLow);
        assert_eq!(reqs.max_latency_ms, 2000);
        assert!(!reqs.requires_reasoning);
    }

    #[test]
    fn test_deep_research_requirements() {
        let reqs = AgentRequirements::deep_research();
        assert!(reqs.max_cost_class >= CostClass::High);
        assert!(reqs.requires_reasoning);
        assert!(reqs.requires_web_search);
    }

    #[test]
    fn test_cost_class_allowed() {
        assert_eq!(CostClass::VeryLow.allowed_classes().len(), 1);
        assert_eq!(CostClass::Low.allowed_classes().len(), 2);
        assert_eq!(CostClass::VeryHigh.allowed_classes().len(), 5);
    }
}