cerno_host/
openai.rs

1//! Adapter for anything that speaks OpenAI's `/v1/chat/completions`: vLLM, llama.cpp's
2//! `llama-server`, LM Studio, OpenAI itself.
3//!
4//! The response shape is the same everywhere, so there is one parser. What differs is which
5//! sampling fields a runtime accepts, and that is all a [`Flavour`] carries.
6
7use crate::{FirstTokenDistribution, FirstTokenRequest, HostCapabilities, HostError, ModelHost};
8use async_trait::async_trait;
9use serde_json::{Map, Value, json};
10use std::collections::HashSet;
11use std::sync::Mutex;
12use std::time::{Duration, Instant};
13
14/// OpenAI rejects more than this, and it is vLLM's default `--max-logprobs`. Past 20 there are
15/// no labels to observe anyway.
16const OPENAI_MAX_TOP_LOGPROBS: usize = 20;
17
18/// Which OpenAI-compatible runtime is on the other end.
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum Flavour {
21    /// Standard fields only. OpenAI answers 400 to a field it does not know, so nothing
22    /// runtime-specific can be sent — `top_k` and `min_p` included.
23    ///
24    /// Measured behind Ollama's `/v1` on `gemma4:26b-a4b-it-q4_K_M`: this body returns the same
25    /// distribution as the native path with its pinned options, to the third decimal.
26    Generic,
27    Vllm,
28    LlamaCpp,
29    LmStudio,
30}
31
32impl Flavour {
33    fn name(self) -> &'static str {
34        match self {
35            Flavour::Generic => "openai",
36            Flavour::Vllm => "vllm",
37            Flavour::LlamaCpp => "llamacpp",
38            Flavour::LmStudio => "lmstudio",
39        }
40    }
41
42    /// Sampling settings on top of the standard ones, for the same reason as Ollama's
43    /// `SamplingOptions::REQUIRED`: a runtime that applies its sampling transform before reporting
44    /// logprobs hands back a truncated, renormalised distribution unless every filter is off.
45    ///
46    /// * vLLM reports raw logprobs by default, but pinning `top_k: -1` (its "disabled") costs
47    ///   nothing and survives a server started with a different `--logprobs-mode`.
48    /// * llama.cpp disables `top_k` with `0`, and `post_sampling_probs: false` asks for the
49    ///   distribution before the sampler chain.
50    fn sampling(self) -> Value {
51        match self {
52            Flavour::Generic => json!({}),
53            Flavour::Vllm => json!({"top_k": -1, "min_p": 0.0}),
54            Flavour::LlamaCpp => json!({"top_k": 0, "min_p": 0.0, "post_sampling_probs": false}),
55            Flavour::LmStudio => json!({"top_k": 0, "min_p": 0.0}),
56        }
57    }
58
59    /// Fields that switch the reasoning preamble off. Without them the first token is a template
60    /// control token — measured on `gemma4:e2b-it-qat` behind Ollama's `/v1`: `<|channel>` at
61    /// `-0.03`, the answer at `-3.5`. `reasoning_effort: "none"` is the standard switch and fixed
62    /// it there; `chat_template_kwargs` is how vLLM and llama.cpp reach Qwen-style templates.
63    fn no_thinking(self) -> Value {
64        match self {
65            Flavour::Vllm | Flavour::LlamaCpp => json!({
66                "reasoning_effort": "none",
67                "chat_template_kwargs": {"enable_thinking": false},
68            }),
69            Flavour::Generic | Flavour::LmStudio => json!({"reasoning_effort": "none"}),
70        }
71    }
72}
73
74pub struct OpenAiCompatHost {
75    client: reqwest::Client,
76    base_url: String,
77    api_key: Option<String>,
78    flavour: Flavour,
79    timeout: Duration,
80    /// Models that have refused the thinking switch once. They are asked without it from then
81    /// on, so the refusal costs one extra round trip per model rather than one per question.
82    refused_thinking: Mutex<HashSet<String>>,
83}
84
85impl OpenAiCompatHost {
86    /// `base_url` includes the version prefix, e.g. `http://localhost:8000/v1`.
87    pub fn new(
88        base_url: impl Into<String>,
89        api_key: Option<String>,
90        flavour: Flavour,
91        timeout: Duration,
92    ) -> Result<Self, HostError> {
93        let client = reqwest::Client::builder()
94            .timeout(timeout)
95            .build()
96            .map_err(|e| HostError::Unavailable(e.to_string()))?;
97        Ok(Self {
98            client,
99            base_url: crate::base_url(&base_url.into())?,
100            api_key,
101            flavour,
102            timeout,
103            refused_thinking: Mutex::default(),
104        })
105    }
106
107    fn body(&self, req: &FirstTokenRequest, think_off: bool) -> Value {
108        let mut messages = Vec::with_capacity(2);
109        if let Some(system) = req.system.as_deref() {
110            messages.push(json!({"role": "system", "content": system}));
111        }
112        messages.push(json!({"role": "user", "content": req.user}));
113
114        let mut body = json!({
115            "model": req.model,
116            "messages": messages,
117            "stream": false,
118            "logprobs": true,
119            "top_logprobs": req.top_logprobs.min(OPENAI_MAX_TOP_LOGPROBS),
120            "max_tokens": 1,
121            "temperature": 1.0,
122            "top_p": 1.0,
123        });
124        let object = body.as_object_mut().expect("body is an object");
125        merge(object, self.flavour.sampling());
126        if think_off {
127            merge(object, self.flavour.no_thinking());
128        }
129        body
130    }
131
132    async fn post(&self, body: &Value) -> Result<(u16, String), HostError> {
133        let mut request = self
134            .client
135            .post(format!("{}/chat/completions", self.base_url))
136            .json(body);
137        if let Some(key) = self.api_key.as_deref() {
138            request = request.bearer_auth(key);
139        }
140
141        let response = request
142            .send()
143            .await
144            .map_err(|e| HostError::transport(&e, self.timeout))?;
145
146        let status = response.status().as_u16();
147        let text = response
148            .text()
149            .await
150            .map_err(|e| HostError::transport(&e, self.timeout))?;
151        Ok((status, text))
152    }
153}
154
155fn merge(into: &mut Map<String, Value>, extra: Value) {
156    if let Value::Object(extra) = extra {
157        into.extend(extra);
158    }
159}
160
161/// Whether a failed response is the runtime refusing the thinking switch — OpenAI does for any
162/// model without a reasoning mode. The request then goes out again without it.
163fn rejects_thinking(body: &str) -> bool {
164    let lower = body.to_ascii_lowercase();
165    lower.contains("reasoning_effort")
166        || lower.contains("chat_template_kwargs")
167        || lower.contains("enable_thinking")
168}
169
170#[async_trait]
171impl ModelHost for OpenAiCompatHost {
172    fn capabilities(&self) -> HostCapabilities {
173        HostCapabilities {
174            max_top_logprobs: OPENAI_MAX_TOP_LOGPROBS,
175        }
176    }
177
178    fn name(&self) -> &str {
179        self.flavour.name()
180    }
181
182    async fn first_token(
183        &self,
184        req: FirstTokenRequest,
185    ) -> Result<FirstTokenDistribution, HostError> {
186        let started = Instant::now();
187
188        let think_off = !self
189            .refused_thinking
190            .lock()
191            .expect("lock is never poisoned")
192            .contains(&req.model);
193
194        let (mut status, mut text) = self.post(&self.body(&req, think_off)).await?;
195
196        if think_off && (400..500).contains(&status) && rejects_thinking(&text) {
197            tracing::debug!(
198                model = %req.model,
199                host = self.flavour.name(),
200                "host rejects the thinking switch; retrying without it, and from now on"
201            );
202            self.refused_thinking
203                .lock()
204                .expect("lock is never poisoned")
205                .insert(req.model.clone());
206            (status, text) = self.post(&self.body(&req, false)).await?;
207        }
208
209        if status >= 400 {
210            return Err(HostError::status(status, text));
211        }
212
213        parse_distribution(&text, &req.model, started.elapsed())
214    }
215}
216
217/// Pull the first position's ranked tokens out of a `/chat/completions` response.
218fn parse_distribution(
219    text: &str,
220    model: &str,
221    latency: Duration,
222) -> Result<FirstTokenDistribution, HostError> {
223    let value: Value =
224        serde_json::from_str(text).map_err(|e| HostError::Protocol(e.to_string()))?;
225    let no_logprobs = || HostError::NoLogprobs {
226        model: model.to_string(),
227    };
228
229    let ranked = value
230        .pointer("/choices/0/logprobs/content/0/top_logprobs")
231        .and_then(Value::as_array)
232        .ok_or_else(no_logprobs)?;
233
234    let mut tokens: Vec<(String, f64)> = ranked
235        .iter()
236        .filter_map(|entry| {
237            let token = entry.get("token")?.as_str()?.to_string();
238            let logprob = entry.get("logprob")?.as_f64()?;
239            Some((token, logprob))
240        })
241        .collect();
242
243    if tokens.is_empty() {
244        return Err(no_logprobs());
245    }
246
247    tokens.sort_by(|a, b| b.1.total_cmp(&a.1));
248
249    let floor = tokens
250        .last()
251        .map(|(_, lp)| *lp)
252        .expect("tokens is non-empty");
253
254    let input_tokens = value
255        .pointer("/usage/prompt_tokens")
256        .and_then(Value::as_u64)
257        .unwrap_or(0) as u32;
258
259    Ok(FirstTokenDistribution {
260        tokens,
261        floor,
262        input_tokens,
263        latency,
264    })
265}
266
267#[cfg(test)]
268mod tests {
269    use super::*;
270
271    /// A real `/v1/chat/completions` response, trimmed of `bytes`. Captured from
272    /// `gemma4:e2b-it-qat` behind Ollama's OpenAI endpoint, with `reasoning_effort: "none"`.
273    const REAL_RESPONSE: &str = r#"{
274        "id": "chatcmpl-147",
275        "object": "chat.completion",
276        "model": "gemma4:e2b-it-qat",
277        "choices": [{
278            "index": 0,
279            "message": {"role": "assistant", "content": "B"},
280            "finish_reason": "length",
281            "logprobs": {"content": [{
282                "token": "B",
283                "logprob": -0.002,
284                "top_logprobs": [
285                    {"token": "B", "logprob": -0.002},
286                    {"token": "b", "logprob": -7.721},
287                    {"token": "**", "logprob": -7.946},
288                    {"token": "", "logprob": -7.976},
289                    {"token": "C", "logprob": -8.674}
290                ]
291            }]}
292        }],
293        "usage": {"prompt_tokens": 40, "completion_tokens": 1, "total_tokens": 41}
294    }"#;
295
296    fn request(top_logprobs: usize) -> FirstTokenRequest {
297        FirstTokenRequest {
298            model: "m".into(),
299            system: Some("sys".into()),
300            user: "usr".into(),
301            top_logprobs,
302            keep_alive: Some("5m".into()),
303            offer: crate::test_offer(),
304        }
305    }
306
307    fn body(flavour: Flavour, think_off: bool) -> Value {
308        OpenAiCompatHost::new(
309            "http://localhost:8000/v1",
310            None,
311            flavour,
312            Duration::from_secs(5),
313        )
314        .unwrap()
315        .body(&request(20), think_off)
316    }
317
318    #[test]
319    fn parses_ranked_tokens_and_floor() {
320        let dist = parse_distribution(REAL_RESPONSE, "m", Duration::from_millis(7)).unwrap();
321
322        assert_eq!(dist.tokens.len(), 5);
323        assert_eq!(dist.tokens[0].0, "B");
324        assert_eq!(dist.input_tokens, 40);
325        assert_eq!(dist.logprob("C"), Some(-8.674));
326        assert_eq!(dist.logprob("A"), None);
327        assert_eq!(dist.floor, -8.674);
328    }
329
330    #[test]
331    fn ranks_tokens_even_when_the_host_reports_them_unsorted() {
332        let shuffled = r#"{"choices":[{"logprobs":{"content":[{"token":"B","top_logprobs":[
333            {"token":"A","logprob":-5.0},{"token":"B","logprob":-0.5},{"token":"C","logprob":-9.0}]}]}}]}"#;
334
335        let dist = parse_distribution(shuffled, "m", Duration::ZERO).unwrap();
336
337        assert_eq!(dist.tokens[0].0, "B");
338        assert_eq!(dist.floor, -9.0);
339    }
340
341    /// A runtime that ignores `logprobs` answers normally with `"logprobs": null`.
342    #[test]
343    fn missing_logprobs_is_an_error_naming_the_model() {
344        let body =
345            r#"{"choices":[{"message":{"role":"assistant","content":"B"},"logprobs":null}]}"#;
346
347        let err = parse_distribution(body, "tiny", Duration::ZERO).unwrap_err();
348
349        assert!(matches!(err, HostError::NoLogprobs { model } if model == "tiny"));
350    }
351
352    #[test]
353    fn empty_top_logprobs_is_treated_as_missing() {
354        let body = r#"{"choices":[{"logprobs":{"content":[{"token":"B","top_logprobs":[]}]}}]}"#;
355
356        assert!(matches!(
357            parse_distribution(body, "m", Duration::ZERO),
358            Err(HostError::NoLogprobs { .. })
359        ));
360    }
361
362    #[test]
363    fn the_standard_fields_are_the_same_for_every_flavour() {
364        for flavour in [
365            Flavour::Generic,
366            Flavour::Vllm,
367            Flavour::LlamaCpp,
368            Flavour::LmStudio,
369        ] {
370            let json = body(flavour, true);
371
372            assert_eq!(json["logprobs"], true, "{flavour:?}");
373            assert_eq!(json["top_logprobs"], 20, "{flavour:?}");
374            assert_eq!(json["max_tokens"], 1, "{flavour:?}");
375            assert_eq!(json["temperature"], 1.0, "{flavour:?}");
376            assert_eq!(json["top_p"], 1.0, "{flavour:?}");
377            assert_eq!(json["stream"], false, "{flavour:?}");
378            assert_eq!(json["reasoning_effort"], "none", "{flavour:?}");
379            assert_eq!(json["messages"][0]["role"], "system");
380            assert_eq!(json["messages"][1]["content"], "usr");
381            // keep_alive is Ollama's; nobody else knows it.
382            assert!(json.get("keep_alive").is_none(), "{flavour:?}");
383        }
384    }
385
386    /// OpenAI answers 400 to a field it does not know, so the generic body carries none.
387    #[test]
388    fn the_generic_flavour_sends_nothing_nonstandard() {
389        let json = body(Flavour::Generic, true);
390
391        for field in [
392            "top_k",
393            "min_p",
394            "post_sampling_probs",
395            "chat_template_kwargs",
396        ] {
397            assert!(json.get(field).is_none(), "{field}");
398        }
399    }
400
401    /// The sampling fields are the reason the distribution is readable; pin them verbatim.
402    #[test]
403    fn runtime_sampling_fields_are_serialised_verbatim() {
404        let vllm = body(Flavour::Vllm, true);
405        assert_eq!(vllm["top_k"], -1);
406        assert_eq!(vllm["min_p"], 0.0);
407        assert_eq!(vllm["chat_template_kwargs"]["enable_thinking"], false);
408
409        let llamacpp = body(Flavour::LlamaCpp, true);
410        assert_eq!(llamacpp["top_k"], 0);
411        assert_eq!(llamacpp["min_p"], 0.0);
412        assert_eq!(llamacpp["post_sampling_probs"], false);
413        assert_eq!(llamacpp["chat_template_kwargs"]["enable_thinking"], false);
414
415        let lmstudio = body(Flavour::LmStudio, true);
416        assert_eq!(lmstudio["top_k"], 0);
417        assert_eq!(lmstudio["min_p"], 0.0);
418    }
419
420    /// The retry drops the thinking switch and nothing else — least of all the sampling fields.
421    #[test]
422    fn the_retry_body_keeps_sampling_and_drops_only_the_thinking_switch() {
423        let json = body(Flavour::Vllm, false);
424
425        assert!(json.get("reasoning_effort").is_none());
426        assert!(json.get("chat_template_kwargs").is_none());
427        assert_eq!(json["top_k"], -1);
428    }
429
430    #[test]
431    fn top_logprobs_is_clamped_to_the_host_ceiling() {
432        let host = OpenAiCompatHost::new(
433            "http://x/v1",
434            None,
435            Flavour::Generic,
436            Duration::from_secs(5),
437        )
438        .unwrap();
439
440        assert_eq!(host.body(&request(99), true)["top_logprobs"], 20);
441    }
442
443    #[test]
444    fn recognises_the_thinking_rejection() {
445        assert!(rejects_thinking(
446            r#"{"error":{"message":"Unsupported parameter: 'reasoning_effort' is not supported with this model.","param":"reasoning_effort"}}"#
447        ));
448        assert!(!rejects_thinking(
449            r#"{"error":{"message":"model not found"}}"#
450        ));
451    }
452
453    #[tokio::test]
454    async fn the_api_key_is_sent_as_a_bearer_token() {
455        let mut server = mockito::Server::new_async().await;
456        let mock = server
457            .mock("POST", "/v1/chat/completions")
458            .match_header("authorization", "Bearer sk-test")
459            .with_body(REAL_RESPONSE)
460            .create_async()
461            .await;
462        let host = OpenAiCompatHost::new(
463            format!("{}/v1/", server.url()),
464            Some("sk-test".into()),
465            Flavour::Generic,
466            Duration::from_secs(5),
467        )
468        .unwrap();
469
470        host.first_token(request(4)).await.unwrap();
471
472        mock.assert_async().await;
473    }
474
475    #[tokio::test]
476    async fn without_a_key_no_authorization_header_is_sent() {
477        let mut server = mockito::Server::new_async().await;
478        let mock = server
479            .mock("POST", "/v1/chat/completions")
480            .match_header("authorization", mockito::Matcher::Missing)
481            .with_body(REAL_RESPONSE)
482            .create_async()
483            .await;
484        let host = OpenAiCompatHost::new(
485            format!("{}/v1", server.url()),
486            None,
487            Flavour::LlamaCpp,
488            Duration::from_secs(5),
489        )
490        .unwrap();
491
492        host.first_token(request(4)).await.unwrap();
493
494        mock.assert_async().await;
495    }
496
497    /// A model without a reasoning mode refuses the switch; the same question goes out again
498    /// without it and is answered, and later questions to that model skip the refusal.
499    #[tokio::test]
500    async fn a_rejected_thinking_switch_is_retried_without_it_and_remembered() {
501        let mut server = mockito::Server::new_async().await;
502        let refused = server
503            .mock("POST", "/v1/chat/completions")
504            .match_body(mockito::Matcher::PartialJson(
505                json!({"reasoning_effort": "none"}),
506            ))
507            .with_status(400)
508            .with_body(r#"{"error":{"message":"Unsupported parameter: 'reasoning_effort'"}}"#)
509            .expect(1)
510            .create_async()
511            .await;
512        let answered = server
513            .mock("POST", "/v1/chat/completions")
514            // Only a body without the switch is answered. Without this guard mockito would hand a
515            // repeated switch to this mock once the refusal had been used up, and the test
516            // would pass whether or not the refusal was remembered.
517            .match_request(|request| {
518                !request
519                    .utf8_lossy_body()
520                    .is_ok_and(|body| body.contains(r#""reasoning_effort""#))
521            })
522            .with_body(REAL_RESPONSE)
523            .expect(2)
524            .create_async()
525            .await;
526        let host = OpenAiCompatHost::new(
527            format!("{}/v1", server.url()),
528            None,
529            Flavour::Generic,
530            Duration::from_secs(5),
531        )
532        .unwrap();
533
534        let dist = host.first_token(request(4)).await.unwrap();
535        assert_eq!(dist.tokens[0].0, "B");
536
537        // The second question to the same model must not pay for the refusal again.
538        host.first_token(request(4)).await.unwrap();
539
540        refused.assert_async().await;
541        answered.assert_async().await;
542    }
543}
cerno_host/openai.rs

cerno_host/
openai.rs
