Skip to content

Commit 115de47

Browse files
Evgenii-KazannikEvgenii-Kazannik
committed
Add Mixedbread AI support
1 parent 5d0c5e0 commit 115de47

29 files changed

+2383
-0
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -329,6 +329,7 @@ static TransportVersion def(int id) {
329329
public static final TransportVersion PROJECT_STATE_REGISTRY_RECORDS_DELETIONS = def(9_113_0_00);
330330
public static final TransportVersion ESQL_SERIALIZE_TIMESERIES_FIELD_TYPE = def(9_114_0_00);
331331
public static final TransportVersion ML_INFERENCE_IBM_WATSONX_COMPLETION_ADDED = def(9_115_0_00);
332+
public static final TransportVersion ML_INFERENCE_MIXEDBREAD_ADDED = def(9_116_0_00);
332333
/*
333334
* STOP! READ THIS FIRST! No, really,
334335
* ____ _____ ___ ____ _ ____ _____ _ ____ _____ _ _ ___ ____ _____ ___ ____ ____ _____ _

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

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,10 @@
106106
import org.elasticsearch.xpack.inference.services.jinaai.rerank.JinaAIRerankTaskSettings;
107107
import org.elasticsearch.xpack.inference.services.mistral.completion.MistralChatCompletionServiceSettings;
108108
import org.elasticsearch.xpack.inference.services.mistral.embeddings.MistralEmbeddingsServiceSettings;
109+
import org.elasticsearch.xpack.inference.services.mixedbread.embeddings.MixedbreadEmbeddingsServiceSettings;
110+
import org.elasticsearch.xpack.inference.services.mixedbread.embeddings.MixedbreadEmbeddingsTaskSettings;
111+
import org.elasticsearch.xpack.inference.services.mixedbread.rerank.MixedbreadRerankServiceSettings;
112+
import org.elasticsearch.xpack.inference.services.mixedbread.rerank.MixedbreadRerankTaskSettings;
109113
import org.elasticsearch.xpack.inference.services.openai.completion.OpenAiChatCompletionServiceSettings;
110114
import org.elasticsearch.xpack.inference.services.openai.completion.OpenAiChatCompletionTaskSettings;
111115
import org.elasticsearch.xpack.inference.services.openai.embeddings.OpenAiEmbeddingsServiceSettings;
@@ -164,6 +168,7 @@ public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
164168
addIbmWatsonxNamedWritables(namedWriteables);
165169
addGoogleVertexAiNamedWriteables(namedWriteables);
166170
addMistralNamedWriteables(namedWriteables);
171+
addMixedbreadNamedWriteables(namedWriteables);
167172
addCustomElandWriteables(namedWriteables);
168173
addAnthropicNamedWritables(namedWriteables);
169174
addAmazonBedrockNamedWriteables(namedWriteables);
@@ -276,6 +281,34 @@ private static void addMistralNamedWriteables(List<NamedWriteableRegistry.Entry>
276281
// note - no task settings for Mistral embeddings...
277282
}
278283

284+
private static void addMixedbreadNamedWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {
285+
namedWriteables.add(
286+
new NamedWriteableRegistry.Entry(
287+
ServiceSettings.class,
288+
MixedbreadEmbeddingsServiceSettings.NAME,
289+
MixedbreadEmbeddingsServiceSettings::new
290+
)
291+
);
292+
namedWriteables.add(
293+
new NamedWriteableRegistry.Entry(
294+
TaskSettings.class,
295+
MixedbreadEmbeddingsTaskSettings.NAME,
296+
MixedbreadEmbeddingsTaskSettings::new
297+
)
298+
);
299+
300+
namedWriteables.add(
301+
new NamedWriteableRegistry.Entry(
302+
ServiceSettings.class,
303+
MixedbreadRerankServiceSettings.NAME,
304+
MixedbreadRerankServiceSettings::new
305+
)
306+
);
307+
namedWriteables.add(
308+
new NamedWriteableRegistry.Entry(TaskSettings.class, MixedbreadRerankTaskSettings.NAME, MixedbreadRerankTaskSettings::new)
309+
);
310+
}
311+
279312
private static void addAzureAiStudioNamedWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {
280313
namedWriteables.add(
281314
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
@@ -133,6 +133,7 @@
133133
import org.elasticsearch.xpack.inference.services.ibmwatsonx.IbmWatsonxService;
134134
import org.elasticsearch.xpack.inference.services.jinaai.JinaAIService;
135135
import org.elasticsearch.xpack.inference.services.mistral.MistralService;
136+
import org.elasticsearch.xpack.inference.services.mixedbread.MixedbreadService;
136137
import org.elasticsearch.xpack.inference.services.openai.OpenAiService;
137138
import org.elasticsearch.xpack.inference.services.sagemaker.SageMakerClient;
138139
import org.elasticsearch.xpack.inference.services.sagemaker.SageMakerService;
@@ -392,6 +393,7 @@ public List<InferenceServiceExtension.Factory> getInferenceServiceFactories() {
392393
context -> new GoogleAiStudioService(httpFactory.get(), serviceComponents.get()),
393394
context -> new GoogleVertexAiService(httpFactory.get(), serviceComponents.get()),
394395
context -> new MistralService(httpFactory.get(), serviceComponents.get()),
396+
context -> new MixedbreadService(httpFactory.get(), serviceComponents.get()),
395397
context -> new AnthropicService(httpFactory.get(), serviceComponents.get()),
396398
context -> new AmazonBedrockService(httpFactory.get(), amazonBedrockFactory.get(), serviceComponents.get()),
397399
context -> new AlibabaCloudSearchService(httpFactory.get(), serviceComponents.get()),

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/mixedbread/MixedbreadConstants.java‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
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.services.mixedbread;
9+
10+
public class MixedbreadConstants {
11+
public static final String EMBEDDINGS_URI_PATH = "/v1/embeddings";
12+
public static final String RERANK_URI_PATH = "/v1/rerank";
13+
14+
// common service settings fields
15+
public static final String API_KEY_FIELD = "api_key";
16+
17+
// embeddings service and request settings
18+
public static final String INPUT_FIELD = "input";
19+
20+
// rerank task settings fields
21+
public static final String QUERY_FIELD = "query";
22+
23+
// embeddings task settings fields
24+
public static final String USER_FIELD = "user";
25+
26+
// rerank task settings fields
27+
public static final String RETURN_DOCUMENTS_FIELD = "return_documents";
28+
public static final String TOP_K_FIELD = "top_k";
29+
30+
private MixedbreadConstants() {}
31+
}

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/mixedbread/MixedbreadModel.java‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
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.services.mixedbread;
9+
10+
import org.elasticsearch.inference.EmptySecretSettings;
11+
import org.elasticsearch.inference.ModelConfigurations;
12+
import org.elasticsearch.inference.ModelSecrets;
13+
import org.elasticsearch.inference.SecretSettings;
14+
import org.elasticsearch.inference.ServiceSettings;
15+
import org.elasticsearch.inference.TaskSettings;
16+
import org.elasticsearch.xpack.inference.services.RateLimitGroupingModel;
17+
import org.elasticsearch.xpack.inference.services.settings.DefaultSecretSettings;
18+
import org.elasticsearch.xpack.inference.services.settings.RateLimitSettings;
19+
20+
import java.net.URI;
21+
import java.net.URISyntaxException;
22+
import java.util.Map;
23+
import java.util.Objects;
24+
25+
/**
26+
* Abstract class representing a Mixedbread model for inference.
27+
* This class extends RateLimitGroupingModel and provides common functionality for Mixedbread models.
28+
*/
29+
public abstract class MixedbreadModel extends RateLimitGroupingModel {
30+
protected String modelId;
31+
protected URI uri;
32+
protected RateLimitSettings rateLimitSettings;
33+
34+
public MixedbreadModel(MixedbreadModel model, TaskSettings taskSettings, RateLimitSettings rateLimitSettings) {
35+
super(model, taskSettings);
36+
this.rateLimitSettings = Objects.requireNonNull(rateLimitSettings);
37+
}
38+
39+
/**
40+
* Constructor for creating a MixedbreadModel with specified configurations and secrets.
41+
*
42+
* @param configurations the model configurations
43+
* @param secrets the secret settings for the model
44+
*/
45+
protected MixedbreadModel(ModelConfigurations configurations, ModelSecrets secrets) {
46+
super(configurations, secrets);
47+
}
48+
49+
/**
50+
* Constructor for creating a MixedbreadModel with specified model, service settings, and secret settings.
51+
* @param model the model configurations
52+
* @param serviceSettings the settings for the inference service
53+
*/
54+
protected MixedbreadModel(RateLimitGroupingModel model, ServiceSettings serviceSettings) {
55+
super(model, serviceSettings);
56+
}
57+
58+
public String model() {
59+
return this.modelId;
60+
}
61+
62+
public URI uri() {
63+
return this.uri;
64+
}
65+
66+
@Override
67+
public RateLimitSettings rateLimitSettings() {
68+
return this.rateLimitSettings;
69+
}
70+
71+
@Override
72+
public int rateLimitGroupingHash() {
73+
return Objects.hash(modelId, uri, getSecretSettings());
74+
}
75+
76+
// Needed for testing only
77+
public void setURI(String newUri) {
78+
try {
79+
this.uri = new URI(newUri);
80+
} catch (URISyntaxException e) {
81+
// swallow any error
82+
}
83+
}
84+
85+
/**
86+
* Retrieves the secret settings from the provided map of secrets.
87+
* If the map is null or empty, it returns an instance of EmptySecretSettings.
88+
* Caused by the fact that Mixedbread model doesn't have out of the box security settings and can be used witout authentication.
89+
*
90+
* @param secrets the map containing secret settings
91+
* @return an instance of SecretSettings
92+
*/
93+
protected static SecretSettings retrieveSecretSettings(Map<String, Object> secrets) {
94+
return (secrets != null && secrets.isEmpty()) ? EmptySecretSettings.INSTANCE : DefaultSecretSettings.fromMap(secrets);
95+
}
96+
97+
@Override
98+
public DefaultSecretSettings getSecretSettings() {
99+
return (DefaultSecretSettings) super.getSecretSettings();
100+
}
101+
}

0 commit comments

Comments
 (0)