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

feat: model selection for chat

Author: chaynabors

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

@@ -109,6 +109,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(
@@ -140,10 +141,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())
                         }
@@ -235,6 +248,7 @@ mod tests {
                     content: "Hello".into(),
                     user_input_message_context: None,
                     user_intent: None,
+                    model_id: Some("model".to_owned()),
                 },
                 history: None,
             })
@@ -261,13 +275,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

@@ -68,6 +68,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 { request_id: Option<String> },
 }
 
 impl ReasonCode for ApiClientError {

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/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,9 @@ 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`.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub current_model_id: Option<String>,
 }
 
 impl ConversationState {
@@ -115,6 +118,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 +161,7 @@ impl ConversationState {
             context_message_length: None,
             latest_summary: None,
             updates,
+            current_model_id,
         }
     }
 
@@ -528,6 +533,7 @@ impl ConversationState {
             context_messages,
             dropped_context_files,
             tools: &self.tools,
+            model_id: self.current_model_id.as_deref(),
         }
     }
 
@@ -599,6 +605,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 +837,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 +854,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 +1068,7 @@ mod tests {
             None,
             None,
             tool_manager,
+            None,
         )
         .await;
 
@@ -1089,6 +1099,7 @@ mod tests {
             None,
             None,
             tool_manager.clone(),
+            None,
         )
         .await;
         conversation_state.set_next_user_message("start".to_string()).await;
@@ -1120,6 +1131,7 @@ mod tests {
             None,
             None,
             tool_manager.clone(),
+            None,
         )
         .await;
         conversation_state.set_next_user_message("start".to_string()).await;
@@ -1165,6 +1177,7 @@ mod tests {
             None,
             None,
             tool_manager,
+            None,
         )
         .await;
 
@@ -1235,6 +1248,7 @@ mod tests {
             None,
             Some(SharedWriter::stdout()),
             tool_manager,
+            None,
         )
         .await;
 

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

@@ -128,6 +128,7 @@ impl UserMessage {
                 ..Default::default()
             }),
             user_intent: None,
+            model_id: None,
         }
     }
 
@@ -158,6 +159,7 @@ impl UserMessage {
                 ..Default::default()
             }),
             user_intent: None,
+            model_id: None,
         }
     }