Skip to content

Commit 2cd928c

Browse files
author
elasticsearchmachine
committed
[CI] Auto commit changes from spotless
1 parent 23b3fda commit 2cd928c

File tree

5 files changed

+99
-129
lines changed

5 files changed

+99
-129
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: 59 additions & 101 deletions
Original file line numberDiff line numberDiff line change
@@ -7,31 +7,29 @@
77

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

10-
import org.elasticsearch.common.bytes.BytesReference;
11-
12-
import org.elasticsearch.common.xcontent.LoggingDeprecationHandler;
13-
import org.elasticsearch.common.xcontent.XContentHelper;
14-
import org.elasticsearch.core.Nullable;
15-
import org.elasticsearch.xcontent.ConstructingObjectParser;
16-
import org.elasticsearch.xcontent.ParseField;
17-
import org.elasticsearch.xcontent.XContentFactory;
18-
import org.elasticsearch.xcontent.XContentParser;
19-
import org.elasticsearch.xcontent.XContentParserConfiguration;
20-
import org.elasticsearch.xcontent.XContentType;
21-
import org.elasticsearch.xpack.core.inference.results.StreamingUnifiedChatCompletionResults;
22-
2310
import software.amazon.awssdk.services.bedrockruntime.model.ContentBlockDeltaEvent;
2411
import software.amazon.awssdk.services.bedrockruntime.model.ContentBlockStartEvent;
2512
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamOutput;
2613
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamResponseHandler;
2714

2815
import org.elasticsearch.ElasticsearchException;
2916
import org.elasticsearch.ExceptionsHelper;
17+
import org.elasticsearch.common.bytes.BytesReference;
3018
import org.elasticsearch.common.util.concurrent.EsExecutors;
19+
import org.elasticsearch.common.xcontent.LoggingDeprecationHandler;
20+
import org.elasticsearch.common.xcontent.XContentHelper;
21+
import org.elasticsearch.core.Nullable;
3122
import org.elasticsearch.core.Strings;
3223
import org.elasticsearch.logging.LogManager;
3324
import org.elasticsearch.logging.Logger;
3425
import org.elasticsearch.threadpool.ThreadPool;
26+
import org.elasticsearch.xcontent.ConstructingObjectParser;
27+
import org.elasticsearch.xcontent.ParseField;
28+
import org.elasticsearch.xcontent.XContentFactory;
29+
import org.elasticsearch.xcontent.XContentParser;
30+
import org.elasticsearch.xcontent.XContentParserConfiguration;
31+
import org.elasticsearch.xcontent.XContentType;
32+
import org.elasticsearch.xpack.core.inference.results.StreamingUnifiedChatCompletionResults;
3533

3634
import java.io.IOException;
3735
import java.util.ArrayDeque;
@@ -51,7 +49,9 @@
5149
import static org.elasticsearch.xpack.inference.external.response.XContentUtils.moveToFirstToken;
5250

