From 365398a7f100451b3ec6465a829d405db7e6b35b Mon Sep 17 00:00:00 2001
From: Peter Tripp <peter@zed.dev>
Date: Tue, 21 Jan 2025 15:05:21 -0500
Subject: [PATCH] open_ai: Move from o1-preview to o1 for OpenAI Assistant
 provider (#23425)

- Closes: https://github.com/zed-industries/zed/issues/22521
- Follow-up to: https://github.com/zed-industries/zed/pull/22376
---
 crates/language_model/src/model/cloud_model.rs |  2 +-
 crates/language_models/src/provider/open_ai.rs |  4 +---
 crates/open_ai/src/open_ai.rs                  | 14 +++++++-------
 3 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/crates/language_model/src/model/cloud_model.rs b/crates/language_model/src/model/cloud_model.rs
index 57ed28d625..d5473d2007 100644
--- a/crates/language_model/src/model/cloud_model.rs
+++ b/crates/language_model/src/model/cloud_model.rs
@@ -80,7 +80,7 @@ impl CloudModel {
                 | open_ai::Model::FourOmni
                 | open_ai::Model::FourOmniMini
                 | open_ai::Model::O1Mini
-                | open_ai::Model::O1Preview
+                | open_ai::Model::O1
                 | open_ai::Model::Custom { .. } => {
                     LanguageModelAvailability::RequiresPlan(Plan::ZedPro)
                 }
diff --git a/crates/language_models/src/provider/open_ai.rs b/crates/language_models/src/provider/open_ai.rs
index 5c740f93e6..2e0c3cffab 100644
--- a/crates/language_models/src/provider/open_ai.rs
+++ b/crates/language_models/src/provider/open_ai.rs
@@ -362,9 +362,7 @@ pub fn count_open_ai_tokens(
                 .collect::<Vec<_>>();
 
             match model {
-                open_ai::Model::Custom { .. }
-                | open_ai::Model::O1Mini
-                | open_ai::Model::O1Preview => {
+                open_ai::Model::Custom { .. } | open_ai::Model::O1Mini | open_ai::Model::O1 => {
                     tiktoken_rs::num_tokens_from_messages("gpt-4", &messages)
                 }
                 _ => tiktoken_rs::num_tokens_from_messages(model.id(), &messages),
diff --git a/crates/open_ai/src/open_ai.rs b/crates/open_ai/src/open_ai.rs
index 1e92841249..bf32910f6f 100644
--- a/crates/open_ai/src/open_ai.rs
+++ b/crates/open_ai/src/open_ai.rs
@@ -72,8 +72,8 @@ pub enum Model {
     FourOmni,
     #[serde(rename = "gpt-4o-mini", alias = "gpt-4o-mini")]
     FourOmniMini,
-    #[serde(rename = "o1-preview", alias = "o1-preview")]
-    O1Preview,
+    #[serde(rename = "o1", alias = "o1-preview")]
+    O1,
     #[serde(rename = "o1-mini", alias = "o1-mini")]
     O1Mini,
 
@@ -96,7 +96,7 @@ impl Model {
             "gpt-4-turbo-preview" => Ok(Self::FourTurbo),
             "gpt-4o" => Ok(Self::FourOmni),
             "gpt-4o-mini" => Ok(Self::FourOmniMini),
-            "o1-preview" => Ok(Self::O1Preview),
+            "o1" => Ok(Self::O1),
             "o1-mini" => Ok(Self::O1Mini),
             _ => Err(anyhow!("invalid model id")),
         }
@@ -109,7 +109,7 @@ impl Model {
             Self::FourTurbo => "gpt-4-turbo",
             Self::FourOmni => "gpt-4o",
             Self::FourOmniMini => "gpt-4o-mini",
-            Self::O1Preview => "o1-preview",
+            Self::O1 => "o1",
             Self::O1Mini => "o1-mini",
             Self::Custom { name, .. } => name,
         }
@@ -122,7 +122,7 @@ impl Model {
             Self::FourTurbo => "gpt-4-turbo",
             Self::FourOmni => "gpt-4o",
             Self::FourOmniMini => "gpt-4o-mini",
-            Self::O1Preview => "o1-preview",
+            Self::O1 => "o1",
             Self::O1Mini => "o1-mini",
             Self::Custom {
                 name, display_name, ..
@@ -137,7 +137,7 @@ impl Model {
             Self::FourTurbo => 128000,
             Self::FourOmni => 128000,
             Self::FourOmniMini => 128000,
-            Self::O1Preview => 128000,
+            Self::O1 => 128000,
             Self::O1Mini => 128000,
             Self::Custom { max_tokens, .. } => *max_tokens,
         }
@@ -475,7 +475,7 @@ pub async fn stream_completion(
     api_key: &str,
     request: Request,
 ) -> Result<BoxStream<'static, Result<ResponseStreamEvent>>> {
-    if request.model == "o1-preview" || request.model == "o1-mini" {
+    if request.model.starts_with("o1") {
         let response = complete(client, api_url, api_key, request).await;
         let response_stream_event = response.map(adapt_response_to_stream);
         return Ok(stream::once(future::ready(response_stream_event)).boxed());