Skip to main content

gproxy_transform/envelope/
synthesize.rs

1use bytes::Bytes;
2use gproxy_protocol::{ContentGenerationKind as Kind, StreamFraming};
3
4use super::SseFrame;
5use crate::TransformError;
6
7/// Convert one complete content-generation response into a strict stream.
8pub fn synthesize_response(
9    kind: Kind,
10    body: Bytes,
11    framing: StreamFraming,
12) -> Result<Vec<Bytes>, TransformError> {
13    match kind {
14        Kind::OpenAiChat => {
15            require_framing(framing, &[StreamFraming::Sse])?;
16            let events = crate::typed::synthesize::openai_chat(serde_json::from_slice(&body)?);
17            encode(events, framing, None, true)
18        }
19        Kind::OpenAiResponses | Kind::OpenAiResponsesWebSocket => {
20            require_framing(framing, &[StreamFraming::Sse, StreamFraming::WebSocket])?;
21            let events = crate::typed::synthesize::openai_responses(serde_json::from_slice(&body)?);
22            encode(events, framing, Some(response_name), false)
23        }
24        Kind::ClaudeMessages => {
25            require_framing(framing, &[StreamFraming::Sse])?;
26            let events = crate::typed::synthesize::claude(serde_json::from_slice(&body)?);
27            encode(events, framing, Some(claude_name), false)
28        }
29        Kind::GeminiGenerateContent => {
30            require_framing(framing, &[StreamFraming::Sse, StreamFraming::JsonArray])?;
31            let events = crate::typed::synthesize::gemini(serde_json::from_slice(&body)?);
32            encode(events, framing, None, false)
33        }
34        #[cfg(not(feature = "exhaustive"))]
35        _ => Err(TransformError::unsupported(
36            "content generation kind",
37            "unrecognized external variant",
38        )),
39    }
40}
41
42/// What keeps a client waiting on a synthesized stream from timing out while
43/// the upstream is still producing the object: `None` when the framing has
44/// no idle frame.
45pub fn synthesize_keepalive(kind: Kind, framing: StreamFraming) -> Option<Bytes> {
46    match framing {
47        StreamFraming::JsonArray => Some(Bytes::from_static(b"\n")),
48        StreamFraming::WebSocket => None,
49        StreamFraming::Sse => Some(match kind {
50            Kind::ClaudeMessages => SseFrame::encode(Some("ping"), r#"{"type":"ping"}"#),
51            Kind::OpenAiChat
52            | Kind::OpenAiResponses
53            | Kind::OpenAiResponsesWebSocket
54            | Kind::GeminiGenerateContent => Bytes::from_static(b": keep-alive\n\n"),
55            #[cfg(not(feature = "exhaustive"))]
56            _ => Bytes::from_static(b": keep-alive\n\n"),
57        }),
58    }
59}
60
61/// The failure a synthesized stream ends with once its headers are already
62/// out: each protocol's own terminal error event.
63pub fn synthesize_error(
64    kind: Kind,
65    framing: StreamFraming,
66    message: &str,
67) -> Result<Vec<Bytes>, TransformError> {
68    let (event, value, done) = match kind {
69        Kind::ClaudeMessages => (
70            Some("error"),
71            serde_json::json!({"type":"error","error":{"type":"api_error","message":message}}),
72            false,
73        ),
74        Kind::OpenAiChat => (
75            None,
76            serde_json::json!({"error":{"type":"upstream_error","message":message}}),
77            true,
78        ),
79        Kind::OpenAiResponses | Kind::OpenAiResponsesWebSocket => (
80            Some("error"),
81            serde_json::json!({"type":"error","code":"upstream_error","message":message,"param":null,"sequence_number":0}),
82            false,
83        ),
84        Kind::GeminiGenerateContent => (
85            None,
86            serde_json::json!({"error":{"code":502,"status":"UNAVAILABLE","message":message}}),
87            false,
88        ),
89        #[cfg(not(feature = "exhaustive"))]
90        _ => {
91            return Err(TransformError::unsupported(
92                "content generation kind",
93                "unrecognized external variant",
94            ));
95        }
96    };
97    match framing {
98        StreamFraming::Sse => {
99            let mut output = vec![SseFrame::typed(event, &value)?];
100            if done {
101                output.push(SseFrame::encode(None, "[DONE]"));
102            }
103            Ok(output)
104        }
105        StreamFraming::JsonArray => Ok(vec![Bytes::from(serde_json::to_vec(&[value])?)]),
106        StreamFraming::WebSocket => Ok(vec![Bytes::from(serde_json::to_vec(&value)?)]),
107    }
108}
109
110fn require_framing(
111    framing: StreamFraming,
112    supported: &[StreamFraming],
113) -> Result<(), TransformError> {
114    if supported.contains(&framing) {
115        Ok(())
116    } else {
117        Err(TransformError::shape(
118            "synthetic stream",
119            "framing is not valid for the target protocol",
120        ))
121    }
122}
123
124fn encode<T: serde::Serialize>(
125    events: Vec<T>,
126    framing: StreamFraming,
127    event_name: Option<fn(&T) -> Option<&str>>,
128    done: bool,
129) -> Result<Vec<Bytes>, TransformError> {
130    match framing {
131        StreamFraming::Sse => {
132            let mut output = events
133                .iter()
134                .map(|event| SseFrame::typed(event_name.and_then(|name| name(event)), event))
135                .collect::<Result<Vec<_>, _>>()?;
136            if done {
137                output.push(SseFrame::encode(None, "[DONE]"));
138            }
139            Ok(output)
140        }
141        StreamFraming::JsonArray => Ok(vec![Bytes::from(serde_json::to_vec(&events)?)]),
142        StreamFraming::WebSocket => events
143            .iter()
144            .map(|event| {
145                serde_json::to_vec(event)
146                    .map(Bytes::from)
147                    .map_err(Into::into)
148            })
149            .collect(),
150    }
151}
152
153fn response_name(event: &gproxy_protocol::openai::ResponseStreamEvent) -> Option<&str> {
154    event.event_name()
155}
156
157fn claude_name(event: &gproxy_protocol::claude::StreamEvent) -> Option<&str> {
158    event.event_name()
159}