loadmaster_probe/
client.rs

1//! OpenAI-compatible client with raw capture and streaming reassembly.
2
3use std::time::{Duration, Instant};
4
5use futures_util::StreamExt;
6use serde::Serialize;
7use serde_json::{Value, json};
8
9/// One request and what came back.
10#[derive(Debug, Clone, Serialize)]
11pub struct Exchange {
12    pub request: Value,
13    pub status: u16,
14    /// The response body exactly as received.
15    #[serde(skip)]
16    pub raw: Vec<u8>,
17    /// The response in non-streamed shape (`choices[0].message`), reassembled from deltas
18    /// for streamed requests; for errors, the error body.
19    pub response: Value,
20    pub stream: Option<StreamInfo>,
21    pub latency_ms: u64,
22    /// Time to the first streamed chunk.
23    pub ttft_ms: Option<u64>,
24}
25
26#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize)]
27pub struct StreamInfo {
28    pub chunks: usize,
29    /// `data: [DONE]` arrived.
30    pub done: bool,
31    /// Lines that were not valid SSE JSON.
32    pub bad_lines: usize,
33}
34
35#[derive(Clone)]
36pub struct Client {
37    http: reqwest::Client,
38    base: String,
39    api_key: Option<String>,
40}
41
42impl Client {
43    pub fn new(base: &str, timeout: Duration, api_key: Option<String>) -> Self {
44        Client {
45            http: reqwest::Client::builder()
46                .timeout(timeout)
47                .build()
48                .expect("default client"),
49            base: base
50                .trim_end_matches('/')
51                .trim_end_matches("/v1")
52                .to_owned(),
53            api_key,
54        }
55    }
56
57    pub fn base(&self) -> &str {
58        &self.base
59    }
60
61    pub async fn get(&self, path: &str) -> Option<(u16, reqwest::header::HeaderMap, Value)> {
62        let mut req = self.http.get(format!("{}{path}", self.base));
63        if let Some(k) = &self.api_key {
64            req = req.bearer_auth(k);
65        }
66        let r = req.timeout(Duration::from_secs(5)).send().await.ok()?;
67        let status = r.status().as_u16();
68        let headers = r.headers().clone();
69        let body = r.json().await.unwrap_or(Value::Null);
70        Some((status, headers, body))
71    }
72
73    /// POST a JSON body; the status (0 when the request failed) and the JSON response.
74    pub async fn post(&self, path: &str, body: Value) -> (u16, Value) {
75        let mut req = self.http.post(format!("{}{path}", self.base)).json(&body);
76        if let Some(k) = &self.api_key {
77            req = req.bearer_auth(k);
78        }
79        match req.send().await {
80            Ok(r) => {
81                let status = r.status().as_u16();
82                (status, r.json().await.unwrap_or(Value::Null))
83            }
84            Err(e) => (0, json!({ "error": { "message": e.to_string() } })),
85        }
86    }
87
88    /// POST `/v1/chat/completions`.
89    pub async fn chat(&self, body: Value) -> Exchange {
90        let stream = body["stream"].as_bool().unwrap_or(false);
91        let start = Instant::now();
92        let mut req = self
93            .http
94            .post(format!("{}/v1/chat/completions", self.base))
95            .json(&body);
96        if let Some(k) = &self.api_key {
97            req = req.bearer_auth(k);
98        }
99        let resp = match req.send().await {
100            Ok(r) => r,
101            Err(e) => {
102                return Exchange {
103                    request: body,
104                    status: 0,
105                    raw: Vec::new(),
106                    response: json!({ "error": { "message": format!("request failed: {e}") } }),
107                    stream: None,
108                    latency_ms: start.elapsed().as_millis() as u64,
109                    ttft_ms: None,
110                };
111            }
112        };
113        let status = resp.status().as_u16();
114        let sse = resp
115            .headers()
116            .get(reqwest::header::CONTENT_TYPE)
117            .and_then(|v| v.to_str().ok())
118            .is_some_and(|v| v.starts_with("text/event-stream"));
119        let mut raw = Vec::new();
120        let mut ttft = None;
121        let mut bytes = resp.bytes_stream();
122        while let Some(chunk) = bytes.next().await {
123            match chunk {
124                Ok(c) => {
125                    if ttft.is_none() {
126                        ttft = Some(start.elapsed().as_millis() as u64);
127                    }
128                    raw.extend_from_slice(&c);
129                }
130                Err(_) => break, // a cut connection is a truncated stream
131            }
132        }
133        let latency_ms = start.elapsed().as_millis() as u64;
134        let (response, info) = if stream && sse {
135            let (r, i) = reassemble(&raw);
136            (r, Some(i))
137        } else {
138            (
139                serde_json::from_slice(&raw)
140                    .unwrap_or_else(|_| json!({ "raw": String::from_utf8_lossy(&raw) })),
141                None,
142            )
143        };
144        Exchange {
145            request: body,
146            status,
147            raw,
148            response,
149            stream: info,
150            latency_ms,
151            ttft_ms: if stream { ttft } else { None },
152        }
153    }
154}
155
156/// Fold SSE chunks into one non-streamed completion: content and reasoning concatenated,
157/// tool calls merged by `index` (id and name from the first delta, arguments appended).
158/// An `error` event (llama-server sends one when its parser rejects the output mid-stream)
159/// is kept as the response's top-level `error`, as in a non-streamed error body.
160pub fn reassemble(raw: &[u8]) -> (Value, StreamInfo) {
161    let text = String::from_utf8_lossy(raw);
162    let mut info = StreamInfo::default();
163    let mut content = String::new();
164    let mut reasoning = String::new();
165    let mut calls: Vec<Value> = Vec::new();
166    let mut finish = Value::Null;
167    let mut usage = Value::Null;
168    let mut model = Value::Null;
169    let mut error = Value::Null;
170    for line in text.lines() {
171        let Some(data) = line.strip_prefix("data:").map(str::trim) else {
172            continue;
173        };
174        if data == "[DONE]" {
175            info.done = true;
176            continue;
177        }
178        let Ok(chunk) = serde_json::from_str::<Value>(data) else {
179            info.bad_lines += 1;
180            continue;
181        };
182        info.chunks += 1;
183        if !chunk["error"].is_null() {
184            error = chunk["error"].clone();
185        }
186        if model.is_null() {
187            model = chunk["model"].clone();
188        }
189        if !chunk["usage"].is_null() {
190            usage = chunk["usage"].clone();
191        }
192        let choice = &chunk["choices"][0];
193        if !choice["finish_reason"].is_null() {
194            finish = choice["finish_reason"].clone();
195        }
196        let delta = &choice["delta"];
197        if let Some(c) = delta["content"].as_str() {
198            content.push_str(c);
199        }
200        for key in ["reasoning_content", "reasoning"] {
201            if let Some(r) = delta[key].as_str() {
202                reasoning.push_str(r);
203            }
204        }
205        for tc in delta["tool_calls"].as_array().into_iter().flatten() {
206            let idx = tc["index"].as_u64().unwrap_or(calls.len() as u64) as usize;
207            while calls.len() <= idx {
208                calls.push(json!({ "id": null, "type": "function", "function": { "name": null, "arguments": "" } }));
209            }
210            let slot = &mut calls[idx];
211            if !tc["id"].is_null() {
212                slot["id"] = tc["id"].clone();
213            }
214            if let Some(n) = tc["function"]["name"].as_str() {
215                slot["function"]["name"] = json!(n);
216            }
217            if let Some(a) = tc["function"]["arguments"].as_str() {
218                let prev = slot["function"]["arguments"]
219                    .as_str()
220                    .unwrap_or("")
221                    .to_owned();
222                slot["function"]["arguments"] = json!(prev + a);
223            }
224        }
225    }
226    let mut message = json!({ "role": "assistant", "content": if content.is_empty() { Value::Null } else { json!(content) } });
227    if !reasoning.is_empty() {
228        message["reasoning_content"] = json!(reasoning);
229    }
230    if !calls.is_empty() {
231        message["tool_calls"] = Value::Array(calls);
232    }
233    let mut response = json!({ "object": "chat.completion", "model": model, "choices": [{ "index": 0, "message": message, "finish_reason": finish }], "usage": usage });
234    if !error.is_null() {
235        response["error"] = error;
236    }
237    (response, info)
238}
239
240#[cfg(test)]
241mod tests {
242    use super::*;
243
244    #[test]
245    fn reassembles_split_tool_calls() {
246        let raw = concat!(
247            "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\"}}]}\n\n",
248            "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"c0\",\"function\":{\"name\":\"get\",\"arguments\":\"\"}}]}}]}\n\n",
249            "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"{\\\"a\\\"\"}}]}}]}\n\n",
250            "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\":1}\"}}]}}]}\n\n",
251            "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"tool_calls\"}]}\n\n",
252            "data: [DONE]\n\n",
253        );
254        let (r, info) = reassemble(raw.as_bytes());
255        assert!(info.done);
256        assert_eq!(info.chunks, 5);
257        let call = &r["choices"][0]["message"]["tool_calls"][0];
258        assert_eq!(call["id"], "c0");
259        assert_eq!(call["function"]["arguments"], "{\"a\":1}");
260        assert_eq!(r["choices"][0]["finish_reason"], "tool_calls");
261        assert!(r.get("error").is_none());
262    }
263
264    #[test]
265    fn keeps_a_streamed_error_event() {
266        // Captured from llama-server b11669 serving llama3.2:1b.
267        let raw = concat!(
268            "data: {\"choices\":[{\"finish_reason\":null,\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":null}}],\"object\":\"chat.completion.chunk\"}\n\n",
269            "data: {\"error\":{\"code\":500,\"message\":\"The model produced output that does not match the expected peg-native format\",\"type\":\"server_error\"}}\n\n",
270        );
271        let (r, info) = reassemble(raw.as_bytes());
272        assert!(!info.done);
273        assert_eq!(r["error"]["code"], 500);
274        assert!(
275            r["error"]["message"]
276                .as_str()
277                .unwrap()
278                .contains("peg-native format")
279        );
280    }
281}
loadmaster_probe/client.rs

loadmaster_probe/
client.rs
