fix:extra error handling for gemini (#1268)

crates/goose/src/providers/utils.rs

@@ -1,4 +1,5 @@
 use super::base::Usage;
+use super::errors::GoogleErrorCode;
 use anyhow::Result;
 use base64::Engine;
 use regex::Regex;
@@ -90,6 +91,97 @@ pub async fn handle_response_openai_compat(response: Response) -> Result<Value,
     }
 }
 
+/// Check if the model is a Google model based on the "model" field in the payload.
+///
+/// ### Arguments
+/// - `payload`: The JSON payload as a `serde_json::Value`.
+///
+/// ### Returns
+/// - `bool`: Returns `true` if the model is a Google model, otherwise `false`.
+pub fn is_google_model(payload: &Value) -> bool {
+    if let Some(model) = payload.get("model").and_then(|m| m.as_str()) {
+        // Check if the model name contains "google"
+        return model.to_lowercase().contains("google");
+    }
+    false
+}
+
+/// Extracts `StatusCode` from response status or payload error code.
+/// This function first checks the status code of the response. If the status is successful (2xx),
+/// it then checks the payload for any error codes and maps them to appropriate `StatusCode`.
+/// If the status is not successful (e.g., 4xx or 5xx), the original status code is returned.
+fn get_google_final_status(status: StatusCode, payload: Option<&Value>) -> StatusCode {
+    // If the status is successful, check for an error in the payload
+    if status.is_success() {
+        if let Some(payload) = payload {
+            if let Some(error) = payload.get("error") {
+                if let Some(code) = error.get("code").and_then(|c| c.as_u64()) {
+                    if let Some(google_error) = GoogleErrorCode::from_code(code) {
+                        return google_error.to_status_code();
+                    }
+                }
+            }
+        }
+    }
+    status
+}
+
+/// Handle response from Google Gemini API-compatible endpoints.
+///
+/// Processes HTTP responses, handling specific statuses and parsing the payload
+/// for error messages. Logs the response payload for debugging purposes.
+///
+/// ### References
+/// - Error Codes: https://ai.google.dev/gemini-api/docs/troubleshooting?lang=python
+///
+/// ### Arguments
+/// - `response`: The HTTP response to process.
+///
+/// ### Returns
+/// - `Ok(Value)`: Parsed JSON on success.
+/// - `Err(ProviderError)`: Describes the failure reason.
+pub async fn handle_response_google_compat(response: Response) -> Result<Value, ProviderError> {
+    let status = response.status();
+    let payload: Option<Value> = response.json().await.ok();
+    let final_status = get_google_final_status(status, payload.as_ref());
+
+    match final_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: {:?}", final_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: {:?}", final_status, payload)
+            );
+            Err(ProviderError::RequestFailed(format!("Request failed with status: {}. Message: {}", final_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: {:?}", final_status, payload)
+            );
+            Err(ProviderError::RequestFailed(format!("Request failed with status: {}", final_status)))
+        }
+    }
+}
+
 pub fn sanitize_function_name(name: &str) -> String {
     let re = Regex::new(r"[^a-zA-Z0-9_-]").unwrap();
     re.replace_all(name, "_").to_string()
@@ -253,8 +345,6 @@ pub fn emit_debug_trace<T: serde::Serialize>(
 mod tests {
     use super::*;
     use serde_json::json;
-    use std::io::Write;
-    use tempfile::NamedTempFile;
 
     #[test]
     fn test_detect_image_path() {
@@ -395,4 +485,70 @@ mod tests {
         let unescaped_value = unescape_json_values(&value);
         assert_eq!(unescaped_value, json!({"text": "Hello World"}));
     }
+
+    #[test]
+    fn test_is_google_model() {
+        // Define the test cases as a vector of tuples
+        let test_cases = vec![
+            // (input, expected_result)
+            (json!({ "model": "google_gemini" }), true),
+            (json!({ "model": "microsoft_bing" }), false),
+            (json!({ "model": "" }), false),
+            (json!({}), false),
+            (json!({ "model": "Google_XYZ" }), true),
+            (json!({ "model": "google_abc" }), true),
+        ];
+
+        // Iterate through each test case and assert the result
+        for (payload, expected_result) in test_cases {
+            assert_eq!(is_google_model(&payload), expected_result);
+        }
+    }
+
+    #[test]
+    fn test_get_google_final_status_success() {
+        let status = StatusCode::OK;
+        let payload = json!({});
+        let result = get_google_final_status(status, Some(&payload));
+        assert_eq!(result, StatusCode::OK);
+    }
+
+    #[test]
+    fn test_get_google_final_status_with_error_code() {
+        // Test error code mappings for different payload error codes
+        let test_cases = vec![
+            // (error code, status, expected status code)
+            (200, None, StatusCode::OK),
+            (429, Some(StatusCode::OK), StatusCode::TOO_MANY_REQUESTS),
+            (400, Some(StatusCode::OK), StatusCode::BAD_REQUEST),
+            (401, Some(StatusCode::OK), StatusCode::UNAUTHORIZED),
+            (403, Some(StatusCode::OK), StatusCode::FORBIDDEN),
+            (404, Some(StatusCode::OK), StatusCode::NOT_FOUND),
+            (500, Some(StatusCode::OK), StatusCode::INTERNAL_SERVER_ERROR),
+            (503, Some(StatusCode::OK), StatusCode::SERVICE_UNAVAILABLE),
+            (999, Some(StatusCode::OK), StatusCode::INTERNAL_SERVER_ERROR),
+            (500, Some(StatusCode::BAD_REQUEST), StatusCode::BAD_REQUEST),
+            (
+                404,
+                Some(StatusCode::INTERNAL_SERVER_ERROR),
+                StatusCode::INTERNAL_SERVER_ERROR,
+            ),
+        ];
+
+        for (error_code, status, expected_status) in test_cases {
+            let payload = if let Some(_status) = status {
+                json!({
+                    "error": {
+                        "code": error_code,
+                        "message": "Error message"
+                    }
+                })
+            } else {
+                json!({})
+            };
+
+            let result = get_google_final_status(status.unwrap_or(StatusCode::OK), Some(&payload));
+            assert_eq!(result, expected_status);
+        }
+    }
 }