fix:extra error handling for gemini (#1268)

This commit is contained in:
Yingjie He
2025-02-18 12:57:29 -08:00
committed by GitHub
parent 7c5aaa1c75
commit d8ca7a3e1b
4 changed files with 211 additions and 44 deletions
+5 -40
View File
@@ -3,11 +3,13 @@ use crate::message::Message;
use crate::model::ModelConfig;
use crate::providers::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage};
use crate::providers::formats::google::{create_request, get_usage, response_to_message};
use crate::providers::utils::{emit_debug_trace, unescape_json_values};
use crate::providers::utils::{
emit_debug_trace, handle_response_google_compat, unescape_json_values,
};
use anyhow::Result;
use async_trait::async_trait;
use mcp_core::tool::Tool;
use reqwest::{Client, StatusCode};
use reqwest::Client;
use serde_json::Value;
use std::time::Duration;
use url::Url;
@@ -84,44 +86,7 @@ impl GoogleProvider {
.send()
.await?;
let status = response.status();
let payload: Option<Value> = response.json().await.ok();
match status {
StatusCode::OK => payload.ok_or_else( || ProviderError::RequestFailed("Response body is not valid JSON".to_string()) ),
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => {
Err(ProviderError::Authentication(format!("Authentication failed. Please ensure your API keys are valid and have the required permissions. \
Status: {}. Response: {:?}", status, payload )))
}
StatusCode::BAD_REQUEST => {
let mut error_msg = "Unknown error".to_string();
if let Some(payload) = &payload {
if let Some(error) = payload.get("error") {
error_msg = error.get("message").and_then(|m| m.as_str()).unwrap_or("Unknown error").to_string();
let error_status = error.get("status").and_then(|s| s.as_str()).unwrap_or("Unknown status");
if error_status == "INVALID_ARGUMENT" && error_msg.to_lowercase().contains("exceeds") {
return Err(ProviderError::ContextLengthExceeded(error_msg.to_string()));
}
}
}
tracing::debug!(
"{}", format!("Provider request failed with status: {}. Payload: {:?}", status, payload)
);
Err(ProviderError::RequestFailed(format!("Request failed with status: {}. Message: {}", status, error_msg)))
}
StatusCode::TOO_MANY_REQUESTS => {
Err(ProviderError::RateLimitExceeded(format!("{:?}", payload)))
}
StatusCode::INTERNAL_SERVER_ERROR | StatusCode::SERVICE_UNAVAILABLE => {
Err(ProviderError::ServerError(format!("{:?}", payload)))
}
_ => {
tracing::debug!(
"{}", format!("Provider request failed with status: {}. Payload: {:?}", status, payload)
);
Err(ProviderError::RequestFailed(format!("Request failed with status: {}", status)))
}
}
handle_response_google_compat(response).await
}
}