sim_lib_openai_server/codec_openai/
streaming.rs

1use serde_json::{Value, json};
2use sim_kernel::{Error, Expr, Result, Symbol};
3use sim_lib_stream_core::StreamPacket;
4
5use crate::{codec_openai::encode_openai_responses_response, objects::GatewayEvent};
6
7/// Selects which OpenAI server-sent-event surface to render gateway events as.
8#[derive(Clone, Copy, Debug, PartialEq, Eq)]
9pub enum OpenAiSseSurface {
10    /// The OpenAI Responses API streaming surface.
11    Responses,
12    /// The OpenAI Chat Completions streaming surface.
13    Chat,
14}
15
16/// Decoded gateway event payload: sequence number, event kind, and body.
17#[derive(Clone, Debug, PartialEq, Eq)]
18pub struct GatewayEventData {
19    sequence: u64,
20    kind: Symbol,
21    payload: Expr,
22}
23
24impl GatewayEventData {
25    /// Extracts the streaming-relevant fields from a [`GatewayEvent`].
26    pub fn from_event(event: &GatewayEvent) -> Self {
27        Self {
28            sequence: event.sequence(),
29            kind: event.kind().clone(),
30            payload: event.payload().clone(),
31        }
32    }
33
34    /// Decodes gateway event data from a stream data packet, erroring if the
35    /// packet is not a data packet of the gateway event kind.
36    pub fn from_packet(packet: &StreamPacket) -> Result<Self> {
37        let StreamPacket::Data(data) = packet else {
38            return Err(Error::TypeMismatch {
39                expected: "OpenAI gateway data packet",
40                found: "non-data packet",
41            });
42        };
43        if data.kind != gateway_event_data_kind() {
44            return Err(Error::Eval(format!(
45                "expected OpenAI gateway data kind {}, found {}",
46                gateway_event_data_kind(),
47                data.kind
48            )));
49        }
50        Ok(Self {
51            sequence: string_field(&data.payload, "sequence")?
52                .parse::<u64>()
53                .map_err(|err| Error::Eval(format!("invalid gateway event sequence: {err}")))?,
54            kind: symbol_value_field(&data.payload, "event-kind")?,
55            payload: map_field(&data.payload, "payload")?.clone(),
56        })
57    }
58
59    /// Returns the event sequence number.
60    pub fn sequence(&self) -> u64 {
61        self.sequence
62    }
63
64    /// Returns the event kind symbol.
65    pub fn kind(&self) -> &Symbol {
66        &self.kind
67    }
68
69    /// Returns the event payload expression.
70    pub fn payload(&self) -> &Expr {
71        &self.payload
72    }
73}
74
75/// Returns the stream data kind symbol `stream/data:openai-gateway-event`.
76pub fn gateway_event_data_kind() -> Symbol {
77    Symbol::qualified("stream/data", "openai-gateway-event")
78}
79
80/// Wraps gateway events as stream data packets of the gateway event kind.
81pub fn gateway_event_data_packets(events: &[GatewayEvent]) -> Vec<StreamPacket> {
82    events
83        .iter()
84        .map(|event| StreamPacket::data(gateway_event_data_kind(), event.to_expr()))
85        .collect()
86}
87
88/// Decodes [`GatewayEventData`] from a stream data packet.
89pub fn gateway_event_data_from_packet(packet: &StreamPacket) -> Result<GatewayEventData> {
90    GatewayEventData::from_packet(packet)
91}
92
93/// Accumulates server-sent-event (SSE) framed bytes for a streamed response.
94#[derive(Clone, Debug, Default)]
95pub struct StreamSink {
96    body: Vec<u8>,
97    done: bool,
98}
99
100impl StreamSink {
101    /// Creates an empty sink.
102    pub fn new() -> Self {
103        Self::default()
104    }
105
106    /// Appends one `data:` SSE frame carrying the JSON `value`.
107    pub fn event(&mut self, value: Value) -> Result<()> {
108        let line = serde_json::to_vec(&value).map_err(|err| {
109            Error::Eval(format!("openai codec failed to encode SSE chunk: {err}"))
110        })?;
111        self.body.extend_from_slice(b"data: ");
112        self.body.extend_from_slice(&line);
113        self.body.extend_from_slice(b"\n\n");
114        Ok(())
115    }
116
117    /// Appends the terminating `data: [DONE]` frame once, if not already done.
118    pub fn done(&mut self) {
119        if !self.done {
120            self.body.extend_from_slice(b"data: [DONE]\n\n");
121            self.done = true;
122        }
123    }
124
125    /// Finalizes the stream (emitting `[DONE]`) and returns the framed bytes.
126    pub fn into_bytes(mut self) -> Vec<u8> {
127        self.done();
128        self.body
129    }
130}
131
132/// Encodes a sequence of gateway events as SSE bytes for the given surface,
133/// using `response_id` and `created_at_ms` to stamp the emitted chunks.
134pub fn encode_gateway_events_sse(
135    events: &[GatewayEvent],
136    surface: OpenAiSseSurface,
137    response_id: &str,
138    created_at_ms: u64,
139) -> Result<Vec<u8>> {
140    let packets = gateway_event_data_packets(events);
141    let events = packets
142        .iter()
143        .map(gateway_event_data_from_packet)
144        .collect::<Result<Vec<_>>>()?;
145    let model = model_from_event_data(&events).unwrap_or_else(|| "fixture/echo".to_owned());
146    let mut sink = StreamSink::new();
147    match surface {
148        OpenAiSseSurface::Responses => {
149            for event in &events {
150                if let Some(chunk) = responses_chunk(event, response_id, created_at_ms, &model)? {
151                    sink.event(chunk)?;
152                }
153            }
154        }
155        OpenAiSseSurface::Chat => {
156            for event in &events {
157                if let Some(chunk) = chat_chunk(event, response_id, created_at_ms, &model)? {
158                    sink.event(chunk)?;
159                }
160            }
161        }
162    }
163    Ok(sink.into_bytes())
164}
165
166fn responses_chunk(
167    event: &GatewayEventData,
168    response_id: &str,
169    created_at_ms: u64,
170    model: &str,
171) -> Result<Option<Value>> {
172    Ok(match event.kind().name.as_ref() {
173        "request-start" => Some(json!({
174            "type": "response.created",
175            "response": response_stub(response_id, created_at_ms, model, "created"),
176        })),
177        "plan-start" => Some(json!({
178            "type": "response.metadata",
179            "sequence": event.sequence(),
180        })),
181        "model-start" => Some(json!({
182            "type": "response.in_progress",
183            "response": response_stub(response_id, created_at_ms, model, "in_progress"),
184        })),
185        "delta" => Some(json!({
186            "type": "response.output_text.delta",
187            "delta": string_payload(event.payload())?,
188        })),
189        "usage" => Some(json!({
190            "type": "response.usage",
191            "usage": usage_json(event.payload())?,
192        })),
193        "error" => Some(json!({
194            "type": "error",
195            "error": error_json(event.payload()),
196        })),
197        "final" => Some(json!({
198            "type": "response.completed",
199            "response": final_response_json(event.payload(), response_id, created_at_ms)?,
200        })),
201        _ => None,
202    })
203}
204
205fn chat_chunk(
206    event: &GatewayEventData,
207    response_id: &str,
208    created_at_ms: u64,
209    model: &str,
210) -> Result<Option<Value>> {
211    Ok(match event.kind().name.as_ref() {
212        "model-start" => Some(chat_choice_chunk(
213            response_id,
214            created_at_ms,
215            model,
216            json!({"role": "assistant"}),
217            Value::Null,
218        )),
219        "delta" => Some(chat_choice_chunk(
220            response_id,
221            created_at_ms,
222            model,
223            json!({"content": string_payload(event.payload())?}),
224            Value::Null,
225        )),
226        "usage" => Some(json!({
227            "id": response_id,
228            "object": "chat.completion.chunk",
229            "created": created_at_ms / 1000,
230            "model": model,
231            "choices": [],
232            "usage": usage_json(event.payload())?,
233        })),
234        "error" => Some(json!({
235            "type": "error",
236            "error": error_json(event.payload()),
237        })),
238        "final" => Some(chat_choice_chunk(
239            response_id,
240            created_at_ms,
241            model,
242            json!({}),
243            json!(finish_reason(event.payload()).unwrap_or_else(|_| "stop".to_owned())),
244        )),
245        _ => None,
246    })
247}
248
249fn response_stub(response_id: &str, created_at_ms: u64, model: &str, status: &str) -> Value {
250    json!({
251        "id": response_id,
252        "object": "response",
253        "created_at": created_at_ms / 1000,
254        "status": status,
255        "model": model,
256    })
257}
258
259fn chat_choice_chunk(
260    response_id: &str,
261    created_at_ms: u64,
262    model: &str,
263    delta: Value,
264    finish_reason: Value,
265) -> Value {
266    json!({
267        "id": response_id,
268        "object": "chat.completion.chunk",
269        "created": created_at_ms / 1000,
270        "model": model,
271        "choices": [{
272            "index": 0,
273            "delta": delta,
274            "finish_reason": finish_reason,
275        }],
276    })
277}
278
279fn final_response_json(expr: &Expr, response_id: &str, created_at_ms: u64) -> Result<Value> {
280    let bytes = encode_openai_responses_response(expr, response_id, created_at_ms)?;
281    serde_json::from_slice(&bytes).map_err(|err| {
282        Error::Eval(format!(
283            "openai codec failed to decode final response chunk: {err}"
284        ))
285    })
286}
287
288fn model_from_event_data(events: &[GatewayEventData]) -> Option<String> {
289    events.iter().find_map(|event| {
290        if event.kind().name.as_ref() == "model-start" {
291            string_payload(event.payload()).ok()
292        } else if event.kind().name.as_ref() == "final" {
293            string_field(event.payload(), "model").ok()
294        } else {
295            None
296        }
297    })
298}
299
300fn finish_reason(expr: &Expr) -> Result<String> {
301    symbol_field(expr, "stop-reason")
302}
303
304fn usage_json(expr: &Expr) -> Result<Value> {
305    let Expr::Map(fields) = expr else {
306        return Err(Error::Eval(
307            "openai SSE usage payload must be a map".to_owned(),
308        ));
309    };
310    let prompt = optional_u64_field(fields, "input-tokens")?.unwrap_or(0);
311    let completion = optional_u64_field(fields, "output-tokens")?.unwrap_or(0);
312    let total = optional_u64_field(fields, "total-tokens")?.unwrap_or(prompt + completion);
313    Ok(json!({
314        "prompt_tokens": prompt,
315        "completion_tokens": completion,
316        "total_tokens": total,
317    }))
318}
319
320fn error_json(expr: &Expr) -> Value {
321    match expr {
322        Expr::String(message) => json!({"message": message}),
323        other => json!({"message": format!("{other:?}")}),
324    }
325}
326
327fn string_payload(expr: &Expr) -> Result<String> {
328    match expr {
329        Expr::String(text) => Ok(text.clone()),
330        other => Err(Error::Eval(format!(
331            "openai SSE event payload must be a string, found {other:?}"
332        ))),
333    }
334}
335
336fn string_field(expr: &Expr, key: &str) -> Result<String> {
337    match map_field(expr, key)? {
338        Expr::String(text) => Ok(text.clone()),
339        _ => Err(Error::Eval(format!(
340            "openai SSE field {key} must be a string"
341        ))),
342    }
343}
344
345fn symbol_field(expr: &Expr, key: &str) -> Result<String> {
346    match map_field(expr, key)? {
347        Expr::Symbol(symbol) => Ok(symbol.name.as_ref().to_owned()),
348        _ => Err(Error::Eval(format!(
349            "openai SSE field {key} must be a symbol"
350        ))),
351    }
352}
353
354fn symbol_value_field(expr: &Expr, key: &str) -> Result<Symbol> {
355    match map_field(expr, key)? {
356        Expr::Symbol(symbol) => Ok(symbol.clone()),
357        _ => Err(Error::Eval(format!(
358            "openai SSE field {key} must be a symbol"
359        ))),
360    }
361}
362
363fn map_field<'a>(expr: &'a Expr, key: &str) -> Result<&'a Expr> {
364    let Expr::Map(entries) = expr else {
365        return Err(Error::Eval("openai SSE payload must be a map".to_owned()));
366    };
367    sim_value::access::entry_field(entries, key)
368        .ok_or_else(|| Error::Eval(format!("openai SSE payload missing {key}")))
369}
370
371fn optional_u64_field(entries: &[(Expr, Expr)], key: &str) -> Result<Option<u64>> {
372    let Some(value) = entries.iter().find_map(|(field, value)| match field {
373        Expr::Symbol(symbol) if symbol.name.as_ref() == key => Some(value),
374        _ => None,
375    }) else {
376        return Ok(None);
377    };
378    match value {
379        Expr::Number(number) => number
380            .canonical
381            .parse::<u64>()
382            .map(Some)
383            .map_err(|err| Error::Eval(format!("openai SSE invalid {key}: {err}"))),
384        Expr::String(text) => text
385            .parse::<u64>()
386            .map(Some)
387            .map_err(|err| Error::Eval(format!("openai SSE invalid {key}: {err}"))),
388        _ => Err(Error::Eval(format!(
389            "openai SSE field {key} must be a number"
390        ))),
391    }
392}
sim_lib_openai_server/codec_openai/streaming.rs

sim_lib_openai_server/codec_openai/
streaming.rs
