From 0ba7dbb000b0d46a8b051bde50d7f1dba6e206d8 Mon Sep 17 00:00:00 2001
From: Tat Dat Duong <david@duong.cz>
Date: Fri, 4 Oct 2024 01:20:41 +0200
Subject: [PATCH] Add missing o1 models

---
 js/src/core.ts                  | 10 +++++++++-
 tiktoken/model_to_encoding.json |  9 ++++++++-
 wasm/src/lib.rs                 | 14 ++++++++++++++
 3 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/js/src/core.ts b/js/src/core.ts
index 7aee83f..5bc298b 100644
--- a/js/src/core.ts
+++ b/js/src/core.ts
@@ -280,7 +280,15 @@ export function getEncodingNameForModel(model: TiktokenModel) {
     case "gpt-4o-2024-05-13":
     case "gpt-4o-2024-08-06":
     case "gpt-4o-mini-2024-07-18":
-    case "gpt-4o-mini": {
+    case "gpt-4o-mini":
+    case "o1-mini":
+    case "o1-preview":
+    case "o1-preview-2024-09-12":
+    case "o1-mini-2024-09-12":
+    case "chatgpt-4o-latest":
+    case "gpt-4o-realtime":
+    case "gpt-4o-realtime-preview-2024-10-01":
+    {
       return "o200k_base";
     }
     default:
diff --git a/tiktoken/model_to_encoding.json b/tiktoken/model_to_encoding.json
index e7890e2..06324b8 100644
--- a/tiktoken/model_to_encoding.json
+++ b/tiktoken/model_to_encoding.json
@@ -59,5 +59,12 @@
     "gpt-4o-2024-05-13": "o200k_base",
     "gpt-4o-2024-08-06":"o200k_base",
     "gpt-4o-mini-2024-07-18": "o200k_base",
-    "gpt-4o-mini": "o200k_base"
+    "gpt-4o-mini": "o200k_base",
+    "o1-mini": "o200k_base",
+    "o1-preview": "o200k_base",
+    "o1-preview-2024-09-12": "o200k_base",
+    "o1-mini-2024-09-12": "o200k_base",
+    "chatgpt-4o-latest": "o200k_base",
+    "gpt-4o-realtime": "o200k_base",
+    "gpt-4o-realtime-preview-2024-10-01": "o200k_base"
 }
diff --git a/wasm/src/lib.rs b/wasm/src/lib.rs
index 72cb06b..8c9f863 100644
--- a/wasm/src/lib.rs
+++ b/wasm/src/lib.rs
@@ -423,6 +423,13 @@ export type TiktokenModel =
     | "gpt-4o-2024-08-06"
     | "gpt-4o-mini-2024-07-18"
     | "gpt-4o-mini"
+    | "o1-mini"
+    | "o1-preview"
+    | "o1-preview-2024-09-12"
+    | "o1-mini-2024-09-12"
+    | "chatgpt-4o-latest"
+    | "gpt-4o-realtime"
+    | "gpt-4o-realtime-preview-2024-10-01"
 
 /**
  * @param {TiktokenModel} encoding
@@ -499,6 +506,13 @@ pub fn encoding_for_model(
         "gpt-4o-2024-08-06" => Ok("o200k_base"),
         "gpt-4o-mini-2024-07-18" => Ok("o200k_base"),
         "gpt-4o-mini" => Ok("o200k_base"),
+        "o1-mini" => Ok("o200k_base"),
+        "o1-preview" => Ok("o200k_base"),
+        "o1-preview-2024-09-12" => Ok("o200k_base"),
+        "o1-mini-2024-09-12" => Ok("o200k_base"),
+        "chatgpt-4o-latest" => Ok("o200k_base"),
+        "gpt-4o-realtime" => Ok("o200k_base"),
+        "gpt-4o-realtime-preview-2024-10-01" => Ok("o200k_base"),
         model => Err(JsError::new(
             format!("Invalid model: {}", model.to_string()).as_str(),
         )),