add task_type validation

jimczi · Mar 20, 2024 · 7b578d1 · 7b578d1
1 parent 38f82fd
commit 7b578d1
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/...nce/src/main/java/org/elasticsearch/xpack/inference/mapper/SemanticTextModelSettings.java b/...nce/src/main/java/org/elasticsearch/xpack/inference/mapper/SemanticTextModelSettings.java
@@ -27,6 +27,9 @@
 import java.util.Map;
 import java.util.Objects;
 
+import static org.elasticsearch.inference.TaskType.SPARSE_EMBEDDING;
+import static org.elasticsearch.inference.TaskType.TEXT_EMBEDDING;
+
 /**
  * Serialization class for specifying the settings of a model from semantic_text inference to field mapper.
  */
@@ -58,6 +61,7 @@ public SemanticTextModelSettings(Model model) {
             model.getServiceSettings().dimensions(),
             model.getServiceSettings().similarity()
         );
+        validate();
     }
 
     public static SemanticTextModelSettings parse(XContentParser parser) throws IOException {
@@ -149,6 +153,18 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         return builder.endObject();
     }
 
+    public void validate() {
+        switch (taskType) {
+            case TEXT_EMBEDDING:
+            case SPARSE_EMBEDDING:
+                break;
+
+            default:
+                throw new IllegalArgumentException("Wrong [" + TASK_TYPE_FIELD.getPreferredName() + "], expected " +
+                        TEXT_EMBEDDING + "or " + SPARSE_EMBEDDING + ", got " + taskType.name());
+        }
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;