Merge pull request #135 from aws/evannliu/model-selection

feat: model selection for chat
merge for getting a bugbash build

Author: evanliu048

crates/chat-cli/src/api_client/clients/client.rs

@@ -78,6 +78,7 @@ impl Client {
         telemetry_event: TelemetryEvent,
         user_context: UserContext,
         telemetry_enabled: bool,
+        model_id: Option<String>,
     ) -> Result<(), ApiClientError> {
         match &self.inner {
             inner::Inner::Codewhisperer(client) => {
@@ -90,6 +91,7 @@ impl Client {
                         false => OptOutPreference::OptOut,
                     })
                     .set_profile_arn(self.profile.as_ref().map(|p| p.arn.clone()))
+                    .set_model_id(model_id)
                     .send()
                     .await;
                 Ok(())
@@ -159,6 +161,7 @@ mod tests {
                     .build()
                     .unwrap(),
                 false,
+                Some("model".to_owned()),
             )
             .await
             .unwrap();

crates/chat-cli/src/api_client/clients/streaming_client.rs

@@ -139,6 +139,7 @@ impl StreamingClient {
 
         match &self.inner {
             inner::Inner::Codewhisperer(client) => {
+                let model_id_opt: Option<String> = user_input_message.model_id.clone();
                 let conversation_state = amzn_codewhisperer_streaming_client::types::ConversationState::builder()
                     .set_conversation_id(conversation_id)
                     .current_message(
@@ -170,10 +171,22 @@ impl StreamingClient {
                                 && err.meta().message() == Some("Input is too long."))
                         });
 
+                        let is_model_unavailable = model_id_opt.is_some()
+                            && e.raw_response().is_some_and(|resp| resp.status().as_u16() == 500)
+                            && e.as_service_error().is_some_and(|err| {
+                                err.meta().message()
+                                == Some("Encountered unexpectedly high load when processing the request, please try again.")
+                            });
                         if is_quota_breach {
                             Err(ApiClientError::QuotaBreach("quota has reached its limit"))
                         } else if is_context_window_overflow {
                             Err(ApiClientError::ContextWindowOverflow)
+                        } else if is_model_unavailable {
+                            let request_id = e
+                                .as_service_error()
+                                .and_then(|err| err.meta().request_id())
+                                .map(|s| s.to_string());
+                            Err(ApiClientError::ModelOverloadedError(request_id))
                         } else {
                             Err(e.into())
                         }
@@ -291,6 +304,7 @@ mod tests {
                     content: "Hello".into(),
                     user_input_message_context: None,
                     user_intent: None,
+                    model_id: Some("model".to_owned()),
                 },
                 history: None,
             })
@@ -317,13 +331,15 @@ mod tests {
                     content: "How about rustc?".into(),
                     user_input_message_context: None,
                     user_intent: None,
+                    model_id: Some("model".to_owned()),
                 },
                 history: Some(vec![
                     ChatMessage::UserInputMessage(UserInputMessage {
                         images: None,
                         content: "What language is the linux kernel written in, and who wrote it?".into(),
                         user_input_message_context: None,
                         user_intent: None,
+                        model_id: None,
                     }),
                     ChatMessage::AssistantResponseMessage(AssistantResponseMessage {
                         content: "It is written in C by Linus Torvalds.".into(),

crates/chat-cli/src/api_client/error.rs

@@ -65,6 +65,11 @@ pub enum ApiClientError {
 
     #[error(transparent)]
     AuthError(#[from] AuthError),
+
+    #[error(
+        "The model you've selected is temporarily unavailable. Please use '/model' to select a different model and try again."
+    )]
+    ModelOverloadedError(Option<String>),
 }
 
 #[cfg(test)]

crates/chat-cli/src/api_client/model.rs

@@ -859,6 +859,7 @@ pub struct UserInputMessage {
     pub user_input_message_context: Option<UserInputMessageContext>,
     pub user_intent: Option<UserIntent>,
     pub images: Option<Vec<ImageBlock>>,
+    pub model_id: Option<String>,
 }
 
 impl From<UserInputMessage> for amzn_codewhisperer_streaming_client::types::UserInputMessage {
@@ -868,6 +869,7 @@ impl From<UserInputMessage> for amzn_codewhisperer_streaming_client::types::User
             .set_images(value.images.map(|images| images.into_iter().map(Into::into).collect()))
             .set_user_input_message_context(value.user_input_message_context.map(Into::into))
             .set_user_intent(value.user_intent.map(Into::into))
+            .set_model_id(value.model_id)
             .origin(amzn_codewhisperer_streaming_client::types::Origin::Cli)
             .build()
             .expect("Failed to build UserInputMessage")
@@ -881,6 +883,7 @@ impl From<UserInputMessage> for amzn_qdeveloper_streaming_client::types::UserInp
             .set_images(value.images.map(|images| images.into_iter().map(Into::into).collect()))
             .set_user_input_message_context(value.user_input_message_context.map(Into::into))
             .set_user_intent(value.user_intent.map(Into::into))
+            .set_model_id(value.model_id)
             .origin(amzn_qdeveloper_streaming_client::types::Origin::Cli)
             .build()
             .expect("Failed to build UserInputMessage")
@@ -976,6 +979,7 @@ mod tests {
                 })]),
             }),
             user_intent: Some(UserIntent::ApplyCommonBestPractices),
+            model_id: Some("model id".to_string()),
         };
 
         let codewhisper_input =
