perspt_sdk/prompt/
budget.rs

1//! Deterministic prompt budget fitting (Definition 4, resolved decision 11).
2//!
3//! Optional sections drop in increasing priority order, ties broken by
4//! section id, until the estimate fits the token budget. Required sections
5//! never drop. If the required sections alone exceed the budget,
6//! compilation fails closed. The dropped-section set is therefore a pure
7//! function of the inputs; it enters the program digest and the ledger.
8
9use crate::error::{Result, SdkError};
10
11use super::accountant::TokenAccountantRef;
12use super::section::{PromptSectionId, RenderedSection};
13
14/// The outcome of one deterministic fit.
15#[derive(Debug, Clone, PartialEq)]
16pub struct BudgetFit {
17    /// Surviving sections, in their original declared order.
18    pub kept: Vec<RenderedSection>,
19    /// Dropped section ids, in drop order (ascending priority, ties by id).
20    pub dropped: Vec<PromptSectionId>,
21}
22
23/// Fit sections to `token_budget` under `accountant`.
24pub fn fit_budget(
25    sections: Vec<RenderedSection>,
26    accountant: &TokenAccountantRef,
27    token_budget: u64,
28) -> Result<BudgetFit> {
29    let required_cost: u64 = sections
30        .iter()
31        .filter(|section| section.required)
32        .map(|section| accountant.count_message(&section.content))
33        .sum();
34    if required_cost > token_budget {
35        return Err(SdkError::Domain(format!(
36            "required sections alone cost {required_cost} tokens over the \
37             {token_budget}-token budget; compilation fails closed"
38        )));
39    }
40
41    // Optional sections in drop order: ascending priority, ties by id.
42    let mut drop_order: Vec<usize> = sections
43        .iter()
44        .enumerate()
45        .filter(|(_, section)| !section.required)
46        .map(|(index, _)| index)
47        .collect();
48    drop_order.sort_by(|&a, &b| {
49        sections[a]
50            .priority
51            .cmp(&sections[b].priority)
52            .then_with(|| sections[a].id.cmp(&sections[b].id))
53    });
54
55    let mut total: u64 = sections
56        .iter()
57        .map(|section| accountant.count_message(&section.content))
58        .sum();
59    let mut dropped_indices = Vec::new();
60    let mut drop_queue = drop_order.into_iter();
61    while total > token_budget {
62        let Some(index) = drop_queue.next() else {
63            // Unreachable: required-only cost already fit.
64            return Err(SdkError::Domain(
65                "budget fitting exhausted optional sections without fitting".into(),
66            ));
67        };
68        total -= accountant.count_message(&sections[index].content);
69        dropped_indices.push(index);
70    }
71
72    let dropped: Vec<PromptSectionId> = dropped_indices
73        .iter()
74        .map(|&index| sections[index].id.clone())
75        .collect();
76    let kept = sections
77        .into_iter()
78        .enumerate()
79        .filter(|(index, _)| !dropped_indices.contains(index))
80        .map(|(_, section)| section)
81        .collect();
82    Ok(BudgetFit { kept, dropped })
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88    use crate::prompt::section::{PromptMessageRole, PromptSectionVersion};
89
90    fn section(id: &str, required: bool, priority: u16, bytes: usize) -> RenderedSection {
91        RenderedSection {
92            id: PromptSectionId(id.into()),
93            version: PromptSectionVersion(1),
94            role: PromptMessageRole::System,
95            required,
96            priority,
97            content_hash: format!("sha256:{id}"),
98            content: "x".repeat(bytes),
99        }
100    }
101
102    /// The spec's worked example: at a 1,800-token budget the priority-40
103    /// optional drops first and the priority-60 optional then fits.
104    #[test]
105    fn the_spec_worked_example_drops_only_repository_hints() {
106        let accountant = TokenAccountantRef::approx_bytes_v1();
107        // Bytes chosen so count_message ≈ the spec's token column.
108        let sections = vec![
109            section("s/role", true, 0, 336),                 // ~120
110            section("s/tool_protocol", true, 0, 1_896),      // ~640
111            section("s/correction", true, 0, 2_076),         // ~700
112            section("s/repository_hints", false, 40, 1_356), // ~460
113            section("s/style_guidance", false, 60, 636),     // ~220
114            section("s/output_contract", true, 0, 306),      // ~110
115        ];
116        let fit = fit_budget(sections, &accountant, 1_800).unwrap();
117        assert_eq!(
118            fit.dropped,
119            vec![PromptSectionId("s/repository_hints".into())]
120        );
121        let kept: Vec<&str> = fit.kept.iter().map(|s| s.id.0.as_str()).collect();
122        assert_eq!(
123            kept,
124            [
125                "s/role",
126                "s/tool_protocol",
127                "s/correction",
128                "s/style_guidance",
129                "s/output_contract"
130            ]
131        );
132    }
133
134    #[test]
135    fn required_overflow_fails_closed_and_ties_break_by_id() {
136        let accountant = TokenAccountantRef::approx_bytes_v1();
137        assert!(fit_budget(vec![section("s/huge", true, 0, 30_000)], &accountant, 100).is_err());
138
139        let sections = vec![
140            section("s/b", false, 10, 3_000),
141            section("s/a", false, 10, 3_000),
142            section("s/base", true, 0, 30),
143        ];
144        let fit = fit_budget(sections, &accountant, 1_200).unwrap();
145        assert_eq!(fit.dropped.first().unwrap().0, "s/a", "ties break by id");
146    }
147
148    #[test]
149    fn fitting_is_deterministic() {
150        let accountant = TokenAccountantRef::approx_bytes_v1();
151        let sections = || {
152            vec![
153                section("s/base", true, 0, 300),
154                section("s/opt1", false, 20, 900),
155                section("s/opt2", false, 30, 900),
156            ]
157        };
158        let first = fit_budget(sections(), &accountant, 450).unwrap();
159        let second = fit_budget(sections(), &accountant, 450).unwrap();
160        assert_eq!(first, second);
161    }
162}
perspt_sdk/prompt/budget.rs

perspt_sdk/prompt/
budget.rs