5351
@SuppressWarnings("checkstyle:LineLength")
54-
class AmazonBedrockUnifiedStreamingChatProcessor implements Flow.Processor<ConverseStreamOutput, StreamingUnifiedChatCompletionResults.Results> {
52+
class AmazonBedrockUnifiedStreamingChatProcessor
53+
implements
54+
Flow.Processor<ConverseStreamOutput, StreamingUnifiedChatCompletionResults.Results> {
5555
private static final Logger logger = LogManager.getLogger(AmazonBedrockStreamingChatProcessor.class);
5656

5757
private final AtomicReference<Throwable> error = new AtomicReference<>(null);
@@ -98,19 +98,19 @@ public void onNext(ConverseStreamOutput item) {
9898
demand.set(0); // reset demand before we fork to another thread
9999
item.accept(ConverseStreamResponseHandler.Visitor.builder().onContentBlockStart(this::handleContentBlockStart).build());
100100

101-
// role := unified.RoleType(msg.Value.Role)
102-
// chunk.Choices[0].Delta.Role = &role
101+
// role := unified.RoleType(msg.Value.Role)
102+
// chunk.Choices[0].Delta.Role = &role
103103

104-
// item.accept(ConverseStreamResponseHandler.Visitor.builder()
105-
// .onMessageStart(this::sendDownstreamOnAnotherThread).build());
104+
// item.accept(ConverseStreamResponseHandler.Visitor.builder()
105+
// .onMessageStart(this::sendDownstreamOnAnotherThread).build());
106106
return;
107107
}
108108

109109
case ConverseStreamOutput.EventType.CONTENT_BLOCK_DELTA -> {
110110
demand.set(0); // reset demand before we fork to another thread
111-
item.accept(ConverseStreamResponseHandler.Visitor.builder()
112-
.onContentBlockDelta(this::sendDownstreamOnAnotherThread)
113-
.build());
111+
item.accept(
112+
ConverseStreamResponseHandler.Visitor.builder().onContentBlockDelta(this::sendDownstreamOnAnotherThread).build()
113+
);
114114
return;
115115
}
116116
default -> {
@@ -137,7 +137,7 @@ private void handleContentBlockStart(ContentBlockStartEvent event) {
137137
} catch (Exception e) {
138138
String eventString = event.start().toString();
139139
logger.warn("Failed to parse event from Amazon Bedrock AI provider: {}", eventString);
140-
// throw errorParser.apply(eventString, e);
140+
// throw errorParser.apply(eventString, e);
141141
}
142142

143143
var results = new StreamingUnifiedChatCompletionResults.Results(result);
@@ -158,7 +158,7 @@ private void sendDownstreamOnAnotherThread(ContentBlockDeltaEvent event) {
158158
} catch (Exception e) {
159159
String eventString = event.delta().text();
160160
logger.warn("Failed to parse event from Amazon Bedrock AI provider: {}", eventString);
161-
// throw errorParser.apply(eventString, e);
161+
// throw errorParser.apply(eventString, e);
162162
}
163163

164164
var results = new StreamingUnifiedChatCompletionResults.Results(result);
@@ -245,16 +245,6 @@ public void cancel() {
245245
}
246246
}
247247

248-
249-
250-
251-
252-
253-
254-
255-
256-
257-
258248
private Iterator<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> parse(
259249
XContentParserConfiguration parserConfig,
260250
String event
@@ -263,29 +253,12 @@ private Iterator<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> pars
263253
moveToFirstToken(jsonParser);
264254
ensureExpectedToken(XContentParser.Token.START_OBJECT, jsonParser.currentToken(), jsonParser);
265255

266-
StreamingUnifiedChatCompletionResults.ChatCompletionChunk chunk = AmazonBedrockUnifiedStreamingChatProcessor.AmazonBedrockChatCompletionChunkParser.parse(jsonParser);
256+
StreamingUnifiedChatCompletionResults.ChatCompletionChunk chunk =
257+
AmazonBedrockUnifiedStreamingChatProcessor.AmazonBedrockChatCompletionChunkParser.parse(jsonParser);
267258
return Collections.singleton(chunk).iterator();
268259
}
269260
}
270261

271-
272-
273-
274-
275-
276-
277-
278-
279-
280-
281-
282-
283-
284-
285-
286-
287-
288-
289262
private static final String CONTENT_BLOCK_START_FIELD = "start";
290263
private static final String START_FIELD = "start";
291264
private static final String TOOL_USE_FIELD = "toolUse";
@@ -308,7 +281,9 @@ private Iterator<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> pars
308281
// --- Nested Parsers for Amazon Bedrock structure ---
309282

310283
private record ContentBlockStart(AmazonBedrockUnifiedStreamingChatProcessor.ToolUseBlockStart toolUse, Type type) {}
284+
311285
private record ToolUseBlockStart(String toolUseId, String name) {}
286+
312287
private record Type(String name) {}
313288

314289
private static class ContentBlockStartParser {
@@ -334,11 +309,12 @@ public static AmazonBedrockUnifiedStreamingChatProcessor.ContentBlockStart parse
334309

335310
private static class ToolUseBlockStartParser {
336311
@SuppressWarnings("unchecked")
337-
private static final ConstructingObjectParser<AmazonBedrockUnifiedStreamingChatProcessor.ToolUseBlockStart, Void> PARSER = new ConstructingObjectParser<>(
338-
START_FIELD,
339-
true,
340-
args -> new AmazonBedrockUnifiedStreamingChatProcessor.ToolUseBlockStart((String) args[0], (String) args[1])
341-
);
312+
private static final ConstructingObjectParser<AmazonBedrockUnifiedStreamingChatProcessor.ToolUseBlockStart, Void> PARSER =
313+
new ConstructingObjectParser<>(
314+
START_FIELD,
315+
true,
316+
args -> new AmazonBedrockUnifiedStreamingChatProcessor.ToolUseBlockStart((String) args[0], (String) args[1])
317+
);
342318

343319
static {
344320
PARSER.declareObject(
@@ -359,11 +335,12 @@ public static AmazonBedrockUnifiedStreamingChatProcessor.ToolUseBlockStart parse
359335
}
360336

361337
private static class ToolUseParser {
362-
private static final ConstructingObjectParser<AmazonBedrockUnifiedStreamingChatProcessor.ToolUseBlockStart, Void> PARSER = new ConstructingObjectParser<>(
363-
"toolUse",
364-
true,
365-
args -> new AmazonBedrockUnifiedStreamingChatProcessor.ToolUseBlockStart((String) args[0], (String) args[1])
366-
);
338+
private static final ConstructingObjectParser<AmazonBedrockUnifiedStreamingChatProcessor.ToolUseBlockStart, Void> PARSER =
339+
new ConstructingObjectParser<>(
340+
"toolUse",
341+
true,
342+
args -> new AmazonBedrockUnifiedStreamingChatProcessor.ToolUseBlockStart((String) args[0], (String) args[1])
343+
);
367344

368345
static {
369346
PARSER.declareString(ConstructingObjectParser.optionalConstructorArg(), new ParseField(TOOL_USE_ID_FIELD));
@@ -378,11 +355,7 @@ public static AmazonBedrockUnifiedStreamingChatProcessor.ToolUseBlockStart parse
378355

379356
private static class TypeParser {
380357
private static final ConstructingObjectParser<AmazonBedrockUnifiedStreamingChatProcessor.Type, Void> PARSER =
381-
new ConstructingObjectParser<>(
382-
"type",
383-
true,
384-
args -> new AmazonBedrockUnifiedStreamingChatProcessor.Type((String) args[0])
385-
);
358+
new ConstructingObjectParser<>("type", true, args -> new AmazonBedrockUnifiedStreamingChatProcessor.Type((String) args[0]));
386359

387360
static {
388361
PARSER.declareString(ConstructingObjectParser.optionalConstructorArg(), new ParseField(TOOL_TYPE_FIELD));
@@ -394,20 +367,6 @@ public static AmazonBedrockUnifiedStreamingChatProcessor.Type parse(XContentPars
394367
}
395368
}
396369

397-
398-
399-
400-
401-
402-
403-
404-
405-
406-
407-
408-
409-
410-
411370
public static class AmazonBedrockChatCompletionChunkParser {
412371
private static @Nullable StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Usage usageMetadataToChunk(
413372
@Nullable AmazonBedrockUnifiedStreamingChatProcessor.UsageMetadata usage
@@ -423,15 +382,17 @@ public static class AmazonBedrockChatCompletionChunkParser {
423382
}
424383

425384
private static StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice blockStartToChoice(
426-
AmazonBedrockUnifiedStreamingChatProcessor.ContentBlockStart blockStart) {
385+
AmazonBedrockUnifiedStreamingChatProcessor.ContentBlockStart blockStart
386+
) {
427387
StringBuilder contentTextBuilder = new StringBuilder();
428388
List<StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice.Delta.ToolCall> toolCalls = new ArrayList<>();
429389

430390
String toolUseId = null;
431391
String toolName = null;
432392

433-
var toolUseIdAndNameAreNotEmpty =
434-
blockStart.toolUse() != null && blockStart.toolUse().toolUseId() != null && blockStart.toolUse().name() != null;
393+
var toolUseIdAndNameAreNotEmpty = blockStart.toolUse() != null
394+
&& blockStart.toolUse().toolUseId() != null
395+
&& blockStart.toolUse().name() != null;
435396

436397
if (toolUseIdAndNameAreNotEmpty) {
437398
toolUseId = blockStart.toolUse().toolUseId();
@@ -448,15 +409,19 @@ private static StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice
448409
@SuppressWarnings("unchecked")
449410
private static final ConstructingObjectParser<StreamingUnifiedChatCompletionResults.ChatCompletionChunk, Void> PARSER =
450411
new ConstructingObjectParser<>("amazon_bedrock_chat_completion_chunk", true, args -> {
451-
List<AmazonBedrockUnifiedStreamingChatProcessor.ContentBlockStart> blockStart = (List<AmazonBedrockUnifiedStreamingChatProcessor.ContentBlockStart>) args[0];
452-
AmazonBedrockUnifiedStreamingChatProcessor.UsageMetadata usage = (AmazonBedrockUnifiedStreamingChatProcessor.UsageMetadata) args[1];
412+
List<AmazonBedrockUnifiedStreamingChatProcessor.ContentBlockStart> blockStart = (List<
413+
AmazonBedrockUnifiedStreamingChatProcessor.ContentBlockStart>) args[0];
414+
AmazonBedrockUnifiedStreamingChatProcessor.UsageMetadata usage =
415+
(AmazonBedrockUnifiedStreamingChatProcessor.UsageMetadata) args[1];
453416
String modelversion = (String) args[2];
454417
String responseId = (String) args[3];
455418

456419
boolean blockStartIsEmpty = blockStart == null || blockStart.isEmpty();
457420
List<StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice> choices = blockStartIsEmpty
458-
? Collections.emptyList()
459-
: blockStart.stream().map(AmazonBedrockUnifiedStreamingChatProcessor.AmazonBedrockChatCompletionChunkParser::blockStartToChoice).toList();
421+
? Collections.emptyList()
422+
: blockStart.stream()
423+
.map(AmazonBedrockUnifiedStreamingChatProcessor.AmazonBedrockChatCompletionChunkParser::blockStartToChoice)
424+
.toList();
460425

461426
return new StreamingUnifiedChatCompletionResults.ChatCompletionChunk(
462427
responseId,
@@ -487,14 +452,11 @@ public static StreamingUnifiedChatCompletionResults.ChatCompletionChunk parse(XC
487452
}
488453
}
489454

490-
491455
private record FunctionCall(String name, String args) {}
492456

493457
private static class FunctionCallParser {
494-
private static final ConstructingObjectParser<AmazonBedrockUnifiedStreamingChatProcessor.FunctionCall, Void> PARSER = new ConstructingObjectParser<>(
495-
FUNCTION_CALL_FIELD,
496-
true,
497-
args -> {
458+
private static final ConstructingObjectParser<AmazonBedrockUnifiedStreamingChatProcessor.FunctionCall, Void> PARSER =
459+
new ConstructingObjectParser<>(FUNCTION_CALL_FIELD, true, args -> {
498460
var name = (String) args[0];
499461

500462
@SuppressWarnings("unchecked")
@@ -510,8 +472,7 @@ private static class FunctionCallParser {
510472
logger.warn("Failed to parse and convert VertexAI function args to json", e);
511473
return new AmazonBedrockUnifiedStreamingChatProcessor.FunctionCall(name, null);
512474
}
513-
}
514-
);
475+
});
515476

516477
static {
517478
PARSER.declareString(ConstructingObjectParser.constructorArg(), new ParseField(FUNCTION_NAME_FIELD));
@@ -526,10 +487,8 @@ public static AmazonBedrockUnifiedStreamingChatProcessor.FunctionCall parse(XCon
526487
private record UsageMetadata(int promptTokenCount, int candidatesTokenCount, int totalTokenCount) {}
527488

528489
private static class UsageMetadataParser {
529-
private static final ConstructingObjectParser<AmazonBedrockUnifiedStreamingChatProcessor.UsageMetadata, Void> PARSER = new ConstructingObjectParser<>(
530-
USAGE_METADATA_FIELD,
531-
true,
532-
args -> {
490+
private static final ConstructingObjectParser<AmazonBedrockUnifiedStreamingChatProcessor.UsageMetadata, Void> PARSER =
491+
new ConstructingObjectParser<>(USAGE_METADATA_FIELD, true, args -> {
533492
if (Objects.isNull(args[0]) && Objects.isNull(args[1]) && Objects.isNull(args[2])) {
534493
return null;
535494
}
@@ -538,8 +497,7 @@ private static class UsageMetadataParser {
538497
args[1] == null ? 0 : (int) args[1],
539498
args[2] == null ? 0 : (int) args[2]
540499
);
541-
}
542-
);
500+
});
543501

544502
static {
545503
PARSER.declareInt(ConstructingObjectParser.optionalConstructorArg(), new ParseField(PROMPT_TOKEN_COUNT_FIELD));

0 commit comments

Comments
 (0)