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#[cfg(test)]
137mod tests {
138    use super::*;
139    use crate::types::{ChatMessage, LlmParams};
140
141    #[test]
142    fn budget_is_max_tokens_when_it_fits() {
143        assert_eq!(plan_budget(100, 50, 2048), Ok(50));
144    }
145
146    #[test]
147    fn budget_is_clamped_to_the_room_left_in_the_context() {
148        assert_eq!(plan_budget(2000, 100, 2048), Ok(48));
149    }
150
151    #[test]
152    fn a_prompt_that_fills_the_context_is_refused() {
153        let err = plan_budget(2048, 10, 2048).unwrap_err();
154        assert_eq!(
155            err,
156            ContextOverflow {
157                prompt_tokens: 2048,
158                n_ctx: 2048
159            }
160        );
161        assert_eq!(
162            err.to_string(),
163            "prompt is 2048 tokens but the context holds 2048; raise contextSize or shorten the prompt"
164        );
165    }
166
167    #[test]
168    fn a_zero_budget_still_generates_one_token() {
169        assert_eq!(plan_budget(10, 0, 2048), Ok(1));
170    }
171
172    #[test]
173    fn context_is_the_configured_size_capped_by_training() {
174        assert_eq!(effective_context(Some(32768), 262144), 32768);
175        assert_eq!(effective_context(None, 262144), DEFAULT_CONTEXT_SIZE);
176        assert_eq!(effective_context(Some(32768), 4096), 4096);
177        assert_eq!(
178            effective_context(Some(32768), 0),
179            32768,
180            "unknown training size"
181        );
182    }
183
184    #[test]
185    fn the_system_prompt_leads_the_messages() {
186        let params = LlmParams {
187            system: Some("be brief".into()),
188            messages: vec![ChatMessage {
189                role: "user".into(),
190                content: "hi".into(),
191            }],
192            ..Default::default()
193        };
194        let msgs = chat_messages(&params);
195        assert_eq!(msgs[0].role, "system");
196        assert_eq!(msgs[0].content, "be brief");
197        assert_eq!(msgs[1].content, "hi");
198        let bare = LlmParams {
199            messages: params.messages.clone(),
200            ..Default::default()
201        };
202        assert_eq!(chat_messages(&bare).len(), 1);
203    }
204
205    #[test]
206    fn bos_is_added_only_when_the_model_wants_it_and_the_template_did_not() {
207        assert!(should_add_bos(true, "hello", "<s>"));
208        assert!(!should_add_bos(true, "<s>hello", "<s>"));
209        assert!(!should_add_bos(false, "hello", "<s>"));
210        assert!(
211            should_add_bos(true, "hello", ""),
212            "an empty BOS text never matches"
213        );
214    }
215
216    #[test]
217    fn finish_is_length_when_the_budget_ran_out() {
218        assert_eq!(finish_for(50, 50, false), Finish::Length);
219        assert_eq!(finish_for(12, 50, true), Finish::Stop);
220        assert_eq!(Finish::Stop.as_str(), "stop");
221        assert_eq!(Finish::Length.as_str(), "length");
222    }
223
224    #[test]
225    fn reasoning_is_split_from_the_answer() {
226        assert_eq!(
227            split_reasoning("<think>\nhmm\n</think>\n\n{\"a\":1}"),
228            ("{\"a\":1}".to_string(), Some("hmm".to_string()))
229        );
230        assert_eq!(split_reasoning("  plain  "), ("plain".to_string(), None));
231        assert_eq!(
232            split_reasoning("still thinking</think>"),
233            (String::new(), Some("still thinking".to_string()))
234        );
235    }
236
237    #[test]
238    fn completion_json_is_openai_shaped_with_real_counts() {
239        let json = completion_json("m", "<think>x</think>answer", 9649, 133, Finish::Stop, 2549);
240        assert_eq!(json["object"], "chat.completion");
241        assert_eq!(json["model"], "m");
242        assert_eq!(json["choices"][0]["message"]["role"], "assistant");
243        assert_eq!(json["choices"][0]["message"]["content"], "answer");
244        assert_eq!(json["choices"][0]["message"]["reasoning_content"], "x");
245        assert_eq!(json["choices"][0]["finish_reason"], "stop");
246        assert_eq!(json["usage"]["prompt_tokens"], 9649);
247        assert_eq!(json["usage"]["completion_tokens"], 133);
248        assert_eq!(json["usage"]["total_tokens"], 9782);
249        assert_eq!(json["elapsed_ms"], 2549);
250    }
251
252    #[test]
253    fn completion_json_omits_absent_reasoning() {
254        let json = completion_json("m", "answer", 1, 1, Finish::Length, 1);
255        assert!(json["choices"][0]["message"]
256            .get("reasoning_content")
257            .is_none());
258        assert_eq!(json["choices"][0]["finish_reason"], "length");
259    }
260}
studio_worker/engine/llm_core.rs

studio_worker/engine/
llm_core.rs
