feat(context): tambah context::tokens dengan tiktoken-rs

Ganti tiga heuristik char-count (/3 di shortsend, /4 di loop turn, /4
di status bar) yang saling tidak konsisten dengan satu BPE tokenizer
nyata. tiktoken-rs membundel vocab lewat include_str! saat build, jadi
tidak ada akses jaringan saat runtime.
This commit is contained in:
asepharyana
2026-07-16 07:49:43 +07:00
parent 8bb697a53f
commit c219b6ec58
5 changed files with 112 additions and 0 deletions
+13
View File
@@ -0,0 +1,13 @@
//! Context management: token counting, cross-call tool-result dedup,
//! per-result compression, budget-based shaping, and shared
//! context-window resolution — replaces `runtime::shortsend`.
//!
//! No facade function here: `dedup`, `shaping`, and `tokens` are called
//! directly from each call site (the per-turn auto-compaction loop in
//! `actions::run_agent_turn`, and `Action::Compact`), matching this
//! codebase's "no DI, call modules directly" convention. An orchestration
//! layer would only serve one of the two callers generically — the
//! auto-loop already needs per-stage control to decide when to emit
//! `TurnEvent::Compacted`.
pub mod tokens;
+69
View File
@@ -0,0 +1,69 @@
//! Unified token-count estimation for context-window budgeting.
//!
//! Flow: text -> `tiktoken_rs::o200k_base_singleton()` (BPE vocab embedded
//! in the binary via `include_str!`, no network access) -> `encode_ordinary`
//! -> token count.
//!
//! Why: replaces three independent char-count heuristics that disagreed
//! with each other (`/3` in the old `shortsend.rs`, `/4` in the turn
//! loop, `/4` again in the status bar) with one real BPE tokenizer.
//! `o200k_base` is an approximation for non-OpenAI providers but is far
//! closer than a flat byte-per-token guess; it's only used for the
//! 85%/95% budget thresholds, not for billing-accurate counts.
use crate::dto::chat::message::ChatMessage;
/// Count tokens in a single string under `o200k_base`.
///
/// Return: the BPE token count for `text`. `encode_ordinary` (not
/// `encode`/`encode_with_special_tokens`) is used deliberately — message
/// content that happens to contain a special-token-shaped substring
/// (e.g. literal text `<|endoftext|>` pasted by a user) must be counted
/// as ordinary text, not interpreted as a control token.
pub fn count_tokens(text: &str) -> usize {
tiktoken_rs::o200k_base_singleton().encode_ordinary(text).len()
}
/// Count tokens in a `ChatMessage`'s text content.
///
/// Return: 0 for a message with no `content` (e.g. an assistant message
/// that only carries `tool_calls`).
pub fn count_message_tokens(msg: &ChatMessage) -> usize {
msg.content.as_deref().map_or(0, count_tokens)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::dto::chat::message::ChatMessage;
#[test]
fn empty_string_has_zero_tokens() {
assert_eq!(count_tokens(""), 0);
}
#[test]
fn known_short_phrase_has_expected_token_count() {
// Verified empirically against tiktoken-rs 0.12's o200k_base:
// "hello world" -> [24912, 2375], i.e. 2 tokens.
assert_eq!(count_tokens("hello world"), 2);
}
#[test]
fn known_code_snippet_has_expected_token_count() {
// Verified empirically: 9 tokens under o200k_base.
assert_eq!(count_tokens("fn main() { println!(\"hi\"); }"), 9);
}
#[test]
fn message_with_no_content_counts_zero() {
let msg = ChatMessage::assistant(None);
assert_eq!(count_message_tokens(&msg), 0);
}
#[test]
fn message_token_count_matches_count_tokens_on_its_content() {
let msg = ChatMessage::user("hello world");
assert_eq!(count_message_tokens(&msg), count_tokens("hello world"));
}
}
+1
View File
@@ -2,5 +2,6 @@
//! and the LLM streaming pipeline.
pub mod actions;
pub mod commands;
pub mod context;
pub mod shortsend;
pub mod stream;