From e6de1d1ce7f1e1e7f9809e4c24b8aed3132fcf10 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9=20Cipriani=20Bandarra?= <andreban@gmail.com>
Date: Sun, 6 Oct 2024 08:18:05 +0100
Subject: [PATCH] Adds error logging when failing o parse response

---
 examples/generate_image.rs |  8 +++++++-
 src/client.rs              | 10 +++++++++-
 src/types/predict_image.rs |  1 +
 3 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/examples/generate_image.rs b/examples/generate_image.rs
index 117c4fd..5343dac 100644
--- a/examples/generate_image.rs
+++ b/examples/generate_image.rs
@@ -1,7 +1,8 @@
 use std::{error::Error, io::Cursor};
 
 use gemini_rs::prelude::{
-    GeminiClient, PredictImageRequest, PredictImageRequestParameters, PredictImageRequestPrompt,
+    GeminiClient, PredictImageRequest, PredictImageRequestParameters,
+    PredictImageRequestParametersOutputOptions, PredictImageRequestPrompt,
 };
 use image::{ImageFormat, ImageReader};
 
@@ -28,6 +29,11 @@ pub async fn main() -> Result<(), Box<dyn Error>> {
         }],
         parameters: PredictImageRequestParameters {
             sample_count: 1,
+            aspect_ratio: Some("1:1".to_string()),
+            output_options: Some(PredictImageRequestParametersOutputOptions {
+                mime_type: Some("image/jpeg".to_string()),
+                compression_quality: Some(75),
+            }),
             ..Default::default()
         },
     };
diff --git a/src/client.rs b/src/client.rs
index 92bfce7..e36ba0b 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -3,6 +3,7 @@ use std::sync::Arc;
 use deadqueue::unlimited::Queue;
 use futures_util::stream::StreamExt;
 use reqwest_eventsource::{Event, EventSource};
+use tracing::error;
 
 use crate::dialogue::{Message, Role};
 use crate::error::{Error, Result};
@@ -274,6 +275,13 @@ impl<T: TokenProvider + Clone> GeminiClient<T> {
             .await?;
 
         let txt_json = resp.text().await?;
-        Ok(serde_json::from_str(&txt_json)?)
+
+        match serde_json::from_str::<PredictImageResponse>(&txt_json) {
+            Ok(response) => return Ok(response),
+            Err(e) => {
+                error!(response = txt_json, error = ?e, "Failed to parse response");
+                return Err(e.into());
+            }
+        }
     }
 }
diff --git a/src/types/predict_image.rs b/src/types/predict_image.rs
index 139052b..9f08ed8 100644
--- a/src/types/predict_image.rs
+++ b/src/types/predict_image.rs
@@ -70,6 +70,7 @@ pub struct PredictImageRequestParameters {
     ///  - "watercolor"
     ///  - "cyberpunk"
     ///  - "pop_art"
+    /// Pre-defined styles is only supported for model imagegeneration@002
     #[serde(skip_serializing_if = "Option::is_none")]
     pub sample_image_style: Option<String>,