llmg_core/
streaming.rs

1//! Streaming types for chat completions
2//!
3//! Provides SSE-compatible chunk types for streaming responses.
4
5use serde::{Deserialize, Serialize};
6
7/// SSE event for streaming chat completions (OpenAI-compatible format)
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct ChatCompletionChunk {
10    /// Unique identifier for the chunk
11    pub id: String,
12    /// Object type (always "chat.completion.chunk")
13    pub object: String,
14    /// Unix timestamp
15    pub created: i64,
16    /// Model used
17    pub model: String,
18    /// Choice deltas
19    pub choices: Vec<ChoiceDelta>,
20    /// Usage statistics (optional, for include_usage streaming)
21    #[serde(skip_serializing_if = "Option::is_none")]
22    pub usage: Option<crate::types::Usage>,
23}
24
25impl ChatCompletionChunk {
26    /// Create a new chunk with a single choice
27    pub fn new(
28        id: String,
29        model: String,
30        index: u32,
31        delta: DeltaContent,
32        finish_reason: Option<String>,
33    ) -> Self {
34        Self {
35            id,
36            object: "chat.completion.chunk".to_string(),
37            created: chrono::Utc::now().timestamp(),
38            model,
39            choices: vec![ChoiceDelta {
40                index,
41                delta,
42                finish_reason,
43            }],
44            usage: None,
45        }
46    }
47
48    /// Create a final chunk with finish_reason
49    pub fn final_chunk(id: String, model: String, finish_reason: &str) -> Self {
50        Self {
51            id,
52            object: "chat.completion.chunk".to_string(),
53            created: chrono::Utc::now().timestamp(),
54            model,
55            choices: vec![ChoiceDelta {
56                index: 0,
57                delta: DeltaContent::default(),
58                finish_reason: Some(finish_reason.to_string()),
59            }],
60            usage: None,
61        }
62    }
63
64    /// Generate a new chunk ID
65    pub fn generate_id() -> String {
66        format!("chatcmpl-{}", uuid::Uuid::new_v4())
67    }
68}
69
70/// Delta in a streaming choice
71#[derive(Debug, Clone, Serialize, Deserialize)]
72pub struct ChoiceDelta {
73    /// Index of the choice
74    pub index: u32,
75    /// Delta object
76    pub delta: DeltaContent,
77    /// Finish reason (if completed)
78    #[serde(skip_serializing_if = "Option::is_none")]
79    pub finish_reason: Option<String>,
80}
81
82/// Content delta for streaming responses
83#[derive(Debug, Clone, Serialize, Deserialize, Default)]
84pub struct DeltaContent {
85    /// Role (only sent in first chunk)
86    #[serde(skip_serializing_if = "Option::is_none")]
87    pub role: Option<String>,
88    /// Content delta
89    #[serde(skip_serializing_if = "Option::is_none")]
90    pub content: Option<String>,
91    /// Tool calls delta
92    #[serde(skip_serializing_if = "Option::is_none")]
93    pub tool_calls: Option<Vec<DeltaToolCall>>,
94}
95
96#[derive(Debug, Clone, Serialize, Deserialize)]
97pub struct DeltaToolCall {
98    pub index: u32,
99    #[serde(skip_serializing_if = "Option::is_none")]
100    pub id: Option<String>,
101    #[serde(skip_serializing_if = "Option::is_none")]
102    pub r#type: Option<String>,
103    #[serde(skip_serializing_if = "Option::is_none")]
104    pub function: Option<DeltaFunctionCall>,
105}
106
107#[derive(Debug, Clone, Serialize, Deserialize)]
108pub struct DeltaFunctionCall {
109    #[serde(skip_serializing_if = "Option::is_none")]
110    pub name: Option<String>,
111    #[serde(skip_serializing_if = "Option::is_none")]
112    pub arguments: Option<String>,
113}
114
115impl DeltaContent {
116    /// Create a role delta (for first chunk)
117    pub fn role() -> Self {
118        Self {
119            role: Some("assistant".to_string()),
120            content: None,
121            tool_calls: None,
122        }
123    }
124
125    /// Create a content delta
126    pub fn content(text: impl Into<String>) -> Self {
127        Self {
128            role: None,
129            content: Some(text.into()),
130            tool_calls: None,
131        }
132    }
133
134    /// Create an empty delta (for final chunk)
135    pub fn empty() -> Self {
136        Self::default()
137    }
138}
139
140#[cfg(test)]
141mod tests {
142    use super::*;
143
144    #[test]
145    fn test_chunk_serialization() {
146        let chunk = ChatCompletionChunk::new(
147            "test-id".to_string(),
148            "gpt-4".to_string(),
149            0,
150            DeltaContent::content("Hello"),
151            None,
152        );
153
154        let json = serde_json::to_string(&chunk).unwrap();
155        assert!(json.contains("chat.completion.chunk"));
156        assert!(json.contains("Hello"));
157    }
158
159    #[test]
160    fn test_final_chunk() {
161        let chunk =
162            ChatCompletionChunk::final_chunk("test-id".to_string(), "gpt-4".to_string(), "stop");
163
164        assert_eq!(chunk.choices[0].finish_reason, Some("stop".to_string()));
165        assert!(chunk.choices[0].delta.content.is_none());
166    }
167}
llmg_core/streaming.rs

llmg_core/
streaming.rs
