zai_rs/model/chat/
data.rs

1//! # Chat Completion Data Models
2//!
3//! This module defines the core data structures for chat completion requests,
4//! implementing type-safe chat interactions with the Zhipu AI API.
5//!
6//! ## Type-State Pattern
7//!
8//! The implementation uses Rust's type system to enforce compile-time guarantees
9//! about streaming capabilities through phantom types (`StreamOn`/`StreamOff`).
10//!
11//! ## Features
12//!
13//! - **Type-safe model binding** - Compile-time verification of model-message compatibility
14//! - **Builder pattern** - Fluent API for request construction
15//! - **Streaming support** - Type-state based streaming capability enforcement
16//! - **Tool integration** - Support for function calling and tool usage
17//! - **Parameter control** - Temperature, top-p, max tokens, and other generation parameters
18
19use super::super::chat_base_request::*;
20use super::super::tools::*;
21use super::super::traits::*;
22use crate::client::http::HttpClient;
23use serde::Serialize;
24use std::marker::PhantomData;
25use validator::Validate;
26
27// Type-state is defined in model::traits::{StreamState, StreamOn, StreamOff}
28
29/// Type-safe chat completion request structure.
30///
31/// This struct represents a chat completion request with compile-time guarantees
32/// for model compatibility and streaming capabilities.
33///
34/// ## Type Parameters
35///
36/// - `N` - The AI model type (must implement `ModelName + Chat`)
37/// - `M` - The message type (must form a valid bound with the model)
38/// - `S` - Stream state (`StreamOn` or `StreamOff`, defaults to `StreamOff`)
39///
40/// ## Examples
41///
42/// ```rust,ignore
43/// let model = GLM4_5_flash {};
44/// let messages = TextMessage::user("Hello, how are you?");
45/// let request = ChatCompletion::new(model, messages, api_key);
46/// ```
47pub struct ChatCompletion<N, M, S = StreamOff>
48where
49    N: ModelName + Chat,
50    (N, M): Bounded,
51    ChatBody<N, M>: Serialize,
52    S: StreamState,
53{
54    /// API key for authentication with the Zhipu AI service.
55    pub key: String,
56
57    /// API endpoint URL for chat completions.
58    /// Defaults to "https://open.bigmodel.cn/api/paas/v4/chat/completions"
59    /// but can be customized using the `with_url()` method.
60    pub url: String,
61
62    /// The request body containing model, messages, and parameters.
63    body: ChatBody<N, M>,
64
65    /// Phantom data to track streaming capability at compile time.
66    _stream: PhantomData<S>,
67}
68
69impl<N, M> ChatCompletion<N, M, StreamOff>
70where
71    N: ModelName + Chat,
72    (N, M): Bounded,
73    ChatBody<N, M>: Serialize,
74{
75    /// Creates a new non-streaming chat completion request.
76    ///
77    /// ## Arguments
78    ///
79    /// * `model` - The AI model to use for completion
80    /// * `messages` - The conversation messages
81    /// * `key` - API key for authentication
82    ///
83    /// ## Returns
84    ///
85    /// A new `ChatCompletion` instance configured for non-streaming requests.
86    pub fn new(model: N, messages: M, key: String) -> ChatCompletion<N, M, StreamOff> {
87        let body = ChatBody::new(model, messages);
88        ChatCompletion {
89            body,
90            key,
91            url: "https://open.bigmodel.cn/api/paas/v4/chat/completions".to_string(),
92            _stream: PhantomData,
93        }
94    }
95
96    /// Gets mutable access to the request body for further customization.
97    ///
98    /// This method allows modification of request parameters after initial creation.
99    pub fn body_mut(&mut self) -> &mut ChatBody<N, M> {
100        &mut self.body
101    }
102
103    /// Adds additional messages to the conversation.
104    ///
105    /// This method provides a fluent interface for building conversation context.
106    ///
107    /// ## Arguments
108    ///
109    /// * `messages` - Additional messages to append to the conversation
110    ///
111    /// ## Returns
112    ///
113    /// Self with the updated message collection, enabling method chaining.
114    pub fn add_messages(mut self, messages: M) -> Self {
115        self.body = self.body.add_messages(messages);
116        self
117    }
118    pub fn with_request_id(mut self, request_id: impl Into<String>) -> Self {
119        self.body = self.body.with_request_id(request_id);
120        self
121    }
122    pub fn with_do_sample(mut self, do_sample: bool) -> Self {
123        self.body = self.body.with_do_sample(do_sample);
124        self
125    }
126
127    pub fn with_temperature(mut self, temperature: f32) -> Self {
128        self.body = self.body.with_temperature(temperature);
129        self
130    }
131    pub fn with_top_p(mut self, top_p: f32) -> Self {
132        self.body = self.body.with_top_p(top_p);
133        self
134    }
135    pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
136        self.body = self.body.with_max_tokens(max_tokens);
137        self
138    }
139    pub fn add_tool(mut self, tool: Tools) -> Self {
140        self.body = self.body.add_tools(tool);
141        self
142    }
143    pub fn add_tools(mut self, tools: Vec<Tools>) -> Self {
144        self.body = self.body.extend_tools(tools);
145        self
146    }
147    pub fn with_user_id(mut self, user_id: impl Into<String>) -> Self {
148        self.body = self.body.with_user_id(user_id);
149        self
150    }
151    pub fn with_stop(mut self, stop: String) -> Self {
152        self.body = self.body.with_stop(stop);
153        self
154    }
155
156    /// Sets a custom API endpoint URL for this chat completion request.
157    ///
158    /// This method allows overriding the default API endpoint with a custom URL,
159    /// enabling support for different deployment environments or proxy configurations.
160    ///
161    /// ## Arguments
162    ///
163    /// * `url` - The custom API endpoint URL
164    ///
165    /// ## Returns
166    ///
167    /// Self with the updated URL, enabling method chaining.
168    ///
169    /// ## Examples
170    ///
171    /// ```rust,ignore
172    /// let request = ChatCompletion::new(model, messages, api_key)
173    ///     .with_url("https://custom-api.example.com/v1/chat/completions");
174    /// ```
175    pub fn with_url(mut self, url: impl Into<String>) -> Self {
176        self.url = url.into();
177        self
178    }
179
180    /// Sets the URL to the coding plan endpoint.
181    ///
182    /// This method configures the chat completion request to use the coding-specific
183    /// API endpoint "https://open.bigmodel.cn/api/coding/paas/v4/chat/completions".
184    ///
185    /// ## Returns
186    ///
187    /// Self with the coding plan URL, enabling method chaining.
188    ///
189    /// ## Examples
190    ///
191    /// ```rust,ignore
192    /// let request = ChatCompletion::new(model, messages, api_key)
193    ///     .with_coding_plan();
194    /// ```
195    pub fn with_coding_plan(mut self) -> Self {
196        self.url = "https://open.bigmodel.cn/api/coding/paas/v4/chat/completions".to_string();
197        self
198    }
199
200    // Optional: only available when model supports thinking
201    pub fn with_thinking(mut self, thinking: ThinkingType) -> Self
202    where
203        N: ThinkEnable,
204    {
205        self.body = self.body.with_thinking(thinking);
206        self
207    }
208
209    /// Enables streaming for this chat completion request.
210    ///
211    /// This method transitions the request to streaming mode, allowing
212    /// real-time response processing through Server-Sent Events (SSE).
213    ///
214    /// ## Returns
215    ///
216    /// A new `ChatCompletion` instance with streaming enabled (`StreamOn`).
217    pub fn enable_stream(mut self) -> ChatCompletion<N, M, StreamOn> {
218        self.body.stream = Some(true);
219        ChatCompletion {
220            key: self.key,
221            url: self.url,
222            body: self.body,
223            _stream: PhantomData,
224        }
225    }
226
227    /// Disables streaming for this chat completion request.
228    ///
229    /// This method ensures the request will receive a complete response
230    /// rather than streaming chunks.
231    ///
232    /// ## Returns
233    ///
234    /// A new `ChatCompletion` instance with streaming disabled (`StreamOff`).
235    pub fn disable_stream(mut self) -> ChatCompletion<N, M, StreamOff> {
236        self.body.stream = Some(false);
237        ChatCompletion {
238            key: self.key,
239            url: self.url,
240            body: self.body,
241            _stream: PhantomData,
242        }
243    }
244    /// Validate request parameters for non-stream chat (StreamOff)
245    pub fn validate(&self) -> anyhow::Result<()> {
246        // Field-level validation from ChatBody (temperature/top_p/max_tokens/user_id/stop...)
247        self.body.validate().map_err(|e| anyhow::anyhow!(e))?;
248        // Ensure not accidentally enabling stream in StreamOff state
249        if matches!(self.body.stream, Some(true)) {
250
251            return Err(anyhow::anyhow!(
252                "stream=true detected; use enable_stream() and streaming APIs instead"
253            ));
254        }
255        Ok(())
256    }
257
258    /// Send the request and parse typed response.
259    /// Automatically runs `validate()` before sending.
260    pub async fn send(
261        &self,
262    ) -> anyhow::Result<crate::model::chat_base_response::ChatCompletionResponse>
263    where
264        N: serde::Serialize,
265        M: serde::Serialize,
266    {
267        self.validate()?;
268        let resp: reqwest::Response = self.post().await?;
269        let parsed = resp
270            .json::<crate::model::chat_base_response::ChatCompletionResponse>()
271            .await?;
272        Ok(parsed)
273    }
274}
275
276impl<N, M> ChatCompletion<N, M, StreamOn>
277where
278    N: ModelName + Chat,
279    (N, M): Bounded,
280    ChatBody<N, M>: Serialize,
281{
282    pub fn with_tool_stream(mut self, tool_stream: bool) -> Self
283    where
284        N: ToolStreamEnable,
285    {
286        self.body = self.body.with_tool_stream(tool_stream);
287        self
288    }
289}
290
291impl<N, M, S> HttpClient for ChatCompletion<N, M, S>
292where
293    N: ModelName + Serialize + Chat,
294    M: Serialize,
295    (N, M): Bounded,
296    S: StreamState,
297{
298    type Body = ChatBody<N, M>;
299    type ApiUrl = String;
300    type ApiKey = String;
301
302    /// Returns the API endpoint URL for chat completions.
303    fn api_url(&self) -> &Self::ApiUrl {
304        &self.url
305    }
306    fn api_key(&self) -> &Self::ApiKey {
307        &self.key
308    }
309    fn body(&self) -> &Self::Body {
310        &self.body
311    }
312}
313
314/// Enables Server-Sent Events (SSE) streaming for streaming-enabled chat completions.
315///
316/// This implementation allows streaming chat completions to be processed
317/// incrementally as responses arrive from the API.
318impl<N, M> crate::model::traits::SseStreamable for ChatCompletion<N, M, StreamOn>
319where
320    N: ModelName + Serialize + Chat,
321    M: Serialize,
322    (N, M): Bounded,
323{
324}
325
326/// Provides streaming extension methods for streaming-enabled chat completions.
327///
328/// This implementation enables the use of streaming-specific methods
329/// for processing chat responses in real-time.
330impl<N, M> crate::model::stream_ext::StreamChatLikeExt for ChatCompletion<N, M, StreamOn>
331where
332    N: ModelName + Serialize + Chat,
333    M: Serialize,
334    (N, M): Bounded,
335{
336}
zai_rs/model/chat/data.rs

zai_rs/model/chat/
data.rs
