Skip to main content

vtcode_llm/providers/
zai.rs

1use reqwest::RequestBuilder;
2use serde_json::{Map, Value};
3use vtcode_config::constants::{env_vars, models, urls};
4
5use super::openai_compat::{OpenAiCompatCore, OpenAiCompatSpec, SystemPromptPlacement, impl_openai_compat_provider};
6use crate::provider::{LLMError, LLMRequest};
7
8pub struct ZaiSpec;
9
10fn resolve_zai_base_url(base_url: Option<String>) -> String {
11    // Semantic variant, deliberately not layered on `override_base_url`:
12    // z.ai has a legacy env layer with lower precedence than the current
13    // env var, and the shared helper's env position would let the legacy
14    // value override an explicit configuration (regression-locked by
15    // `base_url_resolution_chains_explicit_then_legacy_env_then_default`).
16    if let Some(url) = base_url {
17        let trimmed = url.trim();
18        if !trimmed.is_empty() {
19            return trimmed.to_string();
20        }
21    }
22
23    if let Ok(value) = std::env::var(env_vars::ZAI_BASE_URL) {
24        let trimmed = value.trim();
25        if !trimmed.is_empty() {
26            return trimmed.to_string();
27        }
28    }
29
30    if let Ok(legacy) = std::env::var(env_vars::Z_AI_BASE_URL) {
31        let trimmed = legacy.trim();
32        if !trimmed.is_empty() {
33            return trimmed.to_string();
34        }
35    }
36
37    urls::ZAI_API_BASE.to_string()
38}
39
40impl OpenAiCompatSpec for ZaiSpec {
41    const NAME: &'static str = "Z.AI";
42    const KEY: &'static str = "zai";
43    const API_KEY_ENV: &'static str = "ZAI_API_KEY";
44    const DEFAULT_MODEL: &'static str = models::zai::DEFAULT_MODEL;
45    const DEFAULT_BASE_URL: &'static str = urls::ZAI_API_BASE;
46    const BASE_URL_ENV: Option<&'static str> = Some(env_vars::ZAI_BASE_URL);
47    const LISTED_MODELS: &'static [&'static str] = models::zai::SUPPORTED_MODELS;
48    const VALIDATION_ALLOWLIST: Option<&'static [&'static str]> = Some(models::zai::SUPPORTED_MODELS);
49
50    const SYSTEM_PROMPT: SystemPromptPlacement = SystemPromptPlacement::Omitted;
51    const SUPPRESS_SAMPLING_WHEN_REASONING: bool = false;
52    const DELTA_ORDER: super::shared::OpenAiDeltaOrder = super::shared::OpenAiDeltaOrder::ContentFirst;
53
54    fn resolve_base_url(_api_key: &str, base_url: Option<String>) -> String {
55        resolve_zai_base_url(base_url)
56    }
57
58    fn insert_tool_choice(_core: &OpenAiCompatCore<Self>, request: &LLMRequest, payload: &mut Map<String, Value>) {
59        if let Some(choice) = &request.tool_choice {
60            // Z.AI only supports "auto"; any other requested mode is coerced.
61            let tool_choice_value = match choice {
62                crate::provider::ToolChoice::Auto => choice.to_provider_format(Self::KEY),
63                _ => Value::String("auto".to_string()),
64            };
65            payload.insert("tool_choice".to_string(), tool_choice_value);
66        } else if request.tools.as_ref().is_some_and(|tools| !tools.is_empty()) {
67            payload.insert("tool_choice".to_string(), Value::String("auto".to_string()));
68        }
69    }
70
71    fn insert_reasoning(
72        core: &OpenAiCompatCore<Self>,
73        request: &LLMRequest,
74        payload: &mut Map<String, Value>,
75    ) -> Result<(), LLMError> {
76        let has_preserved_reasoning = request.messages.iter().any(|message| {
77            message.role == crate::provider::MessageRole::Assistant
78                && message.reasoning.as_ref().is_some_and(|reasoning| !reasoning.is_empty())
79        });
80
81        if let Some(effort) = request.reasoning_effort {
82            if effort == vtcode_config::types::ReasoningEffortLevel::None {
83                payload.insert("thinking".to_owned(), serde_json::json!({"type": "disabled"}));
84                return Ok(());
85            }
86
87            use crate::rig_adapter::RigProviderCapabilities;
88            use vtcode_config::models::Provider;
89            let supported = crate::provider::catalog_or_explicit_reasoning_efforts(
90                Self::KEY,
91                &request.model,
92                core.model_behavior
93                    .as_ref()
94                    .and_then(|behavior| behavior.model_supports_reasoning_effort)
95                    .unwrap_or(false),
96            );
97            let reasoning_params = RigProviderCapabilities::new(Provider::ZAI, &request.model)
98                .reasoning_parameters_for_supported_efforts(effort, supported)
99                ?
100                .ok_or_else(|| LLMError::InvalidRequest {
101                    message: format!(
102                        "Reasoning effort `{effort}` is unsupported for Z.AI model `{}`; choose one of low, high, or max",
103                        request.model
104                    ),
105                    metadata: None,
106                })?;
107            if let Some(params_obj) = reasoning_params.as_object() {
108                for (k, v) in params_obj {
109                    payload.insert(k.clone(), v.clone());
110                }
111            }
112        }
113
114        if has_preserved_reasoning {
115            if let Some(thinking) = payload.get_mut("thinking").and_then(Value::as_object_mut) {
116                thinking.insert("clear_thinking".to_owned(), Value::Bool(false));
117            } else {
118                payload.insert(
119                    "thinking".to_owned(),
120                    serde_json::json!({
121                        "type": "enabled",
122                        "clear_thinking": false
123                    }),
124                );
125            }
126        }
127
128        Ok(())
129    }
130
131    fn finish_payload(
132        _core: &OpenAiCompatCore<Self>,
133        request: &LLMRequest,
134        payload: &mut Map<String, Value>,
135    ) -> Result<(), LLMError> {
136        if let Some(do_sample) = request.do_sample {
137            payload.insert("do_sample".to_owned(), Value::Bool(do_sample));
138        }
139
140        if request.stream && request.tools.as_ref().is_some_and(|tools| !tools.is_empty()) {
141            payload.insert("tool_stream".to_string(), Value::Bool(true));
142        }
143
144        if request.output_format.is_some() {
145            payload.insert("response_format".to_owned(), serde_json::json!({ "type": "json_object" }));
146        }
147
148        Ok(())
149    }
150
151    fn apply_auth(core: &OpenAiCompatCore<Self>, builder: RequestBuilder) -> RequestBuilder {
152        builder.bearer_auth(&core.api_key).header("Accept-Language", "en-US,en")
153    }
154}
155
156impl_openai_compat_provider!(ZAIProvider, ZaiSpec, {
157    fn supports_reasoning(&self, model: &str) -> bool {
158        // Codex-inspired robustness: Setting model_supports_reasoning to false
159        // does NOT disable it for known reasoning models.
160        model.contains("glm")
161            || self
162                .core
163                .model_behavior
164                .as_ref()
165                .and_then(|b| b.model_supports_reasoning)
166                .unwrap_or(false)
167    }
168
169    fn supports_reasoning_effort(&self, model: &str) -> bool {
170        // Same robustness logic for reasoning effort
171        model.contains("glm")
172            || self
173                .core
174                .model_behavior
175                .as_ref()
176                .and_then(|b| b.model_supports_reasoning_effort)
177                .unwrap_or(false)
178    }
179});
180
181#[cfg(test)]
182mod tests {
183    use super::{ZAIProvider, resolve_zai_base_url};
184    use crate::provider::{LLMRequest, Message, ToolChoice, ToolDefinition};
185    use std::sync::Arc;
186    use vtcode_config::constants::models;
187    use vtcode_config::types::ReasoningEffortLevel;
188
189    /// Explicit config wins over both env layers; blank explicit values keep
190    /// falling to the legacy env var, and the default is the last resort.
191    /// Asymmetric on every layer boundary rather than one happy path.
192    #[test]
193    #[serial_test::serial]
194    fn base_url_resolution_chains_explicit_then_legacy_env_then_default() {
195        use vtcode_config::constants::env_vars;
196        let guard = vtcode_commons::env_lock::lock();
197        let prior_current = std::env::var_os(env_vars::ZAI_BASE_URL);
198        let prior_legacy = std::env::var_os(env_vars::Z_AI_BASE_URL);
199        guard.remove_var(env_vars::ZAI_BASE_URL);
200        guard.remove_var(env_vars::Z_AI_BASE_URL);
201
202        guard.set_var(env_vars::Z_AI_BASE_URL, "https://legacy.example/v4");
203        assert_eq!(resolve_zai_base_url(None), "https://legacy.example/v4");
204        assert_eq!(resolve_zai_base_url(Some("   ".to_string())), "https://legacy.example/v4");
205        assert_eq!(
206            resolve_zai_base_url(Some("  https://explicit.example/v4  ".to_string())),
207            "https://explicit.example/v4",
208            "explicit override wins over the legacy env layer"
209        );
210
211        guard.remove_var(env_vars::Z_AI_BASE_URL);
212        assert_eq!(
213            resolve_zai_base_url(None),
214            vtcode_config::constants::urls::ZAI_API_BASE,
215            "default applies only when both env layers are absent"
216        );
217
218        guard.restore_var(env_vars::ZAI_BASE_URL, prior_current);
219        guard.restore_var(env_vars::Z_AI_BASE_URL, prior_legacy);
220    }
221
222    #[test]
223    fn payload_includes_top_p() {
224        let provider = ZAIProvider::new("test-key".to_string());
225        let request = LLMRequest {
226            model: models::zai::GLM_5_3.to_string(),
227            messages: vec![Message::user("hello".to_string())].into(),
228            top_p: Some(0.95),
229            ..Default::default()
230        };
231
232        let payload = provider.core.convert_request(&request).expect("payload should be valid");
233        let top_p = payload.get("top_p").and_then(|v| v.as_f64()).expect("top_p should be present");
234        assert!((top_p - 0.95).abs() < 1e-6);
235    }
236
237    #[test]
238    fn payload_enables_tool_stream_when_streaming_with_tools() {
239        let provider = ZAIProvider::new("test-key".to_string());
240        let request = LLMRequest {
241            model: models::zai::GLM_5_3.to_string(),
242            messages: vec![Message::user("hello".to_string())].into(),
243            stream: true,
244            tools: Some(Arc::new(vec![ToolDefinition::function(
245                "get_weather".to_string(),
246                "Get weather".to_string(),
247                serde_json::json!({
248                    "type": "object",
249                    "properties": {
250                        "location": {"type": "string"}
251                    },
252                    "required": ["location"]
253                }),
254            )])),
255            ..Default::default()
256        };
257
258        let payload = provider.core.convert_request(&request).expect("payload should be valid");
259        assert_eq!(payload.get("stream").and_then(|v| v.as_bool()), Some(true));
260        assert_eq!(payload.get("tool_stream").and_then(|v| v.as_bool()), Some(true));
261    }
262
263    #[test]
264    fn payload_streaming_without_tools_does_not_set_tool_stream() {
265        let provider = ZAIProvider::new("test-key".to_string());
266        let request = LLMRequest {
267            model: models::zai::GLM_5_3.to_string(),
268            messages: vec![Message::user("hello".to_string())].into(),
269            stream: true,
270            ..Default::default()
271        };
272
273        let payload = provider.core.convert_request(&request).expect("payload should be valid");
274        assert_eq!(payload.get("stream").and_then(|v| v.as_bool()), Some(true));
275        assert!(payload.get("tool_stream").is_none());
276    }
277
278    #[test]
279    fn zai_base_url_uses_explicit_override() {
280        let resolved = resolve_zai_base_url(Some("https://api.z.ai/api/coding/paas/v4".to_string()));
281        assert_eq!(resolved, "https://api.z.ai/api/coding/paas/v4");
282    }
283
284    #[test]
285    fn payload_includes_do_sample() {
286        let provider = ZAIProvider::new("test-key".to_string());
287        let request = LLMRequest {
288            model: models::zai::GLM_5_3.to_string(),
289            messages: vec![Message::user("hello".to_string())].into(),
290            do_sample: Some(false),
291            ..Default::default()
292        };
293
294        let payload = provider.core.convert_request(&request).expect("payload should be valid");
295        assert_eq!(payload.get("do_sample").and_then(|v| v.as_bool()), Some(false));
296    }
297
298    #[test]
299    fn payload_disables_thinking_for_none_effort() {
300        let provider = ZAIProvider::new("test-key".to_string());
301        let request = LLMRequest {
302            model: models::zai::GLM_5_3.to_string(),
303            messages: vec![Message::user("hello".to_string())].into(),
304            reasoning_effort: Some(ReasoningEffortLevel::None),
305            ..Default::default()
306        };
307
308        let payload = provider.core.convert_request(&request).expect("payload should be valid");
309        assert_eq!(payload.get("thinking").and_then(|v| v.get("type")).and_then(|v| v.as_str()), Some("disabled"));
310    }
311
312    #[test]
313    fn payload_enables_thinking_for_low_effort() {
314        let provider = ZAIProvider::new("test-key".to_string());
315        let request = LLMRequest {
316            model: models::zai::GLM_5_3.to_string(),
317            messages: vec![Message::user("hello".to_string())].into(),
318            reasoning_effort: Some(ReasoningEffortLevel::Low),
319            ..Default::default()
320        };
321
322        let payload = provider.core.convert_request(&request).expect("payload should be valid");
323        assert_eq!(payload.get("thinking").and_then(|v| v.get("type")).and_then(|v| v.as_str()), Some("enabled"));
324        assert_eq!(payload.get("reasoning_effort").and_then(|v| v.as_str()), Some("low"));
325    }
326
327    #[test]
328    fn payload_rejects_unsupported_xhigh_and_preserves_max() {
329        let provider = ZAIProvider::new("test-key".to_string());
330        let unsupported = LLMRequest {
331            model: models::zai::GLM_5_3.to_string(),
332            messages: vec![Message::user("hello".to_string())].into(),
333            reasoning_effort: Some(ReasoningEffortLevel::XHigh),
334            ..Default::default()
335        };
336        let error = provider
337            .core
338            .convert_request(&unsupported)
339            .expect_err("unsupported effort must be blocked before transport");
340        assert!(error.to_string().contains("unsupported"));
341
342        let supported = LLMRequest {
343            reasoning_effort: Some(ReasoningEffortLevel::Max),
344            ..unsupported
345        };
346        let payload = provider.core.convert_request(&supported).expect("payload should be valid");
347        assert_eq!(payload.get("thinking").and_then(|v| v.get("type")).and_then(|v| v.as_str()), Some("enabled"));
348        assert_eq!(payload.get("reasoning_effort").and_then(|v| v.as_str()), Some("max"));
349    }
350
351    #[test]
352    fn payload_enables_preserved_thinking_when_reasoning_history_present() {
353        let provider = ZAIProvider::new("test-key".to_string());
354        let mut assistant = Message::assistant("tool planning".to_string());
355        assistant.reasoning = Some("reason step 1".to_string());
356
357        let request = LLMRequest {
358            model: models::zai::GLM_5_3.to_string(),
359            messages: vec![assistant].into(),
360            ..Default::default()
361        };
362
363        let payload = provider.core.convert_request(&request).expect("payload should be valid");
364        assert_eq!(payload.get("thinking").and_then(|v| v.get("type")).and_then(|v| v.as_str()), Some("enabled"));
365        assert_eq!(
366            payload
367                .get("thinking")
368                .and_then(|v| v.get("clear_thinking"))
369                .and_then(|v| v.as_bool()),
370            Some(false)
371        );
372    }
373
374    #[test]
375    fn payload_serializes_assistant_reasoning_content() {
376        let provider = ZAIProvider::new("test-key".to_string());
377        let mut assistant = Message::assistant("answer".to_string());
378        assistant.reasoning = Some("chain".to_string());
379
380        let request = LLMRequest {
381            model: models::zai::GLM_5_3.to_string(),
382            messages: vec![assistant].into(),
383            ..Default::default()
384        };
385
386        let payload = provider.core.convert_request(&request).expect("payload should be valid");
387        let messages = payload
388            .get("messages")
389            .and_then(|v| v.as_array())
390            .expect("messages should be serialized");
391        let first = messages.first().expect("at least one message");
392        assert_eq!(first.get("reasoning_content").and_then(|v| v.as_str()), Some("chain"));
393    }
394
395    #[test]
396    fn payload_serializes_web_search_tool() {
397        let provider = ZAIProvider::new("test-key".to_string());
398        let request = LLMRequest {
399            model: models::zai::GLM_5_3.to_string(),
400            messages: vec![Message::user("latest economic events".to_string())].into(),
401            tools: Some(Arc::new(vec![ToolDefinition::web_search(serde_json::json!({
402                "enable": true,
403                "search_engine": "search-prime",
404                "count": 5
405            }))])),
406            ..Default::default()
407        };
408
409        let payload = provider.core.convert_request(&request).expect("payload should be valid");
410        let tools = payload
411            .get("tools")
412            .and_then(|v| v.as_array())
413            .expect("tools should be serialized");
414        let first = tools.first().expect("at least one tool");
415        assert_eq!(first.get("type").and_then(|v| v.as_str()), Some("web_search"));
416        assert_eq!(
417            first
418                .get("web_search")
419                .and_then(|v| v.get("search_engine"))
420                .and_then(|v| v.as_str()),
421            Some("search-prime")
422        );
423    }
424
425    #[test]
426    fn payload_tool_choice_auto_when_requested() {
427        let provider = ZAIProvider::new("test-key".to_string());
428        let request = LLMRequest {
429            model: models::zai::GLM_5_3.to_string(),
430            messages: vec![Message::user("hello".to_string())].into(),
431            tool_choice: Some(ToolChoice::auto()),
432            ..Default::default()
433        };
434
435        let payload = provider.core.convert_request(&request).expect("payload should be valid");
436        assert_eq!(payload.get("tool_choice").and_then(|v| v.as_str()), Some("auto"));
437    }
438
439    #[test]
440    fn payload_forces_tool_choice_to_auto_for_non_auto_permissions() {
441        let provider = ZAIProvider::new("test-key".to_string());
442        let request = LLMRequest {
443            model: models::zai::GLM_5_3.to_string(),
444            messages: vec![Message::user("hello".to_string())].into(),
445            tool_choice: Some(ToolChoice::none()),
446            ..Default::default()
447        };
448
449        let payload = provider.core.convert_request(&request).expect("payload should be valid");
450        assert_eq!(payload.get("tool_choice").and_then(|v| v.as_str()), Some("auto"));
451    }
452
453    #[test]
454    fn payload_defaults_tool_choice_to_auto_when_tools_provided() {
455        let provider = ZAIProvider::new("test-key".to_string());
456        let request = LLMRequest {
457            model: models::zai::GLM_5_3.to_string(),
458            messages: vec![Message::user("hello".to_string())].into(),
459            tools: Some(Arc::new(vec![ToolDefinition::function(
460                "get_weather".to_string(),
461                "Get weather".to_string(),
462                serde_json::json!({
463                    "type": "object",
464                    "properties": {
465                        "location": {"type": "string"}
466                    },
467                    "required": ["location"]
468                }),
469            )])),
470            ..Default::default()
471        };
472
473        let payload = provider.core.convert_request(&request).expect("payload should be valid");
474        assert_eq!(payload.get("tool_choice").and_then(|v| v.as_str()), Some("auto"));
475    }
476
477    #[test]
478    fn payload_enables_json_mode_when_output_format_requested() {
479        let provider = ZAIProvider::new("test-key".to_string());
480        let request = LLMRequest {
481            model: models::zai::GLM_5_3.to_string(),
482            messages: vec![Message::user("return json".to_string())].into(),
483            output_format: Some(serde_json::json!({
484                "type": "object",
485                "properties": {
486                    "sentiment": {"type": "string"}
487                }
488            })),
489            ..Default::default()
490        };
491
492        let payload = provider.core.convert_request(&request).expect("payload should be valid");
493        assert_eq!(
494            payload
495                .get("response_format")
496                .and_then(|v| v.get("type"))
497                .and_then(|v| v.as_str()),
498            Some("json_object")
499        );
500    }
501
502    #[test]
503    fn payload_keeps_json_mode_when_thinking_disabled() {
504        let provider = ZAIProvider::new("test-key".to_string());
505        let request = LLMRequest {
506            model: models::zai::GLM_5_3.to_string(),
507            messages: vec![Message::user("return json".to_string())].into(),
508            output_format: Some(serde_json::json!({
509                "type": "object",
510                "properties": {
511                    "sentiment": {"type": "string"}
512                }
513            })),
514            reasoning_effort: Some(ReasoningEffortLevel::None),
515            ..Default::default()
516        };
517
518        let payload = provider.core.convert_request(&request).expect("payload should be valid");
519        assert_eq!(
520            payload
521                .get("response_format")
522                .and_then(|v| v.get("type"))
523                .and_then(|v| v.as_str()),
524            Some("json_object")
525        );
526        assert_eq!(payload.get("thinking").and_then(|v| v.get("type")).and_then(|v| v.as_str()), Some("disabled"));
527    }
528}