Skip to content

Commit 3f4df8f

Browse files
author
elasticsearchmachine
committed
[CI] Auto commit changes from spotless
1 parent b677ae7 commit 3f4df8f

File tree

2 files changed

+52
-33
lines changed

2 files changed

+52
-33
lines changed

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

Lines changed: 49 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -7,15 +7,14 @@
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.ContentBlockDelta;
1311
import software.amazon.awssdk.services.bedrockruntime.model.ContentBlockDeltaEvent;
1412
import software.amazon.awssdk.services.bedrockruntime.model.ContentBlockStart;
1513
import software.amazon.awssdk.services.bedrockruntime.model.ContentBlockStartEvent;
1614
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamMetadataEvent;
1715
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamOutput;
1816
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamResponseHandler;
17+
import software.amazon.awssdk.services.bedrockruntime.model.MessageStartEvent;
1918

2019
import org.elasticsearch.ElasticsearchException;
2120
import org.elasticsearch.ExceptionsHelper;
@@ -24,8 +23,7 @@
2423
import org.elasticsearch.logging.LogManager;
2524
import org.elasticsearch.logging.Logger;
2625
import org.elasticsearch.threadpool.ThreadPool;
27-
28-
import software.amazon.awssdk.services.bedrockruntime.model.MessageStartEvent;
26+
import org.elasticsearch.xpack.core.inference.results.StreamingUnifiedChatCompletionResults;
2927

3028
import java.util.ArrayDeque;
3129
import java.util.List;
@@ -38,7 +36,9 @@
3836
import static org.elasticsearch.xpack.inference.InferencePlugin.UTILITY_THREAD_POOL_NAME;
3937

