1use std::collections::HashMap;
2
3use deepseek_config::ProviderKind;
4use serde::{Deserialize, Serialize};
5
6#[derive(Debug, Clone, Serialize, Deserialize)]
7pub struct ModelInfo {
8 pub id: String,
9 pub provider: ProviderKind,
10 pub aliases: Vec<String>,
11 pub supports_tools: bool,
12 pub supports_reasoning: bool,
13}
14
15#[derive(Debug, Clone, Serialize, Deserialize)]
16pub struct ModelResolution {
17 pub requested: Option<String>,
18 pub resolved: ModelInfo,
19 pub used_fallback: bool,
20 pub fallback_chain: Vec<String>,
21}
22
23#[derive(Debug, Clone)]
24pub struct ModelRegistry {
25 models: Vec<ModelInfo>,
26 alias_map: HashMap<String, usize>,
27}
28
29impl Default for ModelRegistry {
30 fn default() -> Self {
31 let models = vec![
32 ModelInfo {
33 id: "deepseek-reasoner".to_string(),
34 provider: ProviderKind::Deepseek,
35 aliases: vec!["deepseek-r1".to_string()],
36 supports_tools: true,
37 supports_reasoning: true,
38 },
39 ModelInfo {
40 id: "deepseek-chat".to_string(),
41 provider: ProviderKind::Deepseek,
42 aliases: vec!["deepseek-v3".to_string(), "deepseek-v3.2".to_string()],
43 supports_tools: true,
44 supports_reasoning: false,
45 },
46 ModelInfo {
47 id: "gpt-4.1".to_string(),
48 provider: ProviderKind::Openai,
49 aliases: vec!["gpt4.1".to_string(), "gpt-4o".to_string()],
50 supports_tools: true,
51 supports_reasoning: true,
52 },
53 ModelInfo {
54 id: "gpt-4.1-mini".to_string(),
55 provider: ProviderKind::Openai,
56 aliases: vec!["gpt-4o-mini".to_string()],
57 supports_tools: true,
58 supports_reasoning: false,
59 },
60 ];
61 Self::new(models)
62 }
63}
64
65impl ModelRegistry {
66 #[must_use]
67 pub fn new(models: Vec<ModelInfo>) -> Self {
68 let mut alias_map = HashMap::new();
69 for (idx, model) in models.iter().enumerate() {
70 alias_map.insert(normalize(&model.id), idx);
71 for alias in &model.aliases {
72 alias_map.insert(normalize(alias), idx);
73 }
74 }
75 Self { models, alias_map }
76 }
77
78 #[must_use]
79 pub fn list(&self) -> Vec<ModelInfo> {
80 self.models.clone()
81 }
82
83 #[must_use]
84 pub fn resolve(
85 &self,
86 requested: Option<&str>,
87 provider_hint: Option<ProviderKind>,
88 ) -> ModelResolution {
89 let mut fallback_chain = Vec::new();
90
91 if let Some(name) = requested {
92 fallback_chain.push(format!("requested:{name}"));
93 if let Some(idx) = self.alias_map.get(&normalize(name)) {
94 return ModelResolution {
95 requested: Some(name.to_string()),
96 resolved: self.models[*idx].clone(),
97 used_fallback: false,
98 fallback_chain,
99 };
100 }
101 }
102
103 let provider = provider_hint.unwrap_or(ProviderKind::Deepseek);
104 fallback_chain.push(format!("provider_default:{}", provider.as_str()));
105 if let Some(model) = self.models.iter().find(|m| m.provider == provider).cloned() {
106 return ModelResolution {
107 requested: requested.map(ToOwned::to_owned),
108 resolved: model,
109 used_fallback: true,
110 fallback_chain,
111 };
112 }
113
114 let final_fallback = self.models.first().cloned().unwrap_or(ModelInfo {
115 id: "deepseek-reasoner".to_string(),
116 provider: ProviderKind::Deepseek,
117 aliases: Vec::new(),
118 supports_tools: true,
119 supports_reasoning: true,
120 });
121 fallback_chain.push("global_default:deepseek-reasoner".to_string());
122 ModelResolution {
123 requested: requested.map(ToOwned::to_owned),
124 resolved: final_fallback,
125 used_fallback: true,
126 fallback_chain,
127 }
128 }
129}
130
131fn normalize(value: &str) -> String {
132 value.trim().to_ascii_lowercase()
133}