Compare commits

..

2 Commits

Author SHA1 Message Date
Slavi Pantaleev
9f1d664319 Capitalize thinking-notice flavor messages and bump 1.24.0 release date
The "thinking…" presets started with a lowercase letter, which reads
oddly as a standalone message. Capitalize the first word of each
literal-led line; lines that open with {{ baibot_name }} are untouched.

Also bump the 1.24.0 CHANGELOG date to the actual release date.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-26 07:02:08 +03:00
Aine
0c78928dab add thinking notice 2026-06-25 19:41:15 +01:00
18 changed files with 61 additions and 321 deletions

View File

@@ -14,31 +14,13 @@ concurrency:
group: ci-${{ github.event.pull_request.number || github.ref }} group: ci-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true cancel-in-progress: true
jobs: jobs:
prek: test-and-clippy:
name: Lint, format & test name: Unit testing and linting
runs-on: ubuntu-latest runs-on: ubuntu-latest
steps: steps:
- uses: actions/checkout@v7 - uses: actions/checkout@v7
- uses: dtolnay/rust-toolchain@1.93.0
# Toolchain version + components come from rust-toolchain.toml. rustflags is
# cleared so plain builds don't fail on warnings; the clippy hook still does.
- uses: actions-rust-lang/setup-rust-toolchain@v1
with:
rustflags: ''
- name: Install SQLite3 - name: Install SQLite3
run: sudo apt-get update && sudo apt-get install -y libsqlite3-dev run: sudo apt-get update && sudo apt-get install -y libsqlite3-dev
- run: cargo test --all-features
# just drives the prek recipes; mise provides the pinned prek (mise.toml). - run: cargo clippy
- uses: taiki-e/install-action@v2
with:
tool: just
- uses: jdx/mise-action@v4
# Run the same prek hooks devs run locally; .pre-commit-config.yaml is the
# source of truth. Tests are a separate step for visible timing.
- name: Lint & format (prek hooks, excluding tests)
run: just prek-run-on-all --skip test-unit
- name: Unit tests
run: just test

View File

