Skip to content

Commit 6a9d678

Browse files
committed
[ML] Move and rename AmazonBedrockSecretSettings (elastic#125323)
In preparation for integrating with SageMaker, we want to reuse the existing SecretSettings. - AmazonBedrockSecretSettings moved from services.amazonbedrock to common.amazon. - AmazonBedrockSecretSettings was renamed to AwsSecretSettings. - accessKey and secretKey are now encapsulated.
1 parent 271b04f commit 6a9d678

File tree

13 files changed

+103
-99
lines changed

13 files changed

+103
-99
lines changed

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

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import org.elasticsearch.xpack.inference.action.task.StreamingTaskManager;
2929
import org.elasticsearch.xpack.inference.chunking.SentenceBoundaryChunkingSettings;
3030
import org.elasticsearch.xpack.inference.chunking.WordBoundaryChunkingSettings;
31+
import org.elasticsearch.xpack.inference.common.amazon.AwsSecretSettings;
3132
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.AlibabaCloudSearchServiceSettings;
3233
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.completion.AlibabaCloudSearchCompletionServiceSettings;
3334
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.completion.AlibabaCloudSearchCompletionTaskSettings;
@@ -37,7 +38,6 @@
3738
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.rerank.AlibabaCloudSearchRerankTaskSettings;
3839
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.sparse.AlibabaCloudSearchSparseServiceSettings;
3940
import org.elasticsearch.xpack.inference.services.alibabacloudsearch.sparse.AlibabaCloudSearchSparseTaskSettings;
40-
import org.elasticsearch.xpack.inference.services.amazonbedrock.AmazonBedrockSecretSettings;
4141
import org.elasticsearch.xpack.inference.services.amazonbedrock.completion.AmazonBedrockChatCompletionServiceSettings;
4242
import org.elasticsearch.xpack.inference.services.amazonbedrock.completion.AmazonBedrockChatCompletionTaskSettings;
4343
import org.elasticsearch.xpack.inference.services.amazonbedrock.embeddings.AmazonBedrockEmbeddingsServiceSettings;
@@ -146,6 +146,7 @@ public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
146146
addCustomElandWriteables(namedWriteables);
147147
addAnthropicNamedWritables(namedWriteables);
148148
addAmazonBedrockNamedWriteables(namedWriteables);
149+
addAwsNamedWriteables(namedWriteables);
149150
addEisNamedWriteables(namedWriteables);
150151
addAlibabaCloudSearchNamedWriteables(namedWriteables);
151152
addJinaAINamedWriteables(namedWriteables);
@@ -165,14 +166,6 @@ private static void addUnifiedNamedWriteables(List<NamedWriteableRegistry.Entry>
165166
}
166167

167168
private static void addAmazonBedrockNamedWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {
168-
namedWriteables.add(
169-
new NamedWriteableRegistry.Entry(
170-
AmazonBedrockSecretSettings.class,
171-
AmazonBedrockSecretSettings.NAME,
172-
AmazonBedrockSecretSettings::new
173-
)
174-
);
175-
176169
namedWriteables.add(
177170
new NamedWriteableRegistry.Entry(
178171
ServiceSettings.class,
@@ -199,6 +192,10 @@ private static void addAmazonBedrockNamedWriteables(List<NamedWriteableRegistry.
199192
);
200193
}
201194

195+
private static void addAwsNamedWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {
196+
namedWriteables.add(new NamedWriteableRegistry.Entry(AwsSecretSettings.class, AwsSecretSettings.NAME, AwsSecretSettings::new));
197+
}
198+
202199
private static void addMistralNamedWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {
203200
namedWriteables.add(
204201
new NamedWriteableRegistry.Entry(
Lines changed: 18 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
* 2.0.
66
*/
77

8-
package org.elasticsearch.xpack.inference.services.amazonbedrock;
8+
package org.elasticsearch.xpack.inference.common.amazon;
99

1010
import org.elasticsearch.TransportVersion;
1111
import org.elasticsearch.TransportVersions;
@@ -33,13 +33,13 @@
3333
import static org.elasticsearch.xpack.inference.services.amazonbedrock.AmazonBedrockConstants.ACCESS_KEY_FIELD;
3434
import static org.elasticsearch.xpack.inference.services.amazonbedrock.AmazonBedrockConstants.SECRET_KEY_FIELD;
3535

36-
public class AmazonBedrockSecretSettings implements SecretSettings {
37-
public static final String NAME = "amazon_bedrock_secret_settings";
36+
public class AwsSecretSettings implements SecretSettings {
37+
public static final String NAME = "aws_secret_settings";
3838

39-
public final SecureString accessKey;
40-
public final SecureString secretKey;
39+
private final SecureString accessKey;
40+
private final SecureString secretKey;
4141

42-
public static AmazonBedrockSecretSettings fromMap(@Nullable Map<String, Object> map) {
42+
public static AwsSecretSettings fromMap(@Nullable Map<String, Object> map) {
4343
if (map == null) {
4444
return null;
4545
}
@@ -62,15 +62,15 @@ public static AmazonBedrockSecretSettings fromMap(@Nullable Map<String, Object>
6262
throw validationException;
6363
}
6464

65-
return new AmazonBedrockSecretSettings(secureAccessKey, secureSecretKey);
65+
return new AwsSecretSettings(secureAccessKey, secureSecretKey);
6666
}
6767

68-
public AmazonBedrockSecretSettings(SecureString accessKey, SecureString secretKey) {
68+
public AwsSecretSettings(SecureString accessKey, SecureString secretKey) {
6969
this.accessKey = Objects.requireNonNull(accessKey);
7070
this.secretKey = Objects.requireNonNull(secretKey);
7171
}
7272

73-
public AmazonBedrockSecretSettings(StreamInput in) throws IOException {
73+
public AwsSecretSettings(StreamInput in) throws IOException {
7474
this.accessKey = in.readSecureString();
7575
this.secretKey = in.readSecureString();
7676
}
@@ -106,7 +106,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
106106
public boolean equals(Object object) {
107107
if (this == object) return true;
108108
if (object == null || getClass() != object.getClass()) return false;
109-
AmazonBedrockSecretSettings that = (AmazonBedrockSecretSettings) object;
109+
AwsSecretSettings that = (AwsSecretSettings) object;
110110
return Objects.equals(accessKey, that.accessKey) && Objects.equals(secretKey, that.secretKey);
111111
}
112112

@@ -120,6 +120,14 @@ public SecretSettings newSecretSettings(Map<String, Object> newSecrets) {
120120
return fromMap(new HashMap<>(newSecrets));
121121
}
122122

123+
public SecureString accessKey() {
124+
return accessKey;
125+
}
126+
127+
public SecureString secretKey() {
128+
return secretKey;
129+
}
130+
123131
public static class Configuration {
124132
public static Map<String, SettingsConfiguration> get() {
125133
return configuration.getOrCompute();

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/amazonbedrock/AmazonBedrockBaseClient.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ protected AmazonBedrockBaseClient(AmazonBedrockModel model, @Nullable TimeValue
2828
public static Integer getModelKeysAndRegionHashcode(AmazonBedrockModel model, @Nullable TimeValue timeout) {
2929
var secretSettings = model.getSecretSettings();
3030
var serviceSettings = model.getServiceSettings();
31-
return Objects.hash(secretSettings.accessKey, secretSettings.secretKey, serviceSettings.region(), timeout);
31+
return Objects.hash(secretSettings.accessKey(), secretSettings.secretKey(), serviceSettings.region(), timeout);
3232
}
3333

3434
public final void setClock(Clock clock) {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/amazonbedrock/AmazonBedrockInferenceClient.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -142,7 +142,7 @@ protected BedrockRuntimeAsyncClient createAmazonBedrockClient(AmazonBedrockModel
142142
try {
143143
SpecialPermission.check();
144144
return AccessController.doPrivileged((PrivilegedExceptionAction<BedrockRuntimeAsyncClient>) () -> {
145-
var credentials = AwsBasicCredentials.create(secretSettings.accessKey.toString(), secretSettings.secretKey.toString());
145+
var credentials = AwsBasicCredentials.create(secretSettings.accessKey().toString(), secretSettings.secretKey().toString());
146146
var credentialsProvider = StaticCredentialsProvider.create(credentials);
147147
var clientConfig = timeout == null
148148
? NettyNioAsyncHttpClient.builder().connectionTimeout(DEFAULT_CLIENT_TIMEOUT_MS)

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/AmazonBedrockRequestManager.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ public static AmazonBedrockRequestManager.RateLimitGrouping of(AmazonBedrockMode
4343

4444
var awsSecretSettings = model.getSecretSettings();
4545

46-
return new RateLimitGrouping(Objects.hash(awsSecretSettings.accessKey, awsSecretSettings.secretKey));
46+
return new RateLimitGrouping(Objects.hash(awsSecretSettings.accessKey(), awsSecretSettings.secretKey()));
4747
}
4848
}
4949

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
import org.elasticsearch.inference.ModelSecrets;
1313
import org.elasticsearch.inference.ServiceSettings;
1414
import org.elasticsearch.inference.TaskSettings;
15+
import org.elasticsearch.xpack.inference.common.amazon.AwsSecretSettings;
1516
import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
1617
import org.elasticsearch.xpack.inference.external.action.amazonbedrock.AmazonBedrockActionVisitor;
1718
import org.elasticsearch.xpack.inference.services.settings.RateLimitSettings;
@@ -81,8 +82,8 @@ public AmazonBedrockServiceSettings getServiceSettings() {
8182
}
8283

8384
@Override
84-
public AmazonBedrockSecretSettings getSecretSettings() {
85-
return (AmazonBedrockSecretSettings) super.getSecretSettings();
85+
public AwsSecretSettings getSecretSettings() {
86+
return (AwsSecretSettings) super.getSecretSettings();
8687
}
8788

8889
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import org.elasticsearch.rest.RestStatus;
3232
import org.elasticsearch.xpack.inference.chunking.ChunkingSettingsBuilder;
3333
import org.elasticsearch.xpack.inference.chunking.EmbeddingRequestChunker;
34+
import org.elasticsearch.xpack.inference.common.amazon.AwsSecretSettings;
3435
import org.elasticsearch.xpack.inference.external.action.amazonbedrock.AmazonBedrockActionCreator;
3536
import org.elasticsearch.xpack.inference.external.amazonbedrock.AmazonBedrockRequestSender;
3637
import org.elasticsearch.xpack.inference.external.http.sender.EmbeddingsInput;
@@ -459,7 +460,7 @@ public static InferenceServiceConfiguration get() {
459460
.build()
460461
);
461462

462-
configurationMap.putAll(AmazonBedrockSecretSettings.Configuration.get());
463+
configurationMap.putAll(AwsSecretSettings.Configuration.get());
463464
configurationMap.putAll(
464465
RateLimitSettings.toSettingsConfigurationWithDescription(
465466
"By default, the amazonbedrock service sets the number of requests allowed per minute to 240.",

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/amazonbedrock/completion/AmazonBedrockChatCompletionModel.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,11 @@
1212
import org.elasticsearch.inference.ModelSecrets;
1313
import org.elasticsearch.inference.TaskSettings;
1414
import org.elasticsearch.inference.TaskType;
15+
import org.elasticsearch.xpack.inference.common.amazon.AwsSecretSettings;
1516
import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
1617
import org.elasticsearch.xpack.inference.external.action.amazonbedrock.AmazonBedrockActionVisitor;
1718
import org.elasticsearch.xpack.inference.services.ConfigurationParseContext;
1819
import org.elasticsearch.xpack.inference.services.amazonbedrock.AmazonBedrockModel;
19-
import org.elasticsearch.xpack.inference.services.amazonbedrock.AmazonBedrockSecretSettings;
2020

2121
import java.util.Map;
2222

@@ -47,7 +47,7 @@ public AmazonBedrockChatCompletionModel(
4747
name,
4848
AmazonBedrockChatCompletionServiceSettings.fromMap(serviceSettings, context),
4949
AmazonBedrockChatCompletionTaskSettings.fromMap(taskSettings),
50-
AmazonBedrockSecretSettings.fromMap(secretSettings)
50+
AwsSecretSettings.fromMap(secretSettings)
5151
);
5252
}
5353

@@ -57,7 +57,7 @@ public AmazonBedrockChatCompletionModel(
5757
String service,
5858
AmazonBedrockChatCompletionServiceSettings serviceSettings,
5959
AmazonBedrockChatCompletionTaskSettings taskSettings,
60-
AmazonBedrockSecretSettings secrets
60+
AwsSecretSettings secrets
6161
) {
6262
super(new ModelConfigurations(inferenceEntityId, taskType, service, serviceSettings, taskSettings), new ModelSecrets(secrets));
6363
}

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/amazonbedrock/embeddings/AmazonBedrockEmbeddingsModel.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,11 @@
1616
import org.elasticsearch.inference.ServiceSettings;
1717
import org.elasticsearch.inference.TaskSettings;
1818
import org.elasticsearch.inference.TaskType;
19+
import org.elasticsearch.xpack.inference.common.amazon.AwsSecretSettings;
1920
import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
2021
import org.elasticsearch.xpack.inference.external.action.amazonbedrock.AmazonBedrockActionVisitor;
2122
import org.elasticsearch.xpack.inference.services.ConfigurationParseContext;
2223
import org.elasticsearch.xpack.inference.services.amazonbedrock.AmazonBedrockModel;
23-
import org.elasticsearch.xpack.inference.services.amazonbedrock.AmazonBedrockSecretSettings;
2424

2525
import java.util.Map;
2626

@@ -54,7 +54,7 @@ public AmazonBedrockEmbeddingsModel(
5454
AmazonBedrockEmbeddingsServiceSettings.fromMap(serviceSettings, context),
5555
new EmptyTaskSettings(),
5656
chunkingSettings,
57-
AmazonBedrockSecretSettings.fromMap(secretSettings)
57+
AwsSecretSettings.fromMap(secretSettings)
5858
);
5959
}
6060

@@ -65,7 +65,7 @@ public AmazonBedrockEmbeddingsModel(
6565
AmazonBedrockEmbeddingsServiceSettings serviceSettings,
6666
TaskSettings taskSettings,
6767
ChunkingSettings chunkingSettings,
68-
AmazonBedrockSecretSettings secrets
68+
AwsSecretSettings secrets
6969
) {
7070
super(
7171
new ModelConfigurations(inferenceEntityId, taskType, service, serviceSettings, new EmptyTaskSettings(), chunkingSettings),
Lines changed: 22 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
* 2.0.
66
*/
77

8-
package org.elasticsearch.xpack.inference.services.amazonbedrock;
8+
package org.elasticsearch.xpack.inference.common.amazon;
99

1010
import org.elasticsearch.TransportVersion;
1111
import org.elasticsearch.common.Strings;
@@ -27,48 +27,46 @@
2727
import static org.hamcrest.Matchers.containsString;
2828
import static org.hamcrest.Matchers.is;
2929

30-
public class AmazonBedrockSecretSettingsTests extends AbstractBWCWireSerializationTestCase<AmazonBedrockSecretSettings> {
30+
public class AwsSecretSettingsTests extends AbstractBWCWireSerializationTestCase<AwsSecretSettings> {
3131

3232
public void testNewSecretSettings() {
33-
AmazonBedrockSecretSettings initialSettings = createRandom();
34-
AmazonBedrockSecretSettings newSettings = createRandom();
33+
AwsSecretSettings initialSettings = createRandom();
34+
AwsSecretSettings newSettings = createRandom();
3535

36-
AmazonBedrockSecretSettings finalSettings = (AmazonBedrockSecretSettings) initialSettings.newSecretSettings(
37-
Map.of(ACCESS_KEY_FIELD, newSettings.accessKey.toString(), SECRET_KEY_FIELD, newSettings.secretKey.toString())
36+
AwsSecretSettings finalSettings = (AwsSecretSettings) initialSettings.newSecretSettings(
37+
Map.of(ACCESS_KEY_FIELD, newSettings.accessKey().toString(), SECRET_KEY_FIELD, newSettings.secretKey().toString())
3838
);
3939

4040
assertEquals(newSettings, finalSettings);
4141
}
4242

4343
public void testIt_CreatesSettings_ReturnsNullFromMap_null() {
44-
var secrets = AmazonBedrockSecretSettings.fromMap(null);
44+
var secrets = AwsSecretSettings.fromMap(null);
4545
assertNull(secrets);
4646
}
4747

4848
public void testIt_CreatesSettings_FromMap_WithValues() {
49-
var secrets = AmazonBedrockSecretSettings.fromMap(
50-
new HashMap<>(Map.of(ACCESS_KEY_FIELD, "accesstest", SECRET_KEY_FIELD, "secrettest"))
51-
);
49+
var secrets = AwsSecretSettings.fromMap(new HashMap<>(Map.of(ACCESS_KEY_FIELD, "accesstest", SECRET_KEY_FIELD, "secrettest")));
5250
assertThat(
5351
secrets,
54-
is(new AmazonBedrockSecretSettings(new SecureString("accesstest".toCharArray()), new SecureString("secrettest".toCharArray())))
52+
is(new AwsSecretSettings(new SecureString("accesstest".toCharArray()), new SecureString("secrettest".toCharArray())))
5553
);
5654
}
5755

5856
public void testIt_CreatesSettings_FromMap_IgnoresExtraKeys() {
59-
var secrets = AmazonBedrockSecretSettings.fromMap(
57+
var secrets = AwsSecretSettings.fromMap(
6058
new HashMap<>(Map.of(ACCESS_KEY_FIELD, "accesstest", SECRET_KEY_FIELD, "secrettest", "extrakey", "extravalue"))
6159
);
6260
assertThat(
6361
secrets,
64-
is(new AmazonBedrockSecretSettings(new SecureString("accesstest".toCharArray()), new SecureString("secrettest".toCharArray())))
62+
is(new AwsSecretSettings(new SecureString("accesstest".toCharArray()), new SecureString("secrettest".toCharArray())))
6563
);
6664
}
6765

6866
public void testIt_FromMap_ThrowsValidationException_AccessKeyMissing() {
6967
var thrownException = expectThrows(
7068
ValidationException.class,
71-
() -> AmazonBedrockSecretSettings.fromMap(new HashMap<>(Map.of(SECRET_KEY_FIELD, "secrettest")))
69+
() -> AwsSecretSettings.fromMap(new HashMap<>(Map.of(SECRET_KEY_FIELD, "secrettest")))
7270
);
7371

7472
assertThat(
@@ -80,7 +78,7 @@ public void testIt_FromMap_ThrowsValidationException_AccessKeyMissing() {
8078
public void testIt_FromMap_ThrowsValidationException_SecretKeyMissing() {
8179
var thrownException = expectThrows(
8280
ValidationException.class,
83-
() -> AmazonBedrockSecretSettings.fromMap(new HashMap<>(Map.of(ACCESS_KEY_FIELD, "accesstest")))
81+
() -> AwsSecretSettings.fromMap(new HashMap<>(Map.of(ACCESS_KEY_FIELD, "accesstest")))
8482
);
8583

8684
assertThat(
@@ -90,9 +88,7 @@ public void testIt_FromMap_ThrowsValidationException_SecretKeyMissing() {
9088
}
9189

9290
public void testToXContent_CreatesProperContent() throws IOException {
93-
var secrets = AmazonBedrockSecretSettings.fromMap(
94-
new HashMap<>(Map.of(ACCESS_KEY_FIELD, "accesstest", SECRET_KEY_FIELD, "secrettest"))
95-
);
91+
var secrets = AwsSecretSettings.fromMap(new HashMap<>(Map.of(ACCESS_KEY_FIELD, "accesstest", SECRET_KEY_FIELD, "secrettest")));
9692

9793
XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON);
9894
secrets.toXContent(builder, null);
@@ -106,26 +102,26 @@ public static Map<String, Object> getAmazonBedrockSecretSettingsMap(String acces
106102
}
107103

108104
@Override
109-
protected AmazonBedrockSecretSettings mutateInstanceForVersion(AmazonBedrockSecretSettings instance, TransportVersion version) {
105+
protected AwsSecretSettings mutateInstanceForVersion(AwsSecretSettings instance, TransportVersion version) {
110106
return instance;
111107
}
112108

113109
@Override
114-
protected Writeable.Reader<AmazonBedrockSecretSettings> instanceReader() {
115-
return AmazonBedrockSecretSettings::new;
110+
protected Writeable.Reader<AwsSecretSettings> instanceReader() {
111+
return AwsSecretSettings::new;
116112
}
117113

118114
@Override
119-
protected AmazonBedrockSecretSettings createTestInstance() {
115+
protected AwsSecretSettings createTestInstance() {
120116
return createRandom();
121117
}
122118

123119
@Override
124-
protected AmazonBedrockSecretSettings mutateInstance(AmazonBedrockSecretSettings instance) throws IOException {
125-
return randomValueOtherThan(instance, AmazonBedrockSecretSettingsTests::createRandom);
120+
protected AwsSecretSettings mutateInstance(AwsSecretSettings instance) throws IOException {
121+
return randomValueOtherThan(instance, AwsSecretSettingsTests::createRandom);
126122
}
127123

128-
private static AmazonBedrockSecretSettings createRandom() {
129-
return new AmazonBedrockSecretSettings(new SecureString(randomAlphaOfLength(10)), new SecureString(randomAlphaOfLength(10)));
124+
private static AwsSecretSettings createRandom() {
125+
return new AwsSecretSettings(new SecureString(randomAlphaOfLength(10)), new SecureString(randomAlphaOfLength(10)));
130126
}
131127
}

0 commit comments

Comments
 (0)