openai_protocol/builders/chat/
stream_response.rs

1//! Builder for ChatCompletionStreamResponse
2//!
3//! Provides an ergonomic fluent API for constructing streaming chat completion responses.
4
5use std::borrow::Cow;
6
7use crate::{
8    chat::*,
9    common::{FunctionCallDelta, ToolCallDelta, Usage},
10};
11
12/// Builder for ChatCompletionStreamResponse
13///
14/// Provides a fluent interface for constructing streaming chat completion chunks with sensible defaults.
15#[must_use = "Builder does nothing until .build() is called"]
16#[derive(Clone, Debug)]
17pub struct ChatCompletionStreamResponseBuilder {
18    id: String,
19    object: String,
20    created: u64,
21    model: String,
22    choices: Vec<ChatStreamChoice>,
23    usage: Option<Usage>,
24    system_fingerprint: Option<String>,
25}
26
27impl ChatCompletionStreamResponseBuilder {
28    /// Create a new builder with required fields
29    ///
30    /// # Arguments
31    /// - `id`: Completion ID (e.g., "chatcmpl_abc123")
32    /// - `model`: Model name used for generation
33    pub fn new(id: impl Into<String>, model: impl Into<String>) -> Self {
34        Self {
35            id: id.into(),
36            object: "chat.completion.chunk".to_string(),
37            created: chrono::Utc::now().timestamp() as u64,
38            model: model.into(),
39            choices: Vec::new(),
40            usage: None,
41            system_fingerprint: None,
42        }
43    }
44
45    /// Copy common fields from a ChatCompletionRequest
46    ///
47    /// This populates the model field from the request.
48    pub fn copy_from_request(mut self, request: &ChatCompletionRequest) -> Self {
49        self.model = request.model.clone();
50        self
51    }
52
53    /// Set the object type (default: "chat.completion.chunk")
54    pub fn object(mut self, object: impl Into<String>) -> Self {
55        self.object = object.into();
56        self
57    }
58
59    /// Set the creation timestamp (default: current time)
60    pub fn created(mut self, timestamp: u64) -> Self {
61        self.created = timestamp;
62        self
63    }
64
65    /// Set the choices
66    pub fn choices(mut self, choices: Vec<ChatStreamChoice>) -> Self {
67        self.choices = choices;
68        self
69    }
70
71    /// Add a single choice (delta)
72    pub fn add_choice(mut self, choice: ChatStreamChoice) -> Self {
73        self.choices.push(choice);
74        self
75    }
76
77    /// Set usage information (typically sent in final chunk)
78    pub fn usage(mut self, usage: Usage) -> Self {
79        self.usage = Some(usage);
80        self
81    }
82
83    /// Set system fingerprint if provided (handles Option)
84    pub fn maybe_system_fingerprint(mut self, fingerprint: Option<impl Into<String>>) -> Self {
85        if let Some(fp) = fingerprint {
86            self.system_fingerprint = Some(fp.into());
87        }
88        self
89    }
90
91    /// Set usage if provided (handles Option)
92    pub fn maybe_usage(mut self, usage: Option<Usage>) -> Self {
93        if let Some(u) = usage {
94            self.usage = Some(u);
95        }
96        self
97    }
98
99    /// Add a choice delta that sets `role` and `content`
100    pub fn add_choice_content(
101        mut self,
102        index: u32,
103        role: impl Into<String>,
104        content: impl Into<String>,
105    ) -> Self {
106        self.choices.push(ChatStreamChoice {
107            index,
108            delta: ChatMessageDelta {
109                role: Some(role.into()),
110                content: Some(content.into()),
111                tool_calls: None,
112                reasoning_content: None,
113            },
114            logprobs: None,
115            finish_reason: None,
116            matched_stop: None,
117        });
118        self
119    }
120
121    /// Add a choice delta that sets `role`, `content`, and `logprobs`
122    pub fn add_choice_content_with_logprobs(
123        mut self,
124        index: u32,
125        role: impl Into<String>,
126        content: impl Into<String>,
127        logprobs: Option<crate::common::ChatLogProbs>,
128    ) -> Self {
129        self.choices.push(ChatStreamChoice {
130            index,
131            delta: ChatMessageDelta {
132                role: Some(role.into()),
133                content: Some(content.into()),
134                tool_calls: None,
135                reasoning_content: None,
136            },
137            logprobs,
138            finish_reason: None,
139            matched_stop: None,
140        });
141        self
142    }
143
144    /// Add a choice delta that only sets `role`
145    pub fn add_choice_role(mut self, index: u32, role: impl Into<String>) -> Self {
146        self.choices.push(ChatStreamChoice {
147            index,
148            delta: ChatMessageDelta {
149                role: Some(role.into()),
150                content: None,
151                tool_calls: None,
152                reasoning_content: None,
153            },
154            logprobs: None,
155            finish_reason: None,
156            matched_stop: None,
157        });
158        self
159    }
160
161    /// Add a choice delta that appends a tool-call *arguments delta*
162    /// Uses `Cow` so you can pass `&str` or `String` without extra clones
163    pub fn add_choice_tool_args(
164        mut self,
165        index: u32,
166        args_delta: impl Into<Cow<'static, str>>,
167    ) -> Self {
168        self.choices.push(ChatStreamChoice {
169            index,
170            delta: ChatMessageDelta {
171                role: Some("assistant".to_string()),
172                content: None,
173                tool_calls: Some(vec![ToolCallDelta {
174                    index: 0,
175                    id: None,
176                    tool_type: None,
177                    function: Some(FunctionCallDelta {
178                        name: None,
179                        arguments: Some(args_delta.into().into_owned()),
180                    }),
181                }]),
182                reasoning_content: None,
183            },
184            logprobs: None,
185            finish_reason: None,
186            matched_stop: None,
187        });
188        self
189    }
190
191    /// Add a choice delta that sets reasoning content (for models that stream reasoning)
192    pub fn add_choice_reasoning(mut self, index: u32, reasoning: impl Into<String>) -> Self {
193        self.choices.push(ChatStreamChoice {
194            index,
195            delta: ChatMessageDelta {
196                role: Some("assistant".to_string()),
197                content: None,
198                tool_calls: None,
199                reasoning_content: Some(reasoning.into()),
200            },
201            logprobs: None,
202            finish_reason: None,
203            matched_stop: None,
204        });
205        self
206    }
207
208    /// Add a choice delta for tool call with function name and ID
209    pub fn add_choice_tool_name(
210        mut self,
211        index: u32,
212        tool_call_id: impl Into<String>,
213        function_name: impl Into<String>,
214    ) -> Self {
215        self.choices.push(ChatStreamChoice {
216            index,
217            delta: ChatMessageDelta {
218                role: Some("assistant".to_string()),
219                content: None,
220                tool_calls: Some(vec![ToolCallDelta {
221                    index: 0,
222                    id: Some(tool_call_id.into()),
223                    tool_type: Some("function".to_string()),
224                    function: Some(FunctionCallDelta {
225                        name: Some(function_name.into()),
226                        arguments: None,
227                    }),
228                }]),
229                reasoning_content: None,
230            },
231            logprobs: None,
232            finish_reason: None,
233            matched_stop: None,
234        });
235        self
236    }
237
238    /// Add a choice delta with a pre-constructed ToolCallDelta
239    /// Useful when you already have a ToolCallDelta object to emit
240    pub fn add_choice_tool_call_delta(
241        mut self,
242        index: u32,
243        tool_call_delta: ToolCallDelta,
244    ) -> Self {
245        self.choices.push(ChatStreamChoice {
246            index,
247            delta: ChatMessageDelta {
248                role: Some("assistant".to_string()),
249                content: None,
250                tool_calls: Some(vec![tool_call_delta]),
251                reasoning_content: None,
252            },
253            logprobs: None,
254            finish_reason: None,
255            matched_stop: None,
256        });
257        self
258    }
259
260    /// Add a choice with finish_reason (final chunk)
261    /// This is used for the last chunk in a stream to signal completion
262    pub fn add_choice_finish_reason(
263        mut self,
264        index: u32,
265        finish_reason: impl Into<String>,
266        matched_stop: Option<serde_json::Value>,
267    ) -> Self {
268        self.choices.push(ChatStreamChoice {
269            index,
270            delta: ChatMessageDelta {
271                role: None,
272                content: None,
273                tool_calls: None,
274                reasoning_content: None,
275            },
276            logprobs: None,
277            finish_reason: Some(finish_reason.into()),
278            matched_stop,
279        });
280        self
281    }
282
283    /// Build the ChatCompletionStreamResponse
284    pub fn build(self) -> ChatCompletionStreamResponse {
285        ChatCompletionStreamResponse {
286            id: self.id,
287            object: self.object,
288            created: self.created,
289            model: self.model,
290            system_fingerprint: self.system_fingerprint,
291            choices: self.choices,
292            usage: self.usage,
293        }
294    }
295}
296
297// ============================================================================
298// Tests
299// ============================================================================
300
301#[cfg(test)]
302mod tests {
303    use super::*;
304
305    #[test]
306    fn test_build_minimal() {
307        let chunk = ChatCompletionStreamResponseBuilder::new("chatcmpl_123", "gpt-4").build();
308
309        assert_eq!(chunk.id, "chatcmpl_123");
310        assert_eq!(chunk.model, "gpt-4");
311        assert_eq!(chunk.object, "chat.completion.chunk");
312        assert!(chunk.choices.is_empty());
313        assert!(chunk.usage.is_none());
314    }
315
316    #[test]
317    fn test_with_content_delta() {
318        let chunk = ChatCompletionStreamResponseBuilder::new("chatcmpl_456", "gpt-4")
319            .add_choice_content(0, "assistant", "Hello")
320            .build();
321
322        assert_eq!(chunk.choices.len(), 1);
323        assert_eq!(chunk.choices[0].index, 0);
324        assert_eq!(chunk.choices[0].delta.content.as_ref().unwrap(), "Hello");
325        assert_eq!(chunk.choices[0].delta.role.as_ref().unwrap(), "assistant");
326        assert!(chunk.choices[0].finish_reason.is_none());
327    }
328
329    #[test]
330    fn test_with_role_delta() {
331        let chunk = ChatCompletionStreamResponseBuilder::new("chatcmpl_789", "gpt-4")
332            .add_choice_role(0, "assistant")
333            .build();
334
335        assert_eq!(chunk.choices.len(), 1);
336        assert_eq!(chunk.choices[0].delta.role.as_ref().unwrap(), "assistant");
337        assert!(chunk.choices[0].delta.content.is_none());
338    }
339
340    #[test]
341    fn test_with_finish_reason() {
342        let chunk = ChatCompletionStreamResponseBuilder::new("chatcmpl_101", "gpt-4")
343            .add_choice_finish_reason(0, "stop", None)
344            .build();
345
346        assert_eq!(chunk.choices.len(), 1);
347        assert_eq!(chunk.choices[0].finish_reason.as_ref().unwrap(), "stop");
348        assert!(chunk.choices[0].delta.content.is_none());
349        assert!(chunk.choices[0].delta.role.is_none());
350    }
351
352    #[test]
353    fn test_multiple_deltas() {
354        let chunk = ChatCompletionStreamResponseBuilder::new("chatcmpl_202", "gpt-4")
355            .add_choice_role(0, "assistant")
356            .add_choice_content(0, "assistant", "Hello")
357            .add_choice_content(0, "assistant", " world")
358            .add_choice_finish_reason(0, "stop", None)
359            .build();
360
361        assert_eq!(chunk.choices.len(), 4); // role + 2 content + finish
362    }
363
364    #[test]
365    fn test_with_usage() {
366        let usage = Usage {
367            prompt_tokens: 10,
368            completion_tokens: 20,
369            total_tokens: 30,
370            completion_tokens_details: None,
371        };
372
373        let chunk = ChatCompletionStreamResponseBuilder::new("chatcmpl_303", "gpt-4")
374            .add_choice_finish_reason(0, "stop", None)
375            .usage(usage)
376            .build();
377
378        assert!(chunk.usage.is_some());
379        assert_eq!(chunk.usage.as_ref().unwrap().total_tokens, 30);
380    }
381
382    #[test]
383    fn test_copy_from_request() {
384        let request = ChatCompletionRequest {
385            messages: vec![],
386            model: "gpt-3.5-turbo".to_string(),
387            ..Default::default()
388        };
389
390        let chunk = ChatCompletionStreamResponseBuilder::new("chatcmpl_404", "gpt-4")
391            .copy_from_request(&request)
392            .add_choice_content(0, "assistant", "test")
393            .build();
394
395        assert_eq!(chunk.model, "gpt-3.5-turbo"); // Copied from request
396    }
397
398    #[test]
399    fn test_add_choice_explicit() {
400        let choice = ChatStreamChoice {
401            index: 0,
402            delta: ChatMessageDelta {
403                role: Some("assistant".to_string()),
404                content: Some("Hello".to_string()),
405                tool_calls: None,
406                reasoning_content: None,
407            },
408            logprobs: None,
409            finish_reason: None,
410            matched_stop: None,
411        };
412
413        let chunk = ChatCompletionStreamResponseBuilder::new("chatcmpl_505", "gpt-4")
414            .add_choice(choice)
415            .build();
416
417        assert_eq!(chunk.choices.len(), 1);
418        assert_eq!(chunk.choices[0].delta.role.as_ref().unwrap(), "assistant");
419        assert_eq!(chunk.choices[0].delta.content.as_ref().unwrap(), "Hello");
420    }
421}
openai_protocol/builders/chat/stream_response.rs

openai_protocol/builders/chat/
stream_response.rs
