Add usage stats for semantic_text fields

docs/changelog/135262.yaml

@@ -0,0 +1,5 @@
+pr: 135262
+summary: Add usage stats for `semantic_text` fields
+area: "Vector Search"
+type: enhancement
+issues: []

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

@@ -21,8 +21,8 @@
 import java.util.Objects;
 
 public enum TaskType implements Writeable {
-    TEXT_EMBEDDING,
-    SPARSE_EMBEDDING,
+    TEXT_EMBEDDING(true),
+    SPARSE_EMBEDDING(true),
     RERANK,
     COMPLETION,
     ANY {
@@ -52,6 +52,16 @@ 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.
@@ -62,6 +72,14 @@ 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);

server/src/main/resources/transport/definitions/referable/inference_telemetry_added_semantic_text_stats.csv

@@ -0,0 +1 @@
+9181000

server/src/main/resources/transport/upper_bounds/9.2.csv

@@ -1 +1 @@
-elastic_reranker_chunking_configuration,9180000
+inference_telemetry_added_semantic_text_stats,9181000

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

@@ -7,9 +7,12 @@
 
 package org.elasticsearch.xpack.core.inference.usage;
 
+import org.elasticsearch.TransportVersion;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.features.NodeFeature;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xcontent.XContentBuilder;
@@ -19,28 +22,42 @@
 
 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"
+    );
+
     private final String service;
     private final TaskType taskType;
     private long count;
+    @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;
         this.count = count;
-    }
-
-    public ModelStats(ModelStats stats) {
-        this(stats.service, stats.taskType, stats.count);
+        this.semanticTextStats = semanticTextStats;
     }
 
     public ModelStats(StreamInput in) throws IOException {
         this.service = in.readString();
         this.taskType = in.readEnum(TaskType.class);
         this.count = in.readLong();
+        if (in.getTransportVersion().supports(INFERENCE_TELEMETRY_ADDED_SEMANTIC_TEXT_STATS)) {
+            this.semanticTextStats = in.readOptional(SemanticTextStats::new);
+        } else {
+            this.semanticTextStats = null;
+        }
     }
 
     public void add() {
@@ -59,6 +76,11 @@ public long count() {
         return count;
     }
 
+    @Nullable
+    public SemanticTextStats semanticTextStats() {
+        return semanticTextStats;
+    }
+
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();
@@ -71,25 +93,34 @@ public void addXContentFragment(XContentBuilder builder, Params params) throws I
         builder.field("service", service);
         builder.field("task_type", taskType.name());
         builder.field("count", count);
+        if (semanticTextStats != null) {
+            builder.field("semantic_text", semanticTextStats);
+        }
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         out.writeString(service);
         out.writeEnum(taskType);
         out.writeLong(count);
+        if (out.getTransportVersion().supports(INFERENCE_TELEMETRY_ADDED_SEMANTIC_TEXT_STATS)) {
+            out.writeOptionalWriteable(semanticTextStats);
+        }
     }
 
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
         ModelStats that = (ModelStats) o;
-        return count == that.count && Objects.equals(service, that.service) && taskType == that.taskType;
+        return count == that.count
+            && Objects.equals(service, that.service)
+            && taskType == that.taskType
+            && Objects.equals(semanticTextStats, that.semanticTextStats);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(service, taskType, count);
+        return Objects.hash(service, taskType, count, semanticTextStats);
     }
 }

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

