Allowing model to be overridden but not working yet

Author: jonathan-buttner

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

@@ -12,7 +12,6 @@
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.inference.InferenceServiceResults;
 import org.elasticsearch.rest.RestStatus;
-import org.elasticsearch.xpack.inference.external.http.sender.DocumentsOnlyInput;
 import org.elasticsearch.xpack.inference.external.http.sender.InferenceInputs;
 import org.elasticsearch.xpack.inference.external.http.sender.RequestManager;
 import org.elasticsearch.xpack.inference.external.http.sender.Sender;
@@ -34,13 +33,7 @@ public SingleInputSenderExecutableAction(
 
     @Override
     public void execute(InferenceInputs inferenceInputs, TimeValue timeout, ActionListener<InferenceServiceResults> listener) {
-        if (inferenceInputs instanceof DocumentsOnlyInput == false) {
-            listener.onFailure(new ElasticsearchStatusException("Invalid inference input type", RestStatus.INTERNAL_SERVER_ERROR));
-            return;
-        }
-
-        var docsOnlyInput = (DocumentsOnlyInput) inferenceInputs;
-        if (docsOnlyInput.getInputs().size() > 1) {
+        if (inferenceInputs.inputSize() > 1) {
             listener.onFailure(
                 new ElasticsearchStatusException(requestTypeForInputValidationError + " only accepts 1 input", RestStatus.BAD_REQUEST)
             );

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

@@ -14,7 +14,7 @@ public class DocumentsOnlyInput extends InferenceInputs {
 
     public static DocumentsOnlyInput of(InferenceInputs inferenceInputs) {
         if (inferenceInputs instanceof DocumentsOnlyInput == false) {
-            throw createUnsupportedTypeException(inferenceInputs);
+            throw createUnsupportedTypeException(inferenceInputs, DocumentsOnlyInput.class);
         }
 
         return (DocumentsOnlyInput) inferenceInputs;
@@ -40,4 +40,8 @@ public List<String> getInputs() {
     public boolean stream() {
         return stream;
     }
+
+    public int inputSize() {
+        return input.size();
+    }
 }

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

@@ -10,15 +10,19 @@
 import org.elasticsearch.common.Strings;
 
 public abstract class InferenceInputs {
-    public static IllegalArgumentException createUnsupportedTypeException(InferenceInputs inferenceInputs) {
-        return new IllegalArgumentException(Strings.format("Unsupported inference inputs type: [%s]", inferenceInputs.getClass()));
+    public static IllegalArgumentException createUnsupportedTypeException(InferenceInputs inferenceInputs, Class<?> clazz) {
+        return new IllegalArgumentException(
+            Strings.format("Unable to convert inference inputs type: [%s] to [%s]", inferenceInputs.getClass(), clazz)
+        );
     }
 
     public <T> T castTo(Class<T> clazz) {
-        if (this.getClass().isInstance(clazz) == false) {
-            throw createUnsupportedTypeException(this);
+        if (clazz.isInstance(this) == false) {
+            throw createUnsupportedTypeException(this, clazz);
         }
 
         return clazz.cast(this);
     }
+
+    public abstract int inputSize();
 }

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

@@ -14,7 +14,7 @@ public class QueryAndDocsInputs extends InferenceInputs {
 
     public static QueryAndDocsInputs of(InferenceInputs inferenceInputs) {
         if (inferenceInputs instanceof QueryAndDocsInputs == false) {
-            throw createUnsupportedTypeException(inferenceInputs);
+            throw createUnsupportedTypeException(inferenceInputs, QueryAndDocsInputs.class);
         }
 
         return (QueryAndDocsInputs) inferenceInputs;
@@ -47,4 +47,7 @@ public boolean stream() {
         return stream;
     }
 
+    public int inputSize() {
+        return chunks.size();
+    }
 }

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

@@ -18,7 +18,7 @@ public class UnifiedChatInput extends InferenceInputs {
 
     public static UnifiedChatInput of(InferenceInputs inferenceInputs) {
         if (inferenceInputs instanceof UnifiedChatInput == false) {
-            throw createUnsupportedTypeException(inferenceInputs);
+            throw createUnsupportedTypeException(inferenceInputs, UnifiedChatInput.class);
         }
 
         return (UnifiedChatInput) inferenceInputs;
@@ -63,4 +63,8 @@ public UnifiedCompletionRequest getRequest() {
     public boolean stream() {
         return stream;
     }
+
+    public int inputSize() {
+        return request.messages().size();
+    }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/openai/OpenAiUnifiedChatCompletionRequestEntity.java

@@ -108,7 +108,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         builder.endArray();
 
         if (unifiedRequest.model() != null) {
-            builder.field(MODEL_FIELD, unifiedRequest.model());
+            builder.field(MODEL_FIELD, model.getServiceSettings().modelId());
         }
         if (unifiedRequest.maxCompletionTokens() != null) {
             builder.field(MAX_COMPLETION_TOKENS_FIELD, unifiedRequest.maxCompletionTokens());

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/completion/OpenAiChatCompletionModel.java

@@ -25,6 +25,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Objects;
 
 import static org.elasticsearch.xpack.inference.services.openai.OpenAiServiceFields.USER;
 
@@ -41,7 +42,24 @@ public static OpenAiChatCompletionModel of(OpenAiChatCompletionModel model, Map<
 
     public static OpenAiChatCompletionModel of(OpenAiChatCompletionModel model, UnifiedCompletionRequest request) {
         var requestTaskSettings = OpenAiChatCompletionRequestTaskSettings.fromUnifiedRequest(request);
-        return new OpenAiChatCompletionModel(model, OpenAiChatCompletionTaskSettings.of(model.getTaskSettings(), requestTaskSettings));
+        var originalModelServiceSettings = model.getServiceSettings();
+        var overriddenServiceSettings = new OpenAiChatCompletionServiceSettings(
+            Objects.requireNonNullElse(request.model(), originalModelServiceSettings.modelId()),
+            originalModelServiceSettings.uri(),
+            originalModelServiceSettings.organizationId(),
+            originalModelServiceSettings.maxInputTokens(),
+            originalModelServiceSettings.rateLimitSettings()
+        );
+
+        var overriddenTaskSettings = OpenAiChatCompletionTaskSettings.of(model.getTaskSettings(), requestTaskSettings);
+        return new OpenAiChatCompletionModel(
+            overriddenServiceSettings.modelId(),
+            model.getTaskType(),
+            model.getConfigurations().getService(),
+            overriddenServiceSettings,
+            overriddenTaskSettings,
+            model.getSecretSettings()
+        );
     }
 
     public OpenAiChatCompletionModel(

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/InferenceInputsTests.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.external.http.sender;
+
+import org.elasticsearch.common.Strings;
+import org.elasticsearch.test.ESTestCase;
+import org.hamcrest.Matchers;
+
+import java.util.List;
+
+public class InferenceInputsTests extends ESTestCase {
+    public void testCastToSucceeds() {
+        InferenceInputs inputs = new DocumentsOnlyInput(List.of(), false);
+        assertThat(inputs.castTo(DocumentsOnlyInput.class), Matchers.instanceOf(DocumentsOnlyInput.class));
+
+        assertThat(UnifiedChatInput.of(List.of(), false).castTo(UnifiedChatInput.class), Matchers.instanceOf(UnifiedChatInput.class));
+        assertThat(
+            new QueryAndDocsInputs("hello", List.of(), false).castTo(QueryAndDocsInputs.class),
+            Matchers.instanceOf(QueryAndDocsInputs.class)
+        );
+    }
+
+    public void testCastToFails() {
+        InferenceInputs inputs = new DocumentsOnlyInput(List.of(), false);
+        var exception = expectThrows(IllegalArgumentException.class, () -> inputs.castTo(QueryAndDocsInputs.class));
+        assertThat(
+            exception.getMessage(),
+            Matchers.containsString(
+                Strings.format("Unable to convert inference inputs type: [%s] to [%s]", DocumentsOnlyInput.class, QueryAndDocsInputs.class)
+            )
+        );
+    }
+}