claude_code_switcher/templates/
zai.rs

1//! ZAI (GLM/Zhipu) AI provider template implementation
2
3use crate::{
4    settings::{ClaudeSettings, Permissions},
5    snapshots::SnapshotScope,
6    templates::Template,
7};
8use anyhow::{Result, anyhow};
9use atty;
10use inquire::Select;
11use std::collections::HashMap;
12
13/// ZAI (GLM/Zhipu) AI provider regions
14#[derive(Debug, Clone)]
15pub enum ZaiRegion {
16    China,
17    International,
18}
19
20impl ZaiRegion {
21    pub fn display_name(&self) -> &'static str {
22        match self {
23            ZaiRegion::China => "ZAI China (智谱AI)",
24            ZaiRegion::International => "ZAI International",
25        }
26    }
27
28    pub fn description(&self) -> &'static str {
29        match self {
30            ZaiRegion::China => {
31                "Zhipu AI GLM-4.6 in China - Fast response with thinking capabilities"
32            }
33            ZaiRegion::International => {
34                "Zhipu AI GLM-4.6 International - Global access with optimized routing"
35            }
36        }
37    }
38
39    pub fn base_url(&self) -> &'static str {
40        match self {
41            ZaiRegion::China => "https://open.bigmodel.cn/api/anthropic",
42            ZaiRegion::International => "https://api.z.ai/api/anthropic",
43        }
44    }
45
46    pub fn model_name(&self) -> &'static str {
47        match self {
48            ZaiRegion::China => "glm-4.6",
49            ZaiRegion::International => "glm-4.6",
50        }
51    }
52
53    pub fn small_fast_model(&self) -> &'static str {
54        match self {
55            ZaiRegion::China => "glm-4.6",
56            ZaiRegion::International => "glm-4.6-air",
57        }
58    }
59
60    pub fn api_key_url(&self) -> &'static str {
61        match self {
62            ZaiRegion::China => "https://open.bigmodel.cn/usercenter/apikeys",
63            ZaiRegion::International => "https://console.z.ai/apikeys",
64        }
65    }
66}
67
68/// ZAI (GLM/Zhipu) AI provider template
69#[derive(Debug, Clone)]
70pub struct ZaiTemplate {
71    region: ZaiRegion,
72}
73
74impl ZaiTemplate {
75    pub fn new(region: ZaiRegion) -> Self {
76        Self { region }
77    }
78
79    pub fn china() -> Self {
80        Self::new(ZaiRegion::China)
81    }
82
83    pub fn international() -> Self {
84        Self::new(ZaiRegion::International)
85    }
86}
87
88impl Template for ZaiTemplate {
89    fn template_type(&self) -> crate::templates::TemplateType {
90        crate::templates::TemplateType::Zai
91    }
92
93    fn env_var_name(&self) -> &'static str {
94        "Z_AI_API_KEY"
95    }
96
97    fn display_name(&self) -> &'static str {
98        self.region.display_name()
99    }
100
101    fn description(&self) -> &'static str {
102        self.region.description()
103    }
104
105    fn api_key_url(&self) -> Option<&'static str> {
106        Some(self.region.api_key_url())
107    }
108
109    fn has_variants(&self) -> bool {
110        true
111    }
112
113    fn get_variants() -> Result<Vec<Self>>
114    where
115        Self: Sized,
116    {
117        Ok(vec![Self::china(), Self::international()])
118    }
119
120    fn create_interactively() -> Result<Self>
121    where
122        Self: Sized,
123    {
124        if !atty::is(atty::Stream::Stdin) {
125            return Err(anyhow!(
126                "ZAI requires interactive mode to select region. Use 'zai-china' or 'zai-international' explicitly if not in interactive mode."
127            ));
128        }
129
130        let regions = [
131            (
132                "ZAI China (智谱AI)",
133                "Fast response with thinking capabilities, optimized for China users",
134            ),
135            (
136                "ZAI International",
137                "Global access with optimized routing for international users",
138            ),
139        ];
140
141        let options: Vec<String> = regions.iter().map(|(name, _)| name.to_string()).collect();
142
143        let choice = Select::new("Select ZAI region:", options)
144            .prompt()
145            .map_err(|e| anyhow!("Failed to get region selection: {}", e))?;
146
147        let template = match choice.as_str() {
148            "ZAI China (智谱AI)" => Self::china(),
149            "ZAI International" => Self::international(),
150            _ => unreachable!(),
151        };
152
153        Ok(template)
154    }
155
156    fn create_settings(&self, api_key: &str, scope: &SnapshotScope) -> ClaudeSettings {
157        let mut settings = ClaudeSettings::new();
158
159        if matches!(scope, SnapshotScope::Common | SnapshotScope::All) {
160            settings.model = Some(self.region.model_name().to_string());
161
162            // Use the new permissions format from the provided version
163            settings.permissions = Some(Permissions {
164                allow: Some(vec![
165                    "Bash".to_string(),
166                    "Read".to_string(),
167                    "Write".to_string(),
168                    "Edit".to_string(),
169                    "MultiEdit".to_string(),
170                    "Glob".to_string(),
171                    "Grep".to_string(),
172                    "WebFetch".to_string(),
173                ]),
174                ask: None,
175                deny: Some(vec!["WebSearch".to_string()]),
176                additional_directories: None,
177                default_mode: None,
178                disable_bypass_permissions_mode: None,
179            });
180        }
181
182        if matches!(scope, SnapshotScope::Env | SnapshotScope::All) {
183            let mut env = HashMap::new();
184            env.insert("ANTHROPIC_AUTH_TOKEN".to_string(), api_key.to_string());
185            env.insert(
186                "ANTHROPIC_BASE_URL".to_string(),
187                self.region.base_url().to_string(),
188            );
189            env.insert("API_TIMEOUT_MS".to_string(), "3000000".to_string());
190            env.insert(
191                "ANTHROPIC_MODEL".to_string(),
192                self.region.model_name().to_string(),
193            );
194            env.insert(
195                "ANTHROPIC_SMALL_FAST_MODEL".to_string(),
196                self.region.small_fast_model().to_string(),
197            );
198            env.insert("ENABLE_THINKING".to_string(), "true".to_string());
199            env.insert("REASONING_EFFORT".to_string(), "ultrathink".to_string());
200            env.insert("MAX_THINKING_TOKENS".to_string(), "32000".to_string());
201            env.insert("ENABLE_STREAMING".to_string(), "true".to_string());
202            env.insert("MAX_OUTPUT_TOKENS".to_string(), "96000".to_string());
203            env.insert("MAX_MCP_OUTPUT_TOKENS".to_string(), "64000".to_string());
204            env.insert("AUTH_HEADER_MODE".to_string(), "x-api-key".to_string());
205            settings.env = Some(env);
206        }
207
208        settings
209    }
210}
211
212/// Create ZAI template settings (legacy compatibility function)
213pub fn create_zai_template(api_key: &str, scope: &SnapshotScope) -> ClaudeSettings {
214    let template = ZaiTemplate::china(); // Default to China for backward compatibility
215    template.create_settings(api_key, scope)
216}