From 1c0aaf57b3fd6c224f8b7c495bcbb3d4cbf2d35b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9=20Cipriani=20Bandarra?= <andreban@gmail.com>
Date: Fri, 1 Mar 2024 13:45:26 +0000
Subject: [PATCH] Allows different models for text-embedding

---
 examples/text-embedding.rs | 4 +++-
 src/client.rs              | 2 +-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/examples/text-embedding.rs b/examples/text-embedding.rs
index e362ac4..4c4f91c 100644
--- a/examples/text-embedding.rs
+++ b/examples/text-embedding.rs
@@ -33,7 +33,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         ],
     };
 
-    let result = gemini.text_embeddings(&embedding_request).await?;
+    let result = gemini
+        .text_embeddings(&embedding_request, "textembedding-gecko@003")
+        .await?;
     println!("Response: {:?}", result);
 
     Ok(())
diff --git a/src/client.rs b/src/client.rs
index bc62a81..2e4b2c3 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -185,8 +185,8 @@ impl<T: TokenProvider + Clone> GeminiClient<T> {
     pub async fn text_embeddings(
         &self,
         request: &TextEmbeddingRequest,
+        model: &str,
     ) -> Result<TextEmbeddingResponse> {
-        let model = "models/embedding-001";
         let endpoint_url = format!(
             "https://{}/v1/projects/{}/locations/{}/publishers/google/models/{}:predict",
             self.api_endpoint, self.project_id, self.location_id, model,