Skip to content

Implemented completion task for Google VertexAI #128694

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 14 commits into from
Jun 9, 2025
Merged

Implemented completion task for Google VertexAI #128694

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
a9df8e3
Google Vertex AI completion model, response entity and tests
leo-hoet May 30, 2025
80af6d3
Fixed GoogleVertexAiServiceTest for Service configuration
leo-hoet May 30, 2025
ca1b6d5
Changelog
leo-hoet May 30, 2025
b6f5e34
Removed downcasting and using `moveToFirstToken`
leo-hoet Jun 3, 2025
ce6d45f
Create GoogleVertexAiChatCompletionResponseHandler for streaming and …
leo-hoet Jun 3, 2025
6cf0c0b
Added unit tests
leo-hoet Jun 4, 2025
7eabd29
PR feedback
leo-hoet Jun 4, 2025
55d8650
Removed googlevertexaicompletion model. Using just GoogleVertexAiChat…
leo-hoet Jun 4, 2025
bf27166
Renamed uri -> nonStreamingUri. Added streamingUri and getters in Goo…
leo-hoet Jun 4, 2025
ab1fe7a
Moved rateLimitGroupHashing to subclasses of GoogleVertexAiModel
leo-hoet Jun 4, 2025
20f1914
Merge branch 'main' into google-vertexai-completion
lhoet-google Jun 5, 2025
30f53cd
Fixed rate limit has of GoogleVertexAiRerankModel and refactored uri …
leo-hoet Jun 6, 2025
5d20b35
Merge branch 'google-vertexai-completion' of github.com:leo-hoet/elas…
leo-hoet Jun 6, 2025
2b865af
Merge branch 'main' into google-vertexai-completion
jonathan-buttner Jun 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 0 additions & 22 deletions .../xpack/inference/services/googlevertexai/GoogleVertexAiChatCompletionResponseHandler.java
View file
Open in desktop
Copy link
Contributor

@lhoet-google lhoet-google Jun 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not sure if I should keep this class or delete it and use the base class instead.

Also, I am using GoogleVertexAiUnifiedChatCompletionResponseHandler.GoogleVertexAiErrorResponse::fromResponse. I preferred to use it that way and avoid putting that in a common class between GoogleVertexAiUnifiedChatCompletionResponseHandler and GoogleVertexAiChatCompletionResponseHandler to avoid extending the class hierarchy, but let me know if you think otherwise

Copy link
Contributor

@jonathan-buttner jonathan-buttner Jun 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not sure if I should keep this class or delete it and use the base class instead.

Up to you, I'd probably use the base class. If you want to keep this one, then I don't think we need to accept the requestType. I think we can set it in this class directly.

Also, I am using GoogleVertexAiUnifiedChatCompletionResponseHandler.GoogleVertexAiErrorResponse::fromResponse. I preferred to use it that way and avoid putting that in a common class between GoogleVertexAiUnifiedChatCompletionResponseHandler and GoogleVertexAiChatCompletionResponseHandler to avoid extending the class hierarchy, but let me know if you think otherwise

Nice! That looks good.

This file was deleted.