@@ -0,0 +1,100 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.core.inference.usage;
+
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.xcontent.ToXContentObject;
+import org.elasticsearch.xcontent.XContentBuilder;
+
+import java.io.IOException;
+import java.util.Objects;
+
+public class SemanticTextStats implements ToXContentObject, Writeable {
+
+    private static final String FIELD_COUNT = "field_count";
+    private static final String INDICES_COUNT = "indices_count";
+    private static final String INFERENCE_ID_COUNT = "inference_id_count";
+
+    private long fieldCount;
+    private long indicesCount;
+    private long inferenceIdCount;
+
+    public SemanticTextStats() {}
+
+    public SemanticTextStats(long fieldCount, long indicesCount, long inferenceIdCount) {
+        this.fieldCount = fieldCount;
+        this.indicesCount = indicesCount;
+        this.inferenceIdCount = inferenceIdCount;
+    }
+
+    public SemanticTextStats(StreamInput in) throws IOException {
+        fieldCount = in.readVLong();
+        indicesCount = in.readVLong();
+        inferenceIdCount = in.readVLong();
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeVLong(fieldCount);
+        out.writeVLong(indicesCount);
+        out.writeVLong(inferenceIdCount);
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startObject();
+        builder.field(FIELD_COUNT, fieldCount);
+        builder.field(INDICES_COUNT, indicesCount);
+        builder.field(INFERENCE_ID_COUNT, inferenceIdCount);
+        builder.endObject();
+        return builder;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        SemanticTextStats that = (SemanticTextStats) o;
+        return fieldCount == that.fieldCount && indicesCount == that.indicesCount && inferenceIdCount == that.inferenceIdCount;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(fieldCount, indicesCount, inferenceIdCount);
+    }
+
+    public long getFieldCount() {
+        return fieldCount;
+    }
+
+    public long getIndicesCount() {
+        return indicesCount;
+    }
+
+    public long getInferenceIdCount() {
+        return inferenceIdCount;
+    }
+
+    public void addFieldCount(long fieldCount) {
+        this.fieldCount += fieldCount;
+    }
+
+    public void incIndicesCount() {
+        this.indicesCount++;
+    }
+
+    public void setInferenceIdCount(long inferenceIdCount) {
+        this.inferenceIdCount = inferenceIdCount;
+    }
+
+    public boolean isEmpty() {
+        return fieldCount == 0 && indicesCount == 0 && inferenceIdCount == 0;
+    }
+}

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

@@ -7,16 +7,17 @@
 
 package org.elasticsearch.xpack.core.inference.usage;
 
+import org.elasticsearch.TransportVersion;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.inference.TaskType;
-import org.elasticsearch.test.AbstractWireSerializingTestCase;
 import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xpack.core.ml.AbstractBWCWireSerializationTestCase;
 
 import java.io.IOException;
 
 import static org.hamcrest.Matchers.equalTo;
 
-public class ModelStatsTests extends AbstractWireSerializingTestCase<ModelStats> {
+public class ModelStatsTests extends AbstractBWCWireSerializationTestCase<ModelStats> {
 
     @Override
     protected Writeable.Reader<ModelStats> instanceReader() {
@@ -33,10 +34,22 @@ protected ModelStats mutateInstance(ModelStats modelStats) throws IOException {
         String service = modelStats.service();
         TaskType taskType = modelStats.taskType();
         long count = modelStats.count();
-        return switch (randomInt(2)) {
-            case 0 -> new ModelStats(randomValueOtherThan(service, ESTestCase::randomIdentifier), taskType, count);
-            case 1 -> new ModelStats(service, randomValueOtherThan(taskType, () -> randomFrom(TaskType.values())), count);
-            case 2 -> new ModelStats(service, taskType, randomValueOtherThan(count, ESTestCase::randomLong));
+        SemanticTextStats semanticTextStats = modelStats.semanticTextStats();
+        return switch (randomInt(3)) {
+            case 0 -> new ModelStats(randomValueOtherThan(service, ESTestCase::randomIdentifier), taskType, count, semanticTextStats);
+            case 1 -> new ModelStats(
+                service,
+                randomValueOtherThan(taskType, () -> randomFrom(TaskType.values())),
+                count,
+                semanticTextStats
+            );
+            case 2 -> new ModelStats(service, taskType, randomValueOtherThan(count, ESTestCase::randomLong), semanticTextStats);
+            case 3 -> new ModelStats(
+                service,
+                taskType,
+                count,
+                randomValueOtherThan(semanticTextStats, SemanticTextStatsTests::createRandomInstance)
+            );
             default -> throw new IllegalArgumentException();
         };
     }
@@ -56,6 +69,20 @@ public void testAdd() {
     }
 
     public static ModelStats createRandomInstance() {
-        return new ModelStats(randomIdentifier(), randomFrom(TaskType.values()), randomLong());
+        TaskType taskType = randomValueOtherThan(TaskType.ANY, () -> randomFrom(TaskType.values()));
+        return new ModelStats(
+            randomIdentifier(),
+            taskType,
+            randomLong(),
+            taskType.isCompatibleWithSemanticText() ? SemanticTextStatsTests.createRandomInstance() : null
+        );
+    }
+
+    @Override
+    protected ModelStats mutateInstanceForVersion(ModelStats instance, TransportVersion version) {
+        if (version.supports(ModelStats.INFERENCE_TELEMETRY_ADDED_SEMANTIC_TEXT_STATS) == false) {
+            return new ModelStats(instance.service(), instance.taskType(), instance.count(), null);
+        }
+        return instance;
     }
 }