Skip to main content

rig_core/providers/
zai.rs

1//! Z.AI API clients and Rig integrations.
2//!
3//! Z.AI exposes OpenAI-compatible APIs for both its general platform and
4//! coding-focused platform, plus an Anthropic-compatible endpoint for tools
5//! like Claude Code.
6//!
7//! # OpenAI-compatible example
8//! ```no_run
9//! use rig_core::client::CompletionClient;
10//! use rig_core::providers::zai;
11//!
12//! let client = zai::Client::new("YOUR_API_KEY").expect("Failed to build client");
13//! let glm_4_6 = client.completion_model(zai::GLM_4_6);
14//! ```
15//!
16//! # Anthropic-compatible example
17//! ```no_run
18//! use rig_core::client::CompletionClient;
19//! use rig_core::providers::zai;
20//!
21//! let client = zai::AnthropicClient::new("YOUR_API_KEY").expect("Failed to build client");
22//! let glm_4_6 = client.completion_model(zai::GLM_4_6);
23//! ```
24
25use crate::client;
26use crate::providers::internal::anthropic_compatible::{
27    AnthropicBaseUrl, impl_dual_dialect_provider,
28};
29
30/// General-purpose OpenAI-compatible base URL.
31pub const GENERAL_API_BASE_URL: &str = "https://api.z.ai/api/paas/v4";
32/// Coding-focused OpenAI-compatible base URL.
33pub const CODING_API_BASE_URL: &str = "https://api.z.ai/api/coding/paas/v4";
34/// Anthropic-compatible base URL.
35pub const ANTHROPIC_API_BASE_URL: &str = "https://api.z.ai/api/anthropic";
36
37/// `glm-4.6`
38pub const GLM_4_6: &str = "glm-4.6";
39/// `glm-4.6-air`
40pub const GLM_4_6_AIR: &str = "glm-4.6-air";
41/// `glm-4.6-x`
42pub const GLM_4_6_X: &str = "glm-4.6-x";
43/// `glm-4.5`
44pub const GLM_4_5: &str = "glm-4.5";
45/// `glm-4.5-air`
46pub const GLM_4_5_AIR: &str = "glm-4.5-air";
47/// `glm-4.5v`
48pub const GLM_4_5V: &str = "glm-4.5v";
49/// `glm-4.5-airx`
50pub const GLM_4_5_AIRX: &str = "glm-4.5-airx";
51
52impl_dual_dialect_provider!(
53    ext = ZAiExt,
54    builder = ZAiBuilder,
55    anthropic_ext = ZAiAnthropicExt,
56    anthropic_builder = ZAiAnthropicBuilder,
57    client_input = client::BearerAuth,
58    api_key_env = "ZAI_API_KEY",
59    base_url = GENERAL_API_BASE_URL,
60    base_url_env = "ZAI_API_BASE",
61    anthropic_provider_name = "z.ai",
62    anthropic_base_url = ANTHROPIC_API_BASE_URL,
63    anthropic_base_url_env = "ZAI_ANTHROPIC_API_BASE",
64);
65
66client::impl_capabilities!(
67    ZAiExt,
68    completion = super::openai::completion::GenericCompletionModel<ZAiExt, H>,
69);
70
71impl super::openai::completion::OpenAICompatibleProvider for ZAiExt {
72    const PROVIDER_NAME: &'static str = "zai";
73
74    type StreamingUsage = super::openai::Usage;
75
76    type Response = super::openai::CompletionResponse;
77}
78
79const ANTHROPIC_BASE_URLS: AnthropicBaseUrl = AnthropicBaseUrl::new(
80    &[
81        (GENERAL_API_BASE_URL, ANTHROPIC_API_BASE_URL),
82        (CODING_API_BASE_URL, ANTHROPIC_API_BASE_URL),
83    ],
84    &[
85        "/api/paas/v4",
86        "/api/paas/v4/",
87        "/api/coding/paas/v4",
88        "/api/coding/paas/v4/",
89    ],
90    "/api/anthropic",
91);
92
93impl<H> ClientBuilder<H> {
94    pub fn general(self) -> Self {
95        self.base_url(GENERAL_API_BASE_URL)
96    }
97
98    pub fn coding(self) -> Self {
99        self.base_url(CODING_API_BASE_URL)
100    }
101}
102
103impl<H> AnthropicClientBuilder<H> {
104    pub fn general(self) -> Self {
105        self.base_url(ANTHROPIC_API_BASE_URL)
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::{
112        ANTHROPIC_API_BASE_URL, ANTHROPIC_BASE_URLS, CODING_API_BASE_URL, GENERAL_API_BASE_URL,
113    };
114
115    #[test]
116    fn test_client_initialization() {
117        let _client = crate::providers::zai::Client::new("dummy-key").expect("Client::new()");
118        let _client_from_builder = crate::providers::zai::Client::builder()
119            .api_key("dummy-key")
120            .build()
121            .expect("Client::builder()");
122        let _anthropic_client = crate::providers::zai::AnthropicClient::new("dummy-key")
123            .expect("AnthropicClient::new()");
124        let _anthropic_client_from_builder = crate::providers::zai::AnthropicClient::builder()
125            .api_key("dummy-key")
126            .build()
127            .expect("AnthropicClient::builder()");
128    }
129
130    #[test]
131    fn normalize_openai_style_bases_to_anthropic_base() {
132        assert_eq!(
133            ANTHROPIC_BASE_URLS
134                .normalize(GENERAL_API_BASE_URL)
135                .as_deref(),
136            Some(ANTHROPIC_API_BASE_URL)
137        );
138        assert_eq!(
139            ANTHROPIC_BASE_URLS
140                .normalize(CODING_API_BASE_URL)
141                .as_deref(),
142            Some(ANTHROPIC_API_BASE_URL)
143        );
144        assert_eq!(
145            ANTHROPIC_BASE_URLS
146                .normalize("https://proxy.example.com/api/paas/v4")
147                .as_deref(),
148            Some("https://proxy.example.com/api/anthropic")
149        );
150        assert_eq!(
151            ANTHROPIC_BASE_URLS
152                .normalize("https://proxy.example.com/api/coding/paas/v4")
153                .as_deref(),
154            Some("https://proxy.example.com/api/anthropic")
155        );
156    }
157
158    #[test]
159    fn normalize_preserves_existing_anthropic_base() {
160        assert_eq!(
161            ANTHROPIC_BASE_URLS
162                .normalize("https://proxy.example.com/api/anthropic")
163                .as_deref(),
164            Some("https://proxy.example.com/api/anthropic")
165        );
166    }
167
168    #[test]
169    fn anthropic_primary_override_wins() {
170        let override_url = ANTHROPIC_BASE_URLS.resolve(
171            Some("https://primary.example.com/api/anthropic"),
172            Some(GENERAL_API_BASE_URL),
173        );
174
175        assert_eq!(
176            override_url.as_deref(),
177            Some("https://primary.example.com/api/anthropic")
178        );
179    }
180}