1use serde::Serialize;
15use validator::*;
16
17use super::{tools::*, traits::*};
18
19#[derive(Debug, Clone, Validate, Serialize)]
48pub struct ChatBody<N, M>
49where
50 N: ModelName,
51 (N, M): Bounded,
52{
53 pub model: N,
55
56 pub messages: Vec<M>,
58
59 #[serde(skip_serializing_if = "Option::is_none")]
62 pub request_id: Option<String>,
63
64 #[serde(skip_serializing_if = "Option::is_none")]
68 pub thinking: Option<ThinkingType>,
69
70 #[serde(skip_serializing_if = "Option::is_none")]
74 pub do_sample: Option<bool>,
75
76 #[serde(skip_serializing_if = "Option::is_none")]
79 pub stream: Option<bool>,
80
81 #[serde(skip_serializing_if = "Option::is_none")]
85 pub tool_stream: Option<bool>,
86
87 #[serde(skip_serializing_if = "Option::is_none")]
91 #[validate(range(min = 0.0, max = 1.0))]
92 pub temperature: Option<f32>,
93
94 #[serde(skip_serializing_if = "Option::is_none")]
98 #[validate(range(min = 0.0, max = 1.0))]
99 pub top_p: Option<f32>,
100
101 #[serde(skip_serializing_if = "Option::is_none")]
104 #[validate(range(min = 1, max = 98304))]
105 pub max_tokens: Option<u32>,
106
107 #[serde(skip_serializing_if = "Option::is_none")]
111 pub tools: Option<Vec<Tools>>,
112
113 #[serde(skip_serializing_if = "Option::is_none")]
117 #[validate(length(min = 6, max = 128))]
118 pub user_id: Option<String>,
119
120 #[serde(skip_serializing_if = "Option::is_none")]
122 #[validate(length(min = 1, max = 1))]
123 pub stop: Option<Vec<String>>,
124
125 #[serde(skip_serializing_if = "Option::is_none")]
128 pub response_format: Option<ResponseFormat>,
129}
130
131impl<N, M> ChatBody<N, M>
132where
133 N: ModelName,
134 (N, M): Bounded,
135{
136 pub fn new(model: N, messages: M) -> Self {
137 Self {
138 model,
139 messages: vec![messages],
140 request_id: None,
141 thinking: None,
142 do_sample: None,
143 stream: None,
144 tool_stream: None,
145 temperature: None,
146 top_p: None,
147 max_tokens: None,
148 tools: None,
149 user_id: None,
150 stop: None,
151 response_format: None,
152 }
153 }
154
155 pub fn add_messages(mut self, messages: M) -> Self {
156 self.messages.push(messages);
157 self
158 }
159 pub fn add_message(mut self, message: M) -> Self {
162 self.messages.push(message);
163 self
164 }
165 pub fn extend_messages(mut self, messages: impl IntoIterator<Item = M>) -> Self {
167 self.messages.extend(messages);
168 self
169 }
170 pub fn with_request_id(mut self, request_id: impl Into<String>) -> Self {
171 self.request_id = Some(request_id.into());
172 self
173 }
174 pub fn with_do_sample(mut self, do_sample: bool) -> Self {
175 self.do_sample = Some(do_sample);
176 self
177 }
178 pub fn with_stream(mut self, stream: bool) -> Self {
179 self.stream = Some(stream);
180 self
181 }
182 pub fn with_temperature(mut self, temperature: f32) -> Self {
183 self.temperature = Some(temperature);
184 self
185 }
186 pub fn with_top_p(mut self, top_p: f32) -> Self {
187 self.top_p = Some(top_p);
188 self
189 }
190 pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
191 self.max_tokens = Some(max_tokens);
192 self
193 }
194 #[deprecated(note = "with_tools is deprecated; use add_tool/add_tools instead")]
198 pub fn with_tools(mut self, tools: impl Into<Vec<Tools>>) -> Self {
199 self.tools = Some(tools.into());
200 self
201 }
202 pub fn add_tools(mut self, tools: Tools) -> Self {
203 self.tools.get_or_insert(Vec::new()).push(tools);
204 self
205 }
206 pub fn extend_tools(mut self, tools: Vec<Tools>) -> Self {
207 self.tools.get_or_insert(Vec::new()).extend(tools);
208 self
209 }
210 pub fn with_user_id(mut self, user_id: impl Into<String>) -> Self {
211 self.user_id = Some(user_id.into());
212 self
213 }
214 pub fn with_stop(mut self, stop: String) -> Self {
215 self.stop.get_or_insert_with(Vec::new).push(stop);
216 self
217 }
218}
219
220impl<N, M> ChatBody<N, M>
221where
222 N: ModelName + ThinkEnable,
223 (N, M): Bounded,
224{
225 pub fn with_thinking(mut self, thinking: ThinkingType) -> Self {
248 self.thinking = Some(thinking);
249 self
250 }
251}
252
253impl<N, M> ChatBody<N, M>
255where
256 N: ModelName + ToolStreamEnable,
257 (N, M): Bounded,
258{
259 pub fn with_tool_stream(mut self, tool_stream: bool) -> Self {
262 self.tool_stream = Some(tool_stream);
263 if tool_stream {
264 self.stream = Some(true);
266 }
267 self
268 }
269}
270
271impl From<Tools> for Vec<Tools> {
273 fn from(tool: Tools) -> Self {
274 vec![tool]
275 }
276}
277
278#[cfg(test)]
279mod tests {
280 use super::*;
281 use crate::model::{chat_message_types::TextMessage, chat_models::GLM4_6};
282
283 #[test]
284 fn test_with_tool_stream_sets_both_fields() {
285 let body: ChatBody<GLM4_6, TextMessage> =
286 ChatBody::new(GLM4_6 {}, TextMessage::user("test"));
287 let body = body.with_tool_stream(true);
288 assert_eq!(body.tool_stream, Some(true));
289 assert_eq!(body.stream, Some(true));
290 }
291
292 #[test]
293 fn test_with_tool_stream_false_does_not_force_stream() {
294 let body: ChatBody<GLM4_6, TextMessage> =
295 ChatBody::new(GLM4_6 {}, TextMessage::user("test"));
296 let body = body.with_tool_stream(false);
297 assert_eq!(body.tool_stream, Some(false));
298 assert_ne!(body.stream, Some(true));
300 }
301
302 #[test]
303 fn test_add_tools_accumulates() {
304 let body: ChatBody<GLM4_6, TextMessage> =
305 ChatBody::new(GLM4_6 {}, TextMessage::user("test"));
306 let tool = crate::model::tools::Function::new(
307 "test_fn",
308 "A test function",
309 serde_json::json!({"type": "object"}),
310 );
311 let body = body.add_tools(crate::model::tools::Tools::Function { function: tool });
312 assert!(body.tools.is_some());
313 assert_eq!(body.tools.unwrap().len(), 1);
314 }
315
316 #[test]
317 fn test_extend_messages() {
318 let body: ChatBody<GLM4_6, TextMessage> =
319 ChatBody::new(GLM4_6 {}, TextMessage::user("first"));
320 let body = body.extend_messages(vec![
321 TextMessage::assistant("second"),
322 TextMessage::user("third"),
323 ]);
324 assert_eq!(body.messages.len(), 3);
325 match &body.messages[0] {
326 TextMessage::User { content } => assert_eq!(content, "first"),
327 _ => panic!("Expected User message"),
328 }
329 }
330
331 #[test]
332 fn test_add_message() {
333 let body: ChatBody<GLM4_6, TextMessage> =
334 ChatBody::new(GLM4_6 {}, TextMessage::user("first"));
335 let body = body.add_message(TextMessage::assistant("second"));
336 assert_eq!(body.messages.len(), 2);
337 }
338}