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

kernel/profiles/
context_budget.rs
