Add Google Model Garden's Meta, Mistral, Hugging Face and Ai21 providers support to Inference Plugin #135701
+624
−347
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
…n requests - Refactor Mistral, Ai21, Llama, and Hugging Face request entities to accept model IDs. - Update GoogleVertexAiActionCreator to handle multiple providers including META, HUGGING_FACE, MISTRAL, and AI21. - Enhance serialization tests for model ID handling in chat completion requests. - Introduce new response handlers for each Google Model Garden provider.
elasticsearchmachine
added
v9.2.0
needs:triage
…garden-openai-providers-integration
…garden-openai-providers-integration
…ai-providers-integration' into feature/google-model-garden-openai-providers-integration
|
Hugging Face (Dedicated Endpoint) Create Completion EndpointPerform Non-Streaming CompletionPerform Streaming CompletionCreate Chat Completion EndpointPerform Basic Chat CompletionPerform Complex Chat Completion |
|
Hugging Face (Shared endpoint) Create Completion EndpointPerform Non-Streaming CompletionPerform Streaming CompletionCreate Chat Completion EndpointPerform Basic Chat CompletionPerform Complex Chat Completion |
|
Mistral (Serverless) Tested with mistral-small-2503 Create Completion EndpointPerform Non-Streaming CompletionPerform Streaming CompletionCreate Chat Completion EndpointPerform Basic Chat CompletionPerform Complex Chat Completion |
|
Mistral (Dedicated endpoint) Create Completion EndpointPerform Non-Streaming CompletionPerform Streaming CompletionCreate Chat Completion EndpointPerform Basic Chat CompletionPerform Complex Chat Completion |
|
Ai21 (Serverless) Create Completion EndpointPerform Non-Streaming CompletionPerform Streaming CompletionCreate Chat Completion EndpointPerform Basic Chat CompletionPerform Complex Chat Completion |
|
Testing of the providers is finished. Other configurations such as:
cannot be performed because endpoints do not return successful results. |
|
@jonathan-buttner @DonalEvans @dan-rubinstein |
elasticsearchmachine
added
Team:ML
|
Pinging @elastic/ml-core (Team:ML) |
2 tasks
DonalEvans
reviewed
Oct 7, 2025
|
|
||
| /** | ||
| * Creates a {@link org.elasticsearch.xcontent.ToXContent.Params} that causes ToXContent to include the key values: | ||
| * - Key: {@link #MAX_TOKENS_FIELD}, Value: {@link #MAX_TOKENS_FIELD} |
There was a problem hiding this comment.
This should be Value: {@link #maxCompletionTokens()}. The existing Javadoc on line 113 in this file has a similar mistake.
There was a problem hiding this comment.
Fixed.
Comment on lines
42
to
48
| if (modelId != null) { | ||
| // Standard Hugging Face models require a model ID | ||
| unifiedRequestEntity.toXContent(builder, UnifiedCompletionRequest.withMaxTokens(modelId, params)); | ||
| } else { | ||
| // Some Hugging Face endpoints may not require a model ID | ||
| unifiedRequestEntity.toXContent(builder, UnifiedCompletionRequest.withMaxTokens(params)); | ||
| } |
There was a problem hiding this comment.
Would it be better to move the null check into the withMaxTokens(String modelId, Params params) method instead of having two different methods and requiring the caller know which one should be called?
There was a problem hiding this comment.
Yeah I agree, let's have withMaxTokens(@Nullable String modelId, Params params) and do the null check there.
Comment on lines
36
to
42
| if (modelId != null) { | ||
| // Some Mistral endpoints require the model ID to be specified in the request body | ||
| unifiedRequestEntity.toXContent(builder, UnifiedCompletionRequest.withMaxTokensAndSkipStreamOptionsField(modelId, params)); | ||
| } else { | ||
| // Some Mistral endpoints do not require the model ID to be specified | ||
| unifiedRequestEntity.toXContent(builder, UnifiedCompletionRequest.withMaxTokensAndSkipStreamOptionsField(params)); | ||
| } |
There was a problem hiding this comment.
Another place where we could move the null check into the method on UnifiedCompletionRequest instead of creating an extra method.
Comment on lines
148
to
154
| testExecute_ThrowsElasticsearchException( | ||
| null, | ||
| null, | ||
| null, | ||
| GoogleModelGardenProvider.META, | ||
| new URI("http://localhost:9200"), | ||
| GOOGLE_MODEL_GARDEN_META_COMPLETION_HANDLER |
There was a problem hiding this comment.
The tests in this class could be simplified a fair bit by removing the irrelevant parameters from the createAction() method signature. projectId, location, modelId and uri can all be null in every test without affecting the behaviour of the test.
There was a problem hiding this comment.
Agreed. Removed irrelevant params.
| GoogleModelGardenProvider googleModelGardenProvider, | ||
| GoogleModelGardenProvider provider, | ||
| URI uri, | ||
| ResponseHandler googleModelGardenAnthropicCompletionHandler |
There was a problem hiding this comment.
This name is a bit misleading, since in most of the tests, it's not specifically an Anthropic handler that's being passed in. This comment also applies to the similarly named parameter in testExecute_ThrowsElasticsearchException_WhenSenderOnFailureIsCalled().
There was a problem hiding this comment.
Missed that. Renamed.
jonathan-buttner
requested changes
Oct 7, 2025
|
|
||
| private GenericRequestManager<ChatCompletionInput> createRequestManager(GoogleVertexAiChatCompletionModel model) { | ||
| switch (model.getServiceSettings().provider()) { | ||
| case GOOGLE -> { |
There was a problem hiding this comment.
How about we try to reduce these switch statements to a map. I'm thinking we could create a Provider class or something that has a static map of provider enums to an internal class that contains functions to constructor the 3 different things that we need these switch cases for.
Ideally we wouldn't need to pass an instance of the Provider class around. If possible we could do something like Provider.createCompletionRequestManager(provider, model)
That would grab the internal class and then call `createRequestManagerWithHandler(model, handler) with the static handler that is appropriate for that provider.
Provider would have a function for each of the switch states that we need.
There was a problem hiding this comment.
Thanks. I had Ideas of moving this logic somewhere but initially decided not to.
Now I came up with an idea of moving the logic of getting response handlers and creating request entities to enum itself.
Of course I'm not a big fan of moving too much logic to enums, but having Provider class when we already have GoogleModelGardenProvider enum alongside is not ideal too IMHO. And that way when new provider is added - we can ideally localize the changes to a single file. To me - looks good. Please let me know what you think of it @jonathan-buttner
…garden-openai-providers-integration
…x unit tests, change remove optional stream_options field from llama requests
…ethod signatures and improve readability
|
@DonalEvans @jonathan-buttner |
DonalEvans
reviewed
Oct 8, 2025
| ); | ||
|
|
||
| private static final ResponseHandler AI21_CHAT_COMPLETION_HANDLER = new Ai21ChatCompletionResponseHandler( | ||
| "Google Model Garden Ai21 chat completions", |
There was a problem hiding this comment.
Minor nitpick, only fix this if there are other changes required, but this should be AI21, with an uppercase i
There was a problem hiding this comment.
but this should be AI21, with an uppercase i
Good catch. Confused it with Ai2 which has Ai with lowercase i. Fixed now.
| URI uri, | ||
| ResponseHandler handler | ||
| ) { | ||
| private ExecutableAction createAction(Sender sender, GoogleModelGardenProvider provider, URI uri, ResponseHandler handler) { |
There was a problem hiding this comment.
The uri method parameter is also irrelevant and can be removed.
There was a problem hiding this comment.
Oops, missed that. Removed now.
DonalEvans
approved these changes
Oct 8, 2025
jonathan-buttner
requested changes
Oct 8, 2025
| * @return the ResponseHandler associated with the provider | ||
| */ | ||
| public ResponseHandler getCompletionResponseHandler() { | ||
| return switch (this) { |
There was a problem hiding this comment.
Let's remove the switches. I think we can do this by having a package private constructor. Something like:
@FunctionalInterface
private interface RequestEntityCreator {
ToXContentObject create(UnifiedChatInput unifiedChatInput,
String modelId,
GoogleVertexAiChatCompletionTaskSettings taskSettings);
}
private final ResponseHandler completionResponseHandler;
private final ResponseHandler chatCompletionResponseHandler;
private final RequestEntityCreator entityCreator;
public enum GoogleModelGardenProvider {
GOOGLE(GOOGLE_VERTEX_AI_COMPLETION_HANDLER, GOOGLE_VERTEX_AI_CHAT_COMPLETION_HANDLER, (unifiedInput, modelId, taskSettings) -> new GoogleVertexAiUnifiedChatCompletionRequestEntity(unifiedChatInput, taskSettings.thinkingConfig()),
...
}
GoogleModelGardenProvider(ResponseHandler a, ResponseHandler b, RequestEntityCreator entityCreator) { ... }
Then the methods in here just return/call the appropriate methods and we won't need the switches anymore.
There was a problem hiding this comment.
Done. Had to move handlers into nested classes because you're not allowed read the value of a field before its definition and you cannot declare any fields before actual enum constants.
…ompletions response handler
…ai-providers-integration' into feature/google-model-garden-openai-providers-integration
|
@jonathan-buttner |
jonathan-buttner
approved these changes
Oct 9, 2025
|
@jonathan-buttner Thanks for the approval! You can merge it when you are ready! |
|
@elasticmachine merge upstream |
|
There are no new commits on the base branch. |
auto-merge was automatically disabled
October 9, 2025 22:09
Head branch was pushed to by a user without write access
|
@jonathan-buttner |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
META (Serverless):
Tested with: meta/llama-3.3-70b-instruct-maas
Create Completion Endpoint
Perform Non-Streaming Completion
Perform Streaming Completion
Create Chat Completion Endpoint
Perform Basic Chat Completion
Perform Complex Chat Completion
META (Dedicated Enpoint):
Tested with: meta/llama-3.3-70b-instruct-maas
Create Completion Endpoint
Perform Non-Streaming Completion
Perform Streaming Completion
Create Chat Completion Endpoint
Perform Basic Chat Completion
Perform Complex Chat Completion
Server Error is caused by the endpoint capabilities and not the integration itself.
gradle check?