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;
15use {super::http::HttpClient, std::sync::Arc};
16
17const GLM_BASE_URL: &str = "https://open.bigmodel.cn";
18const GLM_CHAT_PATH: &str = "/api/paas/v4/chat/completions";
19
20/// Zhipu AI (GLM) client
21pub struct ZhipuClient(OpenAiClient);
22
23impl ZhipuClient {
24    pub fn new(api_key: String, model: String) -> Self {
25        Self(
26            OpenAiClient::new(api_key, model)
27                .with_provider_name("zhipu")
28                .with_base_url(GLM_BASE_URL.to_string())
29                .with_chat_completions_path(GLM_CHAT_PATH),
30        )
31    }
32
33    pub fn with_temperature(mut self, temperature: f32) -> Self {
34        self.0 = self.0.with_temperature(temperature);
35        self
36    }
37
38    pub fn with_max_tokens(mut self, max_tokens: usize) -> Self {
39        self.0 = self.0.with_max_tokens(max_tokens);
40        self
41    }
42
43    pub fn with_logprobs(mut self, enabled: bool) -> Self {
44        self.0 = self.0.with_logprobs(enabled);
45        self
46    }
47
48    pub fn with_top_logprobs(mut self, top_logprobs: usize) -> Self {
49        self.0 = self.0.with_top_logprobs(top_logprobs);
50        self
51    }
52
53    pub fn with_base_url(mut self, base_url: String) -> Self {
54        self.0 = self.0.with_base_url(base_url);
55        self
56    }
57
58    pub fn with_retry_config(mut self, retry_config: RetryConfig) -> Self {
59        self.0 = self.0.with_retry_config(retry_config);
60        self
61    }
62
63    pub fn with_http_client(mut self, http: Arc<dyn HttpClient>) -> Self {
64        self.0 = self.0.with_http_client(http);
65        self
66    }
67}
68
69#[async_trait]
70impl LlmClient for ZhipuClient {
71    async fn complete(
72        &self,
73        messages: &[Message],
74        system: Option<&str>,
75        tools: &[ToolDefinition],
76    ) -> Result<LlmResponse> {
77        self.0.complete(messages, system, tools).await
78    }
79
80    async fn complete_streaming(
81        &self,
82        messages: &[Message],
83        system: Option<&str>,
84        tools: &[ToolDefinition],
85        cancel_token: CancellationToken,
86    ) -> Result<mpsc::Receiver<StreamEvent>> {
87        self.0
88            .complete_streaming(messages, system, tools, cancel_token)
89            .await
90    }
91
92    fn native_structured_support(&self) -> structured::NativeStructuredSupport {
93        self.0.native_structured_support()
94    }
95
96    async fn complete_structured(
97        &self,
98        messages: &[Message],
99        system: Option<&str>,
100        tools: &[ToolDefinition],
101        directive: &structured::StructuredDirective,
102    ) -> Result<LlmResponse> {
103        self.0
104            .complete_structured(messages, system, tools, directive)
105            .await
106    }
107
108    async fn complete_streaming_structured(
109        &self,
110        messages: &[Message],
111        system: Option<&str>,
112        tools: &[ToolDefinition],
113        directive: &structured::StructuredDirective,
114        cancel_token: CancellationToken,
115    ) -> Result<mpsc::Receiver<StreamEvent>> {
116        self.0
117            .complete_streaming_structured(messages, system, tools, directive, cancel_token)
118            .await
119    }
120}