19 changes: 5 additions & 14 deletions ...n/java/org/elasticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiModel.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ public abstract class GoogleVertexAiModel extends RateLimitGroupingModel {

private final GoogleVertexAiRateLimitServiceSettings rateLimitServiceSettings;

protected URI uri;
protected URI nonStreamingUri;

public GoogleVertexAiModel(
ModelConfigurations configurations,
Expand All @@ -39,14 +39,14 @@ public GoogleVertexAiModel(
public GoogleVertexAiModel(GoogleVertexAiModel model, ServiceSettings serviceSettings) {
super(model, serviceSettings);

uri = model.uri();
nonStreamingUri = model.nonStreamingUri();
rateLimitServiceSettings = model.rateLimitServiceSettings();
}

public GoogleVertexAiModel(GoogleVertexAiModel model, TaskSettings taskSettings) {
super(model, taskSettings);

uri = model.uri();
nonStreamingUri = model.nonStreamingUri();
rateLimitServiceSettings = model.rateLimitServiceSettings();
}

Expand All @@ -56,17 +56,8 @@ public GoogleVertexAiRateLimitServiceSettings rateLimitServiceSettings() {
return rateLimitServiceSettings;
}

public URI uri() {
return uri;
}

@Override
public int rateLimitGroupingHash() {
// In VertexAI rate limiting is scoped to the project, region and model. URI already has this information so we are using that.
// API Key does not affect the quota
// https://ai.google.dev/gemini-api/docs/rate-limits
// https://cloud.google.com/vertex-ai/docs/quotas
return Objects.hash(uri);
public URI nonStreamingUri() {
return nonStreamingUri;
}

@Override
Expand Down
15 changes: 2 additions & 13 deletions ...java/org/elasticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,6 @@
import org.elasticsearch.xpack.inference.services.ServiceUtils;
import org.elasticsearch.xpack.inference.services.googlevertexai.action.GoogleVertexAiActionCreator;
import org.elasticsearch.xpack.inference.services.googlevertexai.completion.GoogleVertexAiChatCompletionModel;
import org.elasticsearch.xpack.inference.services.googlevertexai.completion.GoogleVertexAiCompletionModel;
import org.elasticsearch.xpack.inference.services.googlevertexai.embeddings.GoogleVertexAiEmbeddingsModel;
import org.elasticsearch.xpack.inference.services.googlevertexai.embeddings.GoogleVertexAiEmbeddingsServiceSettings;
import org.elasticsearch.xpack.inference.services.googlevertexai.request.GoogleVertexAiUnifiedChatCompletionRequest;
Expand Down Expand Up @@ -89,7 +88,7 @@ public class GoogleVertexAiService extends SenderService {
InputType.INTERNAL_SEARCH
);

private final ResponseHandler COMPLETION_HANDLER = new GoogleVertexAiUnifiedChatCompletionResponseHandler(
public static final ResponseHandler COMPLETION_HANDLER = new GoogleVertexAiUnifiedChatCompletionResponseHandler(
"Google VertexAI chat completion"
);

Expand Down Expand Up @@ -360,17 +359,7 @@ private static GoogleVertexAiModel createModel(
context
);

case CHAT_COMPLETION -> new GoogleVertexAiChatCompletionModel(
inferenceEntityId,
taskType,
NAME,
serviceSettings,
taskSettings,
secretSettings,
context
);

case COMPLETION -> new GoogleVertexAiCompletionModel(
case CHAT_COMPLETION, COMPLETION -> new GoogleVertexAiChatCompletionModel(
inferenceEntityId,
taskType,
NAME,
Expand Down
9 changes: 2 additions & 7 deletions ...asticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiStreamingProcessor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
package org.elasticsearch.xpack.inference.services.googlevertexai;

import org.elasticsearch.ElasticsearchStatusException;
import org.elasticsearch.common.Strings;
import org.elasticsearch.common.xcontent.LoggingDeprecationHandler;
import org.elasticsearch.inference.InferenceServiceResults;
import org.elasticsearch.rest.RestStatus;
Expand All @@ -24,9 +25,6 @@
import java.util.Objects;
import java.util.stream.Stream;

import static org.elasticsearch.common.xcontent.XContentParserUtils.ensureExpectedToken;
import static org.elasticsearch.xpack.inference.external.response.XContentUtils.moveToFirstToken;

public class GoogleVertexAiStreamingProcessor extends DelegatingProcessor<Deque<ServerSentEvent>, InferenceServiceResults.Result> {

@Override
Expand All @@ -44,17 +42,14 @@ protected void next(Deque<ServerSentEvent> item) throws Exception {
public static Stream<StreamingChatCompletionResults.Result> parse(XContentParserConfiguration parserConfig, ServerSentEvent event) {
String data = event.data();
try (XContentParser jsonParser = XContentFactory.xContent(XContentType.JSON).createParser(parserConfig, data)) {
moveToFirstToken(jsonParser);
ensureExpectedToken(XContentParser.Token.START_OBJECT, jsonParser.currentToken(), jsonParser);

var chunk = GoogleVertexAiUnifiedStreamingProcessor.GoogleVertexAiChatCompletionChunkParser.parse(jsonParser);

return chunk.choices()
.stream()
.map(choice -> choice.delta())
.filter(Objects::nonNull)
.map(delta -> delta.content())
.filter(content -> content != null && content.isEmpty() == false)
.filter(content -> Strings.isNullOrEmpty(content) == false)
.map(StreamingChatCompletionResults.Result::new);

} catch (IOException e) {
Expand Down
40 changes: 10 additions & 30 deletions ...ticsearch/xpack/inference/services/googlevertexai/action/GoogleVertexAiActionCreator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,6 @@

package org.elasticsearch.xpack.inference.services.googlevertexai.action;

import org.elasticsearch.ElasticsearchStatusException;
import org.elasticsearch.rest.RestStatus;
import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
import org.elasticsearch.xpack.inference.external.action.SenderExecutableAction;
import org.elasticsearch.xpack.inference.external.action.SingleInputSenderExecutableAction;
Expand All @@ -18,17 +16,16 @@
import org.elasticsearch.xpack.inference.external.http.sender.Sender;
import org.elasticsearch.xpack.inference.external.http.sender.UnifiedChatInput;
import org.elasticsearch.xpack.inference.services.ServiceComponents;
import org.elasticsearch.xpack.inference.services.googlevertexai.GoogleVertexAiChatCompletionResponseHandler;
import org.elasticsearch.xpack.inference.services.googlevertexai.GoogleVertexAiEmbeddingsRequestManager;
import org.elasticsearch.xpack.inference.services.googlevertexai.GoogleVertexAiRerankRequestManager;
import org.elasticsearch.xpack.inference.services.googlevertexai.GoogleVertexAiResponseHandler;
import org.elasticsearch.xpack.inference.services.googlevertexai.GoogleVertexAiUnifiedChatCompletionResponseHandler;
import org.elasticsearch.xpack.inference.services.googlevertexai.completion.GoogleVertexAiChatCompletionModel;
import org.elasticsearch.xpack.inference.services.googlevertexai.completion.GoogleVertexAiCompletionModel;
import org.elasticsearch.xpack.inference.services.googlevertexai.embeddings.GoogleVertexAiEmbeddingsModel;
import org.elasticsearch.xpack.inference.services.googlevertexai.request.GoogleVertexAiUnifiedChatCompletionRequest;
import org.elasticsearch.xpack.inference.services.googlevertexai.rerank.GoogleVertexAiRerankModel;
import org.elasticsearch.xpack.inference.services.googlevertexai.response.GoogleVertexAiCompletionResponseEntity;

import java.net.URISyntaxException;
import java.util.Map;
import java.util.Objects;

Expand All @@ -41,10 +38,13 @@ public class GoogleVertexAiActionCreator implements GoogleVertexAiActionVisitor

private final ServiceComponents serviceComponents;

static final ResponseHandler UNIFIED_CHAT_COMPLETION_HANDLER = new GoogleVertexAiUnifiedChatCompletionResponseHandler(
"Google VertexAI chat completion"
static final ResponseHandler CHAT_COMPLETION_HANDLER = new GoogleVertexAiResponseHandler(
"Google VertexAI completion",
GoogleVertexAiCompletionResponseEntity::fromResponse,
GoogleVertexAiUnifiedChatCompletionResponseHandler.GoogleVertexAiErrorResponse::fromResponse,
true
);
static final ResponseHandler CHAT_COMPLETION_HANDLER = new GoogleVertexAiChatCompletionResponseHandler("Google VertexAI completion");

static final String USER_ROLE = "user";

public GoogleVertexAiActionCreator(Sender sender, ServiceComponents serviceComponents) {
Expand Down Expand Up @@ -73,36 +73,16 @@ public ExecutableAction create(GoogleVertexAiRerankModel model, Map<String, Obje

@Override
public ExecutableAction create(GoogleVertexAiChatCompletionModel model, Map<String, Object> taskSettings) {

var failedToSendRequestErrorMessage = constructFailedToSendRequestMessage(COMPLETION_ERROR_PREFIX);

var manager = new GenericRequestManager<>(
serviceComponents.threadPool(),
model,
UNIFIED_CHAT_COMPLETION_HANDLER,
CHAT_COMPLETION_HANDLER,
inputs -> new GoogleVertexAiUnifiedChatCompletionRequest(new UnifiedChatInput(inputs, USER_ROLE), model),
ChatCompletionInput.class
);

return new SingleInputSenderExecutableAction(sender, manager, failedToSendRequestErrorMessage, COMPLETION_ERROR_PREFIX);
}

@Override
public ExecutableAction create(GoogleVertexAiCompletionModel model, Map<String, Object> taskSettings) {
var failedToSendRequestErrorMessage = constructFailedToSendRequestMessage(COMPLETION_ERROR_PREFIX);

var manager = new GenericRequestManager<>(serviceComponents.threadPool(), model, CHAT_COMPLETION_HANDLER, inputs -> {
try {
model.updateUri(inputs.stream());
} catch (URISyntaxException e) {
throw new ElasticsearchStatusException(
"Error constructing URI for Google VertexAI completion",
RestStatus.INTERNAL_SERVER_ERROR,
e
);
}
return new GoogleVertexAiUnifiedChatCompletionRequest(new UnifiedChatInput(inputs, USER_ROLE), model);
}, ChatCompletionInput.class);

return new SingleInputSenderExecutableAction(sender, manager, failedToSendRequestErrorMessage, COMPLETION_ERROR_PREFIX);
}
}
2 changes: 0 additions & 2 deletions ...ticsearch/xpack/inference/services/googlevertexai/action/GoogleVertexAiActionVisitor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@

import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
import org.elasticsearch.xpack.inference.services.googlevertexai.completion.GoogleVertexAiChatCompletionModel;
import org.elasticsearch.xpack.inference.services.googlevertexai.completion.GoogleVertexAiCompletionModel;
import org.elasticsearch.xpack.inference.services.googlevertexai.embeddings.GoogleVertexAiEmbeddingsModel;
import org.elasticsearch.xpack.inference.services.googlevertexai.rerank.GoogleVertexAiRerankModel;

Expand All @@ -23,5 +22,4 @@ public interface GoogleVertexAiActionVisitor {

ExecutableAction create(GoogleVertexAiChatCompletionModel model, Map<String, Object> taskSettings);

ExecutableAction create(GoogleVertexAiCompletionModel model, Map<String, Object> taskSettings);
}
50 changes: 48 additions & 2 deletions ...xpack/inference/services/googlevertexai/completion/GoogleVertexAiChatCompletionModel.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,9 @@
import static org.elasticsearch.core.Strings.format;

public class GoogleVertexAiChatCompletionModel extends GoogleVertexAiModel {

private final URI streamingURI;

public GoogleVertexAiChatCompletionModel(
String inferenceEntityId,
TaskType taskType,
Expand Down Expand Up @@ -63,7 +66,8 @@ public GoogleVertexAiChatCompletionModel(
serviceSettings
);
try {
this.uri = buildUri(serviceSettings.location(), serviceSettings.projectId(), serviceSettings.modelId());
this.streamingURI = buildUriStreaming(serviceSettings.location(), serviceSettings.projectId(), serviceSettings.modelId());
this.nonStreamingUri = buildUriNonStreaming(serviceSettings.location(), serviceSettings.projectId(), serviceSettings.modelId());
} catch (URISyntaxException e) {
throw new RuntimeException(e);
}
Expand Down Expand Up @@ -114,7 +118,28 @@ public GoogleVertexAiSecretSettings getSecretSettings() {
return (GoogleVertexAiSecretSettings) super.getSecretSettings();
}

public static URI buildUri(String location, String projectId, String model) throws URISyntaxException {
public URI streamingURI() {
return this.streamingURI;
}

public static URI buildUriNonStreaming(String location, String projectId, String model) throws URISyntaxException {
return new URIBuilder().setScheme("https")
.setHost(format("%s%s", location, GoogleVertexAiUtils.GOOGLE_VERTEX_AI_HOST_SUFFIX))
.setPathSegments(
GoogleVertexAiUtils.V1,
GoogleVertexAiUtils.PROJECTS,
projectId,
GoogleVertexAiUtils.LOCATIONS,
GoogleVertexAiUtils.GLOBAL,
GoogleVertexAiUtils.PUBLISHERS,
GoogleVertexAiUtils.PUBLISHER_GOOGLE,
GoogleVertexAiUtils.MODELS,
format("%s:%s", model, GoogleVertexAiUtils.GENERATE_CONTENT)
)
.build();
}

public static URI buildUriStreaming(String location, String projectId, String model) throws URISyntaxException {
return new URIBuilder().setScheme("https")
.setHost(format("%s%s", location, GoogleVertexAiUtils.GOOGLE_VERTEX_AI_HOST_SUFFIX))
.setPathSegments(
Expand All @@ -131,4 +156,25 @@ public static URI buildUri(String location, String projectId, String model) thro
.setCustomQuery(GoogleVertexAiUtils.QUERY_PARAM_ALT_SSE)
.build();
}

@Override
public int rateLimitGroupingHash() {
// In VertexAI rate limiting is scoped to the project, region, model and endpoint.
// API Key does not affect the quota
// https://ai.google.dev/gemini-api/docs/rate-limits
// https://cloud.google.com/vertex-ai/docs/quotas
var projectId = getServiceSettings().projectId();
var location = getServiceSettings().location();
var modelId = getServiceSettings().modelId();

// Since we don't beforehand know which API is going to be used, we take a conservative approach and
// count both endpoint for the rate limit
return Objects.hash(
projectId,
location,
modelId,
GoogleVertexAiUtils.GENERATE_CONTENT,
GoogleVertexAiUtils.STREAM_GENERATE_CONTENT
);
}
}
79 changes: 0 additions & 79 deletions ...rch/xpack/inference/services/googlevertexai/completion/GoogleVertexAiCompletionModel.java
View file
Open in desktop

This file was deleted.

Loading
Loading