cooprate with list avalibale models, remove model options

evanliu048 · evanliu048 · commit 1d6fb475d2fc · 2025-07-24T17:43:21.000-07:00
diff --git a/crates/chat-cli/src/api_client/error.rs b/crates/chat-cli/src/api_client/error.rs
@@ -1,6 +1,7 @@
 use amzn_codewhisperer_client::operation::create_subscription_token::CreateSubscriptionTokenError;
 use amzn_codewhisperer_client::operation::generate_completions::GenerateCompletionsError;
 use amzn_codewhisperer_client::operation::list_available_customizations::ListAvailableCustomizationsError;
+use amzn_codewhisperer_client::operation::list_available_models::ListAvailableModelsError;
 use amzn_codewhisperer_client::operation::list_available_profiles::ListAvailableProfilesError;
 use amzn_codewhisperer_client::operation::send_telemetry_event::SendTelemetryEventError;
 pub use amzn_codewhisperer_streaming_client::operation::generate_assistant_response::GenerateAssistantResponseError;
@@ -93,6 +94,9 @@ pub enum ApiClientError {
     // Credential errors
     #[error("failed to load credentials: {}", .0)]
     Credentials(CredentialsError),
+
+    #[error(transparent)]
+    ListAvailableModelsError(#[from] SdkError<ListAvailableModelsError, HttpResponse>),
 }
 
 impl ApiClientError {
@@ -116,6 +120,7 @@ impl ApiClientError {
             Self::ModelOverloadedError { status_code, .. } => *status_code,
             Self::MonthlyLimitReached { status_code } => *status_code,
             Self::Credentials(_e) => None,
+            Self::ListAvailableModelsError(e) => sdk_status_code(e),
         }
     }
 }
@@ -141,6 +146,7 @@ impl ReasonCode for ApiClientError {
             Self::ModelOverloadedError { .. } => "ModelOverloadedError".to_string(),
             Self::MonthlyLimitReached { .. } => "MonthlyLimitReached".to_string(),
             Self::Credentials(_) => "CredentialsError".to_string(),
+            Self::ListAvailableModelsError(e) => sdk_error_code(e),
         }
     }
 }
@@ -188,6 +194,10 @@ mod tests {
                 ListAvailableCustomizationsError::unhandled("<unhandled>"),
                 response(),
             )),
