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;
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
20pub 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 fn has_distinct_non_streaming_transport(&self) -> bool {
97 self.0.has_distinct_non_streaming_transport()
98 }
99
100 async fn complete_structured(
101 &self,
102 messages: &[Message],
103 system: Option<&str>,
104 tools: &[ToolDefinition],
105 directive: &structured::StructuredDirective,
106 ) -> Result<LlmResponse> {
107 self.0
108 .complete_structured(messages, system, tools, directive)
109 .await
110 }
111
112 async fn complete_streaming_structured(
113 &self,
114 messages: &[Message],
115 system: Option<&str>,
116 tools: &[ToolDefinition],
117 directive: &structured::StructuredDirective,
118 cancel_token: CancellationToken,
119 ) -> Result<mpsc::Receiver<StreamEvent>> {
120 self.0
121 .complete_streaming_structured(messages, system, tools, directive, cancel_token)
122 .await
123 }
124}