Updating inference services API response

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

@@ -25,10 +25,12 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.stream.Collectors;
 
 import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg;
 
@@ -37,75 +39,88 @@
  */
 public class InferenceServiceConfiguration implements Writeable, ToXContentObject {
 
-    private final String provider;
-    private final List<TaskSettingsConfiguration> taskTypes;
-    private final Map<String, SettingsConfiguration> configuration;
+    private final String service;
+    private final String name;
+    private final EnumSet<TaskType> taskTypes;
+    private final Map<String, SettingsConfiguration> configurations;
 
     /**
      * Constructs a new {@link InferenceServiceConfiguration} instance with specified properties.
      *
-     * @param provider       The name of the service provider.
-     * @param taskTypes      A list of {@link TaskSettingsConfiguration} supported by the service provider.
-     * @param configuration  The configuration of the service provider, defined by {@link SettingsConfiguration}.
+     * @param service        The name of the service provider.
+     * @param name           The user-friendly name of the service provider.
+     * @param taskTypes      A list of {@link TaskType} supported by the service provider.
+     * @param configurations  The configuration of the service provider, defined by {@link SettingsConfiguration}.
      */
     private InferenceServiceConfiguration(
-        String provider,
-        List<TaskSettingsConfiguration> taskTypes,
-        Map<String, SettingsConfiguration> configuration
+        String service,
+        String name,
+        EnumSet<TaskType> taskTypes,
+        Map<String, SettingsConfiguration> configurations
     ) {
-        this.provider = provider;
+        this.service = service;
+        this.name = name;
         this.taskTypes = taskTypes;
-        this.configuration = configuration;
+        this.configurations = configurations;
     }
 
     public InferenceServiceConfiguration(StreamInput in) throws IOException {
-        this.provider = in.readString();
-        this.taskTypes = in.readCollectionAsList(TaskSettingsConfiguration::new);
-        this.configuration = in.readMap(SettingsConfiguration::new);
+        this.service = in.readString();
+        this.name = in.readString();
+        this.taskTypes = in.readEnumSet(TaskType.class);
+        this.configurations = in.readMap(SettingsConfiguration::new);
     }
 
-    static final ParseField PROVIDER_FIELD = new ParseField("provider");
+    static final ParseField SERVICE_FIELD = new ParseField("service");
+    static final ParseField NAME_FIELD = new ParseField("name");
     static final ParseField TASK_TYPES_FIELD = new ParseField("task_types");
-    static final ParseField CONFIGURATION_FIELD = new ParseField("configuration");
+    static final ParseField CONFIGURATIONS_FIELD = new ParseField("configurations");
 
     @SuppressWarnings("unchecked")
     private static final ConstructingObjectParser<InferenceServiceConfiguration, Void> PARSER = new ConstructingObjectParser<>(
         "inference_service_configuration",
         true,
         args -> {
-            List<String> taskTypes = (ArrayList<String>) args[1];
-            return new InferenceServiceConfiguration.Builder().setProvider((String) args[0])
-                .setTaskTypes((List<TaskSettingsConfiguration>) args[1])
-                .setConfiguration((Map<String, SettingsConfiguration>) args[2])
+            List<String> taskTypes = (ArrayList<String>) args[2];
+            return new InferenceServiceConfiguration.Builder().setService((String) args[0])
+                .setName((String) args[1])
+                .setTaskTypes(EnumSet.copyOf(taskTypes.stream().map(TaskType::fromString).collect(Collectors.toList())))
+                .setConfigurations((Map<String, SettingsConfiguration>) args[3])
                 .build();
         }
     );
 
     static {
-        PARSER.declareString(constructorArg(), PROVIDER_FIELD);
-        PARSER.declareObjectArray(constructorArg(), (p, c) -> TaskSettingsConfiguration.fromXContent(p), TASK_TYPES_FIELD);
-        PARSER.declareObject(constructorArg(), (p, c) -> p.map(), CONFIGURATION_FIELD);
+        PARSER.declareString(constructorArg(), SERVICE_FIELD);
+        PARSER.declareString(constructorArg(), NAME_FIELD);
+        PARSER.declareStringArray(constructorArg(), TASK_TYPES_FIELD);
+        PARSER.declareObject(constructorArg(), (p, c) -> p.map(), CONFIGURATIONS_FIELD);
     }
 
-    public String getProvider() {
-        return provider;
+    public String getService() {
+        return service;
     }
 
-    public List<TaskSettingsConfiguration> getTaskTypes() {
+    public String getName() {
+        return name;
+    }
+
+    public EnumSet<TaskType> getTaskTypes() {
         return taskTypes;
     }
 
-    public Map<String, SettingsConfiguration> getConfiguration() {
-        return configuration;
+    public Map<String, SettingsConfiguration> getConfigurations() {
+        return configurations;
     }
 
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();
         {
-            builder.field(PROVIDER_FIELD.getPreferredName(), provider);
+            builder.field(SERVICE_FIELD.getPreferredName(), service);
+            builder.field(NAME_FIELD.getPreferredName(), name);
             builder.field(TASK_TYPES_FIELD.getPreferredName(), taskTypes);
-            builder.field(CONFIGURATION_FIELD.getPreferredName(), configuration);
+            builder.field(CONFIGURATIONS_FIELD.getPreferredName(), configurations);
         }
         builder.endObject();
         return builder;
@@ -125,17 +140,19 @@ public static InferenceServiceConfiguration fromXContentBytes(BytesReference sou
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        out.writeString(provider);
+        out.writeString(service);
+        out.writeString(name);
         out.writeCollection(taskTypes);
-        out.writeMapValues(configuration);
+        out.writeMapValues(configurations);
     }
 
     public Map<String, Object> toMap() {
         Map<String, Object> map = new HashMap<>();
 
-        map.put(PROVIDER_FIELD.getPreferredName(), provider);
+        map.put(SERVICE_FIELD.getPreferredName(), service);
+        map.put(NAME_FIELD.getPreferredName(), name);
         map.put(TASK_TYPES_FIELD.getPreferredName(), taskTypes);
-        map.put(CONFIGURATION_FIELD.getPreferredName(), configuration);
+        map.put(CONFIGURATIONS_FIELD.getPreferredName(), configurations);
 
         return map;
     }
@@ -145,39 +162,46 @@ public boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
         InferenceServiceConfiguration that = (InferenceServiceConfiguration) o;
-        return provider.equals(that.provider)
+        return service.equals(that.service)
+            && name.equals(that.name)
             && Objects.equals(taskTypes, that.taskTypes)
-            && Objects.equals(configuration, that.configuration);
+            && Objects.equals(configurations, that.configurations);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(provider, taskTypes, configuration);
+        return Objects.hash(service, name, taskTypes, configurations);
     }
 
     public static class Builder {
 
-        private String provider;
-        private List<TaskSettingsConfiguration> taskTypes;
-        private Map<String, SettingsConfiguration> configuration;
+        private String service;
+        private String name;
+        private EnumSet<TaskType> taskTypes;
+        private Map<String, SettingsConfiguration> configurations;
+
+        public Builder setService(String service) {
+            this.service = service;
+            return this;
+        }
 
-        public Builder setProvider(String provider) {
-            this.provider = provider;
+        public Builder setName(String name) {
+            this.name = name;
             return this;
         }
 
-        public Builder setTaskTypes(List<TaskSettingsConfiguration> taskTypes) {
+        public Builder setTaskTypes(EnumSet<TaskType> taskTypes) {
             this.taskTypes = taskTypes;
             return this;
         }
 
-        public Builder setConfiguration(Map<String, SettingsConfiguration> configuration) {
-            this.configuration = configuration;
+        public Builder setConfigurations(Map<String, SettingsConfiguration> configurations) {
+            this.configurations = configurations;
             return this;
         }
 
         public InferenceServiceConfiguration build() {
-            return new InferenceServiceConfiguration(provider, taskTypes, configuration);
+            return new InferenceServiceConfiguration(service, name, taskTypes, configurations);
         }
     }
 }