Cleaning up

Author: jonathan-buttner

docs/changelog/125679.yaml

@@ -1,5 +1,5 @@
 pr: 125679
-summary: Custom Inference Service
+summary: Adding support for generic Inference services
 area: Machine Learning
 type: enhancement
 issues: []

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

@@ -214,7 +214,6 @@ public static ElasticsearchStatusException unknownSettingsError(Map<String, Obje
     }
 
     public static ElasticsearchStatusException unknownSettingsError(Map<String, Object> config, String field, String scope) {
-        // TODO map as JSON
         return new ElasticsearchStatusException(
             "Model configuration contains unknown settings [{}] while parsing field [{}] for settings [{}]",
             RestStatus.BAD_REQUEST,
@@ -546,15 +545,15 @@ public static List<Tuple<String, String>> extractOptionalListOfStringTuples(
 
             var firstElement = listEntry.get(0);
             var secondElement = listEntry.get(1);
-            validateTuple(firstElement, settingName, scope, "the first element", tuplesIndex, validationException);
-            validateTuple(secondElement, settingName, scope, "the second element", tuplesIndex, validationException);
+            validateString(firstElement, settingName, scope, "the first element", tuplesIndex, validationException);
+            validateString(secondElement, settingName, scope, "the second element", tuplesIndex, validationException);
             tuples.add(new Tuple<>((String) firstElement, (String) secondElement));
         }
 
         return tuples;
     }
 
-    private static void validateTuple(
+    private static void validateString(
         Object tupleValue,
         String settingName,
         String scope,

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

@@ -18,7 +18,7 @@
 import org.elasticsearch.xpack.inference.services.custom.response.ErrorResponseParser;
 
 /**
- * Defines how to handle various errors returned from the custom integration.
+ * Defines how to handle various response types returned from the custom integration.
  */
 public class CustomResponseHandler extends BaseResponseHandler {
     public CustomResponseHandler(String requestType, ResponseParser parseFunction, ErrorResponseParser errorParser) {

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

@@ -197,7 +197,7 @@ public CustomServiceSettings(StreamInput in) throws IOException {
 
     @Override
     public String modelId() {
-        // returning null because the model id is embedded in the url
+        // returning null because the model id is embedded in the url or the request body
         return null;
     }
 

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

@@ -19,15 +19,12 @@
 import org.elasticsearch.inference.SimilarityMeasure;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.test.ESTestCase;
-import org.elasticsearch.test.http.MockResponse;
 import org.elasticsearch.test.http.MockWebServer;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
-import org.elasticsearch.xpack.core.inference.results.TextEmbeddingFloatResults;
 import org.elasticsearch.xpack.inference.external.http.HttpClientManager;
 import org.elasticsearch.xpack.inference.logging.ThrottlerManager;
 import org.elasticsearch.xpack.inference.services.custom.CustomModel;
-import org.elasticsearch.xpack.inference.services.custom.response.TextEmbeddingResponseParser;
 import org.junit.After;
 import org.junit.Assume;
 import org.junit.Before;
@@ -45,12 +42,18 @@
 import static org.elasticsearch.xpack.inference.Utils.getRequestConfigMap;
 import static org.elasticsearch.xpack.inference.Utils.inferenceUtilityPool;
 import static org.elasticsearch.xpack.inference.Utils.mockClusterServiceEmpty;
-import static org.elasticsearch.xpack.inference.external.http.Utils.getUrl;
 import static org.hamcrest.Matchers.containsString;
-import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.is;
 import static org.mockito.Mockito.mock;
 
+/**
+ * Base class for testing inference services.
+ * <p>
+ * This class provides common unit tests for inference services, such as testing the model creation, and calling the infer method.
+ *
+ * To use this class, extend it and pass the constructor a configuration.
+ * </p>
+ */
 public abstract class AbstractServiceTests extends ESTestCase {
 
     protected final MockWebServer webServer = new MockWebServer();
@@ -125,8 +128,6 @@ public CommonConfig(TaskType taskType, @Nullable TaskType unsupportedTaskType) {
 
         protected abstract Map<String, Object> createSecretSettingsMap();
 
-        protected abstract CustomModel createEmbeddingModel(TextEmbeddingResponseParser embeddingResponseParser, String url);
-
         protected abstract void assertModel(Model model, TaskType taskType);
 
         protected abstract EnumSet<TaskType> supportedStreamingTasks();
@@ -434,62 +435,6 @@ public void testParsePersistedConfigWithSecrets_ThrowsWhenAnExtraKeyExistsInSecr
         }
     }
 
-    // infer tests
-
-    public void testInfer_SendsRequest() throws IOException {
-        try (var service = testConfiguration.commonConfig.createService(threadPool, clientManager)) {
-            String responseJson = """
-                {
-                  "object": "list",
-                  "data": [
-                      {
-                          "object": "embedding",
-                          "index": 0,
-                          "embedding": [
-                              0.0123,
-                              -0.0123
-                          ]
-                      }
-                  ],
-                  "model": "text-embedding-ada-002-v2",
-                  "usage": {
-                      "prompt_tokens": 8,
-                      "total_tokens": 8
-                  }
-                }
-                """;
-
-            webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson));
-
-            var model = testConfiguration.commonConfig.createEmbeddingModel(
-                new TextEmbeddingResponseParser("$.data[*].embedding"),
-                getUrl(webServer)
-            );
-            PlainActionFuture<InferenceServiceResults> listener = new PlainActionFuture<>();
-            service.infer(
-                model,
-                null,
-                null,
-                null,
-                List.of("test input"),
-                false,
-                new HashMap<>(),
-                InputType.INTERNAL_SEARCH,
-                InferenceAction.Request.DEFAULT_TIMEOUT,
-                listener
-            );
-
-            InferenceServiceResults results = listener.actionGet(TIMEOUT);
-            assertThat(results, instanceOf(TextEmbeddingFloatResults.class));
-
-            var embeddingResults = (TextEmbeddingFloatResults) results;
-            assertThat(
-                embeddingResults.embeddings(),
-                is(List.of(new TextEmbeddingFloatResults.Embedding(new float[] { 0.0123F, -0.0123F })))
-            );
-        }
-    }
-
     public void testInfer_ThrowsErrorWhenModelIsNotValid() throws IOException {
         try (var service = testConfiguration.commonConfig.createService(threadPool, clientManager)) {
             var listener = new PlainActionFuture<InferenceServiceResults>();

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/custom/CustomModelTests.java

@@ -25,13 +25,12 @@
 import static org.hamcrest.Matchers.is;
 
 public class CustomModelTests extends ESTestCase {
-    public static String taskSettingsKey = "test_taskSettings_key";
-    public static String taskSettingsValue = "test_taskSettings_value";
+    private static final String taskSettingsKey = "test_taskSettings_key";
+    private static final String taskSettingsValue = "test_taskSettings_value";
 
-    public static String secretSettingsKey = "test_secret_key";
-    public static SerializableSecureString secretSettingsValue = new SerializableSecureString("test_secret_value");
-    public static String url = "http://www.abc.com";
-    public static String path = "/endpoint";
+    private static final String secretSettingsKey = "test_secret_key";
+    private static final SerializableSecureString secretSettingsValue = new SerializableSecureString("test_secret_value");
+    private static final String url = "http://www.abc.com";
 
     public void testOverride_DoesNotModifiedFields_TaskSettingsIsEmpty() {
         var model = createModel(