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:
Generated
+28
@@ -196,6 +196,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||||||
checksum = "5cee35f73844aa3014bb606320a6c1f010249dbdf43342fe54b5a4f6a8ed4b79"
|
checksum = "5cee35f73844aa3014bb606320a6c1f010249dbdf43342fe54b5a4f6a8ed4b79"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"memchr",
|
"memchr",
|
||||||
|
"regex-automata",
|
||||||
"serde_core",
|
"serde_core",
|
||||||
]
|
]
|
||||||
|
|
||||||
@@ -807,6 +808,17 @@ dependencies = [
|
|||||||
"regex-syntax",
|
"regex-syntax",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "fancy-regex"
|
||||||
|
version = "0.17.0"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "72cf461f865c862bb7dc573f643dd6a2b6842f7c30b07882b56bd148cc2761b8"
|
||||||
|
dependencies = [
|
||||||
|
"bit-set 0.8.0",
|
||||||
|
"regex-automata",
|
||||||
|
"regex-syntax",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "fast-srgb8"
|
name = "fast-srgb8"
|
||||||
version = "1.0.0"
|
version = "1.0.0"
|
||||||
@@ -3521,6 +3533,21 @@ dependencies = [
|
|||||||
"cfg-if",
|
"cfg-if",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "tiktoken-rs"
|
||||||
|
version = "0.12.0"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "027853bbf8c7763b77c5c595f1c271c7d536ced7d6f83452911b944621e57fc2"
|
||||||
|
dependencies = [
|
||||||
|
"anyhow",
|
||||||
|
"base64",
|
||||||
|
"bstr",
|
||||||
|
"fancy-regex 0.17.0",
|
||||||
|
"lazy_static",
|
||||||
|
"regex",
|
||||||
|
"rustc-hash",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "time"
|
name = "time"
|
||||||
version = "0.3.53"
|
version = "0.3.53"
|
||||||
@@ -4487,6 +4514,7 @@ dependencies = [
|
|||||||
"sha2 0.11.0",
|
"sha2 0.11.0",
|
||||||
"similar",
|
"similar",
|
||||||
"syntect",
|
"syntect",
|
||||||
|
"tiktoken-rs",
|
||||||
"tokio",
|
"tokio",
|
||||||
"tracing",
|
"tracing",
|
||||||
"tracing-subscriber",
|
"tracing-subscriber",
|
||||||
|
|||||||
@@ -58,6 +58,7 @@ chrono = { version = "0.4", features = ["serde"] }
|
|||||||
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
|
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
|
||||||
webbrowser = "1"
|
webbrowser = "1"
|
||||||
lsp-types = "0.97"
|
lsp-types = "0.97"
|
||||||
|
tiktoken-rs = "0.12"
|
||||||
|
|
||||||
[[bin]]
|
[[bin]]
|
||||||
name = "zesdex"
|
name = "zesdex"
|
||||||
|
|||||||
@@ -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;
|
||||||
@@ -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"));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -2,5 +2,6 @@
|
|||||||
//! and the LLM streaming pipeline.
|
//! and the LLM streaming pipeline.
|
||||||
pub mod actions;
|
pub mod actions;
|
||||||
pub mod commands;
|
pub mod commands;
|
||||||
|
pub mod context;
|
||||||
pub mod shortsend;
|
pub mod shortsend;
|
||||||
pub mod stream;
|
pub mod stream;
|
||||||
|
|||||||
Reference in New Issue
Block a user