Skip to main content

gproxy_transform/transform/stream_adapter/
synthesize.rs

1use serde_json::{Value, json};
2
3use super::{ContentGenerationKind, SseFrame};
4
5/// Turn one complete response into the smallest useful SSE event sequence.
6pub fn synthesize_sse(
7    kind: ContentGenerationKind,
8    body: &[u8],
9) -> Result<Vec<u8>, crate::transform::TransformError> {
10    let value: Value = serde_json::from_slice(body).map_err(|error| {
11        crate::transform::TransformError::InvalidInput {
12            reason: format!("synthetic stream response is not JSON: {error}"),
13        }
14    })?;
15    let mut out = String::new();
16    match kind {
17        ContentGenerationKind::OpenAiChatCompletions => synthesize_chat(&value, &mut out),
18        ContentGenerationKind::OpenAiResponses
19        | ContentGenerationKind::OpenAiResponsesWebSocket => synthesize_responses(&value, &mut out),
20        ContentGenerationKind::ClaudeMessages => synthesize_claude(&value, &mut out),
21        ContentGenerationKind::GeminiGenerateContent => {
22            out.push_str(&SseFrame::data(value.to_string()).encode());
23        }
24        _ => {
25            unreachable!("new non-exhaustive protocol variant requires a lockstep transform update")
26        }
27    }
28    Ok(out.into_bytes())
29}
30
31fn synthesize_chat(response: &Value, out: &mut String) {
32    let choices = response
33        .get("choices")
34        .and_then(Value::as_array)
35        .map(|choices| {
36            choices
37                .iter()
38                .map(|choice| {
39                    let mut delta = choice.get("message").cloned().unwrap_or_else(|| json!({}));
40                    if let Some(object) = delta.as_object_mut() {
41                        object.remove("annotations");
42                    }
43                    json!({
44                        "index":choice.get("index").cloned().unwrap_or_else(|| json!(0)),
45                        "delta":delta,
46                        "finish_reason":choice.get("finish_reason").cloned().unwrap_or(Value::Null),
47                        "logprobs":choice.get("logprobs").cloned().unwrap_or(Value::Null),
48                    })
49                })
50                .collect::<Vec<_>>()
51        })
52        .unwrap_or_default();
53    let mut chunk = response.clone();
54    if let Some(object) = chunk.as_object_mut() {
55        object.insert("object".into(), json!("chat.completion.chunk"));
56        object.insert("choices".into(), Value::Array(choices));
57    }
58    out.push_str(&SseFrame::data(chunk.to_string()).encode());
59    out.push_str(&SseFrame::data("[DONE]").encode());
60}
61
62fn synthesize_responses(response: &Value, out: &mut String) {
63    let mut started = response.clone();
64    if let Some(object) = started.as_object_mut() {
65        object.insert("status".into(), json!("in_progress"));
66        object.insert("output".into(), json!([]));
67    }
68    push_named(out, json!({"type":"response.created","response":started}));
69
70    for (output_index, item) in response
71        .get("output")
72        .and_then(Value::as_array)
73        .into_iter()
74        .flatten()
75        .enumerate()
76    {
77        let mut added_item = item.clone();
78        if let Some(object) = added_item.as_object_mut() {
79            object.insert("status".into(), json!("in_progress"));
80        }
81        push_named(
82            out,
83            json!({"type":"response.output_item.added","output_index":output_index,"item":added_item}),
84        );
85        let item_id = item
86            .get("id")
87            .cloned()
88            .unwrap_or_else(|| json!(format!("item_{output_index}")));
89        match item.get("type").and_then(Value::as_str) {
90            Some("message") => synthesize_response_message(out, item, output_index, &item_id),
91            Some("function_call") => {
92                let arguments = item.get("arguments").cloned().unwrap_or_else(|| json!(""));
93                push_named(
94                    out,
95                    json!({"type":"response.function_call_arguments.delta","item_id":item_id,"output_index":output_index,"delta":arguments}),
96                );
97                push_named(
98                    out,
99                    json!({"type":"response.function_call_arguments.done","item_id":item_id,"output_index":output_index,"name":item.get("name"),"arguments":arguments}),
100                );
101            }
102            Some("custom_tool_call") => {
103                let input = item.get("input").cloned().unwrap_or_else(|| json!(""));
104                push_named(
105                    out,
106                    json!({"type":"response.custom_tool_call_input.delta","item_id":item_id,"output_index":output_index,"delta":input}),
107                );
108                push_named(
109                    out,
110                    json!({"type":"response.custom_tool_call_input.done","item_id":item_id,"output_index":output_index,"input":input}),
111                );
112            }
113            Some("reasoning") => synthesize_response_reasoning(out, item, output_index, &item_id),
114            _ => {}
115        }
116        push_named(
117            out,
118            json!({"type":"response.output_item.done","output_index":output_index,"item":item}),
119        );
120    }
121    push_named(
122        out,
123        json!({"type":"response.completed","response":response}),
124    );
125}
126
127fn synthesize_response_message(
128    out: &mut String,
129    item: &Value,
130    output_index: usize,
131    item_id: &Value,
132) {
133    for (content_index, part) in item
134        .get("content")
135        .and_then(Value::as_array)
136        .into_iter()
137        .flatten()
138        .enumerate()
139    {
140        push_named(
141            out,
142            json!({"type":"response.content_part.added","item_id":item_id,"output_index":output_index,"content_index":content_index,"part":part}),
143        );
144        let (delta_type, done_type, field) = match part.get("type").and_then(Value::as_str) {
145            Some("refusal") => ("response.refusal.delta", "response.refusal.done", "refusal"),
146            _ => (
147                "response.output_text.delta",
148                "response.output_text.done",
149                "text",
150            ),
151        };
152        let text = part.get(field).cloned().unwrap_or_else(|| json!(""));
153        push_named(
154            out,
155            json!({"type":delta_type,"item_id":item_id,"output_index":output_index,"content_index":content_index,"delta":text}),
156        );
157        push_named(
158            out,
159            json!({"type":done_type,"item_id":item_id,"output_index":output_index,"content_index":content_index,(field):text}),
160        );
161        push_named(
162            out,
163            json!({"type":"response.content_part.done","item_id":item_id,"output_index":output_index,"content_index":content_index,"part":part}),
164        );
165    }
166}
167
168fn synthesize_response_reasoning(
169    out: &mut String,
170    item: &Value,
171    output_index: usize,
172    item_id: &Value,
173) {
174    for (content_index, part) in item
175        .get("content")
176        .and_then(Value::as_array)
177        .into_iter()
178        .flatten()
179        .enumerate()
180    {
181        if let Some(text) = part.get("text") {
182            push_named(
183                out,
184                json!({"type":"response.reasoning_text.delta","item_id":item_id,"output_index":output_index,"content_index":content_index,"delta":text}),
185            );
186            push_named(
187                out,
188                json!({"type":"response.reasoning_text.done","item_id":item_id,"output_index":output_index,"content_index":content_index,"text":text}),
189            );
190        }
191    }
192}
193
194fn synthesize_claude(response: &Value, out: &mut String) {
195    let mut message = response.clone();
196    if let Some(object) = message.as_object_mut() {
197        object.insert("content".into(), json!([]));
198        object.insert("stop_reason".into(), Value::Null);
199        object.insert("stop_sequence".into(), Value::Null);
200    }
201    push_named(out, json!({"type":"message_start","message":message}));
202    for (index, block) in response
203        .get("content")
204        .and_then(Value::as_array)
205        .into_iter()
206        .flatten()
207        .enumerate()
208    {
209        let mut start = block.clone();
210        if let Some(object) = start.as_object_mut() {
211            match block.get("type").and_then(Value::as_str) {
212                Some("text") => {
213                    object.insert("text".into(), json!(""));
214                }
215                Some("thinking") => {
216                    object.insert("thinking".into(), json!(""));
217                }
218                Some("tool_use") => {
219                    object.insert("input".into(), json!({}));
220                }
221                _ => {}
222            }
223        }
224        push_named(
225            out,
226            json!({"type":"content_block_start","index":index,"content_block":start}),
227        );
228        match block.get("type").and_then(Value::as_str) {
229            Some("text") => push_named(
230                out,
231                json!({"type":"content_block_delta","index":index,"delta":{"type":"text_delta","text":block.get("text").cloned().unwrap_or_else(|| json!(""))}}),
232            ),
233            Some("thinking") => {
234                push_named(
235                    out,
236                    json!({"type":"content_block_delta","index":index,"delta":{"type":"thinking_delta","thinking":block.get("thinking").cloned().unwrap_or_else(|| json!(""))}}),
237                );
238                if let Some(signature) = block.get("signature") {
239                    push_named(
240                        out,
241                        json!({"type":"content_block_delta","index":index,"delta":{"type":"signature_delta","signature":signature}}),
242                    );
243                }
244            }
245            Some("tool_use") => push_named(
246                out,
247                json!({"type":"content_block_delta","index":index,"delta":{"type":"input_json_delta","partial_json":block.get("input").cloned().unwrap_or_else(|| json!({})).to_string()}}),
248            ),
249            _ => {}
250        }
251        push_named(out, json!({"type":"content_block_stop","index":index}));
252    }
253    push_named(
254        out,
255        json!({"type":"message_delta","delta":{"stop_reason":response.get("stop_reason").cloned().unwrap_or(Value::Null),"stop_sequence":response.get("stop_sequence").cloned().unwrap_or(Value::Null)},"usage":response.get("usage").cloned().unwrap_or_else(|| json!({}))}),
256    );
257    push_named(out, json!({"type":"message_stop"}));
258}
259
260fn push_named(out: &mut String, event: Value) {
261    let name = event
262        .get("type")
263        .and_then(Value::as_str)
264        .unwrap_or("message");
265    out.push_str(&SseFrame::event(name, event.to_string()).encode());
266}