Add Google Model Garden's Anthropic support to Inference Plugin (#134080)

* Add Google Model Garden Anthropic integration

* Clean up AnthropicChatCompletionStreamingProcessor

* Enhance GoogleVertexAiChatCompletionServiceSettings to support optional parameters based on transport version

* Add extractOptionalUri method and corresponding tests for URI extraction

* Add GoogleModelGardenProvider support to chat completion models and tests

* Enhance AnthropicChatCompletionStreamingProcessor and related classes to support new content block types and improve parsing logic

* Refactor AnthropicChatCompletionResponseHandler to use a custom error parser and add unit tests for response validation

* Add unit tests for AnthropicChatCompletionStreamingProcessor to validate response parsing and error handling

* Add unit tests for GoogleModelGardenAnthropicChatCompletionRequestEntity to validate serialization of user fields

* Add support for Anthropic provider in Google Vertex AI chat completion model and update related tests

* Add changelog

* Refactor switch case in GoogleVertexAiActionCreator to handle null case

* Validate service settings for Google Vertex AI model configuration

* Enhance Anthropic model tests to validate URI handling and provider requirements

* [CI] Auto commit changes from spotless

* Refactor switch case in GoogleVertexAiService to handle null case

* Simplify version check in GoogleVertexAiChatCompletionServiceSettings

* Make GOOGLE provider default for GoogleModelGarden integration

* Update anthropic_version to vertex-2024-10-22 in request entity and tests

* Refactor Google Vertex AI request handling to improve provider management and error handling

* Enhance validation for Google Model Garden settings to ensure required parameters are provided

* Remove uri streamingUri and provider from rate limit grouping hash calculation

* Refactor null and empty checks for projectId, location, and modelId in GoogleVertexAiChatCompletionServiceSettings

* Refactor Google Model Garden integration to include task settings in request entity and enhance validation for max tokens

* Revert "Update anthropic_version to vertex-2024-10-22 in request entity and tests"

This reverts commit 63ea4b8.

* Refactor Google Vertex AI settings to utilize GoogleVertexAiUtils for model garden support checks

* [CI] Update transport version definitions

* Update anthropic_version in tests and enhance validation logic for Google Vertex AI settings

* Update versions

* Enhance task settings validation in GoogleVertexAiChatCompletionModel

* Address comments regarding anthropic version and configuration

* [CI] Update transport version definitions

* Add nullable annotation for maxTokens parameter in GoogleVertexAiChatCompletionTaskSettings

* [CI] Update transport version definitions

* Clarify URI handling logic in GoogleVertexAiChatCompletionModel comments

* Make maxTokens nullable

* [CI] Update transport version definitions

* Fixed unit tests

* [CI] Update transport version definitions

* Fix validation logic for Google Model Garden and Vertex AI settings

* [CI] Update transport version definitions

* Add validation tests for Google Vertex AI and Model Garden settings

* Refactor validation logic for Google Vertex AI and Model Garden settings

* Add comment

* Update Google Vertex AI Task Settings parsing logic and AnthropicChatCompletionStreamingProcessor readability

---------

Co-authored-by: elasticsearchmachine <[email protected]>
Co-authored-by: Jonathan Buttner <[email protected]>

Author: 3 people

docs/changelog/134080.yaml

@@ -0,0 +1,5 @@
+pr: 134080
+summary: Added Google Model Garden Anthropic Completion and Chat Completion support to the Inference Plugin
+area: Machine Learning
+type: enhancement
+issues: []

server/src/main/java/org/elasticsearch/inference/UnifiedCompletionRequest.java

