Skip to content

Commit a39936f

Browse files
DonalEvansDonalEvans
committed
Apply review feedback
- Move transport version checks from ThinkingConfig to GoogleVertexAiChatCompletionServiceSettings - Remove default value argument from ThinkingConfig.of() - Add test coverage for ServiceUtils.extractOptionalPositiveInteger() and extractOptionalInteger()
1 parent bd109a0 commit a39936f

File tree

5 files changed

+125
-42
lines changed

5 files changed

+125
-42
lines changed

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/completion/GoogleVertexAiChatCompletionServiceSettings.java‎

Lines changed: 14 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,16 @@ public class GoogleVertexAiChatCompletionServiceSettings extends FilteredXConten
5252
private static final ThinkingConfig EMPTY_THINKING_CONFIG = new ThinkingConfig();
5353

5454
public GoogleVertexAiChatCompletionServiceSettings(StreamInput in) throws IOException {
55-
this(in.readString(), in.readString(), in.readString(), new RateLimitSettings(in), new ThinkingConfig(in));
55+
this.projectId = in.readString();
56+
this.location = in.readString();
57+
this.modelId = in.readString();
58+
this.rateLimitSettings = new RateLimitSettings(in);
59+
60+
if (in.getTransportVersion().onOrAfter(TransportVersions.GEMINI_THINKING_BUDGET_ADDED)) {
61+
thinkingConfig = new ThinkingConfig(in);
62+
} else {
63+
thinkingConfig = EMPTY_THINKING_CONFIG;
64+
}
5665
}
5766

5867
@Override
@@ -83,13 +92,7 @@ public static GoogleVertexAiChatCompletionServiceSettings fromMap(Map<String, Ob
8392
);
8493

8594
// Extract optional thinkingConfig settings
86-
ThinkingConfig thinkingConfig = ThinkingConfig.of(
87-
map,
88-
EMPTY_THINKING_CONFIG,
89-
validationException,
90-
GoogleVertexAiService.NAME,
91-
context
92-
);
95+
ThinkingConfig thinkingConfig = ThinkingConfig.of(map, validationException, GoogleVertexAiService.NAME, context);
9396

