Add xAI's Grok Code Fast 1 model (#36959)

Daniel Dye created 2 months ago

Release Notes:

- Add the `grok-code-fast-1` model to xAI's list of available models.

Change summary

crates/language_models/src/provider/x_ai.rs |  2 +-
crates/x_ai/src/x_ai.rs                     | 14 ++++++++++----
2 files changed, 11 insertions(+), 5 deletions(-)

Detailed changes

crates/language_models/src/provider/x_ai.rs

@@ -319,7 +319,7 @@ impl LanguageModel for XAiLanguageModel {
     }
     fn tool_input_format(&self) -> LanguageModelToolSchemaFormat {
         let model_id = self.model.id().trim().to_lowercase();
-        if model_id.eq(x_ai::Model::Grok4.id()) {
+        if model_id.eq(x_ai::Model::Grok4.id()) || model_id.eq(x_ai::Model::GrokCodeFast1.id()) {
             LanguageModelToolSchemaFormat::JsonSchemaSubset
         } else {
             LanguageModelToolSchemaFormat::JsonSchema

crates/x_ai/src/x_ai.rs

@@ -20,6 +20,8 @@ pub enum Model {
     Grok3MiniFast,
     #[serde(rename = "grok-4-latest")]
     Grok4,
+    #[serde(rename = "grok-code-fast-1")]
+    GrokCodeFast1,
     #[serde(rename = "custom")]
     Custom {
         name: String,
@@ -43,6 +45,7 @@ impl Model {
             "grok-3-mini" => Ok(Self::Grok3Mini),
             "grok-3-fast" => Ok(Self::Grok3Fast),
             "grok-3-mini-fast" => Ok(Self::Grok3MiniFast),
+            "grok-code-fast-1" => Ok(Self::GrokCodeFast1),
             _ => anyhow::bail!("invalid model id '{id}'"),
         }
     }
@@ -55,6 +58,7 @@ impl Model {
             Self::Grok3Fast => "grok-3-fast",
             Self::Grok3MiniFast => "grok-3-mini-fast",
             Self::Grok4 => "grok-4",
+            Self::GrokCodeFast1 => "grok-code-fast-1",
             Self::Custom { name, .. } => name,
         }
     }
@@ -67,6 +71,7 @@ impl Model {
             Self::Grok3Fast => "Grok 3 Fast",
             Self::Grok3MiniFast => "Grok 3 Mini Fast",
             Self::Grok4 => "Grok 4",
+            Self::GrokCodeFast1 => "Grok Code Fast 1",
             Self::Custom {
                 name, display_name, ..
             } => display_name.as_ref().unwrap_or(name),
@@ -76,7 +81,7 @@ impl Model {
     pub fn max_token_count(&self) -> u64 {
         match self {
             Self::Grok3 | Self::Grok3Mini | Self::Grok3Fast | Self::Grok3MiniFast => 131_072,
-            Self::Grok4 => 256_000,
+            Self::Grok4 | Self::GrokCodeFast1 => 256_000,
             Self::Grok2Vision => 8_192,
             Self::Custom { max_tokens, .. } => *max_tokens,
         }
@@ -85,7 +90,7 @@ impl Model {
     pub fn max_output_tokens(&self) -> Option<u64> {
         match self {
             Self::Grok3 | Self::Grok3Mini | Self::Grok3Fast | Self::Grok3MiniFast => Some(8_192),
-            Self::Grok4 => Some(64_000),
+            Self::Grok4 | Self::GrokCodeFast1 => Some(64_000),
             Self::Grok2Vision => Some(4_096),
             Self::Custom {
                 max_output_tokens, ..
@@ -101,7 +106,7 @@ impl Model {
             | Self::Grok3Fast
             | Self::Grok3MiniFast
             | Self::Grok4 => true,
-            Model::Custom { .. } => false,
+            Self::GrokCodeFast1 | Model::Custom { .. } => false,
         }
     }
 
@@ -116,7 +121,8 @@ impl Model {
             | Self::Grok3Mini
             | Self::Grok3Fast
             | Self::Grok3MiniFast
-            | Self::Grok4 => true,
+            | Self::Grok4
+            | Self::GrokCodeFast1 => true,
             Model::Custom { .. } => false,
         }
     }