Address some review points

Author: dimitris-athanasiou

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

@@ -11,7 +11,6 @@
 import org.elasticsearch.TransportVersions;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
-import org.elasticsearch.features.NodeFeature;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xpack.core.XPackFeatureUsage;
 import org.elasticsearch.xpack.core.XPackField;
@@ -24,8 +23,6 @@
 
 public class InferenceFeatureSetUsage extends XPackFeatureUsage {
 
-    public static final NodeFeature SEMANTIC_TEXT_USAGE = new NodeFeature("inference.semantic_text_usage");
-
     public static final InferenceFeatureSetUsage EMPTY = new InferenceFeatureSetUsage(List.of());
 
     private final Collection<ModelStats> modelStats;

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

@@ -11,6 +11,7 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.features.NodeFeature;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xcontent.XContentBuilder;
@@ -20,6 +21,8 @@
 
 public class ModelStats implements ToXContentObject, Writeable {
 
+    public static final NodeFeature SEMANTIC_TEXT_USAGE = new NodeFeature("inference.semantic_text_usage");
+
     static final TransportVersion INFERENCE_TELEMETRY_ADDED_SEMANTIC_TEXT_STATS = TransportVersion.fromName(
         "inference_telemetry_added_semantic_text_stats"
     );

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

@@ -9,7 +9,7 @@
 
 import org.elasticsearch.features.FeatureSpecification;
 import org.elasticsearch.features.NodeFeature;
-import org.elasticsearch.xpack.core.inference.InferenceFeatureSetUsage;
+import org.elasticsearch.xpack.core.inference.usage.ModelStats;
 import org.elasticsearch.xpack.inference.mapper.SemanticInferenceMetadataFieldsMapper;
 import org.elasticsearch.xpack.inference.mapper.SemanticTextFieldMapper;
 import org.elasticsearch.xpack.inference.queries.InterceptedInferenceQueryBuilder;
@@ -89,7 +89,7 @@ public Set<NodeFeature> getTestFeatures() {
                 SemanticQueryBuilder.SEMANTIC_QUERY_FILTER_FIELD_CAPS_FIX,
                 InterceptedInferenceQueryBuilder.NEW_SEMANTIC_QUERY_INTERCEPTORS,
                 TEXT_SIMILARITY_RERANKER_SNIPPETS,
-                InferenceFeatureSetUsage.SEMANTIC_TEXT_USAGE
+                ModelStats.SEMANTIC_TEXT_USAGE
             )
         );
         return testFeatures;

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

@@ -35,7 +35,6 @@
 import org.elasticsearch.xpack.inference.registry.ModelRegistry;
 
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -93,7 +92,6 @@ private InferenceFeatureSetUsage collectUsage(List<ModelConfigurations> endpoint
             mapInferenceFieldsByIndexServiceAndTask(indicesMetadata, endpoints);
         Map<String, ModelStats> endpointStats = new TreeMap<>();
         addStatsByServiceAndTask(inferenceFieldsByIndexServiceAndTask, endpoints, endpointStats);
-        addTopLevelSemanticTextStatsByTask(inferenceFieldsByIndexServiceAndTask, endpointStats);
         addStatsForDefaultModels(inferenceFieldsByIndexServiceAndTask, endpoints, endpointStats);
         return new InferenceFeatureSetUsage(endpointStats.values().stream().filter(stats -> stats.count() > 0).toList());
     }
@@ -149,13 +147,17 @@ private static void addStatsByServiceAndTask(
                 endpointStats.get(serviceAndTaskType.toString())
             )
         );