+            ApiClientError::ListAvailableModelsError(SdkError::service_error(
+                ListAvailableModelsError::unhandled("<unhandled>"),
+                response(),
+            )),
             ApiClientError::ListAvailableServices(SdkError::service_error(
                 ListCustomizationsError::unhandled("<unhandled>"),
                 response(),
diff --git a/crates/chat-cli/src/api_client/mod.rs b/crates/chat-cli/src/api_client/mod.rs
@@ -12,7 +12,9 @@ use std::time::Duration;
 
 use amzn_codewhisperer_client::Client as CodewhispererClient;
 use amzn_codewhisperer_client::operation::create_subscription_token::CreateSubscriptionTokenOutput;
+use amzn_codewhisperer_client::types::Origin::Cli;
 use amzn_codewhisperer_client::types::{
+    Model,
     OptOutPreference,
     SubscriptionStatus,
     TelemetryEvent,
@@ -232,6 +234,46 @@ impl ApiClient {
         Ok(profiles)
     }
 
+    pub async fn list_available_models(&self) -> Result<(Vec<Model>, Option<Model>), ApiClientError> {
+        if cfg!(test) {
+            return Ok((
+                vec![
+                    Model::builder()
+                        .model_id("model-1")
+                        .description("Test Model 1")
+                        .build()
+                        .unwrap(),
+                ],
+                Some(
+                    Model::builder()
+                        .model_id("model-1")
+                        .description("Test Model 1")
+                        .build()
+                        .unwrap(),
+                ),
+            ));
+        }
+
+        // todo yifan: add default_model once API is ready
+        let mut models = Vec::new();
+        let default_model = None;
+        let request = self
+            .client
+            .list_available_models()
+            .set_origin(Some(Cli))
+            .set_profile_arn(self.profile.as_ref().map(|p| p.arn.clone()));
+        let mut paginator = request.into_paginator().send();
+
+        while let Some(models_output) = paginator.next().await {
+            models.extend(models_output?.models().iter().cloned());
+            // if default_model.is_none() && output.default_model().is_some() {
+            //     default_model = output.default_model().cloned();
+            // }
+        }
+
+        Ok((models, default_model))
+    }
+
     pub async fn create_subscription_token(&self) -> Result<CreateSubscriptionTokenOutput, ApiClientError> {
         if cfg!(test) {
             return Ok(CreateSubscriptionTokenOutput::builder()
diff --git a/crates/chat-cli/src/cli/chat/cli/mod.rs b/crates/chat-cli/src/cli/chat/cli/mod.rs
@@ -108,7 +108,7 @@ impl SlashCommand {
             Self::Hooks(args) => args.execute(session).await,
             Self::Usage(args) => args.execute(os, session).await,
             Self::Mcp(args) => args.execute(session).await,
-            Self::Model(args) => args.execute(session).await,
+            Self::Model(args) => args.execute(os, session).await,
             Self::Subscribe(args) => args.execute(os, session).await,
             Self::Persist(subcommand) => subcommand.execute(os, session).await,
             // Self::Root(subcommand) => {
diff --git a/crates/chat-cli/src/cli/chat/cli/model.rs b/crates/chat-cli/src/cli/chat/cli/model.rs
@@ -20,44 +20,63 @@ use crate::cli::chat::{
 };
 use crate::os::Os;
 
-pub struct ModelOption {
-    pub name: &'static str,
-    pub model_id: &'static str,
-}
+// pub struct ModelOption {
+//     pub name: &'static str,
+//     pub model_id: &'static str,
+// }
 
-pub const MODEL_OPTIONS: [ModelOption; 2] = [
-    ModelOption {
-        name: "claude-4-sonnet",
-        model_id: "CLAUDE_SONNET_4_20250514_V1_0",
-    },
-    ModelOption {
-        name: "claude-3.7-sonnet",
-        model_id: "CLAUDE_3_7_SONNET_20250219_V1_0",
-    },
-];
+// pub const MODEL_OPTIONS: [ModelOption; 2] = [
+//     ModelOption {
+//         name: "claude-4-sonnet",
+//         model_id: "CLAUDE_SONNET_4_20250514_V1_0",
+//     },
+//     ModelOption {
+//         name: "claude-3.7-sonnet",
+//         model_id: "CLAUDE_3_7_SONNET_20250219_V1_0",
+//     },
+// ];
 
 #[deny(missing_docs)]
 #[derive(Debug, PartialEq, Args)]
 pub struct ModelArgs;
 
 impl ModelArgs {
-    pub async fn execute(self, session: &mut ChatSession) -> Result<ChatState, ChatError> {
-        Ok(select_model(session)?.unwrap_or(ChatState::PromptUser {
+    pub async fn execute(self, os: &mut Os, session: &mut ChatSession) -> Result<ChatState, ChatError> {
+        Ok(select_model(os, session).await?.unwrap_or(ChatState::PromptUser {
             skip_printing_tools: false,
         }))
     }
 }
 
-pub fn select_model(session: &mut ChatSession) -> Result<Option<ChatState>, ChatError> {
+pub async fn select_model(os: &mut Os, session: &mut ChatSession) -> Result<Option<ChatState>, ChatError> {
     queue!(session.stderr, style::Print("\n"))?;
+
+    // Fetch available models from service
+    let (models, _default_model) = os
+        .client
+        .list_available_models()
+        .await
+        .map_err(|e| ChatError::Custom(format!("Failed to fetch available models: {}", e).into()))?;
+
+    if models.is_empty() {
+        queue!(
+            session.stderr,
+            style::SetForegroundColor(Color::Red),
+            style::Print("No models available\n"),
+            style::ResetColor
+        )?;
+        return Ok(None);
+    }
+
     let active_model_id = session.conversation.model.as_deref();
-    let labels: Vec<String> = MODEL_OPTIONS
+
+    let labels: Vec<String> = models
         .iter()
         .map(|opt| {
-            if (opt.model_id.is_empty() && active_model_id.is_none()) || Some(opt.model_id) == active_model_id {
-                format!("{} (active)", opt.name)
+            if Some(opt.model_id.as_str()) == active_model_id {
+                format!("{} (active)", opt.model_id)
             } else {
-                opt.name.to_owned()
+                opt.model_id.to_owned()
             }
         })
         .collect();
@@ -83,14 +102,14 @@ pub fn select_model(session: &mut ChatSession) -> Result<Option<ChatState>, Chat
     queue!(session.stderr, style::ResetColor)?;
 
     if let Some(index) = selection {
-        let selected = &MODEL_OPTIONS[index];
+        let selected = &models[index];
         let model_id_str = selected.model_id.to_string();
         session.conversation.model = Some(model_id_str);
 
         queue!(
             session.stderr,
             style::Print("\n"),
-            style::Print(format!(" Using {}\n\n", selected.name)),
+            style::Print(format!(" Using {}\n\n", selected.model_id)),
             style::ResetColor,
             style::SetForegroundColor(Color::Reset),
             style::SetBackgroundColor(Color::Reset),
@@ -106,21 +125,24 @@ pub fn select_model(session: &mut ChatSession) -> Result<Option<ChatState>, Chat
 
 /// Returns Claude 3.7 for: Amazon IDC users, FRA region users
 /// Returns Claude 4.0 for: Builder ID users, other regions
-pub async fn default_model_id(os: &Os) -> &'static str {
+pub async fn default_model_id(os: &Os) -> String {
+    if let Ok((_, Some(default_model))) = os.client.list_available_models().await {
+        return default_model.model_id().to_string();
+    }
     // Check FRA region first
     if let Ok(Some(profile)) = os.database.get_auth_profile() {
         if profile.arn.split(':').nth(3) == Some("eu-central-1") {
-            return "CLAUDE_3_7_SONNET_20250219_V1_0";
+            return "CLAUDE_3_7_SONNET_20250219_V1_0".to_string();
         }
     }
 
     // Check if Amazon IDC user
     if let Ok(Some(token)) = BuilderIdToken::load(&os.database).await {
         if matches!(token.token_type(), TokenType::IamIdentityCenter) && token.is_amzn_user() {
-            return "CLAUDE_3_7_SONNET_20250219_V1_0";
+            return "CLAUDE_3_7_SONNET_20250219_V1_0".to_string();
         }
     }
 
     // Default to 4.0
-    "CLAUDE_SONNET_4_20250514_V1_0"
+    "CLAUDE_SONNET_4_20250514_V1_0".to_string()
 }
diff --git a/crates/chat-cli/src/cli/chat/mod.rs b/crates/chat-cli/src/cli/chat/mod.rs
@@ -126,10 +126,7 @@ use crate::auth::AuthError;
 use crate::auth::builder_id::is_idc_user;
 use crate::cli::agent::Agents;
 use crate::cli::chat::cli::SlashCommand;
-use crate::cli::chat::cli::model::{
-    MODEL_OPTIONS,
-    default_model_id,
-};
+use crate::cli::chat::cli::model::default_model_id;
 use crate::cli::chat::cli::prompts::{
     GetPromptError,
     PromptsSubcommand,
@@ -267,15 +264,16 @@ impl ChatArgs {
         };
 
         // If modelId is specified, verify it exists before starting the chat
-        let model_id: Option<String> = if let Some(model_name) = self.model {
-            let model_name_lower = model_name.to_lowercase();
-            match MODEL_OPTIONS.iter().find(|opt| opt.name == model_name_lower) {
+        let model_id: Option<String> = if let Some(requested_model_id) = self.model {
+            let requested_model_id_lower = requested_model_id.to_lowercase();
+            let (models, _) = os.client.list_available_models().await?;
+            match models.iter().find(|opt| opt.model_id == requested_model_id_lower) {
                 Some(opt) => Some((opt.model_id).to_string()),
                 None => {
-                    let available_names: Vec<&str> = MODEL_OPTIONS.iter().map(|opt| opt.name).collect();
+                    let available_names: Vec<&str> = models.iter().map(|opt| opt.model_id()).collect();
                     bail!(
                         "Model '{}' does not exist. Available models: {}",
-                        model_name,
+                        requested_model_id,
                         available_names.join(", ")
                     );
                 },
@@ -516,17 +514,19 @@ impl ChatSession {
         tool_config: HashMap<String, ToolSpec>,
         interactive: bool,
     ) -> Result<Self> {
+        let (models, _) = os.client.list_available_models().await?;
+
         let valid_model_id = match model_id {
             Some(id) => id,
             None => {
                 let from_settings = os
                     .database
                     .settings
                     .get_string(Setting::ChatDefaultModel)
-                    .and_then(|model_name| {
-                        MODEL_OPTIONS
+                    .and_then(|model_id| {
+                        models
                             .iter()
-                            .find(|opt| opt.name == model_name)
+                            .find(|opt| opt.model_id == model_id)
                             .map(|opt| opt.model_id.to_owned())
                     });
 
@@ -1133,15 +1133,13 @@ impl ChatSession {
         self.stderr.flush()?;
 
         if let Some(ref id) = self.conversation.model {
-            if let Some(model_option) = MODEL_OPTIONS.iter().find(|option| option.model_id == *id) {
-                execute!(
-                    self.stderr,
-                    style::SetForegroundColor(Color::Cyan),
-                    style::Print(format!("🤖 You are chatting with {}\n", model_option.name)),
-                    style::SetForegroundColor(Color::Reset),
-                    style::Print("\n")
-                )?;
-            }
+            execute!(
+                self.stderr,
+                style::SetForegroundColor(Color::Cyan),
+                style::Print(format!("🤖 You are chatting with {}\n", id)),
+                style::SetForegroundColor(Color::Reset),
+                style::Print("\n")
+            )?;
         }
 
         if let Some(user_input) = self.initial_input.take() {
@@ -2321,7 +2319,7 @@ impl ChatSession {
     }
 
     async fn retry_model_overload(&mut self, os: &mut Os) -> Result<ChatState, ChatError> {
-        match select_model(self) {
+        match select_model(os, self).await {
             Ok(Some(_)) => (),
             Ok(None) => {
                 // User did not select a model, so reset the current request state.
