Improves error handling

src/error.rs

@@ -7,6 +7,7 @@ pub enum Error {
     Env(std::env::VarError),
     HttpClient(reqwest::Error),
     Token(gcp_auth::Error),
+    Serde(serde_json::Error),
 }
 
 impl Display for Error {
@@ -15,6 +16,7 @@ impl Display for Error {
             Error::Env(e) => write!(f, "Environment variable error: {}", e),
             Error::HttpClient(e) => write!(f, "HTTP Client error: {}", e),
             Error::Token(e) => write!(f, "Token error: {}", e),
+            Error::Serde(e) => write!(f, "Serde error: {}", e),
         }
     }
 }
@@ -38,3 +40,9 @@ impl From<gcp_auth::Error> for Error {
         Error::Token(e)
     }
 }
+
+impl From<serde_json::Error> for Error {
+    fn from(e: serde_json::Error) -> Self {
+        Error::Serde(e)
+    }
+}

src/types.rs

@@ -66,8 +66,9 @@ pub struct GenerateContentResponse(pub Vec<ResponseStreamChunk>);
 #[derive(Debug, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct ResponseStreamChunk {
-    pub candidates: Vec<Candidate>,
+    pub candidates: Option<Vec<Candidate>>,
     pub usage_metadata: Option<UsageMetadata>,
+    pub error: Option<Error>,
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -128,3 +129,24 @@ pub struct FunctionParametersProperty {
     pub r#type: String,
     pub description: String,
 }
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct Error {
+    pub code: i32,
+    pub message: String,
+    pub status: String,
+    pub details: Vec<ErrorDetail>,
+}
+
+// TODO: Make ErrorDetail an enum and map to the different types of errors.
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ErrorDetail {
+    #[serde(rename = "@type")]
+    pub r#type: String,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct Link {
+    pub description: String,
+    pub url: String,
+}

src/vertex_client.rs

@@ -66,7 +66,13 @@ impl<T: TokenProvider + Clone> VertexClient<T> {
 
         let txt_json = resp.text().await?;
         tracing::debug!("Vertex API Response: {}", txt_json);
-        Ok(serde_json::from_str(&txt_json).unwrap())
+        match serde_json::from_str(&txt_json) {
+            Ok(response) => Ok(response),
+            Err(e) => {
+                eprintln!("Failed to parse response: {} / {}", txt_json, e);
+                Err(e.into())
+            }
+        }
     }
 
     /// Prompts a conversation to the model.
@@ -92,7 +98,7 @@ impl<T: TokenProvider + Clone> VertexClient<T> {
             .0
             .into_iter()
             .flat_map(|chunk| {
-                chunk.candidates.into_iter().flat_map(|candidate| {
+                chunk.candidates.unwrap().into_iter().flat_map(|candidate| {
                     candidate
                         .content
                         .parts
@@ -133,7 +139,7 @@ impl<T: TokenProvider + Clone> VertexClient<T> {
             .0
             .into_iter()
             .flat_map(|chunk| {
-                chunk.candidates.into_iter().flat_map(|candidate| {
+                chunk.candidates.unwrap().into_iter().flat_map(|candidate| {
                     candidate
                         .content
                         .parts