fix: make idempotency bit correct for retry

Author: georgeh0

Cargo.lock

@@ -133,10 +133,10 @@ time = { version = "0.3", features = ["macros", "serde"] }
 numpy = "0.25.0"
 infer = "0.19.0"
 serde_with = { version = "3.14.0", features = ["base64"] }
-google-cloud-aiplatform-v1 = { version = "1.0.0", default-features = false, features = [
+google-cloud-aiplatform-v1 = { version = "0.4.5", default-features = false, features = [
     "prediction-service",
 ] }
-google-cloud-gax = "1.0.0"
+google-cloud-gax = "0.24.0"
 
 azure_identity = { version = "0.21.0", default-features = false, features = [
     "enable_reqwest_rustls",

Cargo.toml

@@ -7,6 +7,7 @@ use crate::llm::{
 use base64::prelude::*;
 use google_cloud_aiplatform_v1 as vertexai;
 use google_cloud_gax::exponential_backoff::ExponentialBackoff;
+use google_cloud_gax::options::RequestOptionsBuilder;
 use google_cloud_gax::retry_policy::{Aip194Strict, RetryPolicyExt};
 use google_cloud_gax::retry_throttler::{AdaptiveThrottler, SharedRetryThrottler};
 use serde_json::Value;
@@ -251,9 +252,6 @@ impl google_cloud_gax::retry_policy::RetryPolicy for CustomizedGoogleCloudRetryP
     ) -> google_cloud_gax::retry_result::RetryResult {
         use google_cloud_gax::retry_result::RetryResult;
 
-        if !state.idempotent {
-            return RetryResult::Permanent(error);
-        }
         if let Some(status) = error.status() {
             if status.code == google_cloud_gax::error::rpc::Code::ResourceExhausted {
                 return RetryResult::Continue(error);
@@ -350,7 +348,8 @@ impl LlmGenerationClient for VertexAiClient {
             .client
             .generate_content()
             .set_model(self.get_model_path(request.model))
-            .set_contents(contents);
+            .set_contents(contents)
+            .with_idempotency(true);
         if let Some(sys) = system_instruction {
             req = req.set_system_instruction(sys);
         }
@@ -414,6 +413,7 @@ impl LlmEmbeddingClient for VertexAiClient {
             .set_endpoint(self.get_model_path(request.model))
             .set_instances(instances)
             .set_parameters(parameters)
+            .with_idempotency(true)
             .send()
             .await?;