Skip to content

Commit 2a1b433

Browse files
jonathan-buttnerjonathan-buttner
authored
[ML] Adding fields for Inference service configuration API (#121103)
* Adding fields and making some optional * Fixing tests
1 parent b1dcfb5 commit 2a1b433

File tree

15 files changed

+308
-161
lines changed

15 files changed

+308
-161
lines changed

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/alibabacloudsearch/AlibabaCloudSearchService.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -409,7 +409,7 @@ public static InferenceServiceConfiguration get() {
409409
HTTP_SCHEMA_NAME,
410410
new SettingsConfiguration.Builder(supportedTaskTypes).setDescription("")
411411
.setLabel("HTTP Schema")
412-
.setRequired(true)
412+
.setRequired(false)
413413
.setSensitive(false)
414414
.setUpdatable(false)
415415
.setType(SettingsConfigurationFieldType.STRING)

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/amazonbedrock/AmazonBedrockService.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@
5454
import java.util.Map;
5555
import java.util.Set;
5656

57+
import static org.elasticsearch.xpack.inference.services.ServiceFields.DIMENSIONS;
5758
import static org.elasticsearch.xpack.inference.services.ServiceUtils.createInvalidModelException;
5859
import static org.elasticsearch.xpack.inference.services.ServiceUtils.parsePersistedConfigErrorMsg;
5960
import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeFromMap;
@@ -413,6 +414,20 @@ public static InferenceServiceConfiguration get() {
413414
.build()
414415
);
415416

417+
configurationMap.put(
418+
DIMENSIONS,
419+
new SettingsConfiguration.Builder(EnumSet.of(TaskType.TEXT_EMBEDDING)).setDescription(
420+
"The number of dimensions the resulting embeddings should have. For more information refer to "
421+
+ "https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-embed-text.html."
422+
)
423+
.setLabel("Dimensions")
424+
.setRequired(false)
425+
.setSensitive(false)
426+
.setUpdatable(false)
427+
.setType(SettingsConfigurationFieldType.INTEGER)
428+
.build()
429+
);
430+
416431
configurationMap.putAll(AmazonBedrockSecretSettings.Configuration.get());
417432
configurationMap.putAll(
418433
RateLimitSettings.toSettingsConfigurationWithDescription(

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/azureaistudio/AzureAiStudioService.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@
5353
import java.util.Map;
5454
import java.util.Set;
5555

56+
import static org.elasticsearch.xpack.inference.services.ServiceFields.DIMENSIONS;
5657
import static org.elasticsearch.xpack.inference.services.ServiceUtils.createInvalidModelException;
5758
import static org.elasticsearch.xpack.inference.services.ServiceUtils.parsePersistedConfigErrorMsg;
5859
import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeFromMap;
@@ -441,6 +442,20 @@ public static InferenceServiceConfiguration get() {
441442
.build()
442443
);
443444

445+
configurationMap.put(
446+
DIMENSIONS,
447+
new SettingsConfiguration.Builder(EnumSet.of(TaskType.TEXT_EMBEDDING)).setDescription(
448+
"The number of dimensions the resulting embeddings should have. For more information refer to "
449+
+ "https://learn.microsoft.com/en-us/azure/ai-studio/reference/reference-model-inference-embeddings."
450+
)
451+
.setLabel("Dimensions")
452+
.setRequired(false)
453+
.setSensitive(false)
454+
.setUpdatable(false)
455+
.setType(SettingsConfigurationFieldType.INTEGER)
456+
.build()
457+
);
458+
444459
configurationMap.putAll(DefaultSecretSettings.toSettingsConfiguration(supportedTaskTypes));
445460
configurationMap.putAll(RateLimitSettings.toSettingsConfiguration(supportedTaskTypes));
446461

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/azureopenai/AzureOpenAiService.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@
5050
import java.util.Map;
5151
import java.util.Set;
5252

53+
import static org.elasticsearch.xpack.inference.services.ServiceFields.DIMENSIONS;
5354
import static org.elasticsearch.xpack.inference.services.ServiceUtils.createInvalidModelException;
5455
import static org.elasticsearch.xpack.inference.services.ServiceUtils.parsePersistedConfigErrorMsg;
5556
import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeFromMap;
@@ -382,6 +383,20 @@ public static InferenceServiceConfiguration get() {
382383
.build()
383384
);
384385

386+
configurationMap.put(
387+
DIMENSIONS,
388+
new SettingsConfiguration.Builder(EnumSet.of(TaskType.TEXT_EMBEDDING)).setDescription(
389+
"The number of dimensions the resulting embeddings should have. For more information refer to "
390+
+ "https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#request-body-1."
391+
)
392+
.setLabel("Dimensions")
393+
.setRequired(false)
394+
.setSensitive(false)
395+
.setUpdatable(false)
396+
.setType(SettingsConfigurationFieldType.INTEGER)
397+
.build()
398+
);
399+
385400
configurationMap.putAll(AzureOpenAiSecretSettings.Configuration.get());
386401
configurationMap.putAll(
387402
RateLimitSettings.toSettingsConfigurationWithDescription(

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereService.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import org.elasticsearch.inference.SettingsConfiguration;
2626
import org.elasticsearch.inference.SimilarityMeasure;
2727
import org.elasticsearch.inference.TaskType;
28+
import org.elasticsearch.inference.configuration.SettingsConfigurationFieldType;
2829
import org.elasticsearch.rest.RestStatus;
2930
import org.elasticsearch.xpack.inference.chunking.ChunkingSettingsBuilder;
3031
import org.elasticsearch.xpack.inference.chunking.EmbeddingRequestChunker;
@@ -51,6 +52,7 @@
5152
import java.util.Map;
5253
import java.util.Set;
5354

55+
import static org.elasticsearch.xpack.inference.services.ServiceFields.MODEL_ID;
5456
import static org.elasticsearch.xpack.inference.services.ServiceUtils.createInvalidModelException;
5557
import static org.elasticsearch.xpack.inference.services.ServiceUtils.parsePersistedConfigErrorMsg;
5658
import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeFromMap;
@@ -363,6 +365,19 @@ public static InferenceServiceConfiguration get() {
363365
() -> {
364366
var configurationMap = new HashMap<String, SettingsConfiguration>();
365367

368+
configurationMap.put(
369+
MODEL_ID,
370+
new SettingsConfiguration.Builder(supportedTaskTypes).setDescription(
371+
"The name of the model to use for the inference task."
372+
)
373+
.setLabel("Model ID")
374+
.setRequired(false)
375+
.setSensitive(false)
376+
.setUpdatable(false)
377+
.setType(SettingsConfigurationFieldType.STRING)
378+
.build()
379+
);
380+
366381
configurationMap.putAll(DefaultSecretSettings.toSettingsConfiguration(supportedTaskTypes));
367382
configurationMap.putAll(RateLimitSettings.toSettingsConfiguration(supportedTaskTypes));
368383

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiService.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -327,6 +327,8 @@ public static InferenceServiceConfiguration get() {
327327
() -> {
328328
var configurationMap = new HashMap<String, SettingsConfiguration>();
329329

330+
// TODO whether the model ID is required or not depends on the task type
331+
// For rerank it is optional, for text_embedding it is required
330332
configurationMap.put(
331333
MODEL_ID,
332334
new SettingsConfiguration.Builder(supportedTaskTypes).setDescription("ID of the LLM you're using.")

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/jinaai/JinaAIService.java‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import org.elasticsearch.inference.SettingsConfiguration;
2626
import org.elasticsearch.inference.SimilarityMeasure;
2727
import org.elasticsearch.inference.TaskType;
28+
import org.elasticsearch.inference.configuration.SettingsConfigurationFieldType;
2829
import org.elasticsearch.rest.RestStatus;
2930
import org.elasticsearch.xpack.inference.chunking.ChunkingSettingsBuilder;
3031
import org.elasticsearch.xpack.inference.chunking.EmbeddingRequestChunker;
@@ -49,6 +50,7 @@
4950
import java.util.List;
5051
import java.util.Map;
5152

53+
import static org.elasticsearch.xpack.inference.services.ServiceFields.DIMENSIONS;
5254
import static org.elasticsearch.xpack.inference.services.ServiceUtils.createInvalidModelException;
5355
import static org.elasticsearch.xpack.inference.services.ServiceUtils.parsePersistedConfigErrorMsg;
5456
import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeFromMap;
@@ -339,6 +341,33 @@ public static InferenceServiceConfiguration get() {
339341
() -> {
340342
var configurationMap = new HashMap<String, SettingsConfiguration>();
341343

344+
configurationMap.put(
345+
JinaAIServiceSettings.MODEL_ID,
346+
new SettingsConfiguration.Builder(supportedTaskTypes).setDescription(
347+
"The name of the model to use for the inference task."
348+
)
349+
.setLabel("Model ID")
350+
.setRequired(true)
351+
.setSensitive(false)
352+
.setUpdatable(false)
353+
.setType(SettingsConfigurationFieldType.STRING)
354+
.build()
355+
);
356+
357+
configurationMap.put(
358+
DIMENSIONS,
359+
new SettingsConfiguration.Builder(EnumSet.of(TaskType.TEXT_EMBEDDING)).setDescription(
360+
"The number of dimensions the resulting embeddings should have. For more information refer to "
361+
+ "https://api.jina.ai/redoc#tag/embeddings/operation/create_embedding_v1_embeddings_post."
362+
)
363+
.setLabel("Dimensions")
364+
.setRequired(false)
365+
.setSensitive(false)
366+
.setUpdatable(false)
367+
.setType(SettingsConfigurationFieldType.INTEGER)
368+
.build()
369+
);
370+
342371
configurationMap.putAll(DefaultSecretSettings.toSettingsConfiguration(supportedTaskTypes));
343372
configurationMap.putAll(RateLimitSettings.toSettingsConfiguration(supportedTaskTypes));
344373

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/OpenAiService.java‎

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -56,8 +56,8 @@
5656

5757
import static org.elasticsearch.xpack.inference.external.action.ActionUtils.constructFailedToSendRequestMessage;
5858
import static org.elasticsearch.xpack.inference.external.action.openai.OpenAiActionCreator.COMPLETION_ERROR_PREFIX;
59+
import static org.elasticsearch.xpack.inference.services.ServiceFields.DIMENSIONS;
5960
import static org.elasticsearch.xpack.inference.services.ServiceFields.MODEL_ID;
60-
import static org.elasticsearch.xpack.inference.services.ServiceFields.URL;
6161
import static org.elasticsearch.xpack.inference.services.ServiceUtils.createInvalidModelException;
6262
import static org.elasticsearch.xpack.inference.services.ServiceUtils.parsePersistedConfigErrorMsg;
6363
import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeFromMap;
@@ -440,19 +440,16 @@ public static InferenceServiceConfiguration get() {
440440
);
441441

442442
configurationMap.put(
443-
URL,
444-
new SettingsConfiguration.Builder(SUPPORTED_TASK_TYPES_FOR_SERVICES_API).setDefaultValue(
445-
"https://api.openai.com/v1/chat/completions"
443+
DIMENSIONS,
444+
new SettingsConfiguration.Builder(EnumSet.of(TaskType.TEXT_EMBEDDING)).setDescription(
445+
"The number of dimensions the resulting embeddings should have. For more information refer to "
446+
+ "https://platform.openai.com/docs/api-reference/embeddings/create#embeddings-create-dimensions."
446447
)
447-
.setDescription(
448-
"The OpenAI API endpoint URL. For more information on the URL, refer to the "
449-
+ "https://platform.openai.com/docs/api-reference."
450-
)
451-
.setLabel("URL")
452-
.setRequired(true)
448+
.setLabel("Dimensions")
449+
.setRequired(false)
453450
.setSensitive(false)
454451
.setUpdatable(false)
455-
.setType(SettingsConfigurationFieldType.STRING)
452+
.setType(SettingsConfigurationFieldType.INTEGER)
456453
.build()
457454
);
458455

‎x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/alibabacloudsearch/AlibabaCloudSearchServiceTests.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -490,7 +490,7 @@ public void testGetConfiguration() throws Exception {
490490
"http_schema": {
491491
"description": "",
492492
"label": "HTTP Schema",
493-
"required": true,
493+
"required": false,
494494
"sensitive": false,
495495
"updatable": false,
496496
"type": "str",

0 commit comments

Comments
 (0)