Contain task type compatibility in TransportInferenceUsageAction

Author: dimitris-athanasiou

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

@@ -21,8 +21,8 @@
 import java.util.Objects;
 
 public enum TaskType implements Writeable {
-    TEXT_EMBEDDING(true),
-    SPARSE_EMBEDDING(true),
+    TEXT_EMBEDDING,
+    SPARSE_EMBEDDING,
     RERANK,
     COMPLETION,
     ANY {
@@ -52,16 +52,6 @@ public static TaskType fromStringOrStatusException(String name) {
         }
     }
 
-    private final boolean isCompatibleWithSemanticText;
-
-    TaskType(boolean isCompatibleWithSemanticText) {
-        this.isCompatibleWithSemanticText = isCompatibleWithSemanticText;
-    }
-
-    TaskType() {
-        this(false);
-    }
-
     /**
      * Return true if the {@code other} is the {@link #ANY} type
      * or the same as this.
@@ -72,14 +62,6 @@ public boolean isAnyOrSame(TaskType other) {
         return other == TaskType.ANY || other == this;
     }
 
-    /**
-     * Returns true if this task type is compatible with semantic text.
-     * @return True if this task type is compatible with semantic text.
-     */
-    public boolean isCompatibleWithSemanticText() {
-        return isCompatibleWithSemanticText;
-    }
-
     @Override
     public String toString() {
         return name().toLowerCase(Locale.ROOT);

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/usage/ModelStats.java

@@ -34,14 +34,6 @@ public class ModelStats implements ToXContentObject, Writeable {
     @Nullable
     private final SemanticTextStats semanticTextStats;
 
-    public ModelStats(String service, TaskType taskType) {
-        this(service, taskType, 0L);
-    }
-
-    public ModelStats(String service, TaskType taskType, long count) {
-        this(service, taskType, count, taskType.isCompatibleWithSemanticText() ? new SemanticTextStats() : null);
-    }
-
     public ModelStats(String service, TaskType taskType, long count, @Nullable SemanticTextStats semanticTextStats) {
         this.service = service;
         this.taskType = taskType;

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/inference/usage/ModelStatsTests.java

@@ -55,7 +55,7 @@ protected ModelStats mutateInstance(ModelStats modelStats) throws IOException {
     }
 
     public void testAdd() {
-        ModelStats stats = new ModelStats("test_service", randomFrom(TaskType.values()));
+        ModelStats stats = new ModelStats("test_service", randomFrom(TaskType.values()), 0, null);
         assertThat(stats.count(), equalTo(0L));
 
         stats.add();
@@ -74,7 +74,7 @@ public static ModelStats createRandomInstance() {
             randomIdentifier(),
             taskType,
             randomLong(),
-            taskType.isCompatibleWithSemanticText() ? SemanticTextStatsTests.createRandomInstance() : null
+            randomBoolean() ? SemanticTextStatsTests.createRandomInstance() : null
         );
     }
 

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

@@ -32,9 +32,11 @@
 import org.elasticsearch.xpack.core.inference.InferenceFeatureSetUsage;
 import org.elasticsearch.xpack.core.inference.action.GetInferenceModelAction;
 import org.elasticsearch.xpack.core.inference.usage.ModelStats;
+import org.elasticsearch.xpack.core.inference.usage.SemanticTextStats;
 import org.elasticsearch.xpack.inference.registry.ModelRegistry;
 
 import java.util.ArrayList;
+import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -55,6 +57,11 @@ public class TransportInferenceUsageAction extends XPackUsageFeatureTransportAct
     // Some of the default models have optimized variants for linux that will have the following suffix.
     private static final String MODEL_ID_LINUX_SUFFIX = "_linux-x86_64";
 
+    private static final EnumSet<TaskType> TASK_TYPES_WITH_SEMANTIC_TEXT_SUPPORT = EnumSet.of(
+        TaskType.TEXT_EMBEDDING,
+        TaskType.SPARSE_EMBEDDING
+    );
+
     private final ModelRegistry modelRegistry;
     private final Client client;
 
@@ -144,12 +151,12 @@ private static void addStatsByServiceAndTask(
         for (ModelConfigurations model : endpoints) {
             endpointStats.computeIfAbsent(
                 new ServiceAndTaskType(model.getService(), model.getTaskType()).toString(),
-                key -> new ModelStats(model.getService(), model.getTaskType())
+                key -> createEmptyStats(model)
             ).add();
 
             endpointStats.computeIfAbsent(
                 new ServiceAndTaskType(Metadata.ALL, model.getTaskType()).toString(),
-                key -> new ModelStats(Metadata.ALL, model.getTaskType())
+                key -> createEmptyStats(Metadata.ALL, model.getTaskType())
             ).add();
         }
 
@@ -162,6 +169,19 @@ private static void addStatsByServiceAndTask(
         addTopLevelStatsByTask(inferenceFieldsByIndexServiceAndTask, endpointStats);
     }
 
+    private static ModelStats createEmptyStats(ModelConfigurations model) {
+        return createEmptyStats(model.getService(), model.getTaskType());
+    }
+
+    private static ModelStats createEmptyStats(String service, TaskType taskType) {
+        return new ModelStats(
+            service,
+            taskType,
+            0,
+            TASK_TYPES_WITH_SEMANTIC_TEXT_SUPPORT.contains(taskType) ? new SemanticTextStats() : null
+        );
+    }
+
     private static void addTopLevelStatsByTask(
         Map<ServiceAndTaskType, Map<String, List<InferenceFieldMetadata>>> inferenceFieldsByIndexServiceAndTask,
         Map<String, ModelStats> endpointStats
@@ -172,9 +192,9 @@ private static void addTopLevelStatsByTask(
             }
             ModelStats allStatsForTaskType = endpointStats.computeIfAbsent(
                 new ServiceAndTaskType(Metadata.ALL, taskType).toString(),
-                key -> new ModelStats(Metadata.ALL, taskType)
+                key -> createEmptyStats(Metadata.ALL, taskType)
             );
-            if (taskType.isCompatibleWithSemanticText()) {
+            if (TASK_TYPES_WITH_SEMANTIC_TEXT_SUPPORT.contains(taskType)) {
                 Map<String, List<InferenceFieldMetadata>> inferenceFieldsByIndex = inferenceFieldsByIndexServiceAndTask.entrySet()
                     .stream()
                     .filter(e -> e.getKey().taskType == taskType)
@@ -226,7 +246,12 @@ private void addStatsForDefaultModelsCompatibleWithSemanticText(
             // Now that we have all inference fields for this service and task type, we want to keep only the ones that
             // reference the current default model.
             fieldsByIndex = filterFields(fieldsByIndex, f -> statKey.modelId.equals(endpointIdToModelId.get(f.getInferenceId())));
-            ModelStats stats = new ModelStats(statKey.toString(), statKey.taskType, defaultModelStatsKeyToEndpointCount.getValue());
+            ModelStats stats = new ModelStats(
+                statKey.toString(),
+                statKey.taskType,
+                defaultModelStatsKeyToEndpointCount.getValue(),
+                new SemanticTextStats()
+            );
             addSemanticTextStats(fieldsByIndex, stats);
             endpointStats.put(statKey.toString(), stats);
         }
@@ -239,7 +264,7 @@ private Map<DefaultModelStatsKey, Long> createStatsKeysWithEndpointCountsForDefa
         // Note that endpoints could have a null model id, in which case we don't consider them default as this
         // may only happen for external services.
         Set<String> modelIds = endpoints.stream()
-            .filter(endpoint -> endpoint.getTaskType().isCompatibleWithSemanticText())
+            .filter(endpoint -> TASK_TYPES_WITH_SEMANTIC_TEXT_SUPPORT.contains(endpoint.getTaskType()))
             .filter(endpoint -> modelRegistry.containsDefaultConfigId(endpoint.getInferenceEntityId()))
             .filter(endpoint -> endpoint.getServiceSettings().modelId() != null)
             .map(endpoint -> stripLinuxSuffix(endpoint.getServiceSettings().modelId()))

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

@@ -12,8 +12,6 @@
 import org.elasticsearch.test.ESTestCase;
 import org.hamcrest.Matchers;
 
-import static org.hamcrest.core.Is.is;
-
 public class TaskTypeTests extends ESTestCase {
 
     public void testFromStringOrStatusException() {
@@ -26,12 +24,4 @@ public void testFromStringOrStatusException() {
         assertThat(TaskType.fromStringOrStatusException("any"), Matchers.is(TaskType.ANY));
     }
 
-    public void testIsCompatibleWithSemanticText() {
-        assertThat(TaskType.ANY.isCompatibleWithSemanticText(), is(false));
-        assertThat(TaskType.CHAT_COMPLETION.isCompatibleWithSemanticText(), is(false));
-        assertThat(TaskType.COMPLETION.isCompatibleWithSemanticText(), is(false));
-        assertThat(TaskType.RERANK.isCompatibleWithSemanticText(), is(false));
-        assertThat(TaskType.TEXT_EMBEDDING.isCompatibleWithSemanticText(), is(true));
-        assertThat(TaskType.SPARSE_EMBEDDING.isCompatibleWithSemanticText(), is(true));
-    }
 }