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 effective context window for `record`, honoring a per-runtime policy. A
50/// built-in model has a fixed window; other runtimes derive it from the record's
51/// declared context length (and, for Ollama, a caller override).
52pub fn effective_window(
53    record: &ModelRecord,
54    requested_context_length: Option<i64>,
55) -> Option<i64> {
56    if record.source.kind == SourceKind::builtin() {
57        return Some(4096);
58    }
59    let window = record_policy_window(record, requested_context_length)?;
60    (window > 0).then_some(window)
61}
62
63fn record_policy_window(record: &ModelRecord, requested: Option<i64>) -> Option<i64> {
64    let id = record.runtime.id.as_ref()?;
65    if *id == RuntimeId::ollama() {
66        requested.or(record.context_length)
67    } else if *id == RuntimeId::llama_cpp()
68        || *id == RuntimeId::mlx_swift()
69        || *id == RuntimeId::mlx_lm()
70    {
71        record.context_length
72    } else {
73        None
74    }
75}
76
77/// The user-set context length stored in the record's parameter values, if any.
78pub fn stored_context_length(record: &ModelRecord) -> Option<i64> {
79    normalized_param_values(record)
80        .get("context_length")
81        .and_then(JsonValue::as_i64)
82}
83
84/// Count the characters a chat/completion payload will send: the `content` of
85/// each message plus a top-level `prompt` string, if present.
86pub fn prompt_characters(payload: &JsonValue) -> i64 {
87    let JsonValue::Object(object) = payload else {
88        return 0;
89    };
90    let mut total = 0i64;
91    if let Some(JsonValue::Array(messages)) = object.get("messages") {
92        for message in messages {
93            if let JsonValue::Object(fields) = message
94                && let Some(JsonValue::String(content)) = fields.get("content")
95            {
96                total += content.chars().count() as i64;
97            }
98        }
99    }
100    if let Some(JsonValue::String(prompt)) = object.get("prompt") {
101        total += prompt.chars().count() as i64;
102    }
103    total
104}
kernel/profiles/context_budget.rs

kernel/profiles/
context_budget.rs
