From 57424e47434e56bb489eb6a55687f68e4d4f6bea Mon Sep 17 00:00:00 2001
From: Ben Brandt <benjamin.j.brandt@gmail.com>
Date: Mon, 19 May 2025 13:40:36 +0200
Subject: [PATCH] language_models: Update tiktoken-rs to support newer models
 (#30951)

I was able to get this fix in upstream, so now we can have simpler code
paths for our model selection.

I also added a test to catch if this would cause a bug again in the
future.

Release Notes:

- N/A
---
 Cargo.lock                                    |  7 +--
 Cargo.toml                                    |  2 +-
 .../language_models/src/provider/open_ai.rs   | 59 ++++++++++++++++---
 3 files changed, 54 insertions(+), 14 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 6d9ad27f389b70e2dd040fef0995d2d6824df80f..31fab04d58f5ad6ceab16fb1a155e3eaf6c0d59f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -15807,16 +15807,15 @@ dependencies = [
 
 [[package]]
 name = "tiktoken-rs"
-version = "0.6.0"
+version = "0.7.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "44075987ee2486402f0808505dd65692163d243a337fc54363d49afac41087f6"
+checksum = "25563eeba904d770acf527e8b370fe9a5547bacd20ff84a0b6c3bc41288e5625"
 dependencies = [
  "anyhow",
- "base64 0.21.7",
+ "base64 0.22.1",
  "bstr",
  "fancy-regex 0.13.0",
  "lazy_static",
- "parking_lot",
  "regex",
  "rustc-hash 1.1.0",
 ]
diff --git a/Cargo.toml b/Cargo.toml
index 91b63aedd207f010cdc2d7fc846fbacdfc2746ae..65bcec73bf056036467d36b6fbd981eef30bb8c0 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -555,7 +555,7 @@ sysinfo = "0.31.0"
 take-until = "0.2.0"
 tempfile = "3.20.0"
 thiserror = "2.0.12"
-tiktoken-rs = "0.6.0"
+tiktoken-rs = "0.7.0"
 time = { version = "0.3", features = [
     "macros",
     "parsing",
diff --git a/crates/language_models/src/provider/open_ai.rs b/crates/language_models/src/provider/open_ai.rs
index 313224eae9dc357f2822ca41b579adc1ec79898c..f9e749ee6e6725f922292aa104be8c57330f7595 100644
--- a/crates/language_models/src/provider/open_ai.rs
+++ b/crates/language_models/src/provider/open_ai.rs
@@ -628,22 +628,24 @@ pub fn count_open_ai_tokens(
                 };
                 tiktoken_rs::num_tokens_from_messages(model, &messages)
             }
-            // Not currently supported by tiktoken_rs. All use the same tokenizer as gpt-4o (o200k_base)
-            Model::O1
-            | Model::FourPointOne
-            | Model::FourPointOneMini
-            | Model::FourPointOneNano
-            | Model::O3Mini
-            | Model::O3
-            | Model::O4Mini => tiktoken_rs::num_tokens_from_messages("gpt-4o", &messages),
             // Currently supported by tiktoken_rs
+            // Sometimes tiktoken-rs is behind on model support. If that is the case, make a new branch
+            // arm with an override. We enumerate all supported models here so that we can check if new
+            // models are supported yet or not.
             Model::ThreePointFiveTurbo
             | Model::Four
             | Model::FourTurbo
             | Model::FourOmni
             | Model::FourOmniMini
+            | Model::FourPointOne
+            | Model::FourPointOneMini
+            | Model::FourPointOneNano
+            | Model::O1
             | Model::O1Preview
-            | Model::O1Mini => tiktoken_rs::num_tokens_from_messages(model.id(), &messages),
+            | Model::O1Mini
+            | Model::O3
+            | Model::O3Mini
+            | Model::O4Mini => tiktoken_rs::num_tokens_from_messages(model.id(), &messages),
         }
     })
     .boxed()
@@ -839,3 +841,42 @@ impl Render for ConfigurationView {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use gpui::TestAppContext;
+    use language_model::LanguageModelRequestMessage;
+
+    use super::*;
+
+    #[gpui::test]
+    fn tiktoken_rs_support(cx: &TestAppContext) {
+        let request = LanguageModelRequest {
+            thread_id: None,
+            prompt_id: None,
+            mode: None,
+            messages: vec![LanguageModelRequestMessage {
+                role: Role::User,
+                content: vec![MessageContent::Text("message".into())],
+                cache: false,
+            }],
+            tools: vec![],
+            tool_choice: None,
+            stop: vec![],
+            temperature: None,
+        };
+
+        // Validate that all models are supported by tiktoken-rs
+        for model in Model::iter() {
+            let count = cx
+                .executor()
+                .block(count_open_ai_tokens(
+                    request.clone(),
+                    model,
+                    &cx.app.borrow(),
+                ))
+                .unwrap();
+            assert!(count > 0);
+        }
+    }
+}