Tweaks (#42)

- make error handling report response body
- support custom models (so when new ones are out, they can be used)
- allow setting response_schema
- safety_ratings is optional in response
- upgrade reqwest and make default features optional

Author: avastmick

Cargo.toml

@@ -18,8 +18,8 @@ env_logger = { version = "0.11" }
 futures = { version = "0.3" }
 gcp_auth = { version = "0.12" }
 log = { version = "0.4.20" }
-reqwest = { version = "0.11", features = ["json"] }
-reqwest-streams = { version = "0.5.1", features = ["json"] }
+reqwest = { version = "0.12", default-features = false, features = ["json"] }
+reqwest-streams = { version = "0.8.2", default-features = false, features = ["json"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = { version = "1.0" }
 tokio = { version = "1.35", features = ["full"] }

examples/text_request_json.rs

@@ -62,6 +62,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                 max_output_tokens: None,
                 stop_sequences: None,
                 response_mime_type: Some("application/json".to_string()),
+                response_schema: None,
             }),
 
             system_instruction: None,

src/v1/api.rs

@@ -1,6 +1,7 @@
 //! Manages the interaction with the REST API for the Gemini API.
 use futures::prelude::*;
 use futures::stream::StreamExt;
+use reqwest::StatusCode;
 use reqwest_streams::error::StreamBodyError;
 use reqwest_streams::*;
 use serde_json;
@@ -14,7 +15,7 @@ use crate::v1::gemini::request::Request;
 use crate::v1::gemini::response::GeminiResponse;
 use crate::v1::gemini::Model;
 
-use super::gemini::response::{StreamedGeminiResponse, TokenCount};
+use super::gemini::response::{GeminiErrorResponse, StreamedGeminiResponse, TokenCount};
 use super::gemini::{ModelInformation, ModelInformationList, ResponseType};
 
 #[cfg(feature = "beta")]
@@ -155,20 +156,31 @@ impl Client {
             .get_post_response(client, api_request, token_option)
             .await;
 
-        match result {
-            Ok(response) => match response.status() {
-                reqwest::StatusCode::OK => Ok(response.json::<GeminiResponse>().await.map_err(|e|GoogleAPIError {
-                message: format!(
-                        "Failed to deserialize API response into v1::gemini::response::GeminiResponse: {}",
-                        e
-                    ),
-                code: None,
-            })?),
-                _ => Err(self.new_error_from_status_code(response.status())),
-            },
-            Err(e) => Err(self.new_error_from_reqwest_error(e)),
+        if let Ok(result) = result {
+            match result.status() {
+                reqwest::StatusCode::OK => {
+                    Ok(result.json::<GeminiResponse>().await.map_err(|e|GoogleAPIError {
+                    message: format!(
+                            "Failed to deserialize API response into v1::gemini::response::GeminiResponse: {}",
+                            e
+                        ),
+                    code: None,
+                    })?)
+                },
+                _ => {
+                    let status = result.status();
+
+                    match result.json::<GeminiErrorResponse>().await {
+                        Ok(GeminiErrorResponse::Error { message, .. }) => Err(self.new_error_from_api_message(status, message)),
+                        Err(_) => Err(self.new_error_from_status_code(status)),
+                    }
+                },
+            }
+        } else {
+            Err(self.new_error_from_reqwest_error(result.unwrap_err()))
         }
     }
+
     // Define the function that accepts the stream and the consumer
     /// A streamed post request
     async fn get_streamed_post_result(
@@ -410,6 +422,17 @@ impl Client {
             code: Some(code),
         }
     }
+
+    /// Creates a new error from a status code.
+    fn new_error_from_api_message(&self, code: StatusCode, message: String) -> GoogleAPIError {
+        let message = format!("API message: {message}.");
+
+        GoogleAPIError {
+            message,
+            code: Some(code),
+        }
+    }
+
     /// Creates a new error from a reqwest error.
     fn new_error_from_reqwest_error(&self, mut e: reqwest::Error) -> GoogleAPIError {
         if let Some(url) = e.url_mut() {

src/v1/errors.rs

@@ -11,8 +11,8 @@ impl fmt::Display for GoogleAPIError {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         write!(
             f,
-            "GoogleAPIError - code: {} error: {:?}",
-            self.message, self.code
+            "GoogleAPIError - code: {:?} error: {}",
+            self.code, self.message
         )
     }
 }

src/v1/gemini.rs

@@ -81,6 +81,7 @@ pub enum Model {
     #[cfg_attr(docsrs, doc(cfg(feature = "beta")))]
     Gemini1_5Pro,
     GeminiProVision,
+    Custom(String),
     // TODO Embedding001
 }
 impl fmt::Display for Model {
@@ -92,6 +93,8 @@ impl fmt::Display for Model {
             Model::Gemini1_5Pro => write!(f, "gemini-1.5-pro-latest"),
 
             Model::GeminiProVision => write!(f, "gemini-pro-vision"),
+
+            Model::Custom(name) => write!(f, "{}", name),
             // TODO Model::Embedding001 => write!(f, "embedding-001"),
         }
     }
@@ -319,6 +322,10 @@ pub mod request {
         #[cfg(feature = "beta")]
         #[cfg_attr(docsrs, doc(cfg(feature = "beta")))]
         pub response_mime_type: Option<String>,
+
+        #[cfg(feature = "beta")]
+        #[cfg_attr(docsrs, doc(cfg(feature = "beta")))]
+        pub response_schema: Option<serde_json::Value>,
     }
 
     #[cfg(feature = "beta")]
@@ -425,6 +432,16 @@ pub mod response {
         pub prompt_feedback: Option<PromptFeedback>,
         pub usage_metadata: Option<UsageMetadata>,
     }
+    #[derive(Debug, Clone, Deserialize)]
+    #[serde(rename_all = "camelCase")]
+    pub enum GeminiErrorResponse {
+        Error {
+            code: u16,
+            message: String,
+            status: String,
+        },
+    }
+
     impl GeminiResponse {
         /// Returns the total character count of the response as per the Gemini API.
         pub fn get_response_character_count(&self) -> usize {
@@ -448,6 +465,7 @@ pub mod response {
         pub content: Content,
         pub finish_reason: Option<String>,
         pub index: Option<i32>,
+        #[serde(default)]
         pub safety_ratings: Vec<SafetyRating>,
     }
     #[derive(Debug, Clone, Deserialize)]