@@ -1,10 +1,3 @@
# (2026-06-29) Version 1.25.0
- (**Feature**) [♻️ Context management](./docs/configuration/text-generation.md#️-context-management) now works with every provider, not only [OpenAI](./docs/providers.md#openai). Token counting previously went through [tiktoken-rs](https://github.com/zurawiki/tiktoken-rs), which is accurate only for OpenAI models and silently mis-counted everything else (worst of all for non-English text). OpenAI agents keep using tiktoken-rs; every other provider, including the recommended [Venice](./docs/providers.md#venice), now uses a provider-neutral approximation that needs no per-model tokenizer (ASCII counted at about four characters per token, other scripts such as Cyrillic and CJK at about two), landing within roughly 10-20% of the real count. See the [context management docs](./docs/configuration/text-generation.md#️-context-management).
- (**Improvement**) Context management now trims a conversation on whole-turn boundaries for every provider, so an assistant reply is never kept without the user message it answered. This also adjusts how the OpenAI provider trims: a dangling assistant reply at the oldest edge of the kept history is now dropped along with its missing prompt, rather than left in place.
# (2026-06-26) Version 1.24.0 # (2026-06-26) Version 1.24.0
- (**Feature**) Add an opt-in 💭 **thinking notice** for text generation. When enabled, a slow response (for example, from a reasoning model that runs for minutes) posts a "thinking…" placeholder after a short delay, refreshes it periodically with varying flavor text, and then edits that same message into the final answer, so a long wait no longer looks like a stuck bot. The notice is **disabled by default** and configurable per-room or globally via `text-generation set-thinking-notice-enabled true`. Fast responses (under the delay threshold) never show a placeholder. See the [text-generation configuration docs](./docs/configuration/text-generation.md#-thinking-notice). - (**Feature**) Add an opt-in 💭 **thinking notice** for text generation. When enabled, a slow response (for example, from a reasoning model that runs for minutes) posts a "thinking…" placeholder after a short delay, refreshes it periodically with varying flavor text, and then edits that same message into the final answer, so a long wait no longer looks like a stuck bot. The notice is **disabled by default** and configurable per-room or globally via `text-generation set-thinking-notice-enabled true`. Fast responses (under the delay threshold) never show a placeholder. See the [text-generation configuration docs](./docs/configuration/text-generation.md#-thinking-notice).

34
Cargo.lock generated
View File

@@ -95,9 +95,9 @@ dependencies = [
[[package]] [[package]]
name = "anyhow" name = "anyhow"
version = "1.0.103" version = "1.0.102"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c"
[[package]] [[package]]
name = "anymap2" name = "anymap2"
@@ -315,7 +315,7 @@ dependencies = [
[[package]] [[package]]
name = "baibot" name = "baibot"
version = "1.25.0" version = "1.24.0"
dependencies = [ dependencies = [
"anthropic", "anthropic",
"anyhow", "anyhow",
@@ -327,7 +327,7 @@ dependencies = [
"mime_guess", "mime_guess",
"mxidwc", "mxidwc",
"mxlink", "mxlink",
"quick_cache 0.7.0", "quick_cache",
"regex", "regex",
"reqwest 0.13.4", "reqwest 0.13.4",
"serde", "serde",
@@ -1233,12 +1233,6 @@ version = "0.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2"
[[package]]
name = "foldhash"
version = "0.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb"
[[package]] [[package]]
name = "form_urlencoded" name = "form_urlencoded"
version = "1.2.2" version = "1.2.2"
@@ -1463,7 +1457,7 @@ version = "0.15.5"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1"
dependencies = [ dependencies = [
"foldhash 0.1.5", "foldhash",
] ]
[[package]] [[package]]
@@ -2498,7 +2492,7 @@ dependencies = [
"hex", "hex",
"matrix-sdk", "matrix-sdk",
"mime", "mime",
"quick_cache 0.6.24", "quick_cache",
"rand 0.10.1", "rand 0.10.1",
"serde", "serde",
"serde_json", "serde_json",
@@ -2854,9 +2848,9 @@ checksum = "007d8adb5ddab6f8e3f491ac63566a7d5002cc7ed73901f72057943fa71ae1ae"
[[package]] [[package]]
name = "quick_cache" name = "quick_cache"
version = "0.6.24" version = "0.6.23"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b9c6658afe513a3b484e3abfdaa0d03ef3c0bbf017542c178dd55f94eb3051f9" checksum = "3a3db184a8b66cfe87f0263a1de147a6b554c864d1767c6f7fa4eb0e5497b565"
dependencies = [ dependencies = [
"ahash", "ahash",
"equivalent", "equivalent",
@@ -2864,18 +2858,6 @@ dependencies = [
"parking_lot", "parking_lot",
] ]
[[package]]
name = "quick_cache"
version = "0.7.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "403c1a912fec895cafb223201e368234842acb9220aaf08ab042ae89ba5f135c"
dependencies = [
"equivalent",
"foldhash 0.2.0",
"hashbrown 0.17.1",
"parking_lot",
]
[[package]] [[package]]
name = "quinn" name = "quinn"
version = "0.11.9" version = "0.11.9"

View File

@@ -7,7 +7,7 @@ license = "AGPL-3.0-or-later"
readme = "README.md" readme = "README.md"
keywords = ["matrix", "chat", "bot", "AI", "LLM"] keywords = ["matrix", "chat", "bot", "AI", "LLM"]
include = ["/etc/assets/baibot-torso-768.png", "/src", "/README.md", "/CHANGELOG.md", "/LICENSE"] include = ["/etc/assets/baibot-torso-768.png", "/src", "/README.md", "/CHANGELOG.md", "/LICENSE"]
version = "1.25.0" version = "1.24.0"
edition = "2024" edition = "2024"
[lib] [lib]
@@ -26,7 +26,7 @@ mime_guess = "2.0.*"
mxidwc = "1.0.*" mxidwc = "1.0.*"
mxlink = ">=1.15.0" mxlink = ">=1.15.0"
etke_openai_api_rust = "0.1.*" etke_openai_api_rust = "0.1.*"
quick_cache = "0.7.*" quick_cache = "0.6.*"
regex = "1.12.*" regex = "1.12.*"
# HTTP client for the native `venice` provider. rustls only (no extra TLS stack), matching the # HTTP client for the native `venice` provider. rustls only (no extra TLS stack), matching the
# reqwest copy async-openai/matrix-sdk/mxlink already use. # reqwest copy async-openai/matrix-sdk/mxlink already use.

View File

@@ -30,7 +30,7 @@ It's influenced by [chaz](https://github.com/arcuru/chaz), but does **not** use
- 🔒 Supports [encryption](./docs/features.md#-encryption) for Matrix communication and Account-Data-stored configuration - 🔒 Supports [encryption](./docs/features.md#-encryption) for Matrix communication and Account-Data-stored configuration
- ♻️ Supports [context-management](./docs/configuration/text-generation.md#️-context-management) for every [provider](./docs/providers.md) (automatically trimming older messages on whole-turn boundaries once a conversation outgrows the context window) - ♻️ Supports [context-management](./docs/configuration/text-generation.md#️-context-management) handling on some models (automatically adjusting the message history length, etc.)
- 🛠️ Allows **customizing much of the bot's [configuration](./docs/configuration/README.md)** at runtime (using commands sent via chat) - 🛠️ Allows **customizing much of the bot's [configuration](./docs/configuration/README.md)** at runtime (using commands sent via chat)

View File

@@ -50,9 +50,9 @@ Example: `!bai config room text-generation set-auto-usage only_for_voice` (this
### ♻️ Context Management ### ♻️ Context Management
The bot also supports ♻️ **context management**, which automatically trims the oldest messages once a conversation grows past the context window. It drops whole turns at a time, so a reply is never separated from the message it answered. The bot also supports ♻️ **context management**, which automatically adjusts the message history length, etc.
Counting tokens precisely needs the model's own tokenizer. For [OpenAI](../providers.md#openai) models, the bot counts them with the [tiktoken-rs](https://github.com/zurawiki/tiktoken-rs) library. For every other provider, including the recommended [Venice](../providers.md#venice), the bot falls back to a provider-neutral **approximation** that needs no per-model tokenizer: it counts ASCII text at about four characters per token and other scripts (Cyrillic, CJK, and so on) at about two. Treat it as rough, within roughly 10-20% of the real count for typical text, which is plenty for keeping a long conversation inside the context window. This feature relies on [tokenization](https://en.wikipedia.org/wiki/Large_language_model#Tokenization) performed by the [tiktoken-rs](https://github.com/zurawiki/tiktoken-rs) library which is [poorly well-maintained](https://github.com/zurawiki/tiktoken-rs/issues/50) and only works well for [OpenAI](../providers.md#openai) models.
This setting is **disabled by default**, but can be enabled via `!bai config room text-generation set-context-management-enabled true` (this can also be set globally, see [🛠️ Room Settings](./README.md#room-settings)). This setting is **disabled by default**, but can be enabled via `!bai config room text-generation set-context-management-enabled true` (this can also be set globally, see [🛠️ Room Settings](./README.md#room-settings)).

View File

@@ -24,7 +24,7 @@ The list of supported providers is below.
### How to choose a provider ### How to choose a provider
If you're not sure which provider to start with, **we recommend [Venice](#venice)**: it's the most capable provider baibot supports (covering [💬 text-generation](./features.md#-text-generation) with vision, file inputs, prompt caching, and native web search, plus [🖌️ image-generation](./features.md#️-image-creation) incl. editing, [🦻 speech-to-text](./features.md#-speech-to-text), and [🗣️ text-to-speech](./features.md#️-text-to-speech)) and the only one that runs inference with no logging and no training on your data. If you'd rather start with the most widely-used option, [OpenAI](#openai) is a solid, well-supported choice too. If you're not sure which provider to start with, **we recommend [OpenAI](#openai)** as it's the most popular and has the **widest range of capabilities**: [💬 text-generation](./features.md#-text-generation) (incl. vision, incl. [🛠️ tools](./features.md#️-built-in-tools-openai-only)), [🖌️ image-generation](./features.md#️image-generation), [🦻 speech-to-text](./features.md#-speech-to-text), [🗣️ text-to-speech](./features.md#️-text-to-speech).
You don't need to choose just one though. The bot supports [mixing & matching models](./features.md#-mixing--matching-models), so you can use multiple providers at the same time. You don't need to choose just one though. The bot supports [mixing & matching models](./features.md#-mixing--matching-models), so you can use multiple providers at the same time.
@@ -176,10 +176,10 @@ This provider is just as featureful as the [OpenAI](#openai) provider, but is mo
### Venice ### Venice
[Venice AI](https://venice.ai/chat?ref=kpXDe6) _(ref link with a $10 bonus for you)_ runs inference on Venice-controlled GPUs or zero-data-retention partner infrastructure and stores no prompts or responses, so your conversations don't linger anywhere. It serves both frontier proprietary models and the latest open-source ones. [Venice AI](https://venice.ai) runs inference on Venice-controlled GPUs or zero-data-retention partner infrastructure and stores no prompts or responses, so your conversations don't linger anywhere. It serves both frontier proprietary models and the latest open-source ones.
- 🆔 Identifier: `venice` - 🆔 Identifier: `venice`
- 🔗 Links: [🏠 Home page](https://venice.ai/chat?ref=kpXDe6), [👤 Sign up](https://venice.ai/chat?ref=kpXDe6), [📋 Models list](https://docs.venice.ai/models/overview) - 🔗 Links: [🏠 Home page](https://venice.ai), [👤 Sign up](https://venice.ai), [📋 Models list](https://api.venice.ai/api/v1/models)
- 🌟 Capabilities: [🖌️ image-generation](./features.md#️-image-creation) (incl. editing, via the native knob-rich `/image/generate` and `/image/edit` endpoints), [💬 text-generation](./features.md#-text-generation) (incl. vision, file inputs like PDF and DOCX, and prompt caching; native web search via the `venice_parameters` config), [🗣️ text-to-speech](./features.md#️-text-to-speech), [🦻 speech-to-text](./features.md#-speech-to-text) - 🌟 Capabilities: [🖌️ image-generation](./features.md#️-image-creation) (incl. editing, via the native knob-rich `/image/generate` and `/image/edit` endpoints), [💬 text-generation](./features.md#-text-generation) (incl. vision, file inputs like PDF and DOCX, and prompt caching; native web search via the `venice_parameters` config), [🗣️ text-to-speech](./features.md#️-text-to-speech), [🦻 speech-to-text](./features.md#-speech-to-text)
- 🗲 Quick start: - 🗲 Quick start:
- create a room-local agent: `!bai agent create-room-local venice my-venice-agent` - create a room-local agent: `!bai agent create-room-local venice my-venice-agent`

View File

@@ -1,6 +1,6 @@
services: services:
ollama: ollama:
image: docker.io/ollama/ollama:0.30.11 image: docker.io/ollama/ollama:0.30.10
restart: unless-stopped restart: unless-stopped
ports: ports:
- "${SERVICE_OLLAMA_BIND_PORT_HTTP}:11434" - "${SERVICE_OLLAMA_BIND_PORT_HTTP}:11434"

View File

@@ -15,7 +15,7 @@ use crate::agent::provider::{
}; };
use crate::conversation::llm::{ use crate::conversation::llm::{
Author as LLMAuthor, Conversation as LLMConversation, Message as LLMMessage, Author as LLMAuthor, Conversation as LLMConversation, Message as LLMMessage,
MessageContent as LLMMessageContent, TokenEstimate, shorten_messages_list_to_context_size, MessageContent as LLMMessageContent, shorten_messages_list_to_context_size,
}; };
use crate::strings; use crate::strings;
@@ -130,7 +130,7 @@ impl ControllerTrait for Controller {
tracing::trace!("Shortening messages list to context size"); tracing::trace!("Shortening messages list to context size");
conversation_messages = shorten_messages_list_to_context_size( conversation_messages = shorten_messages_list_to_context_size(
TokenEstimate::Approximate, &text_generation_config.model_id,
&prompt_message, &prompt_message,
conversation_messages, conversation_messages,
Some(text_generation_config.max_response_tokens), Some(text_generation_config.max_response_tokens),

View File

@@ -24,7 +24,7 @@ use crate::{
}, },
conversation::llm::{ conversation::llm::{
Author as LLMAuthor, Conversation as LLMConversation, Message as LLMMessage, Author as LLMAuthor, Conversation as LLMConversation, Message as LLMMessage,
MessageContent as LLMMessageContent, TokenEstimate, shorten_messages_list_to_context_size, MessageContent as LLMMessageContent, shorten_messages_list_to_context_size,
}, },
utils::base64::base64_decode, utils::base64::base64_decode,
}; };
@@ -117,7 +117,7 @@ impl ControllerTrait for Controller {
tracing::trace!("Shortening messages list to context size"); tracing::trace!("Shortening messages list to context size");
conversation_messages = shorten_messages_list_to_context_size( conversation_messages = shorten_messages_list_to_context_size(
TokenEstimate::Tiktoken(&text_generation_config.model_id), &text_generation_config.model_id,
&prompt_message, &prompt_message,
conversation_messages, conversation_messages,
text_generation_config.max_response_tokens, text_generation_config.max_response_tokens,

View File

@@ -15,7 +15,7 @@ use crate::{
}, },
conversation::llm::{ conversation::llm::{
Author as LLMAuthor, Conversation as LLMConversation, Message as LLMMessage, Author as LLMAuthor, Conversation as LLMConversation, Message as LLMMessage,
MessageContent as LLMMessageContent, TokenEstimate, shorten_messages_list_to_context_size, MessageContent as LLMMessageContent, shorten_messages_list_to_context_size,
}, },
}; };
use crate::{ use crate::{
@@ -114,7 +114,7 @@ impl ControllerTrait for Controller {
tracing::trace!("Shortening messages list to context size"); tracing::trace!("Shortening messages list to context size");
conversation_messages = shorten_messages_list_to_context_size( conversation_messages = shorten_messages_list_to_context_size(
TokenEstimate::Approximate, &text_generation_config.model_id,
&prompt_message, &prompt_message,
conversation_messages, conversation_messages,
text_generation_config.max_response_tokens, text_generation_config.max_response_tokens,

View File

@@ -8,7 +8,7 @@ use crate::agent::AgentPurpose;
use crate::agent::provider::entity::{TextGenerationParams, TextGenerationResult}; use crate::agent::provider::entity::{TextGenerationParams, TextGenerationResult};
use crate::conversation::llm::{ use crate::conversation::llm::{
Author as LLMAuthor, Conversation as LLMConversation, Message as LLMMessage, Author as LLMAuthor, Conversation as LLMConversation, Message as LLMMessage,
MessageContent as LLMMessageContent, TokenEstimate, shorten_messages_list_to_context_size, MessageContent as LLMMessageContent, shorten_messages_list_to_context_size,
}; };
use crate::strings; use crate::strings;
@@ -64,7 +64,7 @@ pub async fn generate_text(
if params.context_management_enabled { if params.context_management_enabled {
conversation_messages = shorten_messages_list_to_context_size( conversation_messages = shorten_messages_list_to_context_size(
TokenEstimate::Approximate, &text_generation_config.model_id,
&prompt_message, &prompt_message,
conversation_messages, conversation_messages,
text_generation_config.max_response_tokens, text_generation_config.max_response_tokens,

View File

@@ -2,8 +2,7 @@ use mxlink::matrix_sdk::{
Room, Room,
room::edit::EditedContent, room::edit::EditedContent,
ruma::{ ruma::{
EventId, OwnedEventId, EventId, OwnedEventId, api::client::receipt::create_receipt::v3::ReceiptType,
api::client::receipt::create_receipt::v3::ReceiptType,
events::room::message::{ events::room::message::{
OriginalSyncRoomMessageEvent, RoomMessageEventContentWithoutRelation, OriginalSyncRoomMessageEvent, RoomMessageEventContentWithoutRelation,
}, },

View File

@@ -558,9 +558,7 @@ async fn handle_stage_text_generation(
// loop exits the instant generation resolves, so there is no detached task and no late edit can // loop exits the instant generation resolves, so there is no detached task and no late edit can
// ever clobber the real answer. `placeholder` is the event we must finalize in every exit path. // ever clobber the real answer. `placeholder` is the event we must finalize in every exit path.
let (result, placeholder) = if let Some(notice_prompt_variables) = notice_prompt_variables { let (result, placeholder) = if let Some(notice_prompt_variables) = notice_prompt_variables {
let generation = controller let generation = controller.generate_text(conversation, params).instrument(span);
.generate_text(conversation, params)
.instrument(span);
tokio::pin!(generation); tokio::pin!(generation);
let mut placeholder: Option<OwnedEventId> = None; let mut placeholder: Option<OwnedEventId> = None;

View File

@@ -6,5 +6,5 @@ mod utils;
mod tests; mod tests;
pub use entity::*; pub use entity::*;
pub use tokenization::{TokenEstimate, shorten_messages_list_to_context_size}; pub use tokenization::shorten_messages_list_to_context_size;
pub use utils::*; pub use utils::*;

View File

@@ -4,22 +4,6 @@ use tiktoken_rs::tokenizer;
use super::{Author, Message, MessageContent}; use super::{Author, Message, MessageContent};
/// How to count the tokens in a conversation when trimming it to fit the context window.
pub enum TokenEstimate<'a> {
/// Count via the [tiktoken-rs](https://github.com/zurawiki/tiktoken-rs) library.
/// Accurate for OpenAI models; every other model falls back to the gpt-4
/// tokenizer, which misreads it (badly so for non-English text). Use this only
/// for the OpenAI provider.
Tiktoken(&'a str),
/// Provider-neutral approximation that needs no per-model tokenizer. Expect it
/// to land within roughly 10-20% of the real count for typical text, leaning
/// slightly high: over-counting trims a little extra history, while
/// under-counting would overflow the model's real context window. Use this for
/// every non-OpenAI provider.
Approximate,
}
fn get_bpe_for_model(model: &str) -> &'static CoreBPE { fn get_bpe_for_model(model: &str) -> &'static CoreBPE {
let tokenizer = tokenizer::get_tokenizer(model) let tokenizer = tokenizer::get_tokenizer(model)
.or_else(|| tokenizer::get_tokenizer("gpt-4")) .or_else(|| tokenizer::get_tokenizer("gpt-4"))
@@ -29,27 +13,21 @@ fn get_bpe_for_model(model: &str) -> &'static CoreBPE {
} }
pub fn shorten_messages_list_to_context_size( pub fn shorten_messages_list_to_context_size(
estimate: TokenEstimate<'_>, model: &str,
prompt_message: &Option<Message>, prompt_message: &Option<Message>,
mut messages: Vec<Message>, mut messages: Vec<Message>,
max_response_tokens: Option<u32>, max_response_tokens: Option<u32>,
max_context_tokens: u32, max_context_tokens: u32,
) -> Vec<Message> { ) -> Vec<Message> {
// Loading the tiktoken data is expensive, so we resolve the counter once up // Loading the tokenization data is an expensive process, so
// front and reuse it for every message. // se construct the BPE instance once and then use it for all messages.
let tiktoken = match estimate { let bpe = get_bpe_for_model(model);
TokenEstimate::Tiktoken(model) => Some((get_bpe_for_model(model), model)),
TokenEstimate::Approximate => None,
};
let count = |message: &Message| match tiktoken {
Some((bpe, model)) => tiktoken_token_size_for_message(bpe, model, message),
None => approximate_token_size_for_message(message),
};
// We want to retain the prompt in all cases, so we always count it first. // We want to retain the prompt in all cases, so we always count it first.
// We also always reserve enough tokens for the maximum response we expect. // We also always reserve enough tokens for the maximum response we expect.
let mut current_context_length: u32 = if let Some(prompt_message) = prompt_message { let mut current_context_length: u32 = if let Some(prompt_message) = prompt_message {
count(prompt_message) + max_response_tokens.unwrap_or(0) calculate_token_size_for_message(bpe, model, prompt_message)
+ max_response_tokens.unwrap_or(0)
} else { } else {
0 0
}; };
@@ -59,7 +37,7 @@ pub fn shorten_messages_list_to_context_size(
let mut messages_to_keep: Vec<Message> = Vec::new(); let mut messages_to_keep: Vec<Message> = Vec::new();
for message in messages { for message in messages {
let tokens_for_message = count(&message); let tokens_for_message = calculate_token_size_for_message(bpe, model, &message);
if current_context_length + tokens_for_message > max_context_tokens { if current_context_length + tokens_for_message > max_context_tokens {
break; break;
@@ -70,26 +48,14 @@ pub fn shorten_messages_list_to_context_size(
messages_to_keep.push(message); messages_to_keep.push(message);
} }
// Cut on a turn boundary: the loop may stop right after an assistant reply
// whose triggering user message did not fit, which would leave the kept window
// starting on an orphaned reply. `messages_to_keep` is newest-first here, so
// the oldest kept messages are at the end; drop any trailing assistant messages
// until the window begins at the start of a turn (a user message).
while matches!(
messages_to_keep.last().map(|message| &message.author),
Some(Author::Assistant)
) {
messages_to_keep.pop();
}
messages_to_keep.reverse(); messages_to_keep.reverse();
messages_to_keep messages_to_keep
} }
/// Token size of a message via tiktoken, for a preloaded CoreBPE object. /// Calculate the token size of a message for a given model, with a preloaded CoreBPE object.
/// Accurate only for OpenAI models (see [`TokenEstimate::Tiktoken`]). /// Related to `calculate_token_size_for_model_message`.
fn tiktoken_token_size_for_message(bpe: &CoreBPE, model: &str, message: &Message) -> u32 { fn calculate_token_size_for_message(bpe: &CoreBPE, model: &str, message: &Message) -> u32 {
let (tokens_per_message, tokens_per_name) = if model.starts_with("gpt-3.5") { let (tokens_per_message, tokens_per_name) = if model.starts_with("gpt-3.5") {
( (
4, // every message follows <im_start>{role/name}\n{content}<im_end>\n 4, // every message follows <im_start>{role/name}\n{content}<im_end>\n
@@ -114,52 +80,6 @@ fn tiktoken_token_size_for_message(bpe: &CoreBPE, model: &str, message: &Message
(text_length + role_length + tokens_per_message + tokens_per_name) as u32 (text_length + role_length + tokens_per_message + tokens_per_name) as u32
} }
/// ASCII text averages about four characters per token.
const ASCII_TOKENS_PER_CHAR: f32 = 0.25;
/// Non-ASCII scripts (Cyrillic, CJK, and others) pack more information per
/// character: real tokenizers land around two characters per token for them, so
/// each counts as half a token. CJK runs a touch denser than that, so its estimate
/// can read slightly low, still within the tolerance this approximation targets.
const WIDE_TOKENS_PER_CHAR: f32 = 0.5;
/// Structural per-message overhead (role marker plus message framing), mirroring
/// the small constant the tiktoken path adds.
const APPROX_TOKENS_PER_MESSAGE: u32 = 4;
/// Provider-neutral, tokenizer-free token size of a message
/// (see [`TokenEstimate::Approximate`]).
fn approximate_token_size_for_message(message: &Message) -> u32 {
let text_tokens = match &message.content {
MessageContent::Text(text) => approximate_token_size_for_text(text),
// Images and files are not counted as text, matching the tiktoken path.
MessageContent::Image(..) | MessageContent::File(..) => 0,
};
text_tokens + APPROX_TOKENS_PER_MESSAGE
}
/// Rough token estimate for a piece of text, with no tokenizer.
///
/// ASCII characters count as a quarter-token each (~4 chars/token); characters
/// outside ASCII count as half a token each (~2 chars/token), matching how real
/// tokenizers treat Cyrillic and CJK. Weighting non-ASCII up keeps the estimate
/// from badly under-counting non-English text, the case the tiktoken fallback gets
/// most wrong.
fn approximate_token_size_for_text(text: &str) -> u32 {
let mut estimate = 0.0_f32;
for character in text.chars() {
estimate += if character.is_ascii() {
ASCII_TOKENS_PER_CHAR
} else {
WIDE_TOKENS_PER_CHAR
};
}
estimate.ceil() as u32
}
pub mod test { pub mod test {
#[test] #[test]
fn message_size_counting_works() { fn message_size_counting_works() {
@@ -174,7 +94,7 @@ pub mod test {
timestamp: chrono::Utc::now(), timestamp: chrono::Utc::now(),
}; };
let tokens = super::tiktoken_token_size_for_message(bpe, model, &message); let tokens = super::calculate_token_size_for_message(bpe, model, &message);
assert_eq!(8, tokens); assert_eq!(8, tokens);
} }
@@ -198,7 +118,7 @@ pub mod test {
assert_eq!( assert_eq!(
prompt_length, prompt_length,
super::tiktoken_token_size_for_message(bpe, model, &prompt) super::calculate_token_size_for_message(bpe, model, &prompt)
); );
let mut conversation_messages = Vec::new(); let mut conversation_messages = Vec::new();
@@ -213,7 +133,7 @@ pub mod test {
assert_eq!( assert_eq!(
first_length, first_length,
super::tiktoken_token_size_for_message(bpe, model, &first) super::calculate_token_size_for_message(bpe, model, &first)
); );
conversation_messages.push(first); conversation_messages.push(first);
@@ -228,7 +148,7 @@ pub mod test {
assert_eq!( assert_eq!(
second_length, second_length,
super::tiktoken_token_size_for_message(bpe, model, &second) super::calculate_token_size_for_message(bpe, model, &second)
); );
conversation_messages.push(second); conversation_messages.push(second);
@@ -245,7 +165,7 @@ pub mod test {
assert_eq!( assert_eq!(
third_length, third_length,
super::tiktoken_token_size_for_message(bpe, model, &third) super::calculate_token_size_for_message(bpe, model, &third)
); );
conversation_messages.push(third.clone()); conversation_messages.push(third.clone());
@@ -262,7 +182,7 @@ pub mod test {
assert_eq!( assert_eq!(
forth_length, forth_length,
super::tiktoken_token_size_for_message(bpe, model, &forth) super::calculate_token_size_for_message(bpe, model, &forth)
); );
conversation_messages.push(forth.clone()); conversation_messages.push(forth.clone());
@@ -270,7 +190,7 @@ pub mod test {
assert_eq!(4, conversation_messages.len()); assert_eq!(4, conversation_messages.len());
let new_conversation_messages = super::shorten_messages_list_to_context_size( let new_conversation_messages = super::shorten_messages_list_to_context_size(
super::TokenEstimate::Tiktoken(model), model,
&Some(prompt), &Some(prompt),
conversation_messages, conversation_messages,
max_response_tokens, max_response_tokens,
@@ -309,7 +229,7 @@ pub mod test {
assert_eq!( assert_eq!(
prompt_length, prompt_length,
super::tiktoken_token_size_for_message(bpe, model, &prompt) super::calculate_token_size_for_message(bpe, model, &prompt)
); );
let mut conversation_messages = Vec::new(); let mut conversation_messages = Vec::new();
@@ -324,7 +244,7 @@ pub mod test {
assert_eq!( assert_eq!(
first_length, first_length,
super::tiktoken_token_size_for_message(bpe, model, &first) super::calculate_token_size_for_message(bpe, model, &first)
); );
conversation_messages.push(first); conversation_messages.push(first);
@@ -339,7 +259,7 @@ pub mod test {
assert_eq!( assert_eq!(
second_length, second_length,
super::tiktoken_token_size_for_message(bpe, model, &second) super::calculate_token_size_for_message(bpe, model, &second)
); );
conversation_messages.push(second); conversation_messages.push(second);
@@ -356,7 +276,7 @@ pub mod test {
assert_eq!( assert_eq!(
third_length, third_length,
super::tiktoken_token_size_for_message(bpe, model, &third) super::calculate_token_size_for_message(bpe, model, &third)
); );
conversation_messages.push(third.clone()); conversation_messages.push(third.clone());
@@ -373,7 +293,7 @@ pub mod test {
assert_eq!( assert_eq!(
forth_length, forth_length,
super::tiktoken_token_size_for_message(bpe, model, &forth) super::calculate_token_size_for_message(bpe, model, &forth)
); );
conversation_messages.push(forth.clone()); conversation_messages.push(forth.clone());
@@ -381,7 +301,7 @@ pub mod test {
assert_eq!(4, conversation_messages.len()); assert_eq!(4, conversation_messages.len());
let new_conversation_messages = super::shorten_messages_list_to_context_size( let new_conversation_messages = super::shorten_messages_list_to_context_size(
super::TokenEstimate::Tiktoken(model), model,
&Some(prompt), &Some(prompt),
conversation_messages, conversation_messages,
max_response_tokens, max_response_tokens,
@@ -400,126 +320,4 @@ pub mod test {
forth.content forth.content
); );
} }
#[test]
fn approximate_counting_weights_ascii_and_wide_scripts() {
// 12 ASCII characters at ~4 chars/token = 3 text tokens.
assert_eq!(3, super::approximate_token_size_for_text("Hello there!"));
// 5 CJK characters at ~0.5 token/char = 3 text tokens. The ASCII rate would
// have under-counted these to 2, the failure mode this path avoids.
assert_eq!(3, super::approximate_token_size_for_text("こんにちは"));
let message = super::Message {
author: super::Author::User,
sender_id: None,
content: super::MessageContent::Text("Hello there!".to_string()),
timestamp: chrono::Utc::now(),
};
// 3 text tokens plus the per-message overhead (4).
assert_eq!(7, super::approximate_token_size_for_message(&message));
}
#[test]
fn approximate_shortening_trims_to_budget() {
let prompt = super::Message {
author: super::Author::Prompt,
sender_id: None,
content: super::MessageContent::Text("You are a bot!".to_string()),
timestamp: chrono::Utc::now(),
};
let older = super::Message {
author: super::Author::User,
sender_id: None,
content: super::MessageContent::Text("This is the older message.".to_string()),
timestamp: chrono::Utc::now(),
};
let newer = super::Message {
// A user message, so it is a valid window start: keeping a lone
// assistant reply would be an orphan and get trimmed (see
// `shortening_cuts_on_a_turn_boundary`).
author: super::Author::User,
sender_id: None,
content: super::MessageContent::Text("This is the newer message.".to_string()),
timestamp: chrono::Utc::now(),
};
// Budget room for the prompt and only the newest message.
let max_context_tokens = super::approximate_token_size_for_message(&prompt)
+ super::approximate_token_size_for_message(&newer);
let new_conversation_messages = super::shorten_messages_list_to_context_size(
super::TokenEstimate::Approximate,
&Some(prompt),
vec![older, newer.clone()],
None,
max_context_tokens,
);
assert_eq!(1, new_conversation_messages.len());
assert_eq!(
new_conversation_messages.first().unwrap().content,
newer.content
);
}
#[test]
fn shortening_cuts_on_a_turn_boundary() {
// A four-message conversation of two full turns. All four messages are the
// same length, so they cost the same number of tokens.
let prompt = super::Message {
author: super::Author::Prompt,
sender_id: None,
content: super::MessageContent::Text("system".to_string()),
timestamp: chrono::Utc::now(),
};
let user_one = super::Message {
author: super::Author::User,
sender_id: None,
content: super::MessageContent::Text("user msg 1".to_string()),
timestamp: chrono::Utc::now(),
};
let asst_one = super::Message {
author: super::Author::Assistant,
sender_id: None,
content: super::MessageContent::Text("asst msg 1".to_string()),
timestamp: chrono::Utc::now(),
};
let user_two = super::Message {
author: super::Author::User,
sender_id: None,
content: super::MessageContent::Text("user msg 2".to_string()),
timestamp: chrono::Utc::now(),
};
let asst_two = super::Message {
author: super::Author::Assistant,
sender_id: None,
content: super::MessageContent::Text("asst msg 2".to_string()),
timestamp: chrono::Utc::now(),
};
let per_message = super::approximate_token_size_for_message(&user_one);
// Budget fits the prompt plus three messages. By raw token budget the loop
// would keep asst_two, user_two, and asst_one, but asst_one's own user
// message (user_one) does not fit, so it must be dropped too rather than
// left as an orphaned reply.
let max_context_tokens =
super::approximate_token_size_for_message(&prompt) + (per_message * 3);
let kept = super::shorten_messages_list_to_context_size(
super::TokenEstimate::Approximate,
&Some(prompt),
vec![user_one, asst_one, user_two.clone(), asst_two.clone()],
None,
max_context_tokens,
);
// Only the last whole turn survives; the orphaned asst_one is dropped.
assert_eq!(2, kept.len());
assert_eq!(kept.first().unwrap().content, user_two.content);
assert_eq!(kept.last().unwrap().content, asst_two.content);
}
} }

View File

@@ -128,7 +128,7 @@ pub fn text_generation_context_management_intro() -> String {
format!( format!(
"{}\n{}", "{}\n{}",
"Controls the bot's ability to **intelligently drop old messages from the conversation context** when it gets too large.", "Controls the bot's ability to **intelligently drop old messages from the conversation context** when it gets too large.",
"Counting tokens precisely needs the model's own tokenizer. For [OpenAI](./providers.md#openai) models the bot uses the [tiktoken-rs](https://github.com/zurawiki/tiktoken-rs) library; for every other provider (including the recommended [Venice](./providers.md#venice)) it falls back to a provider-neutral **approximation** (ASCII counted at ~4 characters per token, other scripts at ~2), within roughly 10-20% of the real count for typical text.", "This feature relies on [tokenization](https://en.wikipedia.org/wiki/Large_language_model#Tokenization) performed by the [tiktoken-rs](https://github.com/zurawiki/tiktoken-rs) library which is [poorly well-maintained](https://github.com/zurawiki/tiktoken-rs/issues/50) and only works well for [OpenAI](./providers.md#openai) models.",
) )
} }

View File

@@ -84,23 +84,11 @@ mod tests {
// Boundaries: light below 30s, medium [30s, 90s), deep at/after 90s. // Boundaries: light below 30s, medium [30s, 90s), deep at/after 90s.
// The pools have distinct content, so value comparison identifies the tier. // The pools have distinct content, so value comparison identifies the tier.
assert_eq!(messages_for_elapsed(Duration::from_secs(0)), MESSAGES_LIGHT); assert_eq!(messages_for_elapsed(Duration::from_secs(0)), MESSAGES_LIGHT);
assert_eq!( assert_eq!(messages_for_elapsed(Duration::from_secs(29)), MESSAGES_LIGHT);
messages_for_elapsed(Duration::from_secs(29)), assert_eq!(messages_for_elapsed(Duration::from_secs(30)), MESSAGES_MEDIUM);
MESSAGES_LIGHT assert_eq!(messages_for_elapsed(Duration::from_secs(89)), MESSAGES_MEDIUM);
);
assert_eq!(
messages_for_elapsed(Duration::from_secs(30)),
MESSAGES_MEDIUM
);
assert_eq!(
messages_for_elapsed(Duration::from_secs(89)),
MESSAGES_MEDIUM
);
assert_eq!(messages_for_elapsed(Duration::from_secs(90)), MESSAGES_DEEP); assert_eq!(messages_for_elapsed(Duration::from_secs(90)), MESSAGES_DEEP);
assert_eq!( assert_eq!(messages_for_elapsed(Duration::from_secs(600)), MESSAGES_DEEP);
messages_for_elapsed(Duration::from_secs(600)),
MESSAGES_DEEP
);
} }
#[test] #[test]