Fixing various tests

Author: jonathan-buttner

server/src/main/java/org/elasticsearch/inference/UnifiedCompletionRequest.java

@@ -81,8 +81,6 @@ public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
             new NamedWriteableRegistry.Entry(Content.class, ContentString.NAME, ContentString::new),
             new NamedWriteableRegistry.Entry(ToolChoice.class, ToolChoiceObject.NAME, ToolChoiceObject::new),
             new NamedWriteableRegistry.Entry(ToolChoice.class, ToolChoiceString.NAME, ToolChoiceString::new)
-            // new NamedWriteableRegistry.Entry(Stop.class, StopValues.NAME, StopValues::new),
-            // new NamedWriteableRegistry.Entry(Stop.class, StopString.NAME, StopString::new)
         );
     }
 

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

@@ -12,26 +12,20 @@
 
 public class ChatCompletionInput extends InferenceInputs {
     private final List<String> input;
-    private final boolean stream;
 
     public ChatCompletionInput(List<String> input) {
         this(input, false);
     }
 
     public ChatCompletionInput(List<String> input, boolean stream) {
-        super();
+        super(stream);
         this.input = Objects.requireNonNull(input);
-        this.stream = stream;
     }
 
     public List<String> getInputs() {
         return this.input;
     }
 
-    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/DocumentsOnlyInput.java

@@ -21,26 +21,20 @@ public static DocumentsOnlyInput of(InferenceInputs inferenceInputs) {
     }
 
     private final List<String> input;
-    private final boolean stream;
 
     public DocumentsOnlyInput(List<String> input) {
         this(input, false);
     }
 
     public DocumentsOnlyInput(List<String> input, boolean stream) {
-        super();
+        super(stream);
         this.input = Objects.requireNonNull(input);
-        this.stream = stream;
     }
 
     public List<String> getInputs() {
         return this.input;
     }
 
-    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,6 +10,12 @@
 import org.elasticsearch.common.Strings;
 
 public abstract class InferenceInputs {
+    private final boolean stream;
+
+    public InferenceInputs(boolean stream) {
+        this.stream = stream;
+    }
+
     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)
@@ -24,5 +30,9 @@ public <T> T castTo(Class<T> clazz) {
         return clazz.cast(this);
     }
 
+    public boolean stream() {
+        return stream;
+    }
+
     public abstract int inputSize();
 }

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

@@ -22,17 +22,15 @@ public static QueryAndDocsInputs of(InferenceInputs inferenceInputs) {
 
     private final String query;
     private final List<String> chunks;
-    private final boolean stream;
 
     public QueryAndDocsInputs(String query, List<String> chunks) {
         this(query, chunks, false);
     }
 
     public QueryAndDocsInputs(String query, List<String> chunks, boolean stream) {
-        super();
+        super(stream);
         this.query = Objects.requireNonNull(query);
         this.chunks = Objects.requireNonNull(chunks);
-        this.stream = stream;
     }
 
     public String getQuery() {
@@ -43,10 +41,6 @@ public List<String> getChunks() {
         return chunks;
     }
 
-    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

@@ -14,11 +14,10 @@
 
 public class UnifiedChatInput extends InferenceInputs {
     private final UnifiedCompletionRequest request;
-    private final boolean stream;
 
     public UnifiedChatInput(UnifiedCompletionRequest request, boolean stream) {
+        super(stream);
         this.request = Objects.requireNonNull(request);
-        this.stream = stream;
     }
 
     public UnifiedChatInput(ChatCompletionInput completionInput, String roleValue) {
@@ -47,10 +46,6 @@ public UnifiedCompletionRequest getRequest() {
         return request;
     }
 
-    public boolean stream() {
-        return stream;
-    }
-
     public int inputSize() {
         return request.messages().size();
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/SenderService.java

@@ -75,7 +75,7 @@ private static InferenceInputs createInput(Model model, List<String> input, @Nul
         return switch (model.getTaskType()) {
             case COMPLETION -> new ChatCompletionInput(input, stream);
             case RERANK -> new QueryAndDocsInputs(query, input, stream);
-            case TEXT_EMBEDDING -> new DocumentsOnlyInput(input, stream);
+            case TEXT_EMBEDDING, SPARSE_EMBEDDING -> new DocumentsOnlyInput(input, stream);
             default -> throw new ElasticsearchStatusException(
                 Strings.format("Invalid task type received when determining input type: [%s]", model.getTaskType().toString()),
                 RestStatus.BAD_REQUEST

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googleaistudio/GoogleAiStudioService.java

@@ -284,9 +284,8 @@ protected void doInfer(
     ) {
         if (model instanceof GoogleAiStudioCompletionModel completionModel) {
             var requestManager = new GoogleAiStudioCompletionRequestManager(completionModel, getServiceComponents().threadPool());
-            var docsOnly = DocumentsOnlyInput.of(inputs);
             var failedToSendRequestErrorMessage = constructFailedToSendRequestMessage(
-                completionModel.uri(docsOnly.stream()),
+                completionModel.uri(inputs.stream()),
                 "Google AI Studio completion"
             );
             var action = new SingleInputSenderExecutableAction(

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/Utils.java

@@ -19,6 +19,7 @@
 import org.elasticsearch.inference.ModelConfigurations;
 import org.elasticsearch.inference.ModelSecrets;
 import org.elasticsearch.inference.SimilarityMeasure;
+import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.threadpool.ScalingExecutorBuilder;
 import org.elasticsearch.xpack.core.inference.results.ChatCompletionResults;
 import org.elasticsearch.xpack.inference.common.Truncator;
@@ -160,9 +161,11 @@ public static Model getInvalidModel(String inferenceEntityId, String serviceName
         var mockConfigs = mock(ModelConfigurations.class);
         when(mockConfigs.getInferenceEntityId()).thenReturn(inferenceEntityId);
         when(mockConfigs.getService()).thenReturn(serviceName);
+        when(mockConfigs.getTaskType()).thenReturn(TaskType.TEXT_EMBEDDING);
 
         var mockModel = mock(Model.class);
         when(mockModel.getConfigurations()).thenReturn(mockConfigs);
+        when(mockModel.getTaskType()).thenReturn(TaskType.TEXT_EMBEDDING);
 
         return mockModel;
     }

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/SingleInputSenderExecutableActionTests.java

@@ -61,25 +61,11 @@ public void testOneInputIsValid() {
         assertTrue("Test failed to call listener.", testRan.get());
     }
 
-    public void testInvalidInputType() {
-        var badInput = mock(InferenceInputs.class);
-        var actualException = new AtomicReference<Exception>();
-
-        executableAction.execute(
-            badInput,
-            mock(TimeValue.class),
-            ActionListener.wrap(shouldNotSucceed -> fail("Test failed."), actualException::set)
-        );
-
-        assertThat(actualException.get(), notNullValue());
-        assertThat(actualException.get().getMessage(), is("Invalid inference input type"));
-        assertThat(actualException.get(), instanceOf(ElasticsearchStatusException.class));
-        assertThat(((ElasticsearchStatusException) actualException.get()).status(), is(RestStatus.INTERNAL_SERVER_ERROR));
-    }
-
     public void testMoreThanOneInput() {
         var badInput = mock(DocumentsOnlyInput.class);
-        when(badInput.getInputs()).thenReturn(List.of("one", "two"));
+        var input = List.of("one", "two");
+        when(badInput.getInputs()).thenReturn(input);
+        when(badInput.inputSize()).thenReturn(input.size());
         var actualException = new AtomicReference<Exception>();
 
         executableAction.execute(