[ML] Replace test-only SenderService API

prwhelan · Sep 10, 2024 · b00d243 · b00d243
1 parent 72248e3
commit b00d243
Show file tree

Hide file tree

Showing 11 changed files with 13 additions and 14 deletions.
diff --git a/...gin/inference/src/main/java/org/elasticsearch/xpack/inference/services/SenderService.java b/...gin/inference/src/main/java/org/elasticsearch/xpack/inference/services/SenderService.java
@@ -61,19 +61,6 @@ public void infer(
         }
     }
 
-    public void chunkedInfer(
-        Model model,
-        List<String> input,
-        Map<String, Object> taskSettings,
-        InputType inputType,
-        ChunkingOptions chunkingOptions,
-        TimeValue timeout,
-        ActionListener<List<ChunkedInferenceServiceResults>> listener
-    ) {
-        init();
-        chunkedInfer(model, null, input, taskSettings, inputType, chunkingOptions, timeout, listener);
-    }
-
     @Override
     public void chunkedInfer(
         Model model,
@@ -86,7 +73,7 @@ public void chunkedInfer(
         ActionListener<List<ChunkedInferenceServiceResults>> listener
     ) {
         init();
-        doChunkedInfer(model, null, input, taskSettings, inputType, chunkingOptions, timeout, listener);
+        doChunkedInfer(model, query, input, taskSettings, inputType, chunkingOptions, timeout, listener);
     }
 
     protected abstract void doInfer(

diff --git a/...a/org/elasticsearch/xpack/inference/services/amazonbedrock/AmazonBedrockServiceTests.java b/...a/org/elasticsearch/xpack/inference/services/amazonbedrock/AmazonBedrockServiceTests.java
@@ -1072,6 +1072,7 @@ public void testChunkedInfer_CallsInfer_ConvertsFloatResponse_ForEmbeddings() th
                 PlainActionFuture<List<ChunkedInferenceServiceResults>> listener = new PlainActionFuture<>();
                 service.chunkedInfer(
                     model,
+                    null,
                     List.of("abc", "xyz"),
                     new HashMap<>(),
                     InputType.INGEST,

diff --git a/...a/org/elasticsearch/xpack/inference/services/azureaistudio/AzureAiStudioServiceTests.java b/...a/org/elasticsearch/xpack/inference/services/azureaistudio/AzureAiStudioServiceTests.java
@@ -897,6 +897,7 @@ public void testChunkedInfer() throws IOException {
             PlainActionFuture<List<ChunkedInferenceServiceResults>> listener = new PlainActionFuture<>();
             service.chunkedInfer(
                 model,
+                null,
                 List.of("foo", "bar"),
                 new HashMap<>(),
                 InputType.INGEST,

diff --git a/.../java/org/elasticsearch/xpack/inference/services/azureopenai/AzureOpenAiServiceTests.java b/.../java/org/elasticsearch/xpack/inference/services/azureopenai/AzureOpenAiServiceTests.java
@@ -1104,6 +1104,7 @@ public void testChunkedInfer_CallsInfer_ConvertsFloatResponse() throws IOExcepti
             PlainActionFuture<List<ChunkedInferenceServiceResults>> listener = new PlainActionFuture<>();
             service.chunkedInfer(
                 model,
+                null,
                 List.of("foo", "bar"),
                 new HashMap<>(),
                 InputType.INGEST,

diff --git a/...e/src/test/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceTests.java b/...e/src/test/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceTests.java
@@ -1210,6 +1210,7 @@ public void testChunkedInfer_BatchesCalls() throws IOException {
             // 2 inputs
             service.chunkedInfer(
                 model,
+                null,
                 List.of("foo", "bar"),
                 new HashMap<>(),
                 InputType.UNSPECIFIED,
@@ -1301,6 +1302,7 @@ public void testChunkedInfer_BatchesCalls_Bytes() throws IOException {
             // 2 inputs
             service.chunkedInfer(
                 model,
+                null,
                 List.of("foo", "bar"),
                 new HashMap<>(),
                 InputType.UNSPECIFIED,

diff --git a/...java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceTests.java b/...java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceTests.java
@@ -476,6 +476,7 @@ public void testChunkedInfer_PassesThrough() throws IOException {
             PlainActionFuture<List<ChunkedInferenceServiceResults>> listener = new PlainActionFuture<>();
             service.chunkedInfer(
                 model,
+                null,
                 List.of("input text"),
                 new HashMap<>(),
                 InputType.INGEST,

diff --git a/...org/elasticsearch/xpack/inference/services/googleaistudio/GoogleAiStudioServiceTests.java b/...org/elasticsearch/xpack/inference/services/googleaistudio/GoogleAiStudioServiceTests.java
@@ -697,6 +697,7 @@ public void testChunkedInfer_Batches() throws IOException {
             PlainActionFuture<List<ChunkedInferenceServiceResults>> listener = new PlainActionFuture<>();
             service.chunkedInfer(
                 model,
+                null,
                 input,
                 new HashMap<>(),
                 InputType.INGEST,

diff --git a/.../org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceElserServiceTests.java b/.../org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceElserServiceTests.java
@@ -87,6 +87,7 @@ public void testChunkedInfer_CallsInfer_Elser_ConvertsFloatResponse() throws IOE
             PlainActionFuture<List<ChunkedInferenceServiceResults>> listener = new PlainActionFuture<>();
             service.chunkedInfer(
                 model,
+                null,
                 List.of("abc"),
                 new HashMap<>(),
                 InputType.INGEST,

diff --git a/.../java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceServiceTests.java b/.../java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceServiceTests.java
@@ -616,6 +616,7 @@ public void testChunkedInfer_CallsInfer_TextEmbedding_ConvertsFloatResponse() th
             PlainActionFuture<List<ChunkedInferenceServiceResults>> listener = new PlainActionFuture<>();
             service.chunkedInfer(
                 model,
+                null,
                 List.of("abc"),
                 new HashMap<>(),
                 InputType.INGEST,
@@ -671,6 +672,7 @@ public void testChunkedInfer() throws IOException {
             PlainActionFuture<List<ChunkedInferenceServiceResults>> listener = new PlainActionFuture<>();
             service.chunkedInfer(
                 model,
+                null,
                 List.of("abc"),
                 new HashMap<>(),
                 InputType.INGEST,

diff --git a/...src/test/java/org/elasticsearch/xpack/inference/services/mistral/MistralServiceTests.java b/...src/test/java/org/elasticsearch/xpack/inference/services/mistral/MistralServiceTests.java
@@ -464,6 +464,7 @@ public void testChunkedInfer_Embeddings_CallsInfer_ConvertsFloatResponse() throw
             PlainActionFuture<List<ChunkedInferenceServiceResults>> listener = new PlainActionFuture<>();
             service.chunkedInfer(
                 model,
+                null,
                 List.of("abc", "def"),
                 new HashMap<>(),
                 InputType.INGEST,

diff --git a/...e/src/test/java/org/elasticsearch/xpack/inference/services/openai/OpenAiServiceTests.java b/...e/src/test/java/org/elasticsearch/xpack/inference/services/openai/OpenAiServiceTests.java
@@ -1254,6 +1254,7 @@ public void testChunkedInfer_Batches() throws IOException {
             PlainActionFuture<List<ChunkedInferenceServiceResults>> listener = new PlainActionFuture<>();
             service.chunkedInfer(
                 model,
+                null,
                 List.of("foo", "bar"),
                 new HashMap<>(),
                 InputType.INGEST,