Improve token counting for OpenAI models (#30242)

tiktoken_rs is a bit behind (and even upstream tiktoken doesn't have all
of these models)

We were incorrectly using the cl100k tokenizer for some models that
actually use the o200k tokenizers. So that is updated.

I also made the match arms specific so that we do a better job of
catching whether or not tiktoken-rs accurately supports new models we
add in.

I will also do a PR upstream to see if we can move some of this logic
back out if tiktoken better supports the newer models.

Release Notes:

- Improved tokenizer support for openai models.
This commit is contained in:
Ben Brandt 2025-05-08 15:09:29 +02:00 committed by GitHub
parent ee56706d15
commit 3a3d3c05e8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 42 additions and 28 deletions

View file

@ -219,32 +219,24 @@ impl LanguageModel for CopilotChatLanguageModel {
cx: &App,
) -> BoxFuture<'static, Result<usize>> {
match self.model {
CopilotChatModel::Claude3_5Sonnet => count_anthropic_tokens(request, cx),
CopilotChatModel::Claude3_7Sonnet => count_anthropic_tokens(request, cx),
CopilotChatModel::Claude3_7SonnetThinking => count_anthropic_tokens(request, cx),
CopilotChatModel::Claude3_5Sonnet
| CopilotChatModel::Claude3_7Sonnet
| CopilotChatModel::Claude3_7SonnetThinking => count_anthropic_tokens(request, cx),
CopilotChatModel::Gemini20Flash | CopilotChatModel::Gemini25Pro => {
count_google_tokens(request, cx)
}
_ => {
let model = match self.model {
CopilotChatModel::Gpt4o => open_ai::Model::FourOmni,
CopilotChatModel::Gpt4 => open_ai::Model::Four,
CopilotChatModel::Gpt4_1 => open_ai::Model::FourPointOne,
CopilotChatModel::Gpt3_5Turbo => open_ai::Model::ThreePointFiveTurbo,
CopilotChatModel::O1 => open_ai::Model::O1,
CopilotChatModel::O3Mini => open_ai::Model::O3Mini,
CopilotChatModel::O3 => open_ai::Model::O3,
CopilotChatModel::O4Mini => open_ai::Model::O4Mini,
CopilotChatModel::Claude3_5Sonnet
| CopilotChatModel::Claude3_7Sonnet
| CopilotChatModel::Claude3_7SonnetThinking
| CopilotChatModel::Gemini20Flash
| CopilotChatModel::Gemini25Pro => {
unreachable!()
}
};
count_open_ai_tokens(request, model, cx)
CopilotChatModel::Gpt4o => count_open_ai_tokens(request, open_ai::Model::FourOmni, cx),
CopilotChatModel::Gpt4 => count_open_ai_tokens(request, open_ai::Model::Four, cx),
CopilotChatModel::Gpt4_1 => {
count_open_ai_tokens(request, open_ai::Model::FourPointOne, cx)
}
CopilotChatModel::Gpt3_5Turbo => {
count_open_ai_tokens(request, open_ai::Model::ThreePointFiveTurbo, cx)
}
CopilotChatModel::O1 => count_open_ai_tokens(request, open_ai::Model::O1, cx),
CopilotChatModel::O3Mini => count_open_ai_tokens(request, open_ai::Model::O3Mini, cx),
CopilotChatModel::O3 => count_open_ai_tokens(request, open_ai::Model::O3, cx),
CopilotChatModel::O4Mini => count_open_ai_tokens(request, open_ai::Model::O4Mini, cx),
}
}