Skip to content

Commit 5ee867f

Browse files
brandonskiserbrandonskiser
authored
feat: Prompt users to select a different model when at capacity (#2360)
1 parent 3ae90d7 commit 5ee867f

File tree

2 files changed

+115
-53
lines changed

2 files changed

+115
-53
lines changed

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

Lines changed: 54 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -42,60 +42,66 @@ pub struct ModelArgs;
4242

4343
impl ModelArgs {
4444
pub async fn execute(self, session: &mut ChatSession) -> Result<ChatState, ChatError> {
45-
queue!(session.stderr, style::Print("\n"))?;
46-
let active_model_id = session.conversation.model.as_deref();
47-
let labels: Vec<String> = MODEL_OPTIONS
48-
.iter()
49-
.map(|opt| {
50-
if (opt.model_id.is_empty() && active_model_id.is_none()) || Some(opt.model_id) == active_model_id {
51-
format!("{} (active)", opt.name)
52-
} else {
53-
opt.name.to_owned()
54-
}
55-
})
56-
.collect();
57-
58-
let selection: Option<_> = match Select::with_theme(&crate::util::dialoguer_theme())
59-
.with_prompt("Select a model for this chat session")
60-
.items(&labels)
61-
.default(0)
62-
.interact_on_opt(&dialoguer::console::Term::stdout())
63-
{
64-
Ok(sel) => {
65-
let _ = crossterm::execute!(
66-
std::io::stdout(),
67-
crossterm::style::SetForegroundColor(crossterm::style::Color::Magenta)
68-
);
69-
sel
70-
},
71-
// Ctrl‑C -> Err(Interrupted)
72-
Err(dialoguer::Error::IO(ref e)) if e.kind() == std::io::ErrorKind::Interrupted => None,
73-
Err(e) => return Err(ChatError::Custom(format!("Failed to choose model: {e}").into())),
74-
};
45+
Ok(select_model(session)?.unwrap_or(ChatState::PromptUser {
46+
skip_printing_tools: false,
47+
}))
48+
}
49+
}
7550

76-
queue!(session.stderr, style::ResetColor)?;
51+
pub fn select_model(session: &mut ChatSession) -> Result<Option<ChatState>, ChatError> {
52+
queue!(session.stderr, style::Print("\n"))?;
53+
let active_model_id = session.conversation.model.as_deref();
54+
let labels: Vec<String> = MODEL_OPTIONS
55+
.iter()
56+
.map(|opt| {
57+
if (opt.model_id.is_empty() && active_model_id.is_none()) || Some(opt.model_id) == active_model_id {
58+
format!("{} (active)", opt.name)
59+
} else {
60+
opt.name.to_owned()
61+
}
62+
})
63+
.collect();
7764

78-
if let Some(index) = selection {
79-
let selected = &MODEL_OPTIONS[index];
80-
let model_id_str = selected.model_id.to_string();
81-
session.conversation.model = Some(model_id_str);
65+
let selection: Option<_> = match Select::with_theme(&crate::util::dialoguer_theme())
66+
.with_prompt("Select a model for this chat session")
67+
.items(&labels)
68+
.default(0)
69+
.interact_on_opt(&dialoguer::console::Term::stdout())
70+
{
71+
Ok(sel) => {
72+
let _ = crossterm::execute!(
73+
std::io::stdout(),
74+
crossterm::style::SetForegroundColor(crossterm::style::Color::Magenta)
75+
);
76+
sel
77+
},
78+
// Ctrl‑C -> Err(Interrupted)
79+
Err(dialoguer::Error::IO(ref e)) if e.kind() == std::io::ErrorKind::Interrupted => return Ok(None),
80+
Err(e) => return Err(ChatError::Custom(format!("Failed to choose model: {e}").into())),
81+
};
8282

83-
queue!(
84-
session.stderr,
85-
style::Print("\n"),
86-
style::Print(format!(" Using {}\n\n", selected.name)),
87-
style::ResetColor,
88-
style::SetForegroundColor(Color::Reset),
89-
style::SetBackgroundColor(Color::Reset),
90-
)?;
91-
}
83+
queue!(session.stderr, style::ResetColor)?;
9284

93-
execute!(session.stderr, style::ResetColor)?;
85+
if let Some(index) = selection {
86+
let selected = &MODEL_OPTIONS[index];
87+
let model_id_str = selected.model_id.to_string();
88+
session.conversation.model = Some(model_id_str);
9489

95-
Ok(ChatState::PromptUser {
96-
skip_printing_tools: false,
97-
})
90+
queue!(
91+
session.stderr,
92+
style::Print("\n"),
93+
style::Print(format!(" Using {}\n\n", selected.name)),
94+
style::ResetColor,
95+
style::SetForegroundColor(Color::Reset),
96+
style::SetBackgroundColor(Color::Reset),
97+
)?;
9898
}
99+
100+
execute!(session.stderr, style::ResetColor)?;
101+
102+
Ok(Some(ChatState::PromptUser {
103+
skip_printing_tools: false,
104+
}))
99105
}
100106

101107
/// Returns Claude 3.7 for: Amazon IDC users, FRA region users

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

Lines changed: 61 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ use clap::{
3838
Parser,
3939
};
4040
use cli::compact::CompactStrategy;
41+
use cli::model::select_model;
4142
pub use conversation::ConversationState;
4243
use conversation::TokenWarningLevel;
4344
use crossterm::style::{
@@ -641,6 +642,12 @@ impl ChatSession {
641642
Err(ChatError::Interrupted { tool_uses: None })
642643
}
643644
},
645+
ChatState::RetryModelOverload => tokio::select! {
646+
res = self.retry_model_overload(os) => res,
647+
Ok(_) = ctrl_c_stream => {
648+
Err(ChatError::Interrupted { tool_uses: None })
649+
}
650+
},
644651
ChatState::Exit => return Ok(()),
645652
};
646653

@@ -773,12 +780,30 @@ impl ChatSession {
773780
("Amazon Q is having trouble responding right now", eyre!(err), false)
774781
},
775782
ApiClientError::ModelOverloadedError { request_id, .. } => {
776-
let model_instruction = if self.interactive {
777-
"Please use '/model' to select a different model and try again."
778-
} else {
779-
"Please relaunch with '--model <model_id>' to use a different model."
780-
};
783+
if self.interactive {
784+
execute!(
785+
self.stderr,
786+
style::SetAttribute(Attribute::Bold),
787+
style::SetForegroundColor(Color::Red),
788+
style::Print(
789+
"\nThe model you've selected is temporarily unavailable. Please select a different model.\n"
790+
),
791+
style::SetAttribute(Attribute::Reset),
792+
style::SetForegroundColor(Color::Reset),
793+
)?;
794+
795+
if let Some(id) = request_id {
796+
self.conversation
797+
.append_transcript(format!("Model unavailable (Request ID: {})", id));
798+
}
781799

800+
self.inner = Some(ChatState::RetryModelOverload);
801+
802+
return Ok(());
803+
}
804+
805+
// non-interactive throws this error
806+
let model_instruction = "Please relaunch with '--model <model_id>' to use a different model.";
782807
let err = format!(
783808
"The model you've selected is temporarily unavailable. {}{}\n\n",
784809
model_instruction,
@@ -948,6 +973,8 @@ pub enum ChatState {
948973
/// Parameters for how to perform the compaction request.
949974
strategy: CompactStrategy,
950975
},
976+
/// Retry the current request if we encounter a model overloaded error.
977+
RetryModelOverload,
951978
/// Exit the chat.
952979
Exit,
953980
}
@@ -2145,6 +2172,35 @@ impl ChatSession {
21452172
Ok(ChatState::ExecuteTools)
21462173
}
21472174

2175+
async fn retry_model_overload(&mut self, os: &mut Os) -> Result<ChatState, ChatError> {
2176+
match select_model(self) {
2177+
Ok(Some(_)) => (),
2178+
Ok(None) => {
2179+
// User did not select a model, so reset the current request state.
2180+
self.conversation.enforce_conversation_invariants();
2181+
self.conversation.reset_next_user_message();
2182+
self.pending_tool_index = None;
2183+
return Ok(ChatState::PromptUser {
2184+
skip_printing_tools: false,
2185+
});
2186+
},
2187+
Err(err) => return Err(err),
2188+
}
2189+
2190+
let conv_state = self
2191+
.conversation
2192+
.as_sendable_conversation_state(os, &mut self.stderr, true)
2193+
.await?;
2194+
2195+
if self.interactive {
2196+
self.spinner = Some(Spinner::new(Spinners::Dots, "Thinking...".to_owned()));
2197+
}
2198+
2199+
Ok(ChatState::HandleResponseStream(
2200+
os.client.send_message(conv_state).await?,
2201+
))
2202+
}
2203+
21482204
/// Apply program context to tools that Q may not have.
21492205
// We cannot attach this any other way because Tools are constructed by deserializing
21502206
// output from Amazon Q.

0 commit comments

Comments
 (0)