9497
if (validationException.validationErrors().isEmpty() == false) {
9598
throw validationException;
@@ -158,7 +161,9 @@ public void writeTo(StreamOutput out) throws IOException {
158161
out.writeString(location);
159162
out.writeString(modelId);
160163
rateLimitSettings.writeTo(out);
161-
thinkingConfig.writeTo(out);
164+
if (out.getTransportVersion().onOrAfter(TransportVersions.GEMINI_THINKING_BUDGET_ADDED)) {
165+
thinkingConfig.writeTo(out);
166+
}
162167
}
163168

164169
@Override

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/completion/ThinkingConfig.java‎

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77

88
package org.elasticsearch.xpack.inference.services.googlevertexai.completion;
99

10-
import org.elasticsearch.TransportVersions;
1110
import org.elasticsearch.common.Strings;
1211
import org.elasticsearch.common.ValidationException;
1312
import org.elasticsearch.common.io.stream.StreamInput;
@@ -29,6 +28,8 @@
2928
/**
3029
* This class encapsulates the ThinkingConfig object contained within GenerationConfig. Only the thinkingBudget field is currently
3130
* supported, but the includeThoughts field may be added in the future
31+
*
32+
* @see <a href="https://ai.google.dev/gemini-api/docs/thinking"> Gemini Thinking documentation</a>
3233
*/
3334
public class ThinkingConfig implements Writeable, ToXContentFragment {
3435
public static final String THINKING_CONFIG_FIELD = "thinking_config";
@@ -48,16 +49,11 @@ public ThinkingConfig(Integer thinkingBudget) {
4849
}
4950

5051
public ThinkingConfig(StreamInput in) throws IOException {
51-
if (in.getTransportVersion().onOrAfter(TransportVersions.GEMINI_THINKING_BUDGET_ADDED)) {
52-
thinkingBudget = in.readOptionalVInt();
53-
} else {
54-
thinkingBudget = null;
55-
}
52+
thinkingBudget = in.readOptionalVInt();
5653
}
5754

5855
public static ThinkingConfig of(
5956
Map<String, Object> map,
60-
ThinkingConfig defaultValue,
6157
ValidationException validationException,
6258
String serviceName,
6359
ConfigurationParseContext context
@@ -74,7 +70,7 @@ public static ThinkingConfig of(
7470
throwIfNotEmptyMap(thinkingConfigSettings, serviceName);
7571
}
7672

77-
return thinkingBudget == null ? defaultValue : new ThinkingConfig(thinkingBudget);
73+
return new ThinkingConfig(thinkingBudget);
7874
}
7975

8076
public boolean isEmpty() {
@@ -87,9 +83,7 @@ public Integer getThinkingBudget() {
8783

8884
@Override
8985
public void writeTo(StreamOutput out) throws IOException {
90-
if (out.getTransportVersion().onOrAfter(TransportVersions.GEMINI_THINKING_BUDGET_ADDED)) {
91-
out.writeOptionalVInt(thinkingBudget);
92-
}
86+
out.writeOptionalVInt(thinkingBudget);
9387
}
9488

9589
@Override

‎x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/ServiceUtilsTests.java‎

Lines changed: 73 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import static org.elasticsearch.xpack.inference.services.ServiceUtils.convertToUri;
3232
import static org.elasticsearch.xpack.inference.services.ServiceUtils.createUri;
3333
import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalEnum;
34+
import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalInteger;
3435
import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalList;
3536
import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalListOfStringTuples;
3637
import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalMap;
@@ -539,14 +540,85 @@ public void testExtractOptionalList_AddsException_WhenFieldContainsMixedTypeValu
539540
assertTrue(map.isEmpty());
540541
}
541542

542-
public void testExtractOptionalPositiveInt() {
543+
public void testExtractOptionalPositiveInteger_returnsInteger_withPositiveInteger() {
543544
var validation = new ValidationException();
544545
validation.addValidationError("previous error");
545546
Map<String, Object> map = modifiableMap(Map.of("abc", 1));
546547
assertEquals(Integer.valueOf(1), extractOptionalPositiveInteger(map, "abc", "scope", validation));
547548
assertThat(validation.validationErrors(), hasSize(1));
548549
}
549550

551+
public void testExtractOptionalPositiveInteger_returnsNull_whenSettingNotFound() {
552+
var validation = new ValidationException();
553+
validation.addValidationError("previous error");
554+
Map<String, Object> map = modifiableMap(Map.of("abc", 1));
555+
assertThat(extractOptionalPositiveInteger(map, "not_abc", "scope", validation), is(nullValue()));
556+
assertThat(validation.validationErrors(), hasSize(1));
557+
}
558+
559+
public void testExtractOptionalPositiveInteger_returnsNull_addsValidationError_whenObjectIsNotInteger() {
560+
var validation = new ValidationException();
561+
validation.addValidationError("previous error");
562+
String setting = "abc";
563+
Map<String, Object> map = modifiableMap(Map.of(setting, "not_an_int"));
564+
assertThat(extractOptionalPositiveInteger(map, setting, "scope", validation), is(nullValue()));
565+
assertThat(validation.validationErrors(), hasSize(2));
566+
assertThat(validation.validationErrors().getLast(), containsString("cannot be converted to a [Integer]"));
567+
}
568+
569+
public void testExtractOptionalPositiveInteger_returnNull_addsValidationError_withNonPositiveInteger() {
570+
var validation = new ValidationException();
571+
validation.addValidationError("previous error");
572+
String zeroKey = "zero";
573+
String negativeKey = "negative";
574+
Map<String, Object> map = modifiableMap(Map.of(zeroKey, 0, negativeKey, -1));
575+
576+
// Test zero
577+
assertThat(extractOptionalPositiveInteger(map, zeroKey, "scope", validation), is(nullValue()));
578+
assertThat(validation.validationErrors(), hasSize(2));
579+
assertThat(validation.validationErrors().getLast(), containsString("[" + zeroKey + "] must be a positive integer"));
580+
581+
// Test a negative number
582+
assertThat(extractOptionalPositiveInteger(map, negativeKey, "scope", validation), is(nullValue()));
583+
assertThat(validation.validationErrors(), hasSize(3));
584+
assertThat(validation.validationErrors().getLast(), containsString("[" + negativeKey + "] must be a positive integer"));
585+
}
586+
587+
public void testExtractOptionalInteger_returnsInteger() {
588+
var validation = new ValidationException();
589+
validation.addValidationError("previous error");
590+
String positiveKey = "positive";
591+
int positiveValue = 123;
592+
String zeroKey = "zero";
593+
int zeroValue = 0;
594+
String negativeKey = "negative";
595+
int negativeValue = -123;
596+
Map<String, Object> map = modifiableMap(Map.of(positiveKey, positiveValue, zeroKey, zeroValue, negativeKey, negativeValue));
597+
598+
assertThat(extractOptionalInteger(map, positiveKey, "scope", validation), is(positiveValue));
599+
assertThat(extractOptionalInteger(map, zeroKey, "scope", validation), is(zeroValue));
600+
assertThat(extractOptionalInteger(map, negativeKey, "scope", validation), is(negativeValue));
601+
assertThat(validation.validationErrors(), hasSize(1));
602+
}
603+
604+
public void testExtractOptionalInteger_returnsNull_whenSettingNotFound() {
605+
var validation = new ValidationException();
606+
validation.addValidationError("previous error");
607+
Map<String, Object> map = modifiableMap(Map.of("abc", 1));
608+
assertThat(extractOptionalInteger(map, "not_abc", "scope", validation), is(nullValue()));
609+
assertThat(validation.validationErrors(), hasSize(1));
610+
}
611+
612+
public void testExtractOptionalInteger_returnsNull_addsValidationError_whenObjectIsNotInteger() {
613+
var validation = new ValidationException();
614+
validation.addValidationError("previous error");
615+
String setting = "abc";
616+
Map<String, Object> map = modifiableMap(Map.of(setting, "not_an_int"));
617+
assertThat(extractOptionalInteger(map, setting, "scope", validation), is(nullValue()));
618+
assertThat(validation.validationErrors(), hasSize(2));
619+
assertThat(validation.validationErrors().getLast(), containsString("cannot be converted to a [Integer]"));
620+
}
621+
550622
public void testExtractOptionalPositiveLong_IntegerValue() {
551623
var validation = new ValidationException();
552624
validation.addValidationError("previous error");

‎x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googleaistudio/request/completion/ThinkingConfigTests.java‎

Lines changed: 31 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -8,13 +8,15 @@
88
package org.elasticsearch.xpack.inference.services.googleaistudio.request.completion;
99

1010
import org.elasticsearch.ElasticsearchStatusException;
11+
import org.elasticsearch.TransportVersion;
1112
import org.elasticsearch.common.Strings;
1213
import org.elasticsearch.common.ValidationException;
1314
import org.elasticsearch.common.io.stream.Writeable;
14-
import org.elasticsearch.test.AbstractWireSerializingTestCase;
15+
import org.elasticsearch.common.xcontent.XContentHelper;
1516
import org.elasticsearch.xcontent.XContentBuilder;
1617
import org.elasticsearch.xcontent.XContentFactory;
1718
import org.elasticsearch.xcontent.XContentType;
19+
import org.elasticsearch.xpack.core.ml.AbstractBWCWireSerializationTestCase;
1820
import org.elasticsearch.xpack.inference.services.ConfigurationParseContext;
1921
import org.elasticsearch.xpack.inference.services.googlevertexai.completion.ThinkingConfig;
2022

@@ -28,7 +30,7 @@
2830
import static org.hamcrest.Matchers.is;
2931
import static org.hamcrest.Matchers.nullValue;
3032

31-
public class ThinkingConfigTests extends AbstractWireSerializingTestCase<ThinkingConfig> {
33+
public class ThinkingConfigTests extends AbstractBWCWireSerializationTestCase<ThinkingConfig> {
3234

3335
public void testNoArgConstructor_createsEmptyConfig() {
3436
ThinkingConfig thinkingConfig = new ThinkingConfig();
@@ -57,46 +59,45 @@ public void testOf_withThinkingConfigSpecified_andThinkingBudgetSpecified() {
5759
Map.of(THINKING_CONFIG_FIELD, new HashMap<>(Map.of(THINKING_BUDGET_FIELD, thinkingBudget)))
5860
);
5961

60-
ThinkingConfig result = ThinkingConfig.of(settings, new ThinkingConfig(), exception, "test", ConfigurationParseContext.REQUEST);
62+
ThinkingConfig result = ThinkingConfig.of(settings, exception, "test", ConfigurationParseContext.REQUEST);
6163

6264
assertThat(result, is(new ThinkingConfig(thinkingBudget)));
6365
assertThat(exception.validationErrors(), is(empty()));
6466
}
6567

66-
public void testOf_usesDefault_withThinkingConfigSpecified_andThinkingBudgetNotSpecified() {
68+
public void testOf_returnsEmptyThinkingConfig_withThinkingConfigSpecified_andThinkingBudgetNotSpecified() {
6769
ValidationException exception = new ValidationException();
6870
Map<String, Object> settings = new HashMap<>(Map.of(THINKING_CONFIG_FIELD, new HashMap<>()));
69-
ThinkingConfig defaultValue = new ThinkingConfig(123);
7071

71-
ThinkingConfig result = ThinkingConfig.of(settings, defaultValue, exception, "test", ConfigurationParseContext.REQUEST);
72+
ThinkingConfig result = ThinkingConfig.of(settings, exception, "test", ConfigurationParseContext.REQUEST);
7273

73-
assertThat(result, is(defaultValue));
74+
assertThat(result.getThinkingBudget(), is(nullValue()));
75+
assertThat(result.isEmpty(), is(true));
7476
assertThat(exception.validationErrors(), is(empty()));
7577
}
7678

77-
public void testOf_usesDefault_withThinkingConfigNotSpecified_andThinkingBudgetSpecified() {
79+
public void testOf_returnsEmptyThinkingConfig_withThinkingConfigNotSpecified_andThinkingBudgetSpecified() {
7880
ValidationException exception = new ValidationException();
7981
int thinkingBudget = 256;
8082
Map<String, Object> settings = new HashMap<>(
8183
Map.of("not_thinking_config", new HashMap<>(Map.of(THINKING_BUDGET_FIELD, thinkingBudget)))
8284
);
83-
ThinkingConfig defaultValue = new ThinkingConfig(123);
8485

85-
ThinkingConfig result = ThinkingConfig.of(settings, defaultValue, exception, "test", ConfigurationParseContext.REQUEST);
86+
ThinkingConfig result = ThinkingConfig.of(settings, exception, "test", ConfigurationParseContext.REQUEST);
8687

87-
assertThat(result, is(defaultValue));
88+
assertThat(result.getThinkingBudget(), is(nullValue()));
89+
assertThat(result.isEmpty(), is(true));
8890
assertThat(exception.validationErrors(), is(empty()));
8991
}
9092

9193
public void testOf_throwsException_withUnknownField_andRequestContext() {
9294
ValidationException exception = new ValidationException();
9395
int anInt = 42;
9496
Map<String, Object> settings = new HashMap<>(Map.of(THINKING_CONFIG_FIELD, new HashMap<>(Map.of("not_thinking_budget", anInt))));
95-
ThinkingConfig defaultValue = new ThinkingConfig(123);
9697

9798
Exception thrownException = expectThrows(
9899
ElasticsearchStatusException.class,
99-
() -> ThinkingConfig.of(settings, defaultValue, exception, "test", ConfigurationParseContext.REQUEST)
100+
() -> ThinkingConfig.of(settings, exception, "test", ConfigurationParseContext.REQUEST)
100101
);
101102

102103
assertThat(
@@ -105,15 +106,15 @@ public void testOf_throwsException_withUnknownField_andRequestContext() {
105106
);
106107
}
107108

108-
public void testOf_returnsDefault_withUnknownField_andPersistentContext() {
109+
public void testOf_returnsEmptyThinkingConfig_withUnknownField_andPersistentContext() {
109110
ValidationException exception = new ValidationException();
110111
int anInt = 42;
111112
Map<String, Object> settings = new HashMap<>(Map.of(THINKING_CONFIG_FIELD, new HashMap<>(Map.of("not_thinking_budget", anInt))));
112-
ThinkingConfig defaultValue = new ThinkingConfig(123);
113113

114-
ThinkingConfig result = ThinkingConfig.of(settings, defaultValue, exception, "test", ConfigurationParseContext.PERSISTENT);
114+
ThinkingConfig result = ThinkingConfig.of(settings, exception, "test", ConfigurationParseContext.PERSISTENT);
115115

116-
assertThat(result, is(defaultValue));
116+
assertThat(result.getThinkingBudget(), is(nullValue()));
117+
assertThat(result.isEmpty(), is(true));
117118
assertThat(exception.validationErrors(), is(empty()));
118119
}
119120

@@ -126,8 +127,14 @@ public void testToXContent() throws IOException {
126127
builder.endObject();
127128
String xContentResult = Strings.toString(builder);
128129

129-
assertThat(xContentResult, is("""
130-
{"thinking_config":{"thinking_budget":256}}"""));
130+
String expected = XContentHelper.stripWhitespace("""
131+
{
132+
"thinking_config": {
133+
"thinking_budget": 256
134+
}
135+
}
136+
""");
137+
assertThat(xContentResult, is(expected));
131138
}
132139

133140
@Override
@@ -144,4 +151,9 @@ protected ThinkingConfig createTestInstance() {
144151
protected ThinkingConfig mutateInstance(ThinkingConfig instance) throws IOException {
145152
return randomValueOtherThan(instance, () -> new ThinkingConfig(randomInt()));
146153
}
154+
155+
@Override
156+
protected ThinkingConfig mutateInstanceForVersion(ThinkingConfig instance, TransportVersion version) {
157+
return instance;
158+
}
147159
}

‎x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googlevertexai/request/GoogleVertexAiUnifiedChatCompletionRequestEntityTests.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,8 @@ public class GoogleVertexAiUnifiedChatCompletionRequestEntityTests extends ESTes
3131

3232
private static final String USER_ROLE = "user";
3333
private static final String ASSISTANT_ROLE = "assistant";
34-
private final ThinkingConfig thinkingConfig = new ThinkingConfig(256);
35-
private final ThinkingConfig emptyThinkingConfig = new ThinkingConfig();
34+
private static final ThinkingConfig thinkingConfig = new ThinkingConfig(256);
35+
private static final ThinkingConfig emptyThinkingConfig = new ThinkingConfig();
3636

3737
public void testBasicSerialization_SingleMessage() throws IOException {
3838
UnifiedCompletionRequest.Message message = new UnifiedCompletionRequest.Message(

0 commit comments

Comments
 (0)