4038
@SuppressWarnings("checkstyle:LineLength")
41-
class AmazonBedrockUnifiedStreamingChatProcessor implements Flow.Processor<ConverseStreamOutput, StreamingUnifiedChatCompletionResults.Results> {
39+
class AmazonBedrockUnifiedStreamingChatProcessor
40+
implements
41+
Flow.Processor<ConverseStreamOutput, StreamingUnifiedChatCompletionResults.Results> {
4242
private static final Logger logger = LogManager.getLogger(AmazonBedrockStreamingChatProcessor.class);
4343

4444
private final AtomicReference<Throwable> error = new AtomicReference<>(null);
@@ -85,32 +85,37 @@ public void onNext(ConverseStreamOutput item) {
8585
switch (eventType) {
8686
case ConverseStreamOutput.EventType.MESSAGE_START -> {
8787
demand.set(0); // reset demand before we fork to another thread
88-
item.accept(ConverseStreamResponseHandler.Visitor.builder()
89-
.onMessageStart(event -> handleMessageStart(event, chunks)).build());
88+
item.accept(
89+
ConverseStreamResponseHandler.Visitor.builder().onMessageStart(event -> handleMessageStart(event, chunks)).build()
90+
);
9091
return;
9192
}
9293
case ConverseStreamOutput.EventType.CONTENT_BLOCK_START -> {
9394
demand.set(0); // reset demand before we fork to another thread
94-
item.accept(ConverseStreamResponseHandler.Visitor.builder()
95-
.onContentBlockStart(event -> handleContentBlockStart(event, chunks)).build());
95+
item.accept(
96+
ConverseStreamResponseHandler.Visitor.builder()
97+
.onContentBlockStart(event -> handleContentBlockStart(event, chunks))
98+
.build()
99+
);
96100
return;
97101
}
98102
case ConverseStreamOutput.EventType.CONTENT_BLOCK_DELTA -> {
99103
demand.set(0); // reset demand before we fork to another thread
100-
item.accept(ConverseStreamResponseHandler.Visitor.builder()
101-
.onContentBlockDelta(event -> handleContentBlockDelta(event, chunks)).build());
104+
item.accept(
105+
ConverseStreamResponseHandler.Visitor.builder()
106+
.onContentBlockDelta(event -> handleContentBlockDelta(event, chunks))
107+
.build()
108+
);
102109
return;
103110
}
104111
case ConverseStreamOutput.EventType.METADATA -> {
105112
demand.set(0); // reset demand before we fork to another thread
106-
item.accept(ConverseStreamResponseHandler.Visitor.builder()
107-
.onMetadata(event -> handleMetadata(event, chunks)).build());
113+
item.accept(ConverseStreamResponseHandler.Visitor.builder().onMetadata(event -> handleMetadata(event, chunks)).build());
108114
return;
109115
}
110116
case ConverseStreamOutput.EventType.MESSAGE_STOP -> {
111117
demand.set(0); // reset demand before we fork to another thread
112-
item.accept(ConverseStreamResponseHandler.Visitor.builder()
113-
.onMessageStop(event -> Stream.empty()).build());
118+
item.accept(ConverseStreamResponseHandler.Visitor.builder().onMessageStop(event -> Stream.empty()).build());
114119
return;
115120
}
116121
default -> {
@@ -125,9 +130,7 @@ public void onNext(ConverseStreamOutput item) {
125130
}
126131
}
127132

128-
private void handleMessageStart(
129-
MessageStartEvent event,
130-
ArrayDeque<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> chunks) {
133+
private void handleMessageStart(MessageStartEvent event, ArrayDeque<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> chunks) {
131134
runOnUtilityThreadPool(() -> {
132135
try {
133136
var messageStart = handleMessageStart(event);
@@ -143,7 +146,10 @@ private void handleMessageStart(
143146
});
144147
}
145148

146-
private void handleContentBlockStart(ContentBlockStartEvent event, ArrayDeque<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> chunks) {
149+
private void handleContentBlockStart(
150+
ContentBlockStartEvent event,
151+
ArrayDeque<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> chunks
152+
) {
147153
try {
148154
var contentBlockStart = handleContentBlockStart(event);
149155
contentBlockStart.forEach(chunks::offer);
@@ -154,7 +160,10 @@ private void handleContentBlockStart(ContentBlockStartEvent event, ArrayDeque<St
154160
downstream.onNext(results);
155161
}
156162

157-
private void handleContentBlockDelta(ContentBlockDeltaEvent event, ArrayDeque<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> chunks) {
163+
private void handleContentBlockDelta(
164+
ContentBlockDeltaEvent event,
165+
ArrayDeque<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> chunks
166+
) {
158167
runOnUtilityThreadPool(() -> {
159168
try {
160169
var contentBlockDelta = handleContentBlockDelta(event);
@@ -167,7 +176,10 @@ private void handleContentBlockDelta(ContentBlockDeltaEvent event, ArrayDeque<St
167176
});
168177
}
169178

170-
private void handleMetadata(ConverseStreamMetadataEvent event, ArrayDeque<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> chunks) {
179+
private void handleMetadata(
180+
ConverseStreamMetadataEvent event,
181+
ArrayDeque<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> chunks
182+
) {
171183
runOnUtilityThreadPool(() -> {
172184
try {
173185
var messageDelta = handleMetadata(event);
@@ -265,10 +277,10 @@ public void cancel() {
265277
* @return a stream of ChatCompletionChunk
266278
*/
267279
public static Stream<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> handleMessageStart(MessageStartEvent event) {
268-
var delta = new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta(null, null, event.roleAsString(), null);
269-
var choice = new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice(delta, null, 0);
270-
var chunk = new StreamingUnifiedChatCompletionResults.ChatCompletionChunk(null, List.of(choice), null, null, null);
271-
return Stream.of(chunk);
280+
var delta = new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta(null, null, event.roleAsString(), null);
281+
var choice = new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice(delta, null, 0);
282+
var chunk = new StreamingUnifiedChatCompletionResults.ChatCompletionChunk(null, List.of(choice), null, null, null);
283+
return Stream.of(chunk);
272284
}
273285

274286
/**
@@ -278,11 +290,18 @@ public static Stream<StreamingUnifiedChatCompletionResults.ChatCompletionChunk>
278290
* @param start the ContentBlockStart data
279291
* @return a ToolCall
280292
*/
281-
private static StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta.ToolCall handleToolUseStart(ContentBlockStart start) {
293+
private static StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta.ToolCall handleToolUseStart(
294+
ContentBlockStart start
295+
) {
282296
var type = start.type();
283297
var toolUse = start.toolUse();
284298
var function = new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta.ToolCall.Function(null, toolUse.name());
285-
return new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta.ToolCall(0, toolUse.toolUseId(), function, type.name());
299+
return new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta.ToolCall(
300+
0,
301+
toolUse.toolUseId(),
302+
function,
303+
type.name()
304+
);
286305
}
287306

288307
/**
@@ -292,7 +311,9 @@ private static StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.
292311
* @param delta the ContentBlockDelta data
293312
* @return a ToolCall
294313
*/
295-
private static StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta.ToolCall handleToolUseDelta(ContentBlockDelta delta) {
314+
private static StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta.ToolCall handleToolUseDelta(
315+
ContentBlockDelta delta
316+
) {
296317
var type = delta.type();
297318
var toolUse = delta.toolUse();
298319
var function = new StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta.ToolCall.Function(toolUse.input(), null);

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

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,19 +7,17 @@
77

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

10+
import software.amazon.awssdk.services.bedrockruntime.model.*;
11+
1012
import org.elasticsearch.ElasticsearchException;
1113
import org.elasticsearch.common.util.concurrent.EsExecutors;
1214
import org.elasticsearch.test.ESTestCase;
1315
import org.elasticsearch.threadpool.ThreadPool;
1416
import org.elasticsearch.xpack.core.inference.results.StreamingUnifiedChatCompletionResults;
1517
import org.junit.Before;
1618
import org.mockito.ArgumentCaptor;
17-
import software.amazon.awssdk.services.bedrockruntime.model.*;
1819

19-
import java.util.Arrays;
20-
import java.util.concurrent.ExecutorService;
2120
import java.util.concurrent.Flow;
22-
import java.util.concurrent.atomic.AtomicInteger;
2321

2422
import static org.elasticsearch.xpack.inference.InferencePlugin.UTILITY_THREAD_POOL_NAME;
2523
import static org.hamcrest.Matchers.*;
@@ -106,7 +104,7 @@ private void verifyText(Flow.Subscriber<StreamingUnifiedChatCompletionResults.Re
106104
verify(downstream, times(1)).onNext(assertArg(results -> {
107105
assertThat(results, notNullValue());
108106
assertThat(results.chunks().size(), equalTo(1));
109-
// assertThat(results.chunks().getFirst().choices().getFirst(), equalTo(expectedText));
107+
// assertThat(results.chunks().getFirst().choices().getFirst(), equalTo(expectedText));
110108
}));
111109
}
112110

0 commit comments

Comments
 (0)