@@ -58,11 +58,11 @@ public record UnifiedCompletionRequest(
     private static final String ROLE_FIELD = "role";
     private static final String CONTENT_FIELD = "content";
     private static final String STOP_FIELD = "stop";
-    private static final String TEMPERATURE_FIELD = "temperature";
-    private static final String TOOL_CHOICE_FIELD = "tool_choice";
-    private static final String TOOL_FIELD = "tools";
+    public static final String TEMPERATURE_FIELD = "temperature";
+    public static final String TOOL_CHOICE_FIELD = "tool_choice";
+    public static final String TOOL_FIELD = "tools";
     private static final String TEXT_FIELD = "text";
-    private static final String TYPE_FIELD = "type";
+    public static final String TYPE_FIELD = "type";
     private static final String MODEL_FIELD = "model";
     private static final String MAX_COMPLETION_TOKENS_FIELD = "max_completion_tokens";
     private static final String MAX_TOKENS_FIELD = "max_tokens";

server/src/main/resources/transport/definitions/referable/ml_inference_google_model_garden_added.csv

@@ -0,0 +1 @@
+9179000

server/src/main/resources/transport/upper_bounds/9.2.csv

@@ -1 +1 @@
-no_matching_project_exception,9178000
+ml_inference_google_model_garden_added,9179000

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/ServiceUtils.java

@@ -314,6 +314,18 @@ public static URI extractUri(Map<String, Object> map, String fieldName, Validati
         return convertToUri(parsedUrl, fieldName, ModelConfigurations.SERVICE_SETTINGS, validationException);
     }
 
+    /**
+     * Extracts an optional URI from the map. If the field is not present, null is returned. If the field is present but invalid,
+     * @param map the map to extract the URI from
+     * @param fieldName the field name to extract
+     * @param validationException the validation exception to add errors to
+     * @return the extracted URI or null if not present
+     */
+    public static URI extractOptionalUri(Map<String, Object> map, String fieldName, ValidationException validationException) {
+        String parsedUrl = extractOptionalString(map, fieldName, ModelConfigurations.SERVICE_SETTINGS, validationException);
+        return convertToUri(parsedUrl, fieldName, ModelConfigurations.SERVICE_SETTINGS, validationException);
+    }
+
     public static URI convertToUri(@Nullable String url, String settingName, String settingScope, ValidationException validationException) {
         try {
             return createOptionalUri(url);

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/anthropic/AnthropicChatCompletionResponseHandler.java

@@ -0,0 +1,60 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.services.anthropic;
+
+import org.elasticsearch.inference.InferenceServiceResults;
+import org.elasticsearch.xpack.core.inference.results.StreamingUnifiedChatCompletionResults;
+import org.elasticsearch.xpack.core.inference.results.UnifiedChatCompletionException;
+import org.elasticsearch.xpack.inference.external.http.HttpResult;
+import org.elasticsearch.xpack.inference.external.http.retry.ChatCompletionErrorResponseHandler;
+import org.elasticsearch.xpack.inference.external.http.retry.ResponseParser;
+import org.elasticsearch.xpack.inference.external.http.retry.UnifiedChatCompletionErrorParserContract;
+import org.elasticsearch.xpack.inference.external.http.retry.UnifiedChatCompletionErrorResponseUtils;
+import org.elasticsearch.xpack.inference.external.request.Request;
+import org.elasticsearch.xpack.inference.external.response.streaming.ServerSentEventParser;
+import org.elasticsearch.xpack.inference.external.response.streaming.ServerSentEventProcessor;
+import org.elasticsearch.xpack.inference.services.anthropic.response.AnthropicChatCompletionResponseEntity;
+
+import java.util.concurrent.Flow;
+
+/**
+ * Handles streaming chat completion responses and error parsing for Anthropic inference endpoints.
+ * Adapts the AnthropicResponseHandler to support chat completion schema.
+ */
+public class AnthropicChatCompletionResponseHandler extends AnthropicResponseHandler {
+    private static final String ANTHROPIC_ERROR = "anthropic_error";
+    private static final UnifiedChatCompletionErrorParserContract ANTHROPIC_ERROR_PARSER = UnifiedChatCompletionErrorResponseUtils
+        .createErrorParserWithStringify(ANTHROPIC_ERROR);
+
+    private final ChatCompletionErrorResponseHandler chatCompletionErrorResponseHandler;
+
+    public AnthropicChatCompletionResponseHandler(String requestType) {
+        this(requestType, AnthropicChatCompletionResponseEntity::fromResponse);
+    }
+
+    private AnthropicChatCompletionResponseHandler(String requestType, ResponseParser parseFunction) {
+        super(requestType, parseFunction, true);
+        this.chatCompletionErrorResponseHandler = new ChatCompletionErrorResponseHandler(ANTHROPIC_ERROR_PARSER);
+    }
+
+    @Override
+    public InferenceServiceResults parseResult(Request request, Flow.Publisher<HttpResult> flow) {
+        var serverSentEventProcessor = new ServerSentEventProcessor(new ServerSentEventParser());
+        var anthropicProcessor = new AnthropicChatCompletionStreamingProcessor(
+            (m, e) -> chatCompletionErrorResponseHandler.buildMidStreamChatCompletionError(request.getInferenceEntityId(), m, e)
+        );
+        flow.subscribe(serverSentEventProcessor);
+        serverSentEventProcessor.subscribe(anthropicProcessor);
+        return new StreamingUnifiedChatCompletionResults(anthropicProcessor);
+    }
+
+    @Override
+    protected UnifiedChatCompletionException buildError(String message, Request request, HttpResult result) {
+        return chatCompletionErrorResponseHandler.buildChatCompletionError(message, request, result);
+    }
+}