behest_runtime/
token.rs

1//! Token estimation using character-based heuristics.
2//!
3//! Core functions re-exported from `behest-core`.
4//! Store-dependent estimation functions are defined here.
5
6pub use behest_core::token::{
7    estimate_content_part_tokens, estimate_message_tokens, estimate_messages_tokens,
8    estimate_tokens, estimate_tool_call_tokens,
9};
10
11use behest_store::MessageRecord;
12
13/// Estimates the token count for a store [`MessageRecord`].
14#[must_use]
15pub fn estimate_record_tokens(record: &MessageRecord) -> usize {
16    let content_tokens: usize = record
17        .content
18        .iter()
19        .map(estimate_content_part_tokens)
20        .sum();
21
22    let tool_call_tokens: usize = record
23        .tool_calls
24        .iter()
25        .map(estimate_tool_call_tokens)
26        .sum();
27
28    let tool_meta_tokens = match (&record.tool_call_id, &record.tool_name) {
29        (Some(id), Some(name)) => estimate_tokens(id) + estimate_tokens(name),
30        _ => 0,
31    };
32
33    let role_overhead = match record.role {
34        behest_store::MessageRole::System
35        | behest_store::MessageRole::User
36        | behest_store::MessageRole::Assistant => 8,
37        behest_store::MessageRole::Tool => 10,
38        _ => 8,
39    };
40
41    content_tokens + tool_call_tokens + tool_meta_tokens + role_overhead
42}
43
44/// Estimates the total token count for a slice of [`MessageRecord`]s.
45#[must_use]
46pub fn estimate_records_tokens(records: &[MessageRecord]) -> usize {
47    records.iter().map(estimate_record_tokens).sum()
48}
49
50#[cfg(test)]
51#[allow(clippy::unwrap_used)]
52mod tests {
53    use super::*;
54    use behest_core::message::ContentPart;
55    use uuid::Uuid;
56
57    #[test]
58    fn estimate_record() {
59        let record = MessageRecord::new(
60            Uuid::now_v7(),
61            behest_store::MessageRole::User,
62            vec![ContentPart::text("Hello")],
63        );
64        let tokens = estimate_record_tokens(&record);
65        assert_eq!(tokens, 10);
66    }
67
68    #[test]
69    fn estimate_records_slice() {
70        let records = vec![
71            MessageRecord::new(
72                Uuid::now_v7(),
73                behest_store::MessageRole::System,
74                vec![ContentPart::text("System")],
75            ),
76            MessageRecord::new(
77                Uuid::now_v7(),
78                behest_store::MessageRole::User,
79                vec![ContentPart::text("User")],
80            ),
81        ];
82        let total = estimate_records_tokens(&records);
83        assert!(total > 0);
84    }
85}
behest_runtime/token.rs

behest_runtime/
token.rs
