From 9754618153c75cb3c296b491e8cc270441bc89b1 Mon Sep 17 00:00:00 2001
From: Andre Bandarra <andreban@gmail.com>
Date: Thu, 25 Apr 2024 14:45:31 +0100
Subject: [PATCH] Make SafetyRatings option in Candidates

---
 src/client.rs                 | 10 +++++++---
 src/types/generate_content.rs |  2 +-
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index abe7dff..122f0df 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -69,9 +69,13 @@ impl<T: TokenProvider + Clone> GeminiClient<T> {
                 match event {
                     Ok(event) => {
                         if let Event::Message(event) = event {
-                            let response: GenerateContentResponse =
-                                serde_json::from_str(&event.data).unwrap();
-                            cloned_queue.push(Some(response));
+                            let response: serde_json::error::Result<GenerateContentResponse> =
+                                serde_json::from_str(&event.data);
+                            if let Ok(response) = response {
+                                cloned_queue.push(Some(response));
+                            } else {
+                                tracing::error!("Error parsing message: {}", event.data);
+                            };
                         }
                     }
                     Err(e) => {
diff --git a/src/types/generate_content.rs b/src/types/generate_content.rs
index 67aea7c..40796c8 100644
--- a/src/types/generate_content.rs
+++ b/src/types/generate_content.rs
@@ -48,7 +48,7 @@ pub struct GenerationConfig {
 pub struct Candidate {
     pub content: Option<Content>,
     pub citation_metadata: Option<CitationMetadata>,
-    pub safety_ratings: Vec<SafetyRating>,
+    pub safety_ratings: Option<Vec<SafetyRating>>,
     pub finish_reason: Option<String>,
 }