Skip to content

Commit bf27166

Browse files
leo-hoetleo-hoet
committed
Renamed uri -> nonStreamingUri. Added streamingUri and getters in GoogleVertexAiChatCompletionModel
1 parent 55d8650 commit bf27166

File tree

8 files changed

+41
-21
lines changed

8 files changed

+41
-21
lines changed

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

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@ public abstract class GoogleVertexAiModel extends RateLimitGroupingModel {
2424

2525
private final GoogleVertexAiRateLimitServiceSettings rateLimitServiceSettings;
2626

27-
protected URI uri;
2827
protected URI nonStreamingUri;
2928

3029
public GoogleVertexAiModel(
@@ -40,14 +39,14 @@ public GoogleVertexAiModel(
4039
public GoogleVertexAiModel(GoogleVertexAiModel model, ServiceSettings serviceSettings) {
4140
super(model, serviceSettings);
4241

43-
uri = model.uri();
42+
nonStreamingUri = model.nonStreamingUri();
4443
rateLimitServiceSettings = model.rateLimitServiceSettings();
4544
}
4645

4746
public GoogleVertexAiModel(GoogleVertexAiModel model, TaskSettings taskSettings) {
4847
super(model, taskSettings);
4948

50-
uri = model.uri();
49+
nonStreamingUri = model.nonStreamingUri();
5150
rateLimitServiceSettings = model.rateLimitServiceSettings();
5251
}
5352

@@ -57,10 +56,6 @@ public GoogleVertexAiRateLimitServiceSettings rateLimitServiceSettings() {
5756
return rateLimitServiceSettings;
5857
}
5958

60-
public URI uri() {
61-
return uri;
62-
}
63-
6459
public URI nonStreamingUri() {
6560
return nonStreamingUri;
6661
}
@@ -71,7 +66,7 @@ public int rateLimitGroupingHash() {
7166
// API Key does not affect the quota
7267
// https://ai.google.dev/gemini-api/docs/rate-limits
7368
// https://cloud.google.com/vertex-ai/docs/quotas
74-
return Objects.hash(uri);
69+
return Objects.hash(nonStreamingUri);
7570
}
7671

7772
@Override

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

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,9 @@
3030
import static org.elasticsearch.core.Strings.format;
3131

3232
public class GoogleVertexAiChatCompletionModel extends GoogleVertexAiModel {
33+
34+
private final URI streamingURI;
35+
3336
public GoogleVertexAiChatCompletionModel(
3437
String inferenceEntityId,
3538
TaskType taskType,
@@ -63,7 +66,8 @@ public GoogleVertexAiChatCompletionModel(
6366
serviceSettings
6467
);
6568
try {
66-
this.uri = buildUri(serviceSettings.location(), serviceSettings.projectId(), serviceSettings.modelId());
69+
this.streamingURI = buildUriStreaming(serviceSettings.location(), serviceSettings.projectId(), serviceSettings.modelId());
70+
this.nonStreamingUri = buildUriNonStreaming(serviceSettings.location(), serviceSettings.projectId(), serviceSettings.modelId());
6771
} catch (URISyntaxException e) {
6872
throw new RuntimeException(e);
6973
}
@@ -114,7 +118,28 @@ public GoogleVertexAiSecretSettings getSecretSettings() {
114118
return (GoogleVertexAiSecretSettings) super.getSecretSettings();
115119
}
116120

117-
public static URI buildUri(String location, String projectId, String model) throws URISyntaxException {
121+
public URI streamingURI() {
122+
return this.streamingURI;
123+
}
124+
125+
public static URI buildUriNonStreaming(String location, String projectId, String model) throws URISyntaxException {
126+
return new URIBuilder().setScheme("https")
127+
.setHost(format("%s%s", location, GoogleVertexAiUtils.GOOGLE_VERTEX_AI_HOST_SUFFIX))
128+
.setPathSegments(
129+
GoogleVertexAiUtils.V1,
130+
GoogleVertexAiUtils.PROJECTS,
131+
projectId,
132+
GoogleVertexAiUtils.LOCATIONS,
133+
GoogleVertexAiUtils.GLOBAL,
134+
GoogleVertexAiUtils.PUBLISHERS,
135+
GoogleVertexAiUtils.PUBLISHER_GOOGLE,
136+
GoogleVertexAiUtils.MODELS,
137+
format("%s:%s", model, GoogleVertexAiUtils.GENERATE_CONTENT)
138+
)
139+
.build();
140+
}
141+
142+
public static URI buildUriStreaming(String location, String projectId, String model) throws URISyntaxException {
118143
return new URIBuilder().setScheme("https")
119144
.setHost(format("%s%s", location, GoogleVertexAiUtils.GOOGLE_VERTEX_AI_HOST_SUFFIX))
120145
.setPathSegments(

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ public GoogleVertexAiEmbeddingsModel(GoogleVertexAiEmbeddingsModel model, Google
8181
serviceSettings
8282
);
8383
try {
84-
this.uri = buildUri(serviceSettings.location(), serviceSettings.projectId(), serviceSettings.modelId());
84+
this.nonStreamingUri = buildUri(serviceSettings.location(), serviceSettings.projectId(), serviceSettings.modelId());
8585
} catch (URISyntaxException e) {
8686
throw new RuntimeException(e);
8787
}
@@ -103,7 +103,7 @@ protected GoogleVertexAiEmbeddingsModel(
103103
serviceSettings
104104
);
105105
try {
106-
this.uri = new URI(uri);
106+
this.nonStreamingUri = new URI(uri);
107107
} catch (URISyntaxException e) {
108108
throw new RuntimeException(e);
109109
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ public GoogleVertexAiEmbeddingsRequest(
4646

4747
@Override
4848
public HttpRequest createHttpRequest() {
49-
HttpPost httpPost = new HttpPost(model.uri());
49+
HttpPost httpPost = new HttpPost(model.nonStreamingUri());
5050

5151
ByteArrayEntity byteEntity = new ByteArrayEntity(
5252
Strings.toString(new GoogleVertexAiEmbeddingsRequestEntity(truncationResult.input(), inputType, model.getTaskSettings()))
@@ -84,7 +84,7 @@ public String getInferenceEntityId() {
8484

8585
@Override
8686
public URI getURI() {
87-
return model.uri();
87+
return model.nonStreamingUri();
8888
}
8989

9090
@Override

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ public GoogleVertexAiRerankRequest(
5050

5151
@Override
5252
public HttpRequest createHttpRequest() {
53-
HttpPost httpPost = new HttpPost(model.uri());
53+
HttpPost httpPost = new HttpPost(model.nonStreamingUri());
5454

5555
ByteArrayEntity byteEntity = new ByteArrayEntity(
5656
Strings.toString(
@@ -87,7 +87,7 @@ public String getInferenceEntityId() {
8787

8888
@Override
8989
public URI getURI() {
90-
return model.uri();
90+
return model.nonStreamingUri();
9191
}
9292

9393
@Override

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ public GoogleVertexAiUnifiedChatCompletionRequest(UnifiedChatInput unifiedChatIn
3333

3434
@Override
3535
public HttpRequest createHttpRequest() {
36-
var uri = unifiedChatInput.stream() ? model.uri() : model.nonStreamingUri();
36+
var uri = unifiedChatInput.stream() ? model.streamingURI() : model.nonStreamingUri();
3737
HttpPost httpPost = new HttpPost(uri);
3838

3939
var requestEntity = new GoogleVertexAiUnifiedChatCompletionRequestEntity(unifiedChatInput);
@@ -53,7 +53,7 @@ public void decorateWithAuth(HttpPost httpPost) {
5353

5454
@Override
5555
public URI getURI() {
56-
return model.uri();
56+
return model.nonStreamingUri();
5757
}
5858

5959
@Override

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ public GoogleVertexAiRerankModel(GoogleVertexAiRerankModel model, GoogleVertexAi
6565
serviceSettings
6666
);
6767
try {
68-
this.uri = buildUri(serviceSettings.projectId());
68+
this.nonStreamingUri = buildUri(serviceSettings.projectId());
6969
} catch (URISyntaxException e) {
7070
throw new RuntimeException(e);
7171
}
@@ -87,7 +87,7 @@ protected GoogleVertexAiRerankModel(
8787
serviceSettings
8888
);
8989
try {
90-
this.uri = new URI(uri);
90+
this.nonStreamingUri = new URI(uri);
9191
} catch (URISyntaxException e) {
9292
throw new RuntimeException(e);
9393
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,7 @@ public void testOverrideWith_DoesNotOverrideModelUri() {
8383
var model = createModel("model", Boolean.FALSE, InputType.SEARCH);
8484
var overriddenModel = GoogleVertexAiEmbeddingsModel.of(model, Map.of());
8585

86-
MatcherAssert.assertThat(overriddenModel.uri(), is(model.uri()));
86+
MatcherAssert.assertThat(overriddenModel.nonStreamingUri(), is(model.nonStreamingUri()));
8787
}
8888

8989
public static GoogleVertexAiEmbeddingsModel createModel(

0 commit comments

Comments
 (0)