studio_worker/engine/
llm_core.rs

1//! Engine-independent LLM rules: context budget, BOS handling, reasoning
2//! split and the `chat.completion` response shape.  No llama.cpp here, so
3//! every platform and every CI leg tests it.
4
5use crate::types::{ChatMessage, LlmParams};
6
7/// Context window when the catalogue does not set `contextSize`.  Covers
8/// ordinary chat; long-prompt models (summarisers: ~10k-token prompts)
9/// set their own.  Safe range 2048..=the model's training context.
10pub const DEFAULT_CONTEXT_SIZE: u32 = 8192;
11
12/// Why generation stopped, as OpenAI names it.
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum Finish {
15    Stop,
16    Length,
17}
18
19impl Finish {
20    pub fn as_str(self) -> &'static str {
21        match self {
22            Self::Stop => "stop",
23            Self::Length => "length",
24        }
25    }
26}
27
28/// A prompt that leaves no room to answer.  Refused, never truncated:
29/// cutting the operator's prompt would change what they asked.
30#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
31#[error(
32    "prompt is {prompt_tokens} tokens but the context holds {n_ctx}; raise contextSize or \
33     shorten the prompt"
34)]
35pub struct ContextOverflow {
36    pub prompt_tokens: usize,
37    pub n_ctx: u32,
38}
39
40/// Tokens this request may generate: `max_tokens` (at least one), clamped
41/// to the room the prompt leaves in the context.
42pub fn plan_budget(
43    prompt_tokens: usize,
44    max_tokens: u32,
45    n_ctx: u32,
46) -> Result<u32, ContextOverflow> {
47    let room = (n_ctx as usize)
48        .checked_sub(prompt_tokens)
49        .filter(|r| *r > 0);
50    match room {
51        Some(room) => Ok(max_tokens.max(1).min(room.min(u32::MAX as usize) as u32)),
52        None => Err(ContextOverflow {
53            prompt_tokens,
54            n_ctx,
55        }),
56    }
57}
58
59/// The configured context (or the default), never beyond what the model
60/// was trained on; `trained == 0` means unknown.
61pub fn effective_context(configured: Option<u32>, trained: u32) -> u32 {
62    let wanted = configured.unwrap_or(DEFAULT_CONTEXT_SIZE);
63    if trained > 0 {
64        wanted.min(trained)
65    } else {
66        wanted
67    }
68}
69
70/// The request's messages with its separate system prompt first.
71pub fn chat_messages(params: &LlmParams) -> Vec<ChatMessage> {
72    let mut out = Vec::with_capacity(params.messages.len() + 1);
73    if let Some(system) = &params.system {
74        out.push(ChatMessage {
75            role: "system".into(),
76            content: system.clone(),
77        });
78    }
79    out.extend(params.messages.iter().cloned());
80    out
81}
82
83/// Add BOS only when the model asks for it and the rendered template has
84/// not already written it (a double BOS degrades output).
85pub fn should_add_bos(model_adds_bos: bool, prompt: &str, bos_text: &str) -> bool {
86    model_adds_bos && (bos_text.is_empty() || !prompt.starts_with(bos_text))
87}
88
89/// `length` when generation used its whole budget without an end token.
90pub fn finish_for(generated: u32, budget: u32, hit_end: bool) -> Finish {
91    if !hit_end && generated >= budget {
92        Finish::Length
93    } else {
94        Finish::Stop
95    }
96}
97
98/// Split a `<think>...</think>` block (if any) from the answer.
99pub fn split_reasoning(text: &str) -> (String, Option<String>) {
100    match text.split_once("</think>") {
101        Some((thought, answer)) => {
102            let thought = thought.trim().trim_start_matches("<think>").trim();
103            (answer.trim().to_string(), Some(thought.to_string()))
104        }
105        None => (text.trim().to_string(), None),
106    }
107}
108
109/// OpenAI `chat.completion` JSON with real token counts.
110pub fn completion_json(
111    model: &str,
112    text: &str,
113    prompt_tokens: usize,
114    completion_tokens: u32,
115    finish: Finish,
116    elapsed_ms: u64,
117) -> serde_json::Value {
118    let (content, reasoning) = split_reasoning(text);
119    let mut message = serde_json::json!({ "role": "assistant", "content": content });
120    if let Some(reasoning) = reasoning {
121        message["reasoning_content"] = reasoning.into();
122    }
123    serde_json::json!({
124        "object": "chat.completion",
125        "model": model,
126        "choices": [{ "index": 0, "message": message, "finish_reason": finish.as_str() }],
127        "usage": {
128            "prompt_tokens": prompt_tokens,
129            "completion_tokens": completion_tokens,
130            "total_tokens": prompt_tokens + completion_tokens as usize,
131        },
132        "elapsed_ms": elapsed_ms,
133    })
134}
135
136/// Releases generated text as it becomes safe to stream: text that could
137/// still turn into a stop string is held back until it either does (and is
138/// cut) or cannot.
139#[derive(Debug)]
140pub struct StopHold {
141    stops: Vec<String>,
142    out: String,
143    emitted: usize,
144    stopped: Option<usize>,
145}
146
147impl StopHold {
148    pub fn new(stops: &[String]) -> Self {
149        Self {
150            stops: stops.iter().filter(|s| !s.is_empty()).cloned().collect(),
151            out: String::new(),
152            emitted: 0,
153            stopped: None,
154        }
155    }
156
157    /// Add a piece; returns the text now safe to stream.
158    pub fn push(&mut self, piece: &str) -> String {
159        if self.stopped.is_some() {
160            return String::new();
161        }
162        let search_from = self.emitted;
163        self.out.push_str(piece);
164        if let Some(at) = self
165            .stops
166            .iter()
167            .filter_map(|s| {
168                self.out[search_from..]
169                    .find(s.as_str())
170                    .map(|i| i + search_from)
171            })
172            .min()
173        {
174            self.out.truncate(at);
175            self.stopped = Some(at);
176            return self.release(at);
177        }
178        let held = self.held_suffix();
179        self.release(self.out.len() - held)
180    }
181
182    /// Everything still held (generation ended).
183    pub fn finish(&mut self) -> String {
184        self.release(self.out.len())
185    }
186
187    /// Where a stop string cut the text, if one did.
188    pub fn stopped(&self) -> Option<usize> {
189        self.stopped
190    }
191
192    /// The whole text so far (after any stop cut).
193    pub fn text(&self) -> &str {
194        &self.out
195    }
196
197    /// Length of the longest pending suffix that is a proper prefix of a stop.
198    fn held_suffix(&self) -> usize {
199        let pending = &self.out[self.emitted..];
200        pending
201            .char_indices()
202            .map(|(i, _)| i)
203            .find(|&i| {
204                let tail = &pending[i..];
205                self.stops
206                    .iter()
207                    .any(|s| s.len() > tail.len() && s.starts_with(tail))
208            })
209            .map_or(0, |i| pending.len() - i)
210    }
211
212    fn release(&mut self, upto: usize) -> String {
213        let upto = upto.max(self.emitted);
214        let text = self.out[self.emitted..upto].to_string();
215        self.emitted = upto;
216        text
217    }
218}
219
220/// A piece of a streamed answer: the model's thinking or the answer itself.
221#[derive(Debug, Clone, PartialEq, Eq)]
222pub enum Delta {
223    Reasoning(String),
224    Content(String),
225}
226
227#[derive(Debug, Default, PartialEq, Eq)]
228enum ThinkState {
229    /// Not yet known whether the answer opens with a think block.
230    #[default]
231    Undecided,
232    Thinking,
233    Answering,
234}
235
236const THINK_OPEN: &str = "<think>";
237const THINK_CLOSE: &str = "</think>";
238
239/// Streams a `<think>...</think>` block (if the answer opens with one) as
240/// reasoning and the rest as content, the same split `completion_json`
241/// makes on a whole answer.
242#[derive(Debug, Default)]
243pub struct ThinkSplitter {
244    state: ThinkState,
245    buf: String,
246}
247
248impl ThinkSplitter {
249    pub fn push(&mut self, piece: &str) -> Vec<Delta> {
250        self.buf.push_str(piece);
251        let mut out = Vec::new();
252        loop {
253            match self.state {
254                ThinkState::Undecided => {
255                    let trimmed = self.buf.trim_start();
256                    if let Some(rest) = trimmed.strip_prefix(THINK_OPEN) {
257                        self.buf = rest.to_string();
258                        self.state = ThinkState::Thinking;
259                    } else if THINK_OPEN.starts_with(trimmed) {
260                        return out; // could still become <think>
261                    } else {
262                        self.state = ThinkState::Answering;
263                    }
264                }
265                ThinkState::Thinking => match self.buf.find(THINK_CLOSE) {
266                    Some(at) => {
267                        let thought = self.buf[..at].trim().to_string();
268                        if !thought.is_empty() {
269                            out.push(Delta::Reasoning(thought));
270                        }
271                        self.buf = self.buf[at + THINK_CLOSE.len()..].trim_start().to_string();
272                        self.state = ThinkState::Answering;
273                    }
274                    None => return out, // reasoning is sent whole, at its end
275                },
276                ThinkState::Answering => {
277                    if !self.buf.is_empty() {
278                        out.push(Delta::Content(std::mem::take(&mut self.buf)));
279                    }
280                    return out;
281                }
282            }
283        }
284    }
285
286    /// Generation ended: whatever is buffered goes out as what it was.
287    pub fn finish(&mut self) -> Vec<Delta> {
288        let rest = std::mem::take(&mut self.buf);
289        match self.state {
290            ThinkState::Thinking if !rest.trim().is_empty() => {
291                vec![Delta::Reasoning(rest.trim().to_string())]
292            }
293            ThinkState::Undecided | ThinkState::Answering if !rest.is_empty() => {
294                vec![Delta::Content(rest)]
295            }
296            _ => Vec::new(),
297        }
298    }
299}
300
301/// One streamed `chat.completion.chunk`.
302pub fn chunk_frame(model: &str, delta: &Delta) -> serde_json::Value {
303    let delta = match delta {
304        Delta::Content(t) => serde_json::json!({ "content": t }),
305        Delta::Reasoning(t) => serde_json::json!({ "reasoning_content": t }),
306    };
307    serde_json::json!({
308        "object": "chat.completion.chunk",
309        "model": model,
310        "choices": [{ "index": 0, "delta": delta, "finish_reason": null }],
311    })
312}
313
314/// The last chunk: why generation stopped, and the token counts.
315pub fn final_frame(
316    model: &str,
317    finish: Finish,
318    prompt_tokens: usize,
319    completion_tokens: u32,
320) -> serde_json::Value {
321    serde_json::json!({
322        "object": "chat.completion.chunk",
323        "model": model,
324        "choices": [{ "index": 0, "delta": {}, "finish_reason": finish.as_str() }],
325        "usage": {
326            "prompt_tokens": prompt_tokens,
327            "completion_tokens": completion_tokens,
328            "total_tokens": prompt_tokens + completion_tokens as usize,
329        },
330    })
331}
332
333/// A server-sent event carrying `value`.
334pub fn sse_event(value: &serde_json::Value) -> Vec<u8> {
335    format!("data: {value}\n\n").into_bytes()
336}
337
338/// The event that ends an OpenAI stream.
339pub const SSE_DONE: &[u8] = b"data: [DONE]\n\n";
340
341#[cfg(test)]
342mod tests {
343    use super::*;
344    use crate::types::{ChatMessage, LlmParams};
345
346    #[test]
347    fn budget_is_max_tokens_when_it_fits() {
348        assert_eq!(plan_budget(100, 50, 2048), Ok(50));
349    }
350
351    #[test]
352    fn budget_is_clamped_to_the_room_left_in_the_context() {
353        assert_eq!(plan_budget(2000, 100, 2048), Ok(48));
354    }
355
356    #[test]
357    fn a_prompt_that_fills_the_context_is_refused() {
358        let err = plan_budget(2048, 10, 2048).unwrap_err();
359        assert_eq!(
360            err,
361            ContextOverflow {
362                prompt_tokens: 2048,
363                n_ctx: 2048
364            }
365        );
366        assert_eq!(
367            err.to_string(),
368            "prompt is 2048 tokens but the context holds 2048; raise contextSize or shorten the prompt"
369        );
370    }
371
372    #[test]
373    fn a_zero_budget_still_generates_one_token() {
374        assert_eq!(plan_budget(10, 0, 2048), Ok(1));
375    }
376
377    #[test]
378    fn context_is_the_configured_size_capped_by_training() {
379        assert_eq!(effective_context(Some(32768), 262144), 32768);
380        assert_eq!(effective_context(None, 262144), DEFAULT_CONTEXT_SIZE);
381        assert_eq!(effective_context(Some(32768), 4096), 4096);
382        assert_eq!(
383            effective_context(Some(32768), 0),
384            32768,
385            "unknown training size"
386        );
387    }
388
389    #[test]
390    fn the_system_prompt_leads_the_messages() {
391        let params = LlmParams {
392            system: Some("be brief".into()),
393            messages: vec![ChatMessage {
394                role: "user".into(),
395                content: "hi".into(),
396            }],
397            ..Default::default()
398        };
399        let msgs = chat_messages(&params);
400        assert_eq!(msgs[0].role, "system");
401        assert_eq!(msgs[0].content, "be brief");
402        assert_eq!(msgs[1].content, "hi");
403        let bare = LlmParams {
404            messages: params.messages.clone(),
405            ..Default::default()
406        };
407        assert_eq!(chat_messages(&bare).len(), 1);
408    }
409
410    #[test]
411    fn bos_is_added_only_when_the_model_wants_it_and_the_template_did_not() {
412        assert!(should_add_bos(true, "hello", "<s>"));
413        assert!(!should_add_bos(true, "<s>hello", "<s>"));
414        assert!(!should_add_bos(false, "hello", "<s>"));
415        assert!(
416            should_add_bos(true, "hello", ""),
417            "an empty BOS text never matches"
418        );
419    }
420
421    #[test]
422    fn finish_is_length_when_the_budget_ran_out() {
423        assert_eq!(finish_for(50, 50, false), Finish::Length);
424        assert_eq!(finish_for(12, 50, true), Finish::Stop);
425        assert_eq!(Finish::Stop.as_str(), "stop");
426        assert_eq!(Finish::Length.as_str(), "length");
427    }
428
429    #[test]
430    fn reasoning_is_split_from_the_answer() {
431        assert_eq!(
432            split_reasoning("<think>\nhmm\n</think>\n\n{\"a\":1}"),
433            ("{\"a\":1}".to_string(), Some("hmm".to_string()))
434        );
435        assert_eq!(split_reasoning("  plain  "), ("plain".to_string(), None));
436        assert_eq!(
437            split_reasoning("still thinking</think>"),
438            (String::new(), Some("still thinking".to_string()))
439        );
440    }
441
442    #[test]
443    fn completion_json_is_openai_shaped_with_real_counts() {
444        let json = completion_json("m", "<think>x</think>answer", 9649, 133, Finish::Stop, 2549);
445        assert_eq!(json["object"], "chat.completion");
446        assert_eq!(json["model"], "m");
447        assert_eq!(json["choices"][0]["message"]["role"], "assistant");
448        assert_eq!(json["choices"][0]["message"]["content"], "answer");
449        assert_eq!(json["choices"][0]["message"]["reasoning_content"], "x");
450        assert_eq!(json["choices"][0]["finish_reason"], "stop");
451        assert_eq!(json["usage"]["prompt_tokens"], 9649);
452        assert_eq!(json["usage"]["completion_tokens"], 133);
453        assert_eq!(json["usage"]["total_tokens"], 9782);
454        assert_eq!(json["elapsed_ms"], 2549);
455    }
456
457    #[test]
458    fn completion_json_omits_absent_reasoning() {
459        let json = completion_json("m", "answer", 1, 1, Finish::Length, 1);
460        assert!(json["choices"][0]["message"]
461            .get("reasoning_content")
462            .is_none());
463        assert_eq!(json["choices"][0]["finish_reason"], "length");
464    }
465
466    fn drain(hold: &mut StopHold, pieces: &[&str]) -> String {
467        let mut out = String::new();
468        for p in pieces {
469            out.push_str(&hold.push(p));
470        }
471        out
472    }
473
474    #[test]
475    fn stop_hold_releases_text_that_cannot_start_a_stop() {
476        let mut hold = StopHold::new(&["</s>".to_string()]);
477        assert_eq!(hold.push("hello "), "hello ");
478        assert_eq!(hold.push("world"), "world");
479        assert_eq!(hold.finish(), "");
480    }
481
482    #[test]
483    fn stop_hold_never_leaks_a_stop_split_across_pieces() {
484        let mut hold = StopHold::new(&["END".to_string()]);
485        let streamed = drain(&mut hold, &["one E", "N", "D two"]);
486        assert_eq!(hold.stopped(), Some("one ".len()));
487        assert_eq!(streamed + &hold.finish(), "one ");
488    }
489
490    #[test]
491    fn stop_hold_without_stops_holds_nothing() {
492        let mut hold = StopHold::new(&[]);
493        assert_eq!(hold.push("a"), "a");
494        assert_eq!(hold.push("é"), "é");
495        assert_eq!(hold.stopped(), None);
496    }
497
498    #[test]
499    fn stop_hold_keeps_multibyte_characters_whole() {
500        let mut hold = StopHold::new(&["xyz".to_string()]);
501        let mut out = hold.push("ééé");
502        out.push_str(&hold.finish());
503        assert_eq!(out, "ééé");
504    }
505
506    fn split(pieces: &[&str]) -> (String, String) {
507        let mut s = ThinkSplitter::default();
508        let (mut reasoning, mut content) = (String::new(), String::new());
509        for p in pieces.iter().copied().map(Some).chain([None]) {
510            let deltas = match p {
511                Some(p) => s.push(p),
512                None => s.finish(),
513            };
514            for d in deltas {
515                match d {
516                    Delta::Reasoning(t) => reasoning.push_str(&t),
517                    Delta::Content(t) => content.push_str(&t),
518                }
519            }
520        }
521        (reasoning, content)
522    }
523
524    #[test]
525    fn think_splitter_passes_plain_answers_as_content() {
526        assert_eq!(
527            split(&["{\"a\"", ":1}"]),
528            (String::new(), "{\"a\":1}".into())
529        );
530    }
531
532    #[test]
533    fn think_splitter_routes_a_think_block_to_reasoning() {
534        assert_eq!(
535            split(&["<th", "ink>\nhmm", " ok</th", "ink>\n\nanswer"]),
536            ("hmm ok".into(), "answer".into())
537        );
538    }
539
540    #[test]
541    fn think_splitter_treats_text_that_only_looks_like_a_start_as_content() {
542        assert_eq!(split(&["<t", "able>"]), (String::new(), "<table>".into()));
543        assert_eq!(split(&["<th"]), (String::new(), "<th".into()));
544    }
545
546    #[test]
547    fn think_splitter_flushes_an_unfinished_think_as_reasoning() {
548        assert_eq!(split(&["<think>still"]), ("still".into(), String::new()));
549    }
550
551    #[test]
552    fn stream_frames_are_openai_chunks() {
553        let content = chunk_frame("m", &Delta::Content("hi".into()));
554        assert_eq!(
555            content,
556            serde_json::json!({
557                "object": "chat.completion.chunk",
558                "model": "m",
559                "choices": [{ "index": 0, "delta": { "content": "hi" }, "finish_reason": null }],
560            })
561        );
562        let reasoning = chunk_frame("m", &Delta::Reasoning("r".into()));
563        assert_eq!(
564            reasoning["choices"][0]["delta"],
565            serde_json::json!({ "reasoning_content": "r" })
566        );
567        let last = final_frame("m", Finish::Length, 10, 3);
568        assert_eq!(last["choices"][0]["delta"], serde_json::json!({}));
569        assert_eq!(last["choices"][0]["finish_reason"], "length");
570        assert_eq!(last["usage"]["total_tokens"], 13);
571    }
572
573    #[test]
574    fn sse_events_are_data_lines() {
575        assert_eq!(
576            sse_event(&serde_json::json!({ "a": 1 })),
577            b"data: {\"a\":1}\n\n"
578        );
579        assert_eq!(SSE_DONE, b"data: [DONE]\n\n");
580    }
581}
studio_worker/engine/llm_core.rs

studio_worker/engine/
llm_core.rs