+        addTopLevelSemanticTextStatsByTask(inferenceFieldsByIndexServiceAndTask, endpointStats);
     }
 
     private static void addTopLevelSemanticTextStatsByTask(
         Map<ServiceAndTaskType, Map<String, List<InferenceFieldMetadata>>> inferenceFieldsByIndexServiceAndTask,
         Map<String, ModelStats> endpointStats
     ) {
-        for (TaskType taskType : Arrays.stream(TaskType.values()).filter(t -> t != TaskType.ANY).toList()) {
+        for (TaskType taskType : TaskType.values()) {
+            if (taskType == TaskType.ANY) {
+                continue;
+            }
             ModelStats allStatsForTaskType = endpointStats.computeIfAbsent(
                 new ServiceAndTaskType(Metadata.ALL, taskType).toString(),
                 key -> new ModelStats(Metadata.ALL, taskType)

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

@@ -97,7 +97,7 @@ public void close() {
         client.threadPool().shutdown();
     }
 
-    public void testGivenServices_NoSemanticTextFields() throws Exception {
+    public void testGivenServices_NoInferenceFields() throws Exception {
         givenInferenceEndpoints(
             new ModelConfigurations("model-001", TaskType.TEXT_EMBEDDING, "openai", mockServiceSettings("model-id-001")),
             new ModelConfigurations("model-002", TaskType.TEXT_EMBEDDING, "openai", mockServiceSettings("model-id-002")),
@@ -174,7 +174,7 @@ public void testGivenVariousServicesAndInferenceFields() throws Exception {
         assertStats(response, 4, new ModelStats("openai", TaskType.TEXT_EMBEDDING, 2, new SemanticTextStats(6, 2, 2)));
     }
 
-    public void testGivenServices_SemanticTextFieldsDefaultModels() throws Exception {
+    public void testGivenServices_InferenceFieldsReferencingDefaultModels() throws Exception {
         givenInferenceEndpoints(
             new ModelConfigurations(".endpoint-001", TaskType.TEXT_EMBEDDING, "eis", mockServiceSettings(".model-id-001")),
             new ModelConfigurations(".endpoint-002", TaskType.SPARSE_EMBEDDING, "eis", mockServiceSettings(".model-id-002")),

x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/inference/inference_usage.yml

@@ -1,7 +1,7 @@
 setup:
   - requires:
       cluster_features: "inference.semantic_text_usage"
-      reason: tests added after semantic_text usage was added
+      reason: semantic text stats added to inference usage
 
 ---
 "Test usage given default endpoints":
@@ -16,47 +16,47 @@ setup:
   - match: { inference.models.0.service: "_all" }
   - match: { inference.models.0.task_type: "RERANK" }
   - match: { inference.models.0.count: 1 }
-  - is_false: inference.models.0.semantic_text
+  - not_exists: inference.models.0.semantic_text
 
   - match: { inference.models.1.service: "_all" }
   - match: { inference.models.1.task_type: "SPARSE_EMBEDDING" }
   - match: { inference.models.1.count: 1 }
-  - is_false: inference.models.1.semantic_text
+  - not_exists: inference.models.1.semantic_text
 
   - match: { inference.models.2.service: "_all" }
   - match: { inference.models.2.task_type: "TEXT_EMBEDDING" }
   - match: { inference.models.2.count: 1 }
-  - is_false: inference.models.2.semantic_text
+  - not_exists: inference.models.2.semantic_text
 
   - match: { inference.models.3.service: "_elasticsearch__elser_model_2" }
   - match: { inference.models.3.task_type: "SPARSE_EMBEDDING" }
   - match: { inference.models.3.count: 1 }
-  - is_false: inference.models.3.semantic_text
+  - not_exists: inference.models.3.semantic_text
 
   - match: { inference.models.4.service: "_elasticsearch__multilingual-e5-small" }
   - match: { inference.models.4.task_type: "TEXT_EMBEDDING" }
   - match: { inference.models.4.count: 1 }
-  - is_false: inference.models.4.semantic_text
+  - not_exists: inference.models.4.semantic_text
 
   - match: { inference.models.5.service: "_elasticsearch__rerank-v1" }
   - match: { inference.models.5.task_type: "RERANK" }
   - match: { inference.models.5.count: 1 }
-  - is_false: inference.models.5.semantic_text
+  - not_exists: inference.models.5.semantic_text
 
   - match: { inference.models.6.service: "elasticsearch" }
   - match: { inference.models.6.task_type: "RERANK" }
   - match: { inference.models.6.count: 1 }
-  - is_false: inference.models.6.semantic_text
+  - not_exists: inference.models.6.semantic_text
 
   - match: { inference.models.7.service: "elasticsearch" }
   - match: { inference.models.7.task_type: "SPARSE_EMBEDDING" }
   - match: { inference.models.7.count: 1 }
-  - is_false: inference.models.7.semantic_text
+  - not_exists: inference.models.7.semantic_text
 
   - match: { inference.models.8.service: "elasticsearch" }
   - match: { inference.models.8.task_type: "TEXT_EMBEDDING" }
   - match: { inference.models.8.count: 1 }
-  - is_false: inference.models.8.semantic_text
+  - not_exists: inference.models.8.semantic_text
 
 ---
 "Test usage given default endpoints and semantic_text fields":
@@ -91,7 +91,7 @@ setup:
   - match: { inference.models.0.service: "_all" }
   - match: { inference.models.0.task_type: "RERANK" }
   - match: { inference.models.0.count: 1 }
-  - is_false: inference.models.0.semantic_text
+  - not_exists: inference.models.0.semantic_text
 
   - match: { inference.models.1.service: "_all" }
   - match: { inference.models.1.task_type: "SPARSE_EMBEDDING" }
@@ -124,12 +124,12 @@ setup:
   - match: { inference.models.5.service: "_elasticsearch__rerank-v1" }
   - match: { inference.models.5.task_type: "RERANK" }
   - match: { inference.models.5.count: 1 }
-  - is_false: inference.models.5.semantic_text
+  - not_exists: inference.models.5.semantic_text
 
   - match: { inference.models.6.service: "elasticsearch" }
   - match: { inference.models.6.task_type: "RERANK" }
   - match: { inference.models.6.count: 1 }
-  - is_false: inference.models.6.semantic_text
+  - not_exists: inference.models.6.semantic_text
 
   - match: { inference.models.7.service: "elasticsearch" }
   - match: { inference.models.7.task_type: "SPARSE_EMBEDDING" }