Skip to content

Commit c811d35

Browse files
elasticsearchmachineEvgenii-Kazannik
elasticsearchmachine
authored and
Evgenii-Kazannik
committed
[CI] Auto commit changes from spotless
1 parent b2e04a3 commit c811d35

File tree

5 files changed

+85
-57
lines changed

5 files changed

+85
-57
lines changed

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,6 @@
77

88
package org.elasticsearch.xpack.inference.services.amazonbedrock.client;
99

10-
import org.elasticsearch.xpack.core.inference.results.StreamingUnifiedChatCompletionResults;
11-
1210
import software.amazon.awssdk.services.bedrockruntime.model.ConverseRequest;
1311
import software.amazon.awssdk.services.bedrockruntime.model.ConverseResponse;
1412
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamRequest;
@@ -18,6 +16,7 @@
1816
import org.elasticsearch.ElasticsearchException;
1917
import org.elasticsearch.action.ActionListener;
2018
import org.elasticsearch.inference.InferenceServiceResults;
19+
import org.elasticsearch.xpack.core.inference.results.StreamingUnifiedChatCompletionResults;
2120

2221
import java.time.Instant;
2322
import java.util.concurrent.Flow;

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,6 @@
77

88
package org.elasticsearch.xpack.inference.services.amazonbedrock.client;
99

10-
import org.elasticsearch.xpack.core.inference.results.StreamingUnifiedChatCompletionResults;
11-
1210
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
1311
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
1412
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
@@ -33,6 +31,7 @@
3331
import org.elasticsearch.core.TimeValue;
3432
import org.elasticsearch.inference.InferenceServiceResults;
3533
import org.elasticsearch.threadpool.ThreadPool;
34+
import org.elasticsearch.xpack.core.inference.results.StreamingUnifiedChatCompletionResults;
3635
import org.elasticsearch.xpack.inference.services.amazonbedrock.AmazonBedrockModel;
3736
import org.reactivestreams.FlowAdapters;
3837
import org.slf4j.LoggerFactory;

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

Lines changed: 45 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -7,28 +7,26 @@
77

88
package org.elasticsearch.xpack.inference.services.amazonbedrock.client;
99

10-
import org.elasticsearch.common.xcontent.LoggingDeprecationHandler;
11-
import org.elasticsearch.xcontent.XContentFactory;
12-
import org.elasticsearch.xcontent.XContentParser;
13-
import org.elasticsearch.xcontent.XContentParserConfiguration;
14-
import org.elasticsearch.xcontent.XContentType;
15-
import org.elasticsearch.xpack.core.inference.results.StreamingUnifiedChatCompletionResults;
16-
1710
import software.amazon.awssdk.services.bedrockruntime.model.ContentBlockDeltaEvent;
1811
import software.amazon.awssdk.services.bedrockruntime.model.ContentBlockStartEvent;
1912
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamMetadataEvent;
2013
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamOutput;
2114
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamResponseHandler;
15+
import software.amazon.awssdk.services.bedrockruntime.model.MessageStartEvent;
2216

2317
import org.elasticsearch.ElasticsearchException;
2418
import org.elasticsearch.ExceptionsHelper;
2519
import org.elasticsearch.common.util.concurrent.EsExecutors;
20+
import org.elasticsearch.common.xcontent.LoggingDeprecationHandler;
2621
import org.elasticsearch.core.Strings;
2722
import org.elasticsearch.logging.LogManager;
2823
import org.elasticsearch.logging.Logger;
2924
import org.elasticsearch.threadpool.ThreadPool;
30-
31-
import software.amazon.awssdk.services.bedrockruntime.model.MessageStartEvent;
25+
import org.elasticsearch.xcontent.XContentFactory;
26+
import org.elasticsearch.xcontent.XContentParser;
27+
import org.elasticsearch.xcontent.XContentParserConfiguration;
28+
import org.elasticsearch.xcontent.XContentType;
29+
import org.elasticsearch.xpack.core.inference.results.StreamingUnifiedChatCompletionResults;
3230

3331
import java.io.IOException;
3432
import java.util.ArrayDeque;
@@ -45,7 +43,9 @@
4543
import static org.elasticsearch.xpack.inference.external.response.XContentUtils.positionParserAtTokenAfterField;
4644

