Skip to main content

deepseek_agent/
lib.rs

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}