zai_rs/model/chat/
data.rs

1use std::marker::PhantomData;
2use std::pin::Pin;
3use std::task::{Context, Poll};
4
5use futures_util::{Stream, StreamExt};
6use serde::Serialize;
7use validator::Validate;
8
9use super::super::{
10    chat_base_request::*, chat_base_response::ChatCompletionResponse,
11    chat_stream_response::ChatStreamResponse, tools::*, traits::*,
12};
13use crate::client::ZaiClient;
14
15/// Authenticated typed response stream returned by [`ChatCompletion::stream_via`].
16///
17/// [`next`](Self::next) is available without importing an extension trait. The
18/// type also implements [`Stream`] for combinators from `futures-util`.
19pub struct ChatStream {
20    inner: Pin<Box<dyn Stream<Item = crate::ZaiResult<ChatStreamResponse>> + Send + 'static>>,
21}
22
23impl ChatStream {
24    /// Await the next decoded chat chunk, or `None` after `[DONE]`.
25    ///
26    /// An unexpected EOF is yielded as an error before the stream terminates.
27    pub async fn next(&mut self) -> Option<crate::ZaiResult<ChatStreamResponse>> {
28        self.inner.next().await
29    }
30}
31
32impl Stream for ChatStream {
33    type Item = crate::ZaiResult<ChatStreamResponse>;
34
35    fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
36        self.inner.as_mut().poll_next(context)
37    }
38}
39
40/// Typed chat-completion request builder sent through a [`ZaiClient`].
41///
42/// Generic over the model `N`, the message type `M`, and a stream type-state
43/// `S` (`StreamOff` by default, `StreamOn` after
44/// [`enable_stream`](Self::enable_stream)).
45/// Credentials, endpoint selection and transport policy are owned by the
46/// [`ZaiClient`] supplied to [`send_via`](Self::send_via).
47///
48/// Audio chat deliberately has no tool surface:
49///
50/// ```compile_fail
51/// use zai_rs::model::{
52///     chat::ChatCompletion,
53///     chat_message_types::VoiceMessage,
54///     chat_models::GLM4_voice,
55///     tools::Function,
56/// };
57///
58/// ChatCompletion::new(GLM4_voice {}, VoiceMessage::new_user()).add_tool(
59///     Function::new("lookup", "Lookup", serde_json::json!({"type": "object"})),
60/// );
61/// ```
62///
63/// ```compile_fail
64/// use zai_rs::model::{
65///     chat::ChatCompletion,
66///     chat_message_types::VoiceMessage,
67///     chat_models::GLM4_voice,
68///     tools::ToolChoice,
69/// };
70///
71/// ChatCompletion::new(GLM4_voice {}, VoiceMessage::new_user())
72///     .with_tool_choice(ToolChoice::auto());
73/// ```
74///
75/// ```compile_fail
76/// use zai_rs::model::{
77///     chat::ChatCompletion,
78///     chat_message_types::VoiceMessage,
79///     chat_models::GLM4_voice,
80///     tools::ResponseFormat,
81/// };
82///
83/// ChatCompletion::new(GLM4_voice {}, VoiceMessage::new_user())
84///     .with_response_format(ResponseFormat::JsonObject);
85/// ```
86///
87/// Vision chat accepts a bare [`Function`](crate::model::tools::Function), not
88/// the text schema's retrieval, web-search, or MCP tool variants:
89///
90/// ```compile_fail
91/// use zai_rs::model::{
92///     chat::ChatCompletion,
93///     chat_message_types::VisionMessage,
94///     chat_models::GLM4_6v,
95///     tools::{Retrieval, Tools},
96/// };
97///
98/// ChatCompletion::new(GLM4_6v {}, VisionMessage::new_user()).add_tool(
99///     Tools::Retrieval { retrieval: Retrieval::new("kb") },
100/// );
101/// ```
102///
103/// ```compile_fail
104/// use zai_rs::model::{
105///     chat::ChatCompletion,
106///     chat_message_types::VisionMessage,
107///     chat_models::GLM4_6v,
108///     tools::ResponseFormat,
109/// };
110///
111/// ChatCompletion::new(GLM4_6v {}, VisionMessage::new_user())
112///     .with_response_format(ResponseFormat::JsonObject);
113/// ```
114pub struct ChatCompletion<N, M, S = StreamOff>
115where
116    N: ChatRequestModel + Chat,
117    M: Serialize,
118    (N, M): Bounded,
119    ChatBody<N, M>: Serialize,
120    S: StreamState,
121{
122    body: ChatBody<N, M>,
123    _stream: PhantomData<S>,
124}
125
126impl<N, M, S> ChatCompletion<N, M, S>
127where
128    N: ChatRequestModel + Chat,
129    M: Serialize,
130    (N, M): Bounded,
131    ChatBody<N, M>: Serialize,
132    S: StreamState,
133{
134    /// Borrow the frozen request body.
135    pub const fn body(&self) -> &ChatBody<N, M> {
136        &self.body
137    }
138}
139
140impl<N, M> ChatCompletion<N, M, StreamOff>
141where
142    N: ChatRequestModel + Chat,
143    M: Serialize,
144    (N, M): Bounded,
145    ChatBody<N, M>: Serialize,
146{
147    /// Create a new chat request from a model and the first message batch.
148    pub fn new(model: N, messages: M) -> ChatCompletion<N, M, StreamOff> {
149        let body = ChatBody::new(model, messages);
150        ChatCompletion {
151            body,
152            _stream: PhantomData,
153        }
154    }
155
156    /// Append one message to the conversation.
157    pub fn add_message(mut self, message: M) -> Self {
158        self.body = self.body.add_message(message);
159        self
160    }
161    /// Append multiple messages to the conversation.
162    pub fn extend_messages(mut self, messages: impl IntoIterator<Item = M>) -> Self {
163        self.body = self.body.extend_messages(messages);
164        self
165    }
166    /// Set the client-side request id.
167    pub fn with_request_id(mut self, request_id: impl Into<String>) -> Self {
168        self.body = self.body.with_request_id(request_id);
169        self
170    }
171    /// Enable/disable sampling (`do_sample`).
172    pub fn with_do_sample(mut self, do_sample: bool) -> Self {
173        self.body = self.body.with_do_sample(do_sample);
174        self
175    }
176    /// Set the sampling temperature.
177    pub fn with_temperature(mut self, temperature: f64) -> Self {
178        self.body = self.body.with_temperature(temperature);
179        self
180    }
181    /// Set the nucleus-sampling probability (`top_p`).
182    pub fn with_top_p(mut self, top_p: f64) -> Self {
183        self.body = self.body.with_top_p(top_p);
184        self
185    }
186    /// Set the maximum number of tokens to generate.
187    pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
188        self.body = self.body.with_max_tokens(max_tokens);
189        self
190    }
191    /// Add a single tool to the request.
192    pub fn add_tool(mut self, tool: N::Tool) -> Self
193    where
194        N: ChatToolSupport,
195    {
196        self.body = self.body.add_tool(tool);
197        self
198    }
199    /// Add multiple tools to the request at once.
200    pub fn add_tools(mut self, tools: impl IntoIterator<Item = N::Tool>) -> Self
201    where
202        N: ChatToolSupport,
203    {
204        self.body = self.body.add_tools(tools);
205        self
206    }
207    /// Set automatic tool selection.
208    pub fn with_tool_choice(mut self, tool_choice: ToolChoice) -> Self
209    where
210        N: ChatToolSupport,
211    {
212        self.body = self.body.with_tool_choice(tool_choice);
213        self
214    }
215    /// Remove all tools and their selection policy.
216    pub fn clear_tools(mut self) -> Self
217    where
218        N: ChatToolSupport,
219    {
220        self.body = self.body.clear_tools();
221        self
222    }
223    /// Set the text response format.
224    pub fn with_response_format(mut self, format: ResponseFormat) -> Self
225    where
226        N: ResponseFormatEnable,
227    {
228        self.body = self.body.with_response_format(format);
229        self
230    }
231    /// Enable or disable audio-output watermarking.
232    pub fn with_watermark_enabled(mut self, enabled: bool) -> Self
233    where
234        N: WatermarkEnable,
235    {
236        self.body = self.body.with_watermark_enabled(enabled);
237        self
238    }
239    /// Set the end-user id (used for abuse monitoring).
240    pub fn with_user_id(mut self, user_id: impl Into<String>) -> Self {
241        self.body = self.body.with_user_id(user_id);
242        self
243    }
244    /// Add a stop sequence that halts generation when encountered.
245    pub fn with_stop(mut self, stop: impl Into<String>) -> Self {
246        self.body = self.body.with_stop(stop);
247        self
248    }
249
250    /// Enable thinking mode (requires a model that supports it).
251    pub fn with_thinking(mut self, thinking: ThinkingType) -> Self
252    where
253        N: ThinkEnable,
254    {
255        self.body = self.body.with_thinking(thinking);
256        self
257    }
258
259    /// Set the reasoning effort (GLM-5.2+ only).
260    pub fn with_reasoning_effort(mut self, effort: ReasoningEffort) -> Self
261    where
262        N: ReasoningEffortEnable,
263    {
264        self.body = self.body.with_reasoning_effort(effort);
265        self
266    }
267
268    /// Enables streaming for this chat completion request.
269    pub fn enable_stream(mut self) -> ChatCompletion<N, M, StreamOn> {
270        self.body.set_stream(Some(true));
271        ChatCompletion {
272            body: self.body,
273            _stream: PhantomData,
274        }
275    }
276
277    /// Validate request parameters for non-stream chat (StreamOff)
278    pub fn validate(&self) -> crate::ZaiResult<()> {
279        self.body
280            .validate()
281            .map_err(crate::client::error::ZaiError::from)?;
282        if matches!(self.body.stream(), Some(true)) {
283            return Err(crate::client::error::ZaiError::ApiError {
284                code: crate::client::error::codes::SDK_VALIDATION,
285                message: "stream=true detected; use enable_stream() and streaming APIs instead"
286                    .to_string(),
287            });
288        }
289        Ok(())
290    }
291
292    /// Submit the request via a [`ZaiClient`] and await the (non-streaming)
293    /// response.
294    pub async fn send_via(&self, client: &ZaiClient) -> crate::ZaiResult<ChatCompletionResponse>
295    where
296        N: serde::Serialize,
297        M: serde::Serialize,
298    {
299        self.validate()?;
300        let route = crate::client::routes::CHAT_COMPLETE;
301        let url = client.endpoints().resolve_route(route, &[])?;
302        client
303            .send_json::<_, ChatCompletionResponse>(route.method(), url, &self.body)
304            .await
305    }
306
307    /// Submit using a coding-plan endpoint via a [`ZaiClient`].
308    pub async fn send_via_coding_plan(
309        &self,
310        client: &ZaiClient,
311    ) -> crate::ZaiResult<ChatCompletionResponse>
312    where
313        N: serde::Serialize,
314        M: serde::Serialize,
315    {
316        self.validate()?;
317        let route = crate::client::routes::CHAT_COMPLETE_CODING;
318        let url = client.endpoints().resolve_route(route, &[])?;
319        client
320            .send_json::<_, ChatCompletionResponse>(route.method(), url, &self.body)
321            .await
322    }
323}
324
325impl<N, M> ChatCompletion<N, M, StreamOn>
326where
327    N: ChatRequestModel + Chat,
328    M: Serialize,
329    (N, M): Bounded,
330    ChatBody<N, M>: Serialize,
331{
332    /// Enable/disable tool-call streaming (requires a model that supports it).
333    pub fn with_tool_stream(mut self, tool_stream: bool) -> Self
334    where
335        N: ToolStreamEnable,
336    {
337        self.body = self.body.with_tool_stream(tool_stream);
338        self
339    }
340
341    /// Disables streaming for this chat completion request.
342    pub fn disable_stream(mut self) -> ChatCompletion<N, M, StreamOff> {
343        self.body.set_stream(Some(false));
344        self.body.clear_tool_stream();
345        ChatCompletion {
346            body: self.body,
347            _stream: PhantomData,
348        }
349    }
350
351    /// Submit the request and yield parsed chat chunks from its SSE response.
352    ///
353    /// Authentication, response content-type validation, timeouts, and body
354    /// limits remain inside [`ZaiClient`]; the API secret is never exposed to
355    /// the caller. Streaming POST requests are not retried or redirected.
356    pub async fn stream_via(&self, client: &ZaiClient) -> crate::ZaiResult<ChatStream>
357    where
358        N: serde::Serialize,
359        M: serde::Serialize,
360    {
361        self.body
362            .validate()
363            .map_err(crate::client::error::ZaiError::from)?;
364        let url = client
365            .endpoints()
366            .resolve_route(crate::client::routes::CHAT_COMPLETE, &[])?;
367        let raw = client
368            .send_sse_json(
369                crate::client::routes::CHAT_COMPLETE.method(),
370                url,
371                &self.body,
372            )
373            .await?;
374        let stream = crate::model::sse_parser::decode_required_done_stream(raw, |payload| {
375            serde_json::from_slice::<ChatStreamResponse>(payload)
376                .map_err(crate::ZaiError::from)
377                .and_then(validate_stream_finish_reason)
378        });
379        Ok(ChatStream { inner: stream })
380    }
381}
382
383fn validate_stream_finish_reason(
384    chunk: ChatStreamResponse,
385) -> crate::ZaiResult<ChatStreamResponse> {
386    let failure = chunk
387        .choices
388        .iter()
389        .filter_map(|choice| choice.finish_reason.as_deref())
390        .find_map(|reason| match reason {
391            "sensitive" => Some((1301, "stream stopped by content policy")),
392            "network_error" => Some((1234, "stream stopped by an upstream inference error")),
393            "model_context_window_exceeded" => {
394                Some((1261, "stream exceeded the model context window"))
395            },
396            _ => None,
397        });
398
399    match failure {
400        Some((code, message)) => Err(crate::ZaiError::from_api_response(
401            200,
402            code,
403            message.to_string(),
404        )),
405        None => Ok(chunk),
406    }
407}
408
409#[cfg(test)]
410mod tests {
411    use super::*;
412    use crate::model::{chat_message_types::TextMessage, chat_models::GLM5_2};
413
414    #[test]
415    fn stream_field_is_owned_by_the_type_state() {
416        let request = ChatCompletion::new(GLM5_2 {}, TextMessage::user("hello"));
417        let json = serde_json::to_value(request.body()).unwrap();
418        assert!(json.get("stream").is_none());
419        assert!(json.get("tool_stream").is_none());
420
421        let request = request.enable_stream().with_tool_stream(true);
422        let json = serde_json::to_value(request.body()).unwrap();
423        assert_eq!(json["stream"], true);
424        assert_eq!(json["tool_stream"], true);
425    }
426}
zai_rs/model/chat/data.rs

zai_rs/model/chat/
data.rs
