Changes due to the comments

Author: fzowl

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/voyageai/VoyageAIActionCreator.java

@@ -37,21 +37,15 @@ public VoyageAIActionCreator(Sender sender, ServiceComponents serviceComponents)
     @Override
     public ExecutableAction create(VoyageAIEmbeddingsModel model, Map<String, Object> taskSettings, InputType inputType) {
         var overriddenModel = VoyageAIEmbeddingsModel.of(model, taskSettings, inputType);
-        var failedToSendRequestErrorMessage = constructFailedToSendRequestMessage(
-            overriddenModel.getServiceSettings().getCommonSettings().uri(),
-            "VoyageAI embeddings"
-        );
+        var failedToSendRequestErrorMessage = constructFailedToSendRequestMessage(overriddenModel.uri(), "VoyageAI embeddings");
         var requestCreator = VoyageAIEmbeddingsRequestManager.of(overriddenModel, serviceComponents.threadPool());
         return new SenderExecutableAction(sender, requestCreator, failedToSendRequestErrorMessage);
     }
 
     @Override
     public ExecutableAction create(VoyageAIRerankModel model, Map<String, Object> taskSettings) {
         var overriddenModel = VoyageAIRerankModel.of(model, taskSettings);
-        var failedToSendRequestErrorMessage = constructFailedToSendRequestMessage(
-            overriddenModel.getServiceSettings().getCommonSettings().uri(),
-            "VoyageAI rerank"
-        );
+        var failedToSendRequestErrorMessage = constructFailedToSendRequestMessage(overriddenModel.uri(), "VoyageAI rerank");
         var requestCreator = VoyageAIRerankRequestManager.of(overriddenModel, serviceComponents.threadPool());
         return new SenderExecutableAction(sender, requestCreator, failedToSendRequestErrorMessage);
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/VoyageAIRequestManager.java

@@ -10,9 +10,33 @@
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xpack.inference.services.voyageai.VoyageAIModel;
 
+import java.util.Map;
 import java.util.Objects;
 
 abstract class VoyageAIRequestManager extends BaseRequestManager {
+    private static final String DEFAULT_MODEL_FAMILY = "default_model_family";
+    private static final Map<String, String> MODEL_TO_MODEL_FAMILY = Map.of(
+        "voyage-multimodal-3",
+        "embed_multimodal",
+        "voyage-3-large",
+        "embed_large",
+        "voyage-code-3",
+        "embed_large",
+        "voyage-3",
+        "embed_medium",
+        "voyage-3-lite",
+        "embed_small",
+        "voyage-finance-2",
+        "embed_large",
+        "voyage-law-2",
+        "embed_large",
+        "voyage-code-2",
+        "embed_large",
+        "rerank-2",
+        "rerank_large",
+        "rerank-2-lite",
+        "rerank_small"
+    );
 
     protected VoyageAIRequestManager(ThreadPool threadPool, VoyageAIModel model) {
         super(threadPool, model.getInferenceEntityId(), RateLimitGrouping.of(model), model.rateLimitServiceSettings().rateLimitSettings());
@@ -21,8 +45,10 @@ protected VoyageAIRequestManager(ThreadPool threadPool, VoyageAIModel model) {
     record RateLimitGrouping(int apiKeyHash) {
         public static RateLimitGrouping of(VoyageAIModel model) {
             Objects.requireNonNull(model);
+            String modelId = model.getServiceSettings().modelId();
+            String modelFamily = MODEL_TO_MODEL_FAMILY.getOrDefault(modelId, DEFAULT_MODEL_FAMILY);
 
-            return new RateLimitGrouping(model.apiKey().hashCode());
+            return new RateLimitGrouping(modelFamily.hashCode());
         }
     }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/voyageai/VoyageAIEmbeddingsRequest.java

@@ -8,7 +8,6 @@
 package org.elasticsearch.xpack.inference.external.request.voyageai;
 
 import org.apache.http.client.methods.HttpPost;
-import org.apache.http.client.utils.URIBuilder;
 import org.apache.http.entity.ByteArrayEntity;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.xpack.inference.external.request.HttpRequest;
@@ -19,7 +18,6 @@
 import org.elasticsearch.xpack.inference.services.voyageai.embeddings.VoyageAIEmbeddingsTaskSettings;
 
 import java.net.URI;
-import java.net.URISyntaxException;
 import java.nio.charset.StandardCharsets;
 import java.util.List;
 import java.util.Objects;
@@ -36,7 +34,7 @@ public class VoyageAIEmbeddingsRequest extends VoyageAIRequest {
     public VoyageAIEmbeddingsRequest(List<String> input, VoyageAIEmbeddingsModel embeddingsModel) {
         Objects.requireNonNull(embeddingsModel);
 
-        account = VoyageAIAccount.of(embeddingsModel, VoyageAIEmbeddingsRequest::buildDefaultUri);
+        account = VoyageAIAccount.of(embeddingsModel);
         this.input = Objects.requireNonNull(input);
         serviceSettings = embeddingsModel.getServiceSettings();
         taskSettings = embeddingsModel.getTaskSettings();
@@ -54,7 +52,7 @@ public HttpRequest createHttpRequest() {
         );
         httpPost.setEntity(byteEntity);
 
-        decorateWithAuthHeader(httpPost, account);
+        decorateWithHeaders(httpPost, account);
 
         return new HttpRequest(httpPost, getInferenceEntityId());
     }
@@ -86,11 +84,4 @@ public VoyageAIEmbeddingsTaskSettings getTaskSettings() {
     public VoyageAIEmbeddingsServiceSettings getServiceSettings() {
         return serviceSettings;
     }
-
-    public static URI buildDefaultUri() throws URISyntaxException {
-        return new URIBuilder().setScheme("https")
-            .setHost(VoyageAIUtils.HOST)
-            .setPathSegments(VoyageAIUtils.VERSION_1, VoyageAIUtils.EMBEDDINGS_PATH)
-            .build();
-    }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/voyageai/VoyageAIRequest.java

@@ -17,7 +17,7 @@
 
 public abstract class VoyageAIRequest implements Request {
 
-    public static void decorateWithAuthHeader(HttpPost request, VoyageAIAccount account) {
+    public static void decorateWithHeaders(HttpPost request, VoyageAIAccount account) {
         request.setHeader(HttpHeaders.CONTENT_TYPE, XContentType.JSON.mediaType());
         request.setHeader(createAuthBearerHeader(account.apiKey()));
         request.setHeader(VoyageAIUtils.createRequestSourceHeader());

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/voyageai/VoyageAIRerankRequest.java

@@ -8,7 +8,6 @@
 package org.elasticsearch.xpack.inference.external.request.voyageai;
 
 import org.apache.http.client.methods.HttpPost;
-import org.apache.http.client.utils.URIBuilder;
 import org.apache.http.entity.ByteArrayEntity;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.xpack.inference.external.request.HttpRequest;
@@ -18,7 +17,6 @@
 import org.elasticsearch.xpack.inference.services.voyageai.rerank.VoyageAIRerankTaskSettings;
 
 import java.net.URI;
-import java.net.URISyntaxException;
 import java.nio.charset.StandardCharsets;
 import java.util.List;
 import java.util.Objects;
@@ -35,7 +33,7 @@ public class VoyageAIRerankRequest extends VoyageAIRequest {
     public VoyageAIRerankRequest(String query, List<String> input, VoyageAIRerankModel model) {
         Objects.requireNonNull(model);
 
-        this.account = VoyageAIAccount.of(model, VoyageAIRerankRequest::buildDefaultUri);
+        this.account = VoyageAIAccount.of(model);
         this.input = Objects.requireNonNull(input);
         this.query = Objects.requireNonNull(query);
         taskSettings = model.getTaskSettings();
@@ -52,7 +50,7 @@ public HttpRequest createHttpRequest() {
         );
         httpPost.setEntity(byteEntity);
 
-        decorateWithAuthHeader(httpPost, account);
+        decorateWithHeaders(httpPost, account);
 
         return new HttpRequest(httpPost, getInferenceEntityId());
     }
@@ -76,11 +74,4 @@ public Request truncate() {
     public boolean[] getTruncationInfo() {
         return null;
     }
-
-    public static URI buildDefaultUri() throws URISyntaxException {
-        return new URIBuilder().setScheme("https")
-            .setHost(VoyageAIUtils.HOST)
-            .setPathSegments(VoyageAIUtils.VERSION_1, VoyageAIUtils.RERANK_PATH)
-            .build();
-    }
 }