Skip to main content

a3s_code_core/llm/
zhipu.rs

1//! Zhipu AI (GLM) LLM client
2//!
3//! GLM uses an OpenAI-compatible API but with a different endpoint path.
4//! This client wraps `OpenAiClient` with the correct GLM defaults.
5
6use super::openai::OpenAiClient;
7use super::structured;
8use super::types::*;
9use super::LlmClient;
10use crate::retry::RetryConfig;
11use anyhow::Result;
12use async_trait::async_trait;
13use tokio::sync::mpsc;
14use tokio_util::sync::CancellationToken;
15#[cfg(test)]
16use {super::http::HttpClient, std::sync::Arc};
17
18const GLM_BASE_URL: &str = "https://open.bigmodel.cn";
19const GLM_CHAT_PATH: &str = "/api/paas/v4/chat/completions";
20
21/// Zhipu AI (GLM) client
22pub struct ZhipuClient(OpenAiClient);
23
24impl ZhipuClient {
25    pub fn new(api_key: String, model: String) -> Self {
26        Self(
27            OpenAiClient::new(api_key, model)
28                .with_provider_name("zhipu")
29                .with_base_url(GLM_BASE_URL.to_string())
30                .with_chat_completions_path(GLM_CHAT_PATH),
31        )
32    }
33
34    pub fn with_temperature(mut self, temperature: f32) -> Self {
35        self.0 = self.0.with_temperature(temperature);
36        self
37    }
38
39    pub fn with_max_tokens(mut self, max_tokens: usize) -> Self {
40        self.0 = self.0.with_max_tokens(max_tokens);
41        self
42    }
43
44    pub fn with_base_url(mut self, base_url: String) -> Self {
45        self.0 = self.0.with_base_url(base_url);
46        self
47    }
48
49    pub fn with_retry_config(mut self, retry_config: RetryConfig) -> Self {
50        self.0 = self.0.with_retry_config(retry_config);
51        self
52    }
53
54    #[cfg(test)]
55    pub fn with_http_client(mut self, http: Arc<dyn HttpClient>) -> Self {
56        self.0 = self.0.with_http_client(http);
57        self
58    }
59}
60
61#[async_trait]
62impl LlmClient for ZhipuClient {
63    async fn complete(
64        &self,
65        messages: &[Message],
66        system: Option<&str>,
67        tools: &[ToolDefinition],
68    ) -> Result<LlmResponse> {
69        self.0.complete(messages, system, tools).await
70    }
71
72    async fn complete_streaming(
73        &self,
74        messages: &[Message],
75        system: Option<&str>,
76        tools: &[ToolDefinition],
77        cancel_token: CancellationToken,
78    ) -> Result<mpsc::Receiver<StreamEvent>> {
79        self.0
80            .complete_streaming(messages, system, tools, cancel_token)
81            .await
82    }
83
84    fn native_structured_support(&self) -> structured::NativeStructuredSupport {
85        self.0.native_structured_support()
86    }
87
88    async fn complete_structured(
89        &self,
90        messages: &[Message],
91        system: Option<&str>,
92        tools: &[ToolDefinition],
93        directive: &structured::StructuredDirective,
94    ) -> Result<LlmResponse> {
95        self.0
96            .complete_structured(messages, system, tools, directive)
97            .await
98    }
99
100    async fn complete_streaming_structured(
101        &self,
102        messages: &[Message],
103        system: Option<&str>,
104        tools: &[ToolDefinition],
105        directive: &structured::StructuredDirective,
106        cancel_token: CancellationToken,
107    ) -> Result<mpsc::Receiver<StreamEvent>> {
108        self.0
109            .complete_streaming_structured(messages, system, tools, directive, cancel_token)
110            .await
111    }
112}