2525
2626import java .io .IOException ;
2727import java .util .ArrayList ;
28+ import java .util .EnumSet ;
2829import java .util .HashMap ;
2930import java .util .List ;
3031import java .util .Map ;
3132import java .util .Objects ;
33+ import java .util .stream .Collectors ;
3234
3335import static org .elasticsearch .xcontent .ConstructingObjectParser .constructorArg ;
3436
3739 */
3840public class InferenceServiceConfiguration implements Writeable , ToXContentObject {
3941
40- private final String provider ;
41- private final List <TaskSettingsConfiguration > taskTypes ;
42- private final Map <String , SettingsConfiguration > configuration ;
42+ private final String service ;
43+ private final String name ;
44+ private final EnumSet <TaskType > taskTypes ;
45+ private final Map <String , SettingsConfiguration > configurations ;
4346
4447 /**
4548 * Constructs a new {@link InferenceServiceConfiguration} instance with specified properties.
4649 *
47- * @param provider The name of the service provider.
48- * @param taskTypes A list of {@link TaskSettingsConfiguration} supported by the service provider.
49- * @param configuration The configuration of the service provider, defined by {@link SettingsConfiguration}.
50+ * @param service The name of the service provider.
51+ * @param name The user-friendly name of the service provider.
52+ * @param taskTypes A list of {@link TaskType} supported by the service provider.
53+ * @param configurations The configuration of the service provider, defined by {@link SettingsConfiguration}.
5054 */
5155 private InferenceServiceConfiguration (
52- String provider ,
53- List <TaskSettingsConfiguration > taskTypes ,
54- Map <String , SettingsConfiguration > configuration
56+ String service ,
57+ String name ,
58+ EnumSet <TaskType > taskTypes ,
59+ Map <String , SettingsConfiguration > configurations
5560 ) {
56- this .provider = provider ;
61+ this .service = service ;
62+ this .name = name ;
5763 this .taskTypes = taskTypes ;
58- this .configuration = configuration ;
64+ this .configurations = configurations ;
5965 }
6066
6167 public InferenceServiceConfiguration (StreamInput in ) throws IOException {
62- this .provider = in .readString ();
63- this .taskTypes = in .readCollectionAsList (TaskSettingsConfiguration ::new );
64- this .configuration = in .readMap (SettingsConfiguration ::new );
68+ this .service = in .readString ();
69+ this .name = in .readString ();
70+ this .taskTypes = in .readEnumSet (TaskType .class );
71+ this .configurations = in .readMap (SettingsConfiguration ::new );
6572 }
6673
67- static final ParseField PROVIDER_FIELD = new ParseField ("provider" );
74+ static final ParseField SERVICE_FIELD = new ParseField ("service" );
75+ static final ParseField NAME_FIELD = new ParseField ("name" );
6876 static final ParseField TASK_TYPES_FIELD = new ParseField ("task_types" );
69- static final ParseField CONFIGURATION_FIELD = new ParseField ("configuration " );
77+ static final ParseField CONFIGURATIONS_FIELD = new ParseField ("configurations " );
7078
7179 @ SuppressWarnings ("unchecked" )
7280 private static final ConstructingObjectParser <InferenceServiceConfiguration , Void > PARSER = new ConstructingObjectParser <>(
7381 "inference_service_configuration" ,
7482 true ,
7583 args -> {
76- List <String > taskTypes = (ArrayList <String >) args [1 ];
77- return new InferenceServiceConfiguration .Builder ().setProvider ((String ) args [0 ])
78- .setTaskTypes ((List <TaskSettingsConfiguration >) args [1 ])
79- .setConfiguration ((Map <String , SettingsConfiguration >) args [2 ])
84+ List <String > taskTypes = (ArrayList <String >) args [2 ];
85+ return new InferenceServiceConfiguration .Builder ().setService ((String ) args [0 ])
86+ .setName ((String ) args [1 ])
87+ .setTaskTypes (EnumSet .copyOf (taskTypes .stream ().map (TaskType ::fromString ).collect (Collectors .toList ())))
88+ .setConfigurations ((Map <String , SettingsConfiguration >) args [3 ])
8089 .build ();
8190 }
8291 );
8392
8493 static {
85- PARSER .declareString (constructorArg (), PROVIDER_FIELD );
86- PARSER .declareObjectArray (constructorArg (), (p , c ) -> TaskSettingsConfiguration .fromXContent (p ), TASK_TYPES_FIELD );
87- PARSER .declareObject (constructorArg (), (p , c ) -> p .map (), CONFIGURATION_FIELD );
94+ PARSER .declareString (constructorArg (), SERVICE_FIELD );
95+ PARSER .declareString (constructorArg (), NAME_FIELD );
96+ PARSER .declareStringArray (constructorArg (), TASK_TYPES_FIELD );
97+ PARSER .declareObject (constructorArg (), (p , c ) -> p .map (), CONFIGURATIONS_FIELD );
8898 }
8999
90- public String getProvider () {
91- return provider ;
100+ public String getService () {
101+ return service ;
92102 }
93103
94- public List <TaskSettingsConfiguration > getTaskTypes () {
104+ public String getName () {
105+ return name ;
106+ }
107+
108+ public EnumSet <TaskType > getTaskTypes () {
95109 return taskTypes ;
96110 }
97111
98- public Map <String , SettingsConfiguration > getConfiguration () {
99- return configuration ;
112+ public Map <String , SettingsConfiguration > getConfigurations () {
113+ return configurations ;
100114 }
101115
102116 @ Override
103117 public XContentBuilder toXContent (XContentBuilder builder , Params params ) throws IOException {
104118 builder .startObject ();
105119 {
106- builder .field (PROVIDER_FIELD .getPreferredName (), provider );
120+ builder .field (SERVICE_FIELD .getPreferredName (), service );
121+ builder .field (NAME_FIELD .getPreferredName (), name );
107122 builder .field (TASK_TYPES_FIELD .getPreferredName (), taskTypes );
108- builder .field (CONFIGURATION_FIELD .getPreferredName (), configuration );
123+ builder .field (CONFIGURATIONS_FIELD .getPreferredName (), configurations );
109124 }
110125 builder .endObject ();
111126 return builder ;
@@ -125,17 +140,19 @@ public static InferenceServiceConfiguration fromXContentBytes(BytesReference sou
125140
126141 @ Override
127142 public void writeTo (StreamOutput out ) throws IOException {
128- out .writeString (provider );
143+ out .writeString (service );
144+ out .writeString (name );
129145 out .writeCollection (taskTypes );
130- out .writeMapValues (configuration );
146+ out .writeMapValues (configurations );
131147 }
132148
133149 public Map <String , Object > toMap () {
134150 Map <String , Object > map = new HashMap <>();
135151
136- map .put (PROVIDER_FIELD .getPreferredName (), provider );
152+ map .put (SERVICE_FIELD .getPreferredName (), service );
153+ map .put (NAME_FIELD .getPreferredName (), name );
137154 map .put (TASK_TYPES_FIELD .getPreferredName (), taskTypes );
138- map .put (CONFIGURATION_FIELD .getPreferredName (), configuration );
155+ map .put (CONFIGURATIONS_FIELD .getPreferredName (), configurations );
139156
140157 return map ;
141158 }
@@ -145,39 +162,46 @@ public boolean equals(Object o) {
145162 if (this == o ) return true ;
146163 if (o == null || getClass () != o .getClass ()) return false ;
147164 InferenceServiceConfiguration that = (InferenceServiceConfiguration ) o ;
148- return provider .equals (that .provider )
165+ return service .equals (that .service )
166+ && name .equals (that .name )
149167 && Objects .equals (taskTypes , that .taskTypes )
150- && Objects .equals (configuration , that .configuration );
168+ && Objects .equals (configurations , that .configurations );
151169 }
152170
153171 @ Override
154172 public int hashCode () {
155- return Objects .hash (provider , taskTypes , configuration );
173+ return Objects .hash (service , name , taskTypes , configurations );
156174 }
157175
158176 public static class Builder {
159177
160- private String provider ;
161- private List <TaskSettingsConfiguration > taskTypes ;
162- private Map <String , SettingsConfiguration > configuration ;
178+ private String service ;
179+ private String name ;
180+ private EnumSet <TaskType > taskTypes ;
181+ private Map <String , SettingsConfiguration > configurations ;
182+
183+ public Builder setService (String service ) {
184+ this .service = service ;
185+ return this ;
186+ }
163187
164- public Builder setProvider (String provider ) {
165- this .provider = provider ;
188+ public Builder setName (String name ) {
189+ this .name = name ;
166190 return this ;
167191 }
168192
169- public Builder setTaskTypes (List < TaskSettingsConfiguration > taskTypes ) {
193+ public Builder setTaskTypes (EnumSet < TaskType > taskTypes ) {
170194 this .taskTypes = taskTypes ;
171195 return this ;
172196 }
173197
174- public Builder setConfiguration (Map <String , SettingsConfiguration > configuration ) {
175- this .configuration = configuration ;
198+ public Builder setConfigurations (Map <String , SettingsConfiguration > configurations ) {
199+ this .configurations = configurations ;
176200 return this ;
177201 }
178202
179203 public InferenceServiceConfiguration build () {
180- return new InferenceServiceConfiguration (provider , taskTypes , configuration );
204+ return new InferenceServiceConfiguration (service , name , taskTypes , configurations );
181205 }
182206 }
183207}
0 commit comments