Skip to main content

oxicode_ai/router/
profiles.rs

1//! Profile management for the router.
2
3#![allow(missing_docs)]
4
5use super::types::{RoutedTierConfig, RouterConfig, RouterProfile, RouterTier};
6use std::collections::HashMap;
7
8/// Indexed collection of routing profiles loaded from config.
9#[derive(Debug, Clone)]
10pub struct RouterProfiles {
11    profiles: HashMap<String, RouterProfile>,
12    default_name: String,
13}
14
15impl RouterProfiles {
16    /// Build a profile index from a [`RouterConfig`].
17    pub fn from_config(config: &RouterConfig) -> Self {
18        Self {
19            profiles: config.profiles.clone(),
20            default_name: config.default_profile.clone(),
21        }
22    }
23
24    /// Get a profile by exact name (no fallback).
25    pub fn get(&self, name: &str) -> Option<&RouterProfile> {
26        self.profiles.get(name)
27    }
28
29    /// Get a profile by name, falling back to the default profile.
30    pub fn get_with_fallback(&self, name: &str) -> Option<&RouterProfile> {
31        self.profiles
32            .get(name)
33            .or_else(|| self.profiles.get(&self.default_name))
34    }
35
36    /// Get the default profile.
37    pub fn default_profile(&self) -> Option<&RouterProfile> {
38        self.profiles.get(&self.default_name)
39    }
40
41    /// Get the [`RoutedTierConfig`] for a specific profile and tier.
42    pub fn tier_config(&self, profile_name: &str, tier: RouterTier) -> Option<&RoutedTierConfig> {
43        self.get_with_fallback(profile_name)
44            .map(|p| p.tier_config(tier))
45    }
46
47    /// List all profile names.
48    pub fn profile_names(&self) -> Vec<&str> {
49        self.profiles.keys().map(|s| s.as_str()).collect()
50    }
51}
52
53/// Parsed `"provider/model-id"` pair.
54#[derive(Debug, Clone, PartialEq, Eq)]
55pub struct ProviderModel {
56    /// Provider name (e.g. `"anthropic"`).
57    pub provider: String,
58    /// Model identifier (e.g. `"claude-sonnet-4"`).
59    pub model_id: String,
60}
61
62impl ProviderModel {
63    /// Parse a `"provider/model-id"` string.
64    pub fn parse(s: &str) -> Option<Self> {
65        let (provider, model_id) = s.split_once('/')?;
66        let provider = provider.trim().to_string();
67        let model_id = model_id.trim().to_string();
68        if provider.is_empty() || model_id.is_empty() {
69            return None;
70        }
71        Some(Self { provider, model_id })
72    }
73}
74
75/// Parse the model field from a [`RoutedTierConfig`].
76pub fn parse_tier_model(config: &RoutedTierConfig) -> Option<ProviderModel> {
77    ProviderModel::parse(&config.model)
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83    use crate::ThinkingLevel;
84
85    fn make_config() -> RouterConfig {
86        let mut config = RouterConfig::default();
87        config.profiles.insert(
88            "auto".to_string(),
89            RouterProfile {
90                high: RoutedTierConfig {
91                    model: "anthropic/claude-sonnet-4".to_string(),
92                    thinking: Some(ThinkingLevel::High),
93                    fallbacks: vec!["openai/gpt-4o".to_string()],
94                },
95                medium: RoutedTierConfig {
96                    model: "anthropic/claude-haiku-4".to_string(),
97                    thinking: None,
98                    fallbacks: vec![],
99                },
100                low: RoutedTierConfig {
101                    model: "google/gemini-2.0-flash".to_string(),
102                    thinking: None,
103                    fallbacks: vec![],
104                },
105            },
106        );
107        config
108    }
109
110    #[test]
111    fn profiles_exact_lookup() {
112        let profiles = RouterProfiles::from_config(&make_config());
113        assert!(profiles.get("auto").is_some());
114        assert!(profiles.get("nonexistent").is_none());
115    }
116
117    #[test]
118    fn profiles_fallback_lookup() {
119        let profiles = RouterProfiles::from_config(&make_config());
120        assert!(profiles.get_with_fallback("nonexistent").is_some());
121    }
122
123    #[test]
124    fn default_profile() {
125        let profiles = RouterProfiles::from_config(&make_config());
126        assert!(profiles.default_profile().is_some());
127    }
128
129    #[test]
130    fn tier_config_lookup() {
131        let profiles = RouterProfiles::from_config(&make_config());
132        let tc = profiles.tier_config("auto", RouterTier::High).unwrap();
133        assert_eq!(tc.model, "anthropic/claude-sonnet-4");
134    }
135
136    #[test]
137    fn parse_provider_model() {
138        let pm = ProviderModel::parse("anthropic/claude-sonnet-4").unwrap();
139        assert_eq!(pm.provider, "anthropic");
140        assert_eq!(pm.model_id, "claude-sonnet-4");
141    }
142
143    #[test]
144    fn parse_provider_model_no_slash() {
145        assert!(ProviderModel::parse("just-a-model").is_none());
146    }
147
148    #[test]
149    fn parse_provider_model_empty() {
150        assert!(ProviderModel::parse("/model").is_none());
151        assert!(ProviderModel::parse("provider/").is_none());
152    }
153
154    #[test]
155    fn profile_names() {
156        let profiles = RouterProfiles::from_config(&make_config());
157        let names = profiles.profile_names();
158        assert!(names.contains(&"auto"));
159    }
160}