Fixed rate limit has of GoogleVertexAiRerankModel and refactored uri for GoogleVertexAiUnifiedChatCompletionRequest

Author: leo-hoet

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/request/GoogleVertexAiUnifiedChatCompletionRequest.java

@@ -25,15 +25,16 @@ public class GoogleVertexAiUnifiedChatCompletionRequest implements GoogleVertexA
 
     private final GoogleVertexAiChatCompletionModel model;
     private final UnifiedChatInput unifiedChatInput;
+    private final URI uri;
 
     public GoogleVertexAiUnifiedChatCompletionRequest(UnifiedChatInput unifiedChatInput, GoogleVertexAiChatCompletionModel model) {
         this.model = Objects.requireNonNull(model);
         this.unifiedChatInput = Objects.requireNonNull(unifiedChatInput);
+        this.uri = unifiedChatInput.stream() ? model.streamingURI() : model.nonStreamingUri();
     }
 
     @Override
     public HttpRequest createHttpRequest() {
-        var uri = unifiedChatInput.stream() ? model.streamingURI() : model.nonStreamingUri();
         HttpPost httpPost = new HttpPost(uri);
 
         var requestEntity = new GoogleVertexAiUnifiedChatCompletionRequestEntity(unifiedChatInput);
@@ -53,7 +54,7 @@ public void decorateWithAuth(HttpPost httpPost) {
 
     @Override
     public URI getURI() {
-        return model.nonStreamingUri();
+        return this.uri;
     }
 
     @Override

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/rerank/GoogleVertexAiRerankModel.java

@@ -27,6 +27,7 @@
 import static org.elasticsearch.core.Strings.format;
 
 public class GoogleVertexAiRerankModel extends GoogleVertexAiModel {
+    private static final String RERANK_RATE_LIMIT_ENDPOINT_ID = "rerank";
 
     public GoogleVertexAiRerankModel(
         String inferenceEntityId,
@@ -143,7 +144,6 @@ public int rateLimitGroupingHash() {
         var projectId = getServiceSettings().projectId();
         var modelId = getServiceSettings().modelId();
 
-        // In this case, we don't include the location since it's a global endpoint
-        return Objects.hash(projectId, modelId);
+        return Objects.hash(projectId, modelId, RERANK_RATE_LIMIT_ENDPOINT_ID);
     }
 }