Skip to main content

zai_rs/model/async_chat/
data.rs

1use std::{marker::PhantomData, sync::Arc};
2
3use serde::Serialize;
4use validator::Validate;
5
6use super::super::{chat_base_request::*, tools::*, traits::*};
7use crate::client::{
8    endpoints::{ApiBase, EndpointConfig, paths},
9    http::{HttpClient, HttpClientConfig, parse_typed_response},
10};
11
12pub struct AsyncChatCompletion<N, M, S = StreamOff>
13where
14    N: ModelName + AsyncChat,
15    (N, M): Bounded,
16    ChatBody<N, M>: Serialize,
17    S: StreamState,
18{
19    pub key: String,
20    url: String,
21    endpoint_config: EndpointConfig,
22    api_base: ApiBase,
23    http_config: Arc<HttpClientConfig>,
24    body: ChatBody<N, M>,
25    _stream: PhantomData<S>,
26}
27
28impl<N, M> AsyncChatCompletion<N, M, StreamOff>
29where
30    N: ModelName + AsyncChat,
31    (N, M): Bounded,
32    ChatBody<N, M>: Serialize,
33{
34    pub fn new(model: N, messages: M, key: String) -> Self {
35        let body = ChatBody::new(model, messages);
36        let endpoint_config = EndpointConfig::default();
37        let api_base = ApiBase::PaasV4;
38        let url = endpoint_config.url(&api_base, paths::ASYNC_CHAT_COMPLETIONS);
39        Self {
40            body,
41            key,
42            url,
43            endpoint_config,
44            api_base,
45            http_config: Arc::new(HttpClientConfig::default()),
46            _stream: PhantomData,
47        }
48    }
49
50    pub fn body_mut(&mut self) -> &mut ChatBody<N, M> {
51        &mut self.body
52    }
53
54    // Fluent, builder-style forwarding methods to mutate inner ChatBody and return
55    // Self
56    pub fn add_messages(mut self, messages: M) -> Self {
57        self.body = self.body.add_messages(messages);
58        self
59    }
60    pub fn with_request_id(mut self, request_id: impl Into<String>) -> Self {
61        self.body = self.body.with_request_id(request_id);
62        self
63    }
64    pub fn with_do_sample(mut self, do_sample: bool) -> Self {
65        self.body = self.body.with_do_sample(do_sample);
66        self
67    }
68    #[deprecated(note = "Use enable_stream()/disable_stream() for compile-time guarantees")]
69    pub fn with_stream(mut self, stream: bool) -> Self {
70        self.body = self.body.with_stream(stream);
71        self
72    }
73    pub fn with_tool_stream(mut self, tool_stream: bool) -> Self
74    where
75        N: ToolStreamEnable,
76    {
77        self.body = self.body.with_tool_stream(tool_stream);
78        self
79    }
80
81    pub fn with_temperature(mut self, temperature: f64) -> Self {
82        self.body = self.body.with_temperature(temperature);
83        self
84    }
85    pub fn with_top_p(mut self, top_p: f64) -> Self {
86        self.body = self.body.with_top_p(top_p);
87        self
88    }
89    pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
90        self.body = self.body.with_max_tokens(max_tokens);
91        self
92    }
93    pub fn add_tool(mut self, tool: Tools) -> Self {
94        self.body = self.body.add_tools(tool);
95        self
96    }
97    pub fn add_tools(mut self, tools: Vec<Tools>) -> Self {
98        self.body = self.body.extend_tools(tools);
99        self
100    }
101    pub fn with_user_id(mut self, user_id: impl Into<String>) -> Self {
102        self.body = self.body.with_user_id(user_id);
103        self
104    }
105    pub fn with_stop(mut self, stop: String) -> Self {
106        self.body = self.body.with_stop(stop);
107        self
108    }
109
110    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
111        self.api_base = ApiBase::Custom(base_url.into());
112        self.url = self
113            .endpoint_config
114            .url(&self.api_base, paths::ASYNC_CHAT_COMPLETIONS);
115        self
116    }
117
118    pub fn with_endpoint_config(mut self, endpoint_config: EndpointConfig) -> Self {
119        self.endpoint_config = endpoint_config;
120        self.url = self
121            .endpoint_config
122            .url(&self.api_base, paths::ASYNC_CHAT_COMPLETIONS);
123        self
124    }
125
126    pub fn with_http_config(mut self, config: HttpClientConfig) -> Self {
127        self.http_config = Arc::new(config);
128        self
129    }
130
131    // Optional: only available when model supports thinking
132    pub fn with_thinking(mut self, thinking: ThinkingType) -> Self
133    where
134        N: ThinkEnable,
135    {
136        self.body = self.body.with_thinking(thinking);
137        self
138    }
139
140    // Optional: only available for GLM-5.2+ (reasoning_effort support)
141    pub fn with_reasoning_effort(mut self, effort: ReasoningEffort) -> Self
142    where
143        N: ReasoningEffortEnable,
144    {
145        self.body = self.body.with_reasoning_effort(effort);
146        self
147    }
148
149    // Type-state toggles
150    pub fn enable_stream(mut self) -> AsyncChatCompletion<N, M, StreamOn> {
151        self.body.stream = Some(true);
152        AsyncChatCompletion {
153            key: self.key,
154            url: self.url,
155            endpoint_config: self.endpoint_config,
156            api_base: self.api_base,
157            http_config: self.http_config,
158            body: self.body,
159            _stream: PhantomData,
160        }
161    }
162
163    /// Validate request parameters for non-stream async chat (StreamOff)
164    pub fn validate(&self) -> crate::ZaiResult<()> {
165        self.body
166            .validate()
167            .map_err(crate::client::error::ZaiError::from)?;
168        if matches!(self.body.stream, Some(true)) {
169            return Err(crate::client::error::ZaiError::ApiError {
170                code: 1200,
171                message: "stream=true detected; use enable_stream() and streaming APIs instead"
172                    .to_string(),
173            });
174        }
175
176        Ok(())
177    }
178
179    pub async fn send(
180        &self,
181    ) -> crate::ZaiResult<crate::model::chat_base_response::ChatCompletionResponse>
182    where
183        N: serde::Serialize,
184        M: serde::Serialize,
185    {
186        self.validate()?;
187
188        let resp: reqwest::Response = self.post().await?;
189
190        let parsed =
191            parse_typed_response::<crate::model::chat_base_response::ChatCompletionResponse>(resp)
192                .await?;
193        Ok(parsed)
194    }
195}
196
197impl<N, M> AsyncChatCompletion<N, M, StreamOn>
198where
199    N: ModelName + AsyncChat,
200    (N, M): Bounded,
201    ChatBody<N, M>: Serialize,
202{
203    pub fn with_tool_stream(mut self, tool_stream: bool) -> Self
204    where
205        N: ToolStreamEnable,
206    {
207        self.body = self.body.with_tool_stream(tool_stream);
208        self
209    }
210
211    pub fn disable_stream(mut self) -> AsyncChatCompletion<N, M, StreamOff> {
212        self.body.stream = Some(false);
213        // Reset tool_stream when disabling streaming since tool_stream depends on
214        // stream
215        self.body.tool_stream = None;
216        AsyncChatCompletion {
217            key: self.key,
218            url: self.url,
219            endpoint_config: self.endpoint_config,
220            api_base: self.api_base,
221            http_config: self.http_config,
222            body: self.body,
223            _stream: PhantomData,
224        }
225    }
226}
227
228impl<N, M, S> HttpClient for AsyncChatCompletion<N, M, S>
229where
230    N: ModelName + Serialize + AsyncChat,
231    M: Serialize,
232    (N, M): Bounded,
233    S: StreamState,
234{
235    type Body = ChatBody<N, M>;
236    type ApiUrl = String;
237    type ApiKey = String;
238
239    fn api_url(&self) -> &Self::ApiUrl {
240        &self.url
241    }
242    fn api_key(&self) -> &Self::ApiKey {
243        &self.key
244    }
245    fn body(&self) -> &Self::Body {
246        &self.body
247    }
248
249    fn http_config(&self) -> Arc<HttpClientConfig> {
250        self.http_config.clone()
251    }
252}
253
254impl<N, M> crate::model::traits::SseStreamable for AsyncChatCompletion<N, M, StreamOn>
255where
256    N: ModelName + Serialize + AsyncChat,
257    M: Serialize,
258    (N, M): Bounded,
259{
260}
261
262// Enable typed streaming extension methods for AsyncChatCompletion<...,
263// StreamOn>
264impl<N, M> crate::model::stream_ext::StreamChatLikeExt for AsyncChatCompletion<N, M, StreamOn>
265where
266    N: ModelName + Serialize + AsyncChat,
267    M: Serialize,
268    (N, M): Bounded,
269{
270}