Skip to main content

TokenCounter

Trait TokenCounter 

Source
pub trait TokenCounter: Send + Sync {
    // Required method
    fn count<'life0, 'life1, 'async_trait>(
        &'life0 self,
        text: &'life1 str,
    ) -> Pin<Box<dyn Future<Output = Result<usize, MemoryError>> + Send + 'async_trait>>
       where 'life0: 'async_trait,
             'life1: 'async_trait,
             Self: 'async_trait;
}
Expand description

Counts the token number of a text.

Memory does not know the model; “exact” token counts are user-side knowledge, so this trait carries no model or provider information:

  • single-model setups: inject a counter built for that model when constructing WindowMemory (e.g., wiring in tiktoken-rs); the counting convention is bound at construction;
  • multi-model dynamic routing: the custom implementation holds shared state internally and switches conventions when the application switches models.

The default implementation CharTokenCounter is a heuristic approximation that depends on no model.

The trait is async: it supports remote counting (e.g., calling a vendor’s counting API); local implementations just return Ok(approx). All call sites (record / context / trim) are already async, so remote counting costs nothing extra.

§Example

Inject a custom counter (e.g., trimming by message count):

use molo::memory::{Memory, MemoryError, TokenCounter, WindowMemory};
use molo::Message;

// Each message counts as exactly 1 token: the token budget degenerates to
// a message-count limit.
#[derive(Default)]
struct OnePerMessage;

#[molo::async_trait]
impl TokenCounter for OnePerMessage {
    async fn count(&self, _text: &str) -> Result<usize, MemoryError> {
        Ok(1)
    }
}

#[tokio::main]
async fn main() -> Result<(), MemoryError> {
    let mut memory =
        WindowMemory::new(2).with_token_counter(Box::new(OnePerMessage));
    memory.record(Message::user("u1")).await?;
    memory.record(Message::assistant("a1")).await?;
    memory.record(Message::user("u2")).await?;
    memory.record(Message::assistant("a2")).await?;

    // 4 messages > 2 tokens: trimmed to the most recent round.
    assert_eq!(memory.context().await?.len(), 2);
    Ok(())
}

Required Methods§

Source

fn count<'life0, 'life1, 'async_trait>( &'life0 self, text: &'life1 str, ) -> Pin<Box<dyn Future<Output = Result<usize, MemoryError>> + Send + 'async_trait>>
where 'life0: 'async_trait, 'life1: 'async_trait, Self: 'async_trait,

Counts the token number of a text.

§Errors

Returns MemoryError::TokenCount when counting fails (e.g., a remote counting API is unavailable).

Dyn Compatibility§

This trait is dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors§