Skip to content

Commit 5e54128

Browse files
committed
change model throttled error checking condition
1 parent 6ed98e3 commit 5e54128

File tree

1 file changed

+29
-30
lines changed
  • crates/chat-cli/src/api_client

1 file changed

+29
-30
lines changed

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

Lines changed: 29 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -285,36 +285,34 @@ impl ApiClient {
285285
{
286286
Ok(response) => Ok(SendMessageOutput::Codewhisperer(response)),
287287
Err(err) => {
288-
use amzn_codewhisperer_streaming_client::operation::generate_assistant_response::GenerateAssistantResponseError::ThrottlingError as OperationThrottlingError;
289-
use amzn_codewhisperer_streaming_client::types::ThrottlingExceptionReason;
290-
use amzn_codewhisperer_streaming_client::types::error::ThrottlingError;
291-
292288
let status_code = err.raw_response().map(|res| res.status().as_u16());
293289
let is_quota_breach = status_code.is_some_and(|status| status == 429);
294290
let is_context_window_overflow = err.as_service_error().is_some_and(|err| {
295291
matches!(err, err if err.meta().code() == Some("ValidationException") && err.meta().message() == Some("Input is too long."))
296292
});
297293

298-
let is_model_unavailable =
299-
// Handling the updated error response
300-
err.as_service_error().is_some_and(|err| {
301-
matches!(
302-
err,
303-
OperationThrottlingError(ThrottlingError {
304-
reason: Some(ThrottlingExceptionReason::InsufficientModelCapacity),
305-
..
306-
})
307-
)
308-
})
309-
// Legacy error response
294+
let is_model_unavailable = {
295+
// check if ThrottlingException
296+
let is_throttling_exception = err
297+
.as_service_error()
298+
.is_some_and(|service_err| service_err.meta().code() == Some("ThrottlingException"));
299+
300+
// check if the response contains INSUFFICIENT_MODEL_CAPACITY
301+
let has_insufficient_capacity = err
302+
.raw_response()
303+
.and_then(|resp| resp.body().bytes())
304+
.and_then(|bytes| String::from_utf8(bytes.to_vec()).ok())
305+
.is_some_and(|body| body.contains("INSUFFICIENT_MODEL_CAPACITY"));
306+
307+
(is_throttling_exception && has_insufficient_capacity)
308+
// Legacy error response fallback
310309
|| (model_id_opt.is_some()
311-
&& status_code.is_some_and(|status| status == 500)
312-
&& err.as_service_error().is_some_and(|err| {
313-
err.meta().message()
314-
== Some(
315-
"Encountered unexpectedly high load when processing the request, please try again.",
316-
)
317-
}));
310+
&& status_code.is_some_and(|status| status == 500)
311+
&& err.as_service_error().is_some_and(|err| {
312+
err.meta().message() == Some(
313+
"Encountered unexpectedly high load when processing the request, please try again.",
314+
)}))
315+
};
318316

319317
let is_monthly_limit_err = err
320318
.raw_response()
@@ -325,17 +323,12 @@ impl ApiClient {
325323
})
326324
.unwrap_or(false);
327325

328-
if is_quota_breach {
329-
return Err(ApiClientError::QuotaBreach {
330-
message: "quota has reached its limit",
331-
status_code,
332-
});
333-
}
334-
335326
if is_context_window_overflow {
336327
return Err(ApiClientError::ContextWindowOverflow { status_code });
337328
}
338329

330+
// Both ModelOverloadedError and QuotaBreach return 429,
331+
// so check is_model_unavailable first.
339332
if is_model_unavailable {
340333
return Err(ApiClientError::ModelOverloadedError {
341334
request_id: err
@@ -345,6 +338,12 @@ impl ApiClient {
345338
status_code,
346339
});
347340
}
341+
if is_quota_breach {
342+
return Err(ApiClientError::QuotaBreach {
343+
message: "quota has reached its limit",
344+
status_code,
345+
});
346+
}
348347

349348
if is_monthly_limit_err {
350349
return Err(ApiClientError::MonthlyLimitReached { status_code });

0 commit comments

Comments
 (0)