4745
@SuppressWarnings("checkstyle:LineLength")
48-
class AmazonBedrockUnifiedStreamingChatProcessor implements Flow.Processor<ConverseStreamOutput, StreamingUnifiedChatCompletionResults.Results> {
46+
class AmazonBedrockUnifiedStreamingChatProcessor
47+
implements
48+
Flow.Processor<ConverseStreamOutput, StreamingUnifiedChatCompletionResults.Results> {
4949
private static final Logger logger = LogManager.getLogger(AmazonBedrockStreamingChatProcessor.class);
5050
private static final String FAILED_TO_FIND_FIELD_TEMPLATE = "Failed to find required field [%s] in Anthropic chat completions response";
5151

@@ -94,32 +94,41 @@ public void onNext(ConverseStreamOutput item) {
9494
switch (eventType) {
9595
case ConverseStreamOutput.EventType.MESSAGE_START -> {
9696
demand.set(0); // reset demand before we fork to another thread
97-
item.accept(ConverseStreamResponseHandler.Visitor.builder()
98-
.onMessageStart(event -> handleMessageStart(event, chunks, parserConfig)).build());
97+
item.accept(
98+
ConverseStreamResponseHandler.Visitor.builder()
99+
.onMessageStart(event -> handleMessageStart(event, chunks, parserConfig))
100+
.build()
101+
);
99102
return;
100103
}
101104
case ConverseStreamOutput.EventType.CONTENT_BLOCK_START -> {
102105
demand.set(0); // reset demand before we fork to another thread
103-
item.accept(ConverseStreamResponseHandler.Visitor.builder()
104-
.onContentBlockStart(event -> handleContentBlockStart(event, chunks, parserConfig)).build());
106+
item.accept(
107+
ConverseStreamResponseHandler.Visitor.builder()
108+
.onContentBlockStart(event -> handleContentBlockStart(event, chunks, parserConfig))
109+
.build()
110+
);
105111
return;
106112
}
107113
case ConverseStreamOutput.EventType.CONTENT_BLOCK_DELTA -> {
108114
demand.set(0); // reset demand before we fork to another thread
109-
item.accept(ConverseStreamResponseHandler.Visitor.builder()
110-
.onContentBlockDelta(event -> handleContentBlockDelta(event, chunks, parserConfig)).build());
115+
item.accept(
116+
ConverseStreamResponseHandler.Visitor.builder()
117+
.onContentBlockDelta(event -> handleContentBlockDelta(event, chunks, parserConfig))
118+
.build()
119+
);
111120
return;
112121
}
113122
case ConverseStreamOutput.EventType.METADATA -> {
114123
demand.set(0); // reset demand before we fork to another thread
115-
item.accept(ConverseStreamResponseHandler.Visitor.builder()
116-
.onMetadata(event -> handleMetadata(event, chunks, parserConfig)).build());
124+
item.accept(
125+
ConverseStreamResponseHandler.Visitor.builder().onMetadata(event -> handleMetadata(event, chunks, parserConfig)).build()
126+
);
117127
return;
118128
}
119129
case ConverseStreamOutput.EventType.MESSAGE_STOP -> {
120130
demand.set(0); // reset demand before we fork to another thread
121-
item.accept(ConverseStreamResponseHandler.Visitor.builder()
122-
.onMessageStop(event -> Stream.empty()).build());
131+
item.accept(ConverseStreamResponseHandler.Visitor.builder().onMessageStop(event -> Stream.empty()).build());
123132
return;
124133
}
125134
default -> {
@@ -136,7 +145,9 @@ public void onNext(ConverseStreamOutput item) {
136145

137146
private void handleMessageStart(
138147
MessageStartEvent event,
139-
ArrayDeque<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> chunks, XContentParserConfiguration parserConfig) {
148+
ArrayDeque<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> chunks,
149+
XContentParserConfiguration parserConfig
150+
) {
140151
runOnUtilityThreadPool(() -> {
141152
var data = event.role().name();
142153
try {
@@ -150,8 +161,11 @@ private void handleMessageStart(
150161
});
151162
}
152163

153-
private void handleContentBlockStart(ContentBlockStartEvent event,
154-
ArrayDeque<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> chunks, XContentParserConfiguration parserConfig) {
164+
private void handleContentBlockStart(
165+
ContentBlockStartEvent event,
166+
ArrayDeque<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> chunks,
167+
XContentParserConfiguration parserConfig
168+
) {
155169
var data = event.start().toString();
156170
try {
157171
var contentBlockStart = parseContentBlockStart(parserConfig, data);
@@ -163,8 +177,11 @@ private void handleContentBlockStart(ContentBlockStartEvent event,
163177
downstream.onNext(results);
164178
}
165179

166-
private void handleContentBlockDelta(ContentBlockDeltaEvent event,
167-
ArrayDeque<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> chunks, XContentParserConfiguration parserConfig) {
180+
private void handleContentBlockDelta(
181+
ContentBlockDeltaEvent event,
182+
ArrayDeque<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> chunks,
183+
XContentParserConfiguration parserConfig
184+
) {
168185
runOnUtilityThreadPool(() -> {
169186
var data = event.delta().toString();
170187
try {
@@ -180,7 +197,9 @@ private void handleContentBlockDelta(ContentBlockDeltaEvent event,
180197

181198
private void handleMetadata(
182199
ConverseStreamMetadataEvent event,
183-
ArrayDeque<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> chunks, XContentParserConfiguration parserConfig) {
200+
ArrayDeque<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> chunks,
201+
XContentParserConfiguration parserConfig
202+
) {
184203
runOnUtilityThreadPool(() -> {
185204
var data = event.toString();
186205
try {
@@ -282,7 +301,6 @@ public void cancel() {
282301
public static final String ROLE_FIELD = "role";
283302
public static final String CONTENT_FIELD = "content";
284303

285-
286304
public static final String INDEX_FIELD = "index";
287305
public static final String NAME_FIELD = "name";
288306
public static final String INPUT_TOKENS_FIELD = "input_tokens";

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

Lines changed: 37 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -7,20 +7,18 @@
77

88
package org.elasticsearch.xpack.inference.services.amazonbedrock.request.completion;
99

10-
import org.elasticsearch.inference.UnifiedCompletionRequest;
11-
1210
import software.amazon.awssdk.core.document.Document;
1311
import software.amazon.awssdk.services.bedrockruntime.model.ContentBlock;
1412
import software.amazon.awssdk.services.bedrockruntime.model.InferenceConfiguration;
1513
import software.amazon.awssdk.services.bedrockruntime.model.Message;
16-
17-
import org.elasticsearch.core.Nullable;
18-
import org.elasticsearch.core.Strings;
19-
2014
import software.amazon.awssdk.services.bedrockruntime.model.Tool;
2115
import software.amazon.awssdk.services.bedrockruntime.model.ToolInputSchema;
2216
import software.amazon.awssdk.services.bedrockruntime.model.ToolSpecification;
2317

18+
import org.elasticsearch.core.Nullable;
19+
import org.elasticsearch.core.Strings;
20+
import org.elasticsearch.inference.UnifiedCompletionRequest;
21+
2422
import java.util.List;
2523
import java.util.Map;
2624
import java.util.Optional;
@@ -50,27 +48,44 @@ public static List<Message> getUnifiedConverseMessageList(List<UnifiedCompletion
5048
public static Tool getUnifiedConverseTool(UnifiedCompletionRequest.Tool tool) {
5149
var t = ToolInputSchema.builder();
5250
return Tool.builder()
53-
.toolSpec(ToolSpecification.builder()
54-
.name(tool.function().name())
55-
.description(tool.function().description())
56-
.inputSchema(getInputSchema())
57-
.build())
51+
.toolSpec(
52+
ToolSpecification.builder()
53+
.name(tool.function().name())
54+
.description(tool.function().description())
55+
.inputSchema(getInputSchema())
56+
.build()
57+
)
5858
.build();
5959
}
6060

6161
static ToolInputSchema getInputSchema() {
62-
return ToolInputSchema.fromJson(Document.fromMap(Map.of(
63-
"type", Document.fromString("object"), "properties", Document.fromMap(Map.of(
64-
"sign", Document.fromMap(Map.of(
62+
return ToolInputSchema.fromJson(
63+
Document.fromMap(
64+
Map.of(
6565
"type",
66-
Document.fromString("string"),
67-
"description",
68-
Document.fromString(
69-
"The call sign for the radio station for which you want the most popular song. "
70-
+ "Example calls signs are WZPZ and WKRP.")
71-
))
72-
)), "required", Document.fromList(List.of(Document.fromString("sign")))
73-
)));
66+
Document.fromString("object"),
67+
"properties",
68+
Document.fromMap(
69+
Map.of(
70+
"sign",
71+
Document.fromMap(
72+
Map.of(
73+
"type",
74+
Document.fromString("string"),
75+
"description",
76+
Document.fromString(
77+
"The call sign for the radio station for which you want the most popular song. "
78+
+ "Example calls signs are WZPZ and WKRP."
79+
)
80+
)
81+
)
82+
)
83+
),
84+
"required",
85+
Document.fromList(List.of(Document.fromString("sign")))
86+
)
87+
)
88+
);
7489
}
7590

7691
public static Optional<InferenceConfiguration> inferenceConfig(AmazonBedrockConverseRequestEntity request) {

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

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,20 +7,17 @@
77

88
package org.elasticsearch.xpack.inference.services.amazonbedrock.request.completion;
99

10-
import org.elasticsearch.xpack.core.inference.results.StreamingUnifiedChatCompletionResults;
11-
12-
import software.amazon.awssdk.core.document.Document;
1310
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamRequest;
1411
import software.amazon.awssdk.services.bedrockruntime.model.SpecificToolChoice;
1512
import software.amazon.awssdk.services.bedrockruntime.model.Tool;
1613
import software.amazon.awssdk.services.bedrockruntime.model.ToolChoice;
1714
import software.amazon.awssdk.services.bedrockruntime.model.ToolConfiguration;
18-
import software.amazon.awssdk.services.bedrockruntime.model.ToolInputSchema;
1915
import software.amazon.awssdk.services.bedrockruntime.model.ToolSpecification;
2016

2117
import org.elasticsearch.core.Nullable;
2218
import org.elasticsearch.core.TimeValue;
2319
import org.elasticsearch.inference.TaskType;
20+
import org.elasticsearch.xpack.core.inference.results.StreamingUnifiedChatCompletionResults;
2421
import org.elasticsearch.xpack.inference.services.amazonbedrock.client.AmazonBedrockBaseClient;
2522
import org.elasticsearch.xpack.inference.services.amazonbedrock.completion.AmazonBedrockChatCompletionModel;
2623
import org.elasticsearch.xpack.inference.services.amazonbedrock.request.AmazonBedrockRequest;

0 commit comments

Comments
 (0)