cherenkov/server/
response.rs

1//! Endpoint-specific JSON and SSE framing, independent of token generation.
2
3use super::{
4    ApiKind, MODEL,
5    http::{respond, sse},
6    tool_call::WireToolCall,
7};
8use anyhow::Result;
9use serde_json::{Value, json};
10use std::io::Write;
11
12pub(super) struct Response<'a, W: Write> {
13    out: &'a mut W,
14    kind: ApiKind,
15    id: String,
16    created: u64,
17    streaming: bool,
18    include_usage: bool,
19}
20
21impl<'a, W: Write> Response<'a, W> {
22    pub(super) fn for_writer(
23        out: &'a mut W,
24        kind: ApiKind,
25        id: &str,
26        created: u64,
27        streaming: bool,
28        include_usage: bool,
29    ) -> Self {
30        Self {
31            out,
32            kind,
33            id: id.to_owned(),
34            created,
35            streaming,
36            include_usage,
37        }
38    }
39
40    pub(super) fn start(&mut self) -> Result<()> {
41        if self.streaming {
42            write!(
43                self.out,
44                "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nCache-Control: no-cache\r\nX-Request-ID: {}\r\nConnection: close\r\n\r\n",
45                self.id
46            )?;
47
48            if self.kind == ApiKind::Chat {
49                let mut choice = self.delta_choice(Some(""), None);
50                choice["delta"]["role"] = json!("assistant");
51
52                self.send_chunk(choice)?;
53            }
54        }
55
56        Ok(())
57    }
58
59    pub(super) fn fail(&mut self, message: &str) -> Result<()> {
60        let body = json!({"error":{"message":message,"type":"server_error"}});
61
62        if self.streaming {
63            sse(self.out, &body)?;
64            self.out.write_all(b"data: [DONE]\n\n")?;
65
66            return Ok(());
67        }
68
69        respond(self.out, 500, &body)
70    }
71
72    pub(super) fn text(&mut self, text: &str) -> Result<()> {
73        if self.streaming {
74            self.send_chunk(self.delta_choice(Some(text), None))?;
75        }
76
77        Ok(())
78    }
79
80    pub(super) fn finish(
81        &mut self,
82        text: &str,
83        tool_calls: Option<&[WireToolCall]>,
84        finish_reason: &str,
85        usage: Value,
86    ) -> Result<()> {
87        if self.streaming {
88            // Tool calls arrive as one indexed delta chunk, then an empty
89            // delta carries the terminal reason.
90            if self.kind == ApiKind::Chat {
91                if let Some(calls) = tool_calls.filter(|calls| !calls.is_empty()) {
92                    let items = calls
93                        .iter()
94                        .enumerate()
95                        .map(|(index, call)| {
96                            json!({
97                                "id": call.id,
98                                "type": "function",
99                                "function": {"name": call.name, "arguments": call.arguments},
100                                "index": index,
101                            })
102                        })
103                        .collect::<Vec<Value>>();
104
105                    let mut choice = self.delta_choice(None, None);
106                    choice["delta"]["tool_calls"] = Value::Array(items);
107
108                    self.send_chunk(choice)?;
109                }
110            }
111
112            self.send_chunk(self.delta_choice(None, Some(finish_reason)))?;
113
114            if self.include_usage {
115                let mut body = self.envelope(vec![]);
116                body["usage"] = usage;
117
118                sse(self.out, &body)?;
119            }
120
121            self.out.write_all(b"data: [DONE]\n\n")?;
122            self.out.flush()?;
123        } else {
124            let mut choice = self.choice(Some(finish_reason));
125
126            match self.kind {
127                ApiKind::Chat => {
128                    let calls = tool_calls.filter(|calls| !calls.is_empty());
129                    let mut message = json!({"role":"assistant", "content":text});
130
131                    if let Some(calls) = calls {
132                        if text.is_empty() {
133                            message["content"] = Value::Null;
134                        }
135
136                        message["tool_calls"] = Value::Array(
137                            calls
138                                .iter()
139                                .map(|call| {
140                                    json!({
141                                        "id": call.id,
142                                        "type": "function",
143                                        "function": {"name": call.name, "arguments": call.arguments},
144                                    })
145                                })
146                                .collect::<Vec<Value>>(),
147                        );
148                    }
149
150                    choice["message"] = message;
151                }
152                ApiKind::Completion => choice["text"] = json!(text),
153            }
154
155            let mut body = self.envelope(vec![choice]);
156            body["usage"] = usage;
157
158            respond(self.out, 200, &body)?;
159        }
160
161        Ok(())
162    }
163
164    fn choice(&self, finish_reason: Option<&str>) -> Value {
165        let mut choice = json!({"index":0, "finish_reason":finish_reason});
166
167        if self.kind == ApiKind::Completion {
168            choice["logprobs"] = Value::Null;
169        }
170
171        choice
172    }
173
174    // A final chat chunk has an empty delta, not an empty content field.
175    fn delta_choice(&self, text: Option<&str>, finish_reason: Option<&str>) -> Value {
176        let mut choice = self.choice(finish_reason);
177
178        match self.kind {
179            ApiKind::Chat => {
180                let mut delta = json!({});
181
182                if let Some(text) = text {
183                    delta["content"] = json!(text);
184                }
185
186                choice["delta"] = delta;
187            }
188            ApiKind::Completion => choice["text"] = json!(text.unwrap_or("")),
189        }
190
191        choice
192    }
193
194    fn envelope(&self, choices: Vec<Value>) -> Value {
195        let object = match (self.kind, self.streaming) {
196            (ApiKind::Chat, true) => "chat.completion.chunk",
197            (ApiKind::Chat, false) => "chat.completion",
198            (ApiKind::Completion, _) => "text_completion",
199        };
200
201        json!({"id":self.id, "object":object, "created":self.created, "model":MODEL, "choices":choices})
202    }
203
204    fn send_chunk(&mut self, choice: Value) -> Result<()> {
205        sse(self.out, &self.envelope(vec![choice]))
206    }
207}
208
209#[cfg(test)]
210#[path = "../../tests/unit/server/response.rs"]
211mod tests;
cherenkov/server/response.rs

cherenkov/server/
response.rs
