mirror of
https://github.com/openai/codex.git
synced 2026-04-24 14:45:27 +00:00
use shared tokenizer
This commit is contained in:
@@ -59,6 +59,7 @@ use crate::config::Config;
|
||||
use crate::config_types::McpServerTransportConfig;
|
||||
use crate::config_types::ShellEnvironmentPolicy;
|
||||
use crate::conversation_history::ConversationHistory;
|
||||
use crate::conversation_history::prefetch_tokenizer_in_background;
|
||||
use crate::environment_context::EnvironmentContext;
|
||||
use crate::error::CodexErr;
|
||||
use crate::error::Result as CodexResult;
|
||||
@@ -159,6 +160,8 @@ impl Codex {
|
||||
conversation_history: InitialHistory,
|
||||
session_source: SessionSource,
|
||||
) -> CodexResult<CodexSpawnOk> {
|
||||
// Start loading the tokenizer in the background so we don't block later.
|
||||
prefetch_tokenizer_in_background();
|
||||
let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY);
|
||||
let (tx_event, rx_event) = async_channel::unbounded();
|
||||
|
||||
|
||||
@@ -6,8 +6,11 @@ use codex_protocol::models::ResponseItem;
|
||||
use codex_protocol::protocol::TokenUsage;
|
||||
use codex_protocol::protocol::TokenUsageInfo;
|
||||
use codex_utils_tokenizer::Tokenizer;
|
||||
use tokio::task;
|
||||
use tracing::error;
|
||||
|
||||
static TOKENIZER: OnceLock<Option<Arc<Tokenizer>>> = OnceLock::new();
|
||||
|
||||
use crate::error::CodexErr;
|
||||
|
||||
/// Transcript of conversation history
|
||||
@@ -16,17 +19,13 @@ pub(crate) struct ConversationHistory {
|
||||
/// The oldest items are at the beginning of the vector.
|
||||
items: Vec<ResponseItem>,
|
||||
token_info: Option<TokenUsageInfo>,
|
||||
tokenizer: Option<Arc<Tokenizer>>,
|
||||
}
|
||||
|
||||
impl ConversationHistory {
|
||||
pub(crate) fn new() -> Self {
|
||||
// load the tokenizer once and share it across all instances. It takes ~500 ms to load.
|
||||
let tokenizer = shared_tokenizer();
|
||||
Self {
|
||||
items: Vec::new(),
|
||||
token_info: None,
|
||||
tokenizer,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -121,8 +120,9 @@ impl ConversationHistory {
|
||||
let Some(context_window) = info.model_context_window else {
|
||||
return Ok(());
|
||||
};
|
||||
let Some(tokenizer) = self.tokenizer.as_deref() else {
|
||||
return Ok(());
|
||||
let tokenizer = match shared_tokenizer() {
|
||||
Some(t) => t,
|
||||
None => return Ok(()),
|
||||
};
|
||||
|
||||
let mut input_tokens: i64 = 0;
|
||||
@@ -413,17 +413,33 @@ fn error_or_panic(message: String) {
|
||||
}
|
||||
|
||||
fn shared_tokenizer() -> Option<Arc<Tokenizer>> {
|
||||
static TOKENIZER: OnceLock<Option<Arc<Tokenizer>>> = OnceLock::new();
|
||||
TOKENIZER
|
||||
.get_or_init(|| match Tokenizer::try_default() {
|
||||
Ok(tokenizer) => Some(Arc::new(tokenizer)),
|
||||
Err(error) => {
|
||||
error!("failed to create tokenizer: {error}");
|
||||
None
|
||||
TOKENIZER.get().and_then(|opt| opt.as_ref().map(Arc::clone))
|
||||
}
|
||||
|
||||
/// Kick off background initialization of the shared tokenizer without blocking the caller.
|
||||
pub(crate) fn prefetch_tokenizer_in_background() {
|
||||
if TOKENIZER.get().is_some() {
|
||||
return;
|
||||
}
|
||||
|
||||
// Spawn a background task to initialize the tokenizer. Use spawn_blocking in case
|
||||
// initialization performs CPU-heavy work or file I/O.
|
||||
tokio::spawn(async {
|
||||
let result = task::spawn_blocking(Tokenizer::try_default).await;
|
||||
match result {
|
||||
Ok(Ok(tokenizer)) => {
|
||||
let _ = TOKENIZER.set(Some(Arc::new(tokenizer)));
|
||||
}
|
||||
})
|
||||
.as_ref()
|
||||
.map(Arc::clone)
|
||||
Ok(Err(error)) => {
|
||||
error!("failed to create tokenizer: {error}");
|
||||
let _ = TOKENIZER.set(None);
|
||||
}
|
||||
Err(join_error) => {
|
||||
error!("failed to join tokenizer init task: {join_error}");
|
||||
let _ = TOKENIZER.set(None);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// Anything that is not a system message or "reasoning" message is considered
|
||||
|
||||
Reference in New Issue
Block a user