claude_code_switcher/templates/
kimi.rs

1//! Kimi/Moonshot AI provider templates implementation
2//!
3//! This module provides unified support for Moonshot's various services:
4//! - K2: General-purpose conversational AI
5//! - K2 Thinking: High-speed reasoning with extended context
6//! - Kimi: Specialized coding AI
7
8use crate::{
9    settings::{ClaudeSettings, Permissions},
10    snapshots::SnapshotScope,
11    templates::Template,
12};
13use anyhow::{Result, anyhow};
14use atty;
15use inquire::Select;
16use std::collections::HashMap;
17
18/// Kimi/Moonshot service variants
19#[derive(Debug, Clone)]
20pub enum KimiVariant {
21    K2,
22    K2Thinking,
23    KimiForCoding,
24}
25
26impl KimiVariant {
27    pub fn display_name(&self) -> &'static str {
28        match self {
29            KimiVariant::K2 => "K2 (Moonshot)",
30            KimiVariant::K2Thinking => "K2 Thinking (Moonshot)",
31            KimiVariant::KimiForCoding => "Kimi For Coding",
32        }
33    }
34
35    pub fn description(&self) -> &'static str {
36        match self {
37            KimiVariant::K2 => "Moonshot K2 API - Advanced conversational AI with large context",
38            KimiVariant::K2Thinking => {
39                "Moonshot K2 Thinking API - High-speed reasoning with 256K context"
40            }
41            KimiVariant::KimiForCoding => "Kimi For Coding API - Specialized for coding tasks",
42        }
43    }
44
45    pub fn model_name(&self) -> &'static str {
46        match self {
47            KimiVariant::K2 => "kimi-k2-0905-preview",
48            KimiVariant::K2Thinking => "kimi-k2-thinking",
49            KimiVariant::KimiForCoding => "kimi-for-coding",
50        }
51    }
52
53    pub fn api_base(&self) -> &'static str {
54        match self {
55            KimiVariant::K2 => "https://api.moonshot.cn/v1",
56            KimiVariant::K2Thinking => "https://api.moonshot.cn/anthropic",
57            KimiVariant::KimiForCoding => "https://api.kimi.com/coding/",
58        }
59    }
60
61    pub fn api_key_url(&self) -> &'static str {
62        match self {
63            KimiVariant::K2 | KimiVariant::K2Thinking => {
64                "https://platform.moonshot.cn/console/api-keys"
65            }
66            KimiVariant::KimiForCoding => "https://kimi.moonshot.cn/user-center/apikeys",
67        }
68    }
69
70    pub fn env_var_name(&self) -> &'static str {
71        match self {
72            KimiVariant::K2 | KimiVariant::K2Thinking => "MOONSHOT_API_KEY",
73            KimiVariant::KimiForCoding => "KIMI_API_KEY",
74        }
75    }
76}
77
78/// Kimi/Moonshot AI provider template
79#[derive(Debug, Clone)]
80pub struct KimiTemplate {
81    variant: KimiVariant,
82}
83
84impl KimiTemplate {
85    pub fn new(variant: KimiVariant) -> Self {
86        Self { variant }
87    }
88
89    pub fn k2() -> Self {
90        Self::new(KimiVariant::K2)
91    }
92
93    pub fn k2_thinking() -> Self {
94        Self::new(KimiVariant::K2Thinking)
95    }
96
97    pub fn kimi_for_coding() -> Self {
98        Self::new(KimiVariant::KimiForCoding)
99    }
100}
101
102impl Template for KimiTemplate {
103    fn template_type(&self) -> crate::templates::TemplateType {
104        crate::templates::TemplateType::Kimi
105    }
106
107    fn env_var_name(&self) -> &'static str {
108        self.variant.env_var_name()
109    }
110
111    fn display_name(&self) -> &'static str {
112        self.variant.display_name()
113    }
114
115    fn description(&self) -> &'static str {
116        self.variant.description()
117    }
118
119    fn api_key_url(&self) -> Option<&'static str> {
120        Some(self.variant.api_key_url())
121    }
122
123    fn has_variants(&self) -> bool {
124        true
125    }
126
127    fn get_variants() -> Result<Vec<Self>>
128    where
129        Self: Sized,
130    {
131        Ok(vec![
132            Self::k2(),
133            Self::k2_thinking(),
134            Self::kimi_for_coding(),
135        ])
136    }
137
138    fn create_interactively() -> Result<Self>
139    where
140        Self: Sized,
141    {
142        if !atty::is(atty::Stream::Stdin) {
143            return Err(anyhow!(
144                "Kimi/Moonshot requires interactive mode to select service. Use 'k2', 'k2-thinking', or 'kimi' explicitly if not in interactive mode."
145            ));
146        }
147
148        let variants = [
149            ("Kimi For Coding", "Kimi - Specialized for coding tasks"),
150            ("K2", "Moonshot K2 - Advanced conversational AI"),
151            ("K2 Thinking", "Moonshot K2 Thinking - High-speed reasoning"),
152        ];
153
154        let options: Vec<String> = variants.iter().map(|(name, _)| name.to_string()).collect();
155
156        let choice = Select::new("Select Moonshot service:", options)
157            .prompt()
158            .map_err(|e| anyhow!("Failed to get service selection: {}", e))?;
159
160        let template = match choice.as_str() {
161            "K2" => Self::k2(),
162            "K2 Thinking" => Self::k2_thinking(),
163            "Kimi For Coding" => Self::kimi_for_coding(),
164            _ => unreachable!(),
165        };
166
167        Ok(template)
168    }
169
170    fn create_settings(&self, api_key: &str, scope: &SnapshotScope) -> ClaudeSettings {
171        let mut settings = ClaudeSettings::new();
172
173        if matches!(scope, SnapshotScope::Common | SnapshotScope::All) {
174            settings.model = Some(self.variant.model_name().to_string());
175
176            settings.permissions = Some(Permissions {
177                allow: Some(vec![
178                    "Bash".to_string(),
179                    "Read".to_string(),
180                    "Write".to_string(),
181                    "Edit".to_string(),
182                    "MultiEdit".to_string(),
183                    "Glob".to_string(),
184                    "Grep".to_string(),
185                    "WebFetch".to_string(),
186                ]),
187                ask: None,
188                deny: Some(vec!["WebSearch".to_string()]),
189                additional_directories: None,
190                default_mode: None,
191                disable_bypass_permissions_mode: None,
192            });
193        }
194
195        if matches!(scope, SnapshotScope::Env | SnapshotScope::All) {
196            let mut env = HashMap::new();
197
198            // Different authentication for different services
199            match self.variant {
200                KimiVariant::K2 => {
201                    env.insert("ANTHROPIC_API_KEY".to_string(), api_key.to_string());
202                    env.insert("ANTHROPIC_AUTH_TOKEN".to_string(), api_key.to_string());
203                }
204                KimiVariant::K2Thinking | KimiVariant::KimiForCoding => {
205                    env.insert("ANTHROPIC_AUTH_TOKEN".to_string(), api_key.to_string());
206                }
207            }
208
209            env.insert(
210                "ANTHROPIC_BASE_URL".to_string(),
211                self.variant.api_base().to_string(),
212            );
213            env.insert(
214                "ANTHROPIC_MODEL".to_string(),
215                self.variant.model_name().to_string(),
216            );
217            env.insert(
218                "ANTHROPIC_SMALL_FAST_MODEL".to_string(),
219                self.variant.model_name().to_string(),
220            );
221
222            // Additional models for K2
223            if matches!(self.variant, KimiVariant::K2) {
224                env.insert(
225                    "ANTHROPIC_DEFAULT_SONNET_MODEL".to_string(),
226                    self.variant.model_name().to_string(),
227                );
228            }
229
230            env.insert("API_TIMEOUT_MS".to_string(), "600000".to_string());
231            env.insert(
232                "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC".to_string(),
233                "1".to_string(),
234            );
235            settings.env = Some(env);
236        }
237
238        settings
239    }
240}
241
242/// Legacy compatibility functions
243pub fn create_k2_template(api_key: &str, scope: &SnapshotScope) -> ClaudeSettings {
244    let template = KimiTemplate::k2();
245    template.create_settings(api_key, scope)
246}
247
248pub fn create_k2_thinking_template(api_key: &str, scope: &SnapshotScope) -> ClaudeSettings {
249    let template = KimiTemplate::k2_thinking();
250    template.create_settings(api_key, scope)
251}
252
253pub fn create_kimi_template(api_key: &str, scope: &SnapshotScope) -> ClaudeSettings {
254    let template = KimiTemplate::kimi_for_coding();
255    template.create_settings(api_key, scope)
256}