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