a3s_code_core/llm/
zhipu.rs1use 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
21pub 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}