Skip to content

Commit cffbfd7

Browse files
weizijuncbuescher
authored andcommitted
[Inference API] Add Alibaba Cloud AI Search Model support to Inference API (elastic#111181)
Add Alibaba Cloud AI Search Model support to Inference API. Supports the text_embedding, sparse_embedding and rerank tasks. Requires an Alibaba Cloud Account with Alibaba Cloud Opensearch access and an api key used to access Alibaba Cloud AI Search Model.
1 parent 053a886 commit cffbfd7

File tree

59 files changed

+4756
-0
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

59 files changed

+4756
-0
lines changed

docs/changelog/111181.yaml

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
pr: 111181
2+
summary: "[Inference API] Add Alibaba Cloud AI Search Model support to Inference API"
3+
area: Machine Learning
4+
type: enhancement
5+
issues: [ ]

server/src/main/java/org/elasticsearch/TransportVersions.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,7 @@ static TransportVersion def(int id) {
200200
public static final TransportVersion ESQL_ES_FIELD_CACHED_SERIALIZATION = def(8_730_00_0);
201201
public static final TransportVersion ADD_MANAGE_ROLES_PRIVILEGE = def(8_731_00_0);
202202
public static final TransportVersion REPOSITORIES_TELEMETRY = def(8_732_00_0);
203+
public static final TransportVersion ML_INFERENCE_ALIBABACLOUD_SEARCH_ADDED = def(8_733_00_0);
203204

204205
/*
205206
* STOP! READ THIS FIRST! No, really,

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferenceNamedWriteablesProvider.java

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,13 @@
2525
import org.elasticsearch.xpack.core.inference.results.LegacyTextEmbeddingResults;
2626
import org.elasticsearch.xpack.core.inference.results.RankedDocsResults;
2727
import org.elasticsearch.xpack.core.inference.results.SparseEmbeddingResults;
28+
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.AlibabaCloudSearchServiceSettings;
29+
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.embeddings.AlibabaCloudSearchEmbeddingsServiceSettings;
30+
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.embeddings.AlibabaCloudSearchEmbeddingsTaskSettings;
31+
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.rerank.AlibabaCloudSearchRerankServiceSettings;
32+
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.rerank.AlibabaCloudSearchRerankTaskSettings;
33+
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.sparse.AlibabaCloudSearchSparseServiceSettings;
34+
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.sparse.AlibabaCloudSearchSparseTaskSettings;
2835
import org.elasticsearch.xpack.inference.services.amazonbedrock.AmazonBedrockSecretSettings;
2936
import org.elasticsearch.xpack.inference.services.amazonbedrock.completion.AmazonBedrockChatCompletionServiceSettings;
3037
import org.elasticsearch.xpack.inference.services.amazonbedrock.completion.AmazonBedrockChatCompletionTaskSettings;
@@ -117,6 +124,7 @@ public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
117124
addAnthropicNamedWritables(namedWriteables);
118125
addAmazonBedrockNamedWriteables(namedWriteables);
119126
addEisNamedWriteables(namedWriteables);
127+
addAlibabaCloudSearchNamedWriteables(namedWriteables);
120128

121129
return namedWriteables;
122130
}
@@ -482,6 +490,59 @@ private static void addAnthropicNamedWritables(List<NamedWriteableRegistry.Entry
482490
);
483491
}
484492

493+
private static void addAlibabaCloudSearchNamedWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {
494+
namedWriteables.add(
495+
new NamedWriteableRegistry.Entry(
496+
ServiceSettings.class,
497+
AlibabaCloudSearchServiceSettings.NAME,
498+
AlibabaCloudSearchServiceSettings::new
499+
)
500+
);
501+
namedWriteables.add(
502+
new NamedWriteableRegistry.Entry(
503+
ServiceSettings.class,
504+
AlibabaCloudSearchEmbeddingsServiceSettings.NAME,
505+
AlibabaCloudSearchEmbeddingsServiceSettings::new
506+
)
507+
);
508+
namedWriteables.add(
509+
new NamedWriteableRegistry.Entry(
510+
TaskSettings.class,
511+
AlibabaCloudSearchEmbeddingsTaskSettings.NAME,
512+
AlibabaCloudSearchEmbeddingsTaskSettings::new
513+
)
514+
);
515+
namedWriteables.add(
516+
new NamedWriteableRegistry.Entry(
517+
ServiceSettings.class,
518+
AlibabaCloudSearchSparseServiceSettings.NAME,
519+
AlibabaCloudSearchSparseServiceSettings::new
520+
)
521+
);
522+
namedWriteables.add(
523+
new NamedWriteableRegistry.Entry(
524+
TaskSettings.class,
525+
AlibabaCloudSearchSparseTaskSettings.NAME,
526+
AlibabaCloudSearchSparseTaskSettings::new
527+
)
528+
);
529+
namedWriteables.add(
530+
new NamedWriteableRegistry.Entry(
531+
ServiceSettings.class,
532+
AlibabaCloudSearchRerankServiceSettings.NAME,
533+
AlibabaCloudSearchRerankServiceSettings::new
534+
)
535+
);
536+
namedWriteables.add(
537+
new NamedWriteableRegistry.Entry(
538+
TaskSettings.class,
539+
AlibabaCloudSearchRerankTaskSettings.NAME,
540+
AlibabaCloudSearchRerankTaskSettings::new
541+
)
542+
);
543+
544+
}
545+
485546
private static void addEisNamedWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {
486547
namedWriteables.add(
487548
new NamedWriteableRegistry.Entry(

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,7 @@
7474
import org.elasticsearch.xpack.inference.rest.RestInferenceAction;
7575
import org.elasticsearch.xpack.inference.rest.RestPutInferenceModelAction;
7676
import org.elasticsearch.xpack.inference.services.ServiceComponents;
77+
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.AlibabaCloudSearchService;
7778
import org.elasticsearch.xpack.inference.services.amazonbedrock.AmazonBedrockService;
7879
import org.elasticsearch.xpack.inference.services.anthropic.AnthropicService;
7980
import org.elasticsearch.xpack.inference.services.azureaistudio.AzureAiStudioService;
@@ -237,6 +238,7 @@ public List<InferenceServiceExtension.Factory> getInferenceServiceFactories() {
237238
context -> new MistralService(httpFactory.get(), serviceComponents.get()),
238239
context -> new AnthropicService(httpFactory.get(), serviceComponents.get()),
239240
context -> new AmazonBedrockService(httpFactory.get(), amazonBedrockFactory.get(), serviceComponents.get()),
241+
context -> new AlibabaCloudSearchService(httpFactory.get(), serviceComponents.get()),
240242
ElasticsearchInternalService::new
241243
);
242244
}
Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.inference.external.action.alibabacloudsearch;
9+
10+
import org.elasticsearch.inference.InputType;
11+
import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
12+
import org.elasticsearch.xpack.inference.external.http.sender.Sender;
13+
import org.elasticsearch.xpack.inference.services.ServiceComponents;
14+
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.embeddings.AlibabaCloudSearchEmbeddingsModel;
15+
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.rerank.AlibabaCloudSearchRerankModel;
16+
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.sparse.AlibabaCloudSearchSparseModel;
17+
18+
import java.util.Map;
19+
import java.util.Objects;
20+
21+
/**
22+
* Provides a way to construct an {@link ExecutableAction} using the visitor pattern based on the alibaba cloud search model type.
23+
*/
24+
public class AlibabaCloudSearchActionCreator implements AlibabaCloudSearchActionVisitor {
25+
private final Sender sender;
26+
private final ServiceComponents serviceComponents;
27+
28+
public AlibabaCloudSearchActionCreator(Sender sender, ServiceComponents serviceComponents) {
29+
this.sender = Objects.requireNonNull(sender);
30+
this.serviceComponents = Objects.requireNonNull(serviceComponents);
31+
}
32+
33+
@Override
34+
public ExecutableAction create(AlibabaCloudSearchEmbeddingsModel model, Map<String, Object> taskSettings, InputType inputType) {
35+
var overriddenModel = AlibabaCloudSearchEmbeddingsModel.of(model, taskSettings, inputType);
36+
37+
return new AlibabaCloudSearchEmbeddingsAction(sender, overriddenModel, serviceComponents);
38+
}
39+
40+
@Override
41+
public ExecutableAction create(AlibabaCloudSearchSparseModel model, Map<String, Object> taskSettings, InputType inputType) {
42+
var overriddenModel = AlibabaCloudSearchSparseModel.of(model, taskSettings, inputType);
43+
44+
return new AlibabaCloudSearchSparseAction(sender, overriddenModel, serviceComponents);
45+
}
46+
47+
@Override
48+
public ExecutableAction create(AlibabaCloudSearchRerankModel model, Map<String, Object> taskSettings) {
49+
var overriddenModel = AlibabaCloudSearchRerankModel.of(model, taskSettings);
50+
51+
return new AlibabaCloudSearchRerankAction(sender, overriddenModel, serviceComponents);
52+
}
53+
}
Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.inference.external.action.alibabacloudsearch;
9+
10+
import org.elasticsearch.inference.InputType;
11+
import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
12+
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.embeddings.AlibabaCloudSearchEmbeddingsModel;
13+
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.rerank.AlibabaCloudSearchRerankModel;
14+
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.sparse.AlibabaCloudSearchSparseModel;
15+
16+
import java.util.Map;
17+
18+
public interface AlibabaCloudSearchActionVisitor {
19+
ExecutableAction create(AlibabaCloudSearchEmbeddingsModel model, Map<String, Object> taskSettings, InputType inputType);
20+
21+
ExecutableAction create(AlibabaCloudSearchSparseModel model, Map<String, Object> taskSettings, InputType inputType);
22+
23+
ExecutableAction create(AlibabaCloudSearchRerankModel model, Map<String, Object> taskSettings);
24+
}
Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.inference.external.action.alibabacloudsearch;
9+
10+
import org.elasticsearch.ElasticsearchException;
11+
import org.elasticsearch.action.ActionListener;
12+
import org.elasticsearch.core.TimeValue;
13+
import org.elasticsearch.inference.InferenceServiceResults;
14+
import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
15+
import org.elasticsearch.xpack.inference.external.alibabacloudsearch.AlibabaCloudSearchAccount;
16+
import org.elasticsearch.xpack.inference.external.http.sender.AlibabaCloudSearchEmbeddingsRequestManager;
17+
import org.elasticsearch.xpack.inference.external.http.sender.InferenceInputs;
18+
import org.elasticsearch.xpack.inference.external.http.sender.Sender;
19+
import org.elasticsearch.xpack.inference.services.ServiceComponents;
20+
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.embeddings.AlibabaCloudSearchEmbeddingsModel;
21+
22+
import java.util.Objects;
23+
24+
import static org.elasticsearch.xpack.inference.external.action.ActionUtils.constructFailedToSendRequestMessage;
25+
import static org.elasticsearch.xpack.inference.external.action.ActionUtils.createInternalServerError;
26+
import static org.elasticsearch.xpack.inference.external.action.ActionUtils.wrapFailuresInElasticsearchException;
27+
28+
public class AlibabaCloudSearchEmbeddingsAction implements ExecutableAction {
29+
private final AlibabaCloudSearchAccount account;
30+
private final AlibabaCloudSearchEmbeddingsModel model;
31+
private final String failedToSendRequestErrorMessage;
32+
private final Sender sender;
33+
private final AlibabaCloudSearchEmbeddingsRequestManager requestCreator;
34+
35+
public AlibabaCloudSearchEmbeddingsAction(Sender sender, AlibabaCloudSearchEmbeddingsModel model, ServiceComponents serviceComponents) {
36+
this.model = Objects.requireNonNull(model);
37+
this.sender = Objects.requireNonNull(sender);
38+
this.account = new AlibabaCloudSearchAccount(this.model.getSecretSettings().apiKey());
39+
this.failedToSendRequestErrorMessage = constructFailedToSendRequestMessage(null, "AlibabaCloud Search text embeddings");
40+
this.requestCreator = AlibabaCloudSearchEmbeddingsRequestManager.of(account, model, serviceComponents.threadPool());
41+
}
42+
43+
@Override
44+
public void execute(InferenceInputs inferenceInputs, TimeValue timeout, ActionListener<InferenceServiceResults> listener) {
45+
try {
46+
ActionListener<InferenceServiceResults> wrappedListener = wrapFailuresInElasticsearchException(
47+
failedToSendRequestErrorMessage,
48+
listener
49+
);
50+
sender.send(requestCreator, inferenceInputs, timeout, wrappedListener);
51+
} catch (ElasticsearchException e) {
52+
listener.onFailure(e);
53+
} catch (Exception e) {
54+
listener.onFailure(createInternalServerError(e, failedToSendRequestErrorMessage));
55+
}
56+
}
57+
}
Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.inference.external.action.alibabacloudsearch;
9+
10+
import org.apache.logging.log4j.LogManager;
11+
import org.apache.logging.log4j.Logger;
12+
import org.elasticsearch.ElasticsearchException;
13+
import org.elasticsearch.action.ActionListener;
14+
import org.elasticsearch.core.TimeValue;
15+
import org.elasticsearch.inference.InferenceServiceResults;
16+
import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
17+
import org.elasticsearch.xpack.inference.external.alibabacloudsearch.AlibabaCloudSearchAccount;
18+
import org.elasticsearch.xpack.inference.external.http.sender.AlibabaCloudSearchRerankRequestManager;
19+
import org.elasticsearch.xpack.inference.external.http.sender.InferenceInputs;
20+
import org.elasticsearch.xpack.inference.external.http.sender.Sender;
21+
import org.elasticsearch.xpack.inference.services.ServiceComponents;
22+
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.rerank.AlibabaCloudSearchRerankModel;
23+
24+
import java.util.Objects;
25+
26+
import static org.elasticsearch.xpack.inference.external.action.ActionUtils.constructFailedToSendRequestMessage;
27+
import static org.elasticsearch.xpack.inference.external.action.ActionUtils.createInternalServerError;
28+
import static org.elasticsearch.xpack.inference.external.action.ActionUtils.wrapFailuresInElasticsearchException;
29+
30+
public class AlibabaCloudSearchRerankAction implements ExecutableAction {
31+
private static final Logger logger = LogManager.getLogger(AlibabaCloudSearchRerankAction.class);
32+
33+
private final AlibabaCloudSearchAccount account;
34+
private final AlibabaCloudSearchRerankModel model;
35+
private final String failedToSendRequestErrorMessage;
36+
private final Sender sender;
37+
private final AlibabaCloudSearchRerankRequestManager requestCreator;
38+
39+
public AlibabaCloudSearchRerankAction(Sender sender, AlibabaCloudSearchRerankModel model, ServiceComponents serviceComponents) {
40+
this.model = Objects.requireNonNull(model);
41+
this.account = new AlibabaCloudSearchAccount(this.model.getSecretSettings().apiKey());
42+
this.failedToSendRequestErrorMessage = constructFailedToSendRequestMessage(null, "AlibabaCloud Search rerank");
43+
this.sender = Objects.requireNonNull(sender);
44+
this.requestCreator = AlibabaCloudSearchRerankRequestManager.of(account, model, serviceComponents.threadPool());
45+
}
46+
47+
@Override
48+
public void execute(InferenceInputs inferenceInputs, TimeValue timeout, ActionListener<InferenceServiceResults> listener) {
49+
try {
50+
ActionListener<InferenceServiceResults> wrappedListener = wrapFailuresInElasticsearchException(
51+
failedToSendRequestErrorMessage,
52+
listener
53+
);
54+
sender.send(requestCreator, inferenceInputs, timeout, wrappedListener);
55+
} catch (ElasticsearchException e) {
56+
listener.onFailure(e);
57+
} catch (Exception e) {
58+
listener.onFailure(createInternalServerError(e, failedToSendRequestErrorMessage));
59+
}
60+
}
61+
}
Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.inference.external.action.alibabacloudsearch;
9+
10+
import org.apache.logging.log4j.LogManager;
11+
import org.apache.logging.log4j.Logger;
12+
import org.elasticsearch.ElasticsearchException;
13+
import org.elasticsearch.action.ActionListener;
14+
import org.elasticsearch.core.TimeValue;
15+
import org.elasticsearch.inference.InferenceServiceResults;
16+
import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
17+
import org.elasticsearch.xpack.inference.external.alibabacloudsearch.AlibabaCloudSearchAccount;
18+
import org.elasticsearch.xpack.inference.external.http.sender.AlibabaCloudSearchSparseRequestManager;
19+
import org.elasticsearch.xpack.inference.external.http.sender.InferenceInputs;
20+
import org.elasticsearch.xpack.inference.external.http.sender.Sender;
21+
import org.elasticsearch.xpack.inference.services.ServiceComponents;
22+
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.sparse.AlibabaCloudSearchSparseModel;
23+
24+
import java.util.Objects;
25+
26+
import static org.elasticsearch.xpack.inference.external.action.ActionUtils.constructFailedToSendRequestMessage;
27+
import static org.elasticsearch.xpack.inference.external.action.ActionUtils.createInternalServerError;
28+
import static org.elasticsearch.xpack.inference.external.action.ActionUtils.wrapFailuresInElasticsearchException;
29+
30+
public class AlibabaCloudSearchSparseAction implements ExecutableAction {
31+
private static final Logger logger = LogManager.getLogger(AlibabaCloudSearchSparseAction.class);
32+
33+
private final AlibabaCloudSearchAccount account;
34+
private final AlibabaCloudSearchSparseModel model;
35+
private final String failedToSendRequestErrorMessage;
36+
private final Sender sender;
37+
private final AlibabaCloudSearchSparseRequestManager requestCreator;
38+
39+
public AlibabaCloudSearchSparseAction(Sender sender, AlibabaCloudSearchSparseModel model, ServiceComponents serviceComponents) {
40+
this.model = Objects.requireNonNull(model);
41+
this.account = new AlibabaCloudSearchAccount(this.model.getSecretSettings().apiKey());
42+
this.failedToSendRequestErrorMessage = constructFailedToSendRequestMessage(null, "AlibabaCloud Search sparse embeddings");
43+
this.sender = Objects.requireNonNull(sender);
44+
requestCreator = AlibabaCloudSearchSparseRequestManager.of(account, model, serviceComponents.threadPool());
45+
}
46+
47+
@Override
48+
public void execute(InferenceInputs inferenceInputs, TimeValue timeout, ActionListener<InferenceServiceResults> listener) {
49+
try {
50+
ActionListener<InferenceServiceResults> wrappedListener = wrapFailuresInElasticsearchException(
51+
failedToSendRequestErrorMessage,
52+
listener
53+
);
54+
sender.send(requestCreator, inferenceInputs, timeout, wrappedListener);
55+
} catch (ElasticsearchException e) {
56+
listener.onFailure(e);
57+
} catch (Exception e) {
58+
listener.onFailure(createInternalServerError(e, failedToSendRequestErrorMessage));
59+
}
60+
}
61+
}
Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.inference.external.alibabacloudsearch;
9+
10+
import org.elasticsearch.common.settings.SecureString;
11+
12+
import java.util.Objects;
13+
14+
public record AlibabaCloudSearchAccount(SecureString apiKey) {
15+
16+
public AlibabaCloudSearchAccount {
17+
Objects.requireNonNull(apiKey);
18+
}
19+
}

0 commit comments

Comments
 (0)