@@ -989,6 +993,7 @@ mod tests {
             content: "test content".to_string(),
             user_input_message_context: None,
             user_intent: None,
+            model_id: Some("model id".to_string()),
         };
 
         let codewhisper_minimal =

crates/chat-cli/src/cli/chat/cli.rs

@@ -0,0 +1,151 @@
+use std::collections::HashMap;
+
+use clap::{
+    Args,
+    Parser,
+    Subcommand,
+    ValueEnum,
+};
+
+#[derive(Debug, Clone, PartialEq, Eq, Default, Parser)]
+pub struct Chat {
+    /// (Deprecated, use --trust-all-tools) Enabling this flag allows the model to execute
+    /// all commands without first accepting them.
+    #[arg(short, long, hide = true)]
+    pub accept_all: bool,
+    /// Print the first response to STDOUT without interactive mode. This will fail if the
+    /// prompt requests permissions to use a tool, unless --trust-all-tools is also used.
+    #[arg(long)]
+    pub no_interactive: bool,
+    /// Resumes the previous conversation from this directory.
+    #[arg(short, long)]
+    pub resume: bool,
+    /// The first question to ask
+    pub input: Option<String>,
+    /// Context profile to use
+    #[arg(long = "profile")]
+    pub profile: Option<String>,
+    /// Current model to use
+    #[arg(long = "model")]
+    pub model: Option<String>,
+    /// Allows the model to use any tool to run commands without asking for confirmation.
+    #[arg(long)]
+    pub trust_all_tools: bool,
+    /// Trust only this set of tools. Example: trust some tools:
+    /// '--trust-tools=fs_read,fs_write', trust no tools: '--trust-tools='
+    #[arg(long, value_delimiter = ',', value_name = "TOOL_NAMES")]
+    pub trust_tools: Option<Vec<String>>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Subcommand)]
+pub enum Mcp {
+    /// Add or replace a configured server
+    Add(McpAdd),
+    /// Remove a server from the MCP configuration
+    #[command(alias = "rm")]
+    Remove(McpRemove),
+    /// List configured servers
+    List(McpList),
+    /// Import a server configuration from another file
+    Import(McpImport),
+    /// Get the status of a configured server
+    Status {
+        #[arg(long)]
+        name: String,
+    },
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Args)]
+pub struct McpAdd {
+    /// Name for the server
+    #[arg(long)]
+    pub name: String,
+    /// The command used to launch the server
+    #[arg(long)]
+    pub command: String,
+    /// Where to add the server to.
+    #[arg(long, value_enum)]
+    pub scope: Option<Scope>,
+    /// Environment variables to use when launching the server
+    #[arg(long, value_parser = parse_env_vars)]
+    pub env: Vec<HashMap<String, String>>,
+    /// Server launch timeout, in milliseconds
+    #[arg(long)]
+    pub timeout: Option<u64>,
+    /// Overwrite an existing server with the same name
+    #[arg(long, default_value_t = false)]
+    pub force: bool,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Args)]
+pub struct McpRemove {
+    #[arg(long)]
+    pub name: String,
+    #[arg(long, value_enum)]
+    pub scope: Option<Scope>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Args)]
+pub struct McpList {
+    #[arg(value_enum)]
+    pub scope: Option<Scope>,
+    #[arg(long, hide = true)]
+    pub profile: Option<String>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Args)]
+pub struct McpImport {
+    #[arg(long)]
+    pub file: String,
+    #[arg(value_enum)]
+    pub scope: Option<Scope>,
+    /// Overwrite an existing server with the same name
+    #[arg(long, default_value_t = false)]
+    pub force: bool,
+}
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq, ValueEnum)]
+pub enum Scope {
+    Workspace,
+    Global,
+}
+
+impl std::fmt::Display for Scope {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Scope::Workspace => write!(f, "workspace"),
+            Scope::Global => write!(f, "global"),
+        }
+    }
+}
+
+#[derive(Debug)]
+struct EnvVarParseError(String);
+
+impl std::fmt::Display for EnvVarParseError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "Failed to parse environment variables: {}", self.0)
+    }
+}
+
+impl std::error::Error for EnvVarParseError {}
+
+fn parse_env_vars(arg: &str) -> Result<HashMap<String, String>, EnvVarParseError> {
+    let mut vars = HashMap::new();
+
+    for pair in arg.split(",") {
+        match pair.split_once('=') {
+            Some((key, value)) => {
+                vars.insert(key.trim().to_string(), value.trim().to_string());
+            },
+            None => {
+                return Err(EnvVarParseError(format!(
+                    "Invalid environment variable '{}'. Expected 'name=value'",
+                    pair
+                )));
+            },
+        }
+    }
+
+    Ok(vars)
+}

crates/chat-cli/src/cli/chat/command.rs

@@ -59,6 +59,7 @@ pub enum Command {
         force: bool,
     },
     Mcp,
+    Model,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -839,6 +840,7 @@ impl Command {
                     Self::Save { path, force }
                 },
                 "mcp" => Self::Mcp,
+                "model" => Self::Model,
                 unknown_command => {
                     let looks_like_path = {
                         let after_slash_command_str = parts[1..].join(" ");

crates/chat-cli/src/cli/chat/conversation_state.rs

@@ -105,6 +105,10 @@ pub struct ConversationState {
     latest_summary: Option<String>,
     #[serde(skip)]
     pub updates: Option<SharedWriter>,
+    /// Model explicitly selected by the user in this conversation state via `/model`. (`None` ==
+    /// auto)
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub current_model_id: Option<String>,
 }
 
 impl ConversationState {
@@ -115,6 +119,7 @@ impl ConversationState {
         profile: Option<String>,
         updates: Option<SharedWriter>,
         tool_manager: ToolManager,
+        current_model_id: Option<String>,
     ) -> Self {
         // Initialize context manager
         let context_manager = match ContextManager::new(ctx, None).await {
@@ -157,6 +162,7 @@ impl ConversationState {
             context_message_length: None,
             latest_summary: None,
             updates,
+            current_model_id,
         }
     }
 
@@ -528,6 +534,7 @@ impl ConversationState {
             context_messages,
             dropped_context_files,
             tools: &self.tools,
+            model_id: self.current_model_id.as_deref(),
         }
     }
 
@@ -599,6 +606,7 @@ impl ConversationState {
             user_input_message_context: None,
             user_intent: None,
             images: None,
+            model_id: self.current_model_id.clone(),
         };
 
         // If the last message contains tool uses, then add cancelled tool results to the summary
@@ -830,6 +838,7 @@ pub struct BackendConversationStateImpl<'a, T, U> {
     pub context_messages: U,
     pub dropped_context_files: Vec<(String, String)>,
     pub tools: &'a HashMap<ToolOrigin, Vec<Tool>>,
+    pub model_id: Option<&'a str>,
 }
 
 impl
@@ -846,6 +855,7 @@ impl
             .cloned()
             .map(UserMessage::into_user_input_message)
             .ok_or(eyre::eyre!("next user message is not set"))?;
+        user_input_message.model_id = self.model_id.map(str::to_string);
         if let Some(ctx) = user_input_message.user_input_message_context.as_mut() {
             ctx.tools = Some(self.tools.values().flatten().cloned().collect::<Vec<_>>());
         }
@@ -1059,6 +1069,7 @@ mod tests {
             None,
             None,
             tool_manager,
+            None,
         )
         .await;
 
@@ -1089,6 +1100,7 @@ mod tests {
             None,
             None,
             tool_manager.clone(),
+            None,
         )
         .await;
         conversation_state.set_next_user_message("start".to_string()).await;
@@ -1120,6 +1132,7 @@ mod tests {
             None,
             None,
             tool_manager.clone(),
+            None,
         )
         .await;
         conversation_state.set_next_user_message("start".to_string()).await;
@@ -1165,6 +1178,7 @@ mod tests {
             None,
             None,
             tool_manager,
+            None,
         )
         .await;
 
@@ -1235,6 +1249,7 @@ mod tests {
             None,
             Some(SharedWriter::stdout()),
             tool_manager,
+            None,
         )
         .await;