From e056edc2c467bfba2a84aa2ac47c507393312865 Mon Sep 17 00:00:00 2001
From: Andre Bandarra <andreban@gmail.com>
Date: Tue, 14 May 2024 22:33:12 +0100
Subject: [PATCH] Implements result conversation for embedding results

---
 examples/text-embedding.rs   |  3 ++-
 src/types/text_embeddings.rs | 30 ++++++++++++++++++++++++------
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/examples/text-embedding.rs b/examples/text-embedding.rs
index 7c087e8..6b2f01e 100644
--- a/examples/text-embedding.rs
+++ b/examples/text-embedding.rs
@@ -35,7 +35,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let result = gemini
         .text_embeddings(&embedding_request, "textembedding-gecko@003")
-        .await?;
+        .await?
+        .into_result()?;
     println!("Response: {:?}", result);
 
     Ok(())
diff --git a/src/types/text_embeddings.rs b/src/types/text_embeddings.rs
index 09be2fe..d1d3173 100644
--- a/src/types/text_embeddings.rs
+++ b/src/types/text_embeddings.rs
@@ -1,5 +1,7 @@
 use serde::{Deserialize, Serialize};
 
+use crate::error::{Error, Result};
+
 use super::VertexApiError;
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
@@ -17,12 +19,28 @@ pub struct TextEmbeddingRequestInstance {
 #[derive(Clone, Debug, Serialize, Deserialize)]
 #[serde(untagged)]
 pub enum TextEmbeddingResponse {
-    Ok {
-        predictions: Vec<TextEmbeddingPrediction>,
-    },
-    Error {
-        error: VertexApiError,
-    },
+    Ok(TextEmbeddingResponseOk),
+    Error { error: VertexApiError },
+}
+
+impl TextEmbeddingResponse {
+    pub fn into_result(self) -> Result<TextEmbeddingResponseOk> {
+        self.into()
+    }
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct TextEmbeddingResponseOk {
+    pub predictions: Vec<TextEmbeddingPrediction>,
+}
+
+impl From<TextEmbeddingResponse> for Result<TextEmbeddingResponseOk> {
+    fn from(value: TextEmbeddingResponse) -> Self {
+        match value {
+            TextEmbeddingResponse::Ok(ok) => Ok(ok),
+            TextEmbeddingResponse::Error { error } => Err(Error::VertexError(error)),
+        }
+    }
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]