language_models: Update tiktoken-rs to support newer models (#30951)

Ben Brandt created 6 months ago

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

Change summary

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

Detailed changes

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",
 ]

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",

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);
+        }
+    }
+}