kernel/profiles/
context_budget.rs

1//! Context-window budgeting: estimating prompt token cost, deciding whether a
2//! request fits, and clamping the completion length to what remains.
3
4use crate::profiles::configuration::normalized_param_values;
5use crate::records::{JsonValue, ModelRecord, RuntimeId, SourceKind};
6
7/// The number of tokens reserved for the completion when deciding if a prompt
8/// fits the window, wherever the window can afford it. A window too small for
9/// it reserves half of itself instead, so a model declaring fewer than twice
10/// this is still usable.
11pub const COMPLETION_FLOOR: i64 = 256;
12
13/// The outcome of assessing a prompt against a context window.
14#[derive(Debug, Clone, PartialEq, Eq)]
15pub enum Verdict {
16    /// The prompt fits; `clamped_max_tokens` is the completion length capped to
17    /// what the window leaves free.
18    Fits {
19        /// The completion length after clamping to the remaining window.
20        clamped_max_tokens: Option<i64>,
21    },
22    /// The prompt does not fit the window.
23    Exceeds {
24        /// The estimated token cost of the prompt.
25        estimated: i64,
26        /// The window it exceeded.
27        window: i64,
28    },
29}
30
31/// A rough token estimate for `characters` of text (about four characters per
32/// token).
33pub fn estimated_tokens(characters: i64) -> i64 {
34    (characters + 3) / 4
35}
36
37/// The tokens held back for the reply out of a window of `window`: the
38/// [`COMPLETION_FLOOR`], or half the window when the window cannot afford the
39/// floor. A model declaring fewer than twice the floor would otherwise be
40/// refused every prompt it was ever given, however short.
41///
42/// At least one token is always held back, so a prompt only ever fits a window
43/// that has room to answer it.
44fn reserved_for_completion(window: i64) -> i64 {
45    COMPLETION_FLOOR.min(window / 2).max(1)
46}
47
48/// Decide whether a prompt of `prompt_characters` fits `window`, and clamp the
49/// completion length to the space that remains.
50pub fn assess(prompt_characters: i64, window: i64, requested_max_tokens: Option<i64>) -> Verdict {
51    let estimated = estimated_tokens(prompt_characters);
52    if window <= 0 || estimated + reserved_for_completion(window) > window {
53        return Verdict::Exceeds { estimated, window };
54    }
55    let available = window - estimated;
56    let clamped = requested_max_tokens.unwrap_or(available).min(available);
57    Verdict::Fits {
58        clamped_max_tokens: Some(clamped),
59    }
60}
61
62/// The context window Apple's built-in model serves: fixed at 4096 tokens per
63/// session. The single source for every consumer — the runtime facade's
64/// budget, the apple-foundation scanner's record hint, and this module's
65/// policy — so the value cannot drift apart.
66pub const BUILTIN_CONTEXT_WINDOW: i64 = 4096;
67
68/// The effective context window for `record`, honoring a per-runtime policy. A
69/// built-in model has a fixed window; other runtimes derive it from the record's
70/// declared context length (and, for Ollama, a caller override).
71pub fn effective_window(
72    record: &ModelRecord,
73    requested_context_length: Option<i64>,
74) -> Option<i64> {
75    if record.source.kind == SourceKind::builtin() {
76        return Some(BUILTIN_CONTEXT_WINDOW);
77    }
78    let window = record_policy_window(record, requested_context_length)?;
79    (window > 0).then_some(window)
80}
81
82fn record_policy_window(record: &ModelRecord, requested: Option<i64>) -> Option<i64> {
83    let id = record.runtime.id.as_ref()?;
84    if *id == RuntimeId::ollama() {
85        requested.or(record.context_length)
86    } else if *id == RuntimeId::llama_cpp()
87        || *id == RuntimeId::mlx_swift()
88        || *id == RuntimeId::mlx_lm()
89    {
90        record.context_length
91    } else {
92        None
93    }
94}
95
96/// The user-set context length stored in the record's parameter values, if any.
97pub fn stored_context_length(record: &ModelRecord) -> Option<i64> {
98    normalized_param_values(record)
99        .get("context_length")
100        .and_then(JsonValue::as_i64)
101}
102
103/// Count the characters a chat/completion payload will send: the `content` of
104/// each message plus a top-level `prompt` string, if present.
105pub fn prompt_characters(payload: &JsonValue) -> i64 {
106    let JsonValue::Object(object) = payload else {
107        return 0;
108    };
109    let mut total = 0i64;
110    if let Some(JsonValue::Array(messages)) = object.get("messages") {
111        for message in messages {
112            if let JsonValue::Object(fields) = message
113                && let Some(JsonValue::String(content)) = fields.get("content")
114            {
115                total += content.chars().count() as i64;
116            }
117        }
118    }
119    if let Some(JsonValue::String(prompt)) = object.get("prompt") {
120        total += prompt.chars().count() as i64;
121    }
122    total
123}
kernel/profiles/context_budget.rs

kernel/profiles/
context_budget.rs
