cortiq_server/
streaming.rs

1//! SSE (Server-Sent Events) streaming for OpenAI-compatible chat completions.
2
3use axum::response::sse::{Event, Sse};
4use futures::stream::Stream;
5use serde::Serialize;
6use std::convert::Infallible;
7use std::pin::Pin;
8use std::task::{Context, Poll};
9use tokio::sync::mpsc;
10
11/// A streaming chat completion response.
12pub struct ChatStream {
13    rx: mpsc::Receiver<StreamChunk>,
14    state: StreamState,
15}
16
17/// SSE termination protocol: after the finish_reason chunk (or channel
18/// close) exactly one `data: [DONE]` is emitted, then the stream ends.
19#[derive(PartialEq)]
20enum StreamState {
21    Open,
22    Finishing,
23    Done,
24}
25
26/// A single chunk in the SSE stream.
27#[derive(Debug, Clone, Serialize)]
28pub struct StreamChunk {
29    pub id: String,
30    pub object: String,
31    pub created: u64,
32    pub model: String,
33    pub choices: Vec<StreamChoice>,
34}
35
36#[derive(Debug, Clone, Serialize)]
37pub struct StreamChoice {
38    pub index: u32,
39    pub delta: StreamDelta,
40    pub finish_reason: Option<String>,
41}
42
43#[derive(Debug, Clone, Serialize)]
44pub struct StreamDelta {
45    #[serde(skip_serializing_if = "Option::is_none")]
46    pub role: Option<String>,
47    #[serde(skip_serializing_if = "Option::is_none")]
48    pub content: Option<String>,
49}
50
51impl ChatStream {
52    /// Create a new stream pair (sender, SSE response).
53    pub fn new(buffer: usize) -> (mpsc::Sender<StreamChunk>, Self) {
54        let (tx, rx) = mpsc::channel(buffer);
55        (
56            tx,
57            Self {
58                rx,
59                state: StreamState::Open,
60            },
61        )
62    }
63
64    /// Convert to axum SSE response.
65    pub fn into_sse(self) -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
66        Sse::new(SseStream { inner: self })
67    }
68}
69
70struct SseStream {
71    inner: ChatStream,
72}
73
74impl Stream for SseStream {
75    type Item = Result<Event, Infallible>;
76
77    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
78        match self.inner.state {
79            StreamState::Done => return Poll::Ready(None),
80            StreamState::Finishing => {
81                self.inner.state = StreamState::Done;
82                return Poll::Ready(Some(Ok(Event::default().data("[DONE]"))));
83            }
84            StreamState::Open => {}
85        }
86
87        match self.inner.rx.poll_recv(cx) {
88            Poll::Ready(Some(chunk)) => {
89                let is_finish = chunk
90                    .choices
91                    .first()
92                    .and_then(|c| c.finish_reason.as_deref())
93                    .is_some();
94                if is_finish {
95                    self.inner.state = StreamState::Finishing;
96                }
97                let json = serde_json::to_string(&chunk).unwrap_or_default();
98                Poll::Ready(Some(Ok(Event::default().data(json))))
99            }
100            Poll::Ready(None) => {
101                // Channel closed without a finish chunk — still terminate
102                // the protocol correctly.
103                self.inner.state = StreamState::Done;
104                Poll::Ready(Some(Ok(Event::default().data("[DONE]"))))
105            }
106            Poll::Pending => Poll::Pending,
107        }
108    }
109}
110
111/// A content-delta chunk (built synchronously from the generation thread).
112pub fn token_chunk(id: &str, model: &str, token: &str, created: u64) -> StreamChunk {
113    StreamChunk {
114        id: id.to_string(),
115        object: "chat.completion.chunk".to_string(),
116        created,
117        model: model.to_string(),
118        choices: vec![StreamChoice {
119            index: 0,
120            delta: StreamDelta {
121                role: None,
122                content: Some(token.to_string()),
123            },
124            finish_reason: None,
125        }],
126    }
127}
128
129/// The terminal chunk carrying the real finish_reason.
130pub fn finish_chunk(id: &str, model: &str, reason: &str, created: u64) -> StreamChunk {
131    StreamChunk {
132        id: id.to_string(),
133        object: "chat.completion.chunk".to_string(),
134        created,
135        model: model.to_string(),
136        choices: vec![StreamChoice {
137            index: 0,
138            delta: StreamDelta {
139                role: None,
140                content: None,
141            },
142            finish_reason: Some(reason.to_string()),
143        }],
144    }
145}
cortiq_server/streaming.rs

cortiq_server/
streaming.rs
