tea_context/
budget.rs

1use crate::{ContextError, ContextErrorCode};
2
3/// Fixed separator between accepted prompt segments.
4pub const PROMPT_SEPARATOR: &str = "\n\n";
5/// Fixed marker appended to deterministically shortened segments.
6pub const TRUNCATION_MARKER: &str = "[truncated]";
7
8/// Exact byte and conservative token limits for one compiled prompt.
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub struct PromptBudget {
11    max_bytes: usize,
12    max_estimated_tokens: usize,
13}
14
15impl PromptBudget {
16    /// Creates non-zero bounded prompt limits.
17    ///
18    /// # Errors
19    ///
20    /// Returns an error for zero values or limits above 16 MiB / safe integer.
21    pub fn new(max_bytes: usize, max_estimated_tokens: usize) -> Result<Self, ContextError> {
22        if max_bytes == 0
23            || max_bytes > 16 * 1024 * 1024
24            || max_estimated_tokens == 0
25            || u64::try_from(max_estimated_tokens)
26                .map_or(true, |value| value > tea_protocol::MAX_SAFE_INTEGER)
27        {
28            return Err(ContextError::new(
29                ContextErrorCode::InvalidValue,
30                "prompt budget is invalid",
31            ));
32        }
33        Ok(Self {
34            max_bytes,
35            max_estimated_tokens,
36        })
37    }
38    /// Returns exact output byte limit.
39    #[must_use]
40    pub const fn max_bytes(self) -> usize {
41        self.max_bytes
42    }
43    /// Returns conservative estimated-token limit.
44    #[must_use]
45    pub const fn max_estimated_tokens(self) -> usize {
46        self.max_estimated_tokens
47    }
48}
49
50/// Deterministic conservative token estimate `ceil(utf8_bytes / 3)`.
51#[must_use]
52pub const fn estimate_tokens(bytes: usize) -> usize {
53    bytes.saturating_add(2) / 3
54}
55
56pub(crate) fn effective_remaining_bytes(budget: PromptBudget, used_bytes: usize) -> usize {
57    let byte_remaining = budget.max_bytes.saturating_sub(used_bytes);
58    let token_capacity = budget
59        .max_estimated_tokens
60        .saturating_mul(3)
61        .saturating_sub(used_bytes);
62    byte_remaining.min(token_capacity)
63}
64
65pub(crate) fn truncate(content: &str, maximum: usize) -> Option<String> {
66    if maximum < TRUNCATION_MARKER.len() {
67        return None;
68    }
69    let content_limit = maximum - TRUNCATION_MARKER.len();
70    let mut boundary = content_limit.min(content.len());
71    while boundary > 0 && !content.is_char_boundary(boundary) {
72        boundary -= 1;
73    }
74    let mut output = content[..boundary].to_owned();
75    output.push_str(TRUNCATION_MARKER);
76    Some(output)
77}
tea_context/budget.rs

tea_context/
budget.rs
