anthropic vertex ai support

Signed-off-by: omar <[email protected]>

Author: puertomontt

api/v1alpha1/ai_backend.go

@@ -254,10 +254,13 @@ type GeminiConfig struct {
 	ApiVersion string `json:"apiVersion"`
 }
 
-// Publisher configures the type of publisher model to use for VertexAI. Currently, only Google is supported.
+// Publisher configures the type of publisher model to use for VertexAI. Google and Anthropic are supported.
 type Publisher string
 
-const GOOGLE Publisher = "GOOGLE"
+const (
+	GOOGLE    Publisher = "GOOGLE"
+	ANTHROPIC Publisher = "ANTHROPIC"
+)
 
 // VertexAIConfig settings for the [Vertex AI](https://cloud.google.com/vertex-ai/docs) LLM provider.
 // To find the values for the project ID, project location, and publisher, you can check the fields of an API request, such as
@@ -293,8 +296,8 @@ type VertexAIConfig struct {
 	// Optional: The model path to route to. Defaults to the Gemini model path, `generateContent`.
 	ModelPath *string `json:"modelPath,omitempty"`
 
-	// The type of publisher model to use. Currently, only Google is supported.
-	// +kubebuilder:validation:Enum=GOOGLE
+	// The type of publisher model to use. Google and Anthropic are supported.
+	// +kubebuilder:validation:Enum=GOOGLE;ANTHROPIC
 	Publisher Publisher `json:"publisher"`
 }
 

install/helm/kgateway-crds/templates/gateway.kgateway.dev_backends.yaml

@@ -542,10 +542,11 @@ spec:
                             minLength: 1
                             type: string
                           publisher:
-                            description: The type of publisher model to use. Currently,
-                              only Google is supported.
+                            description: The type of publisher model to use. Google
+                              and Anthropic are supported.
                             enum:
                             - GOOGLE
+                            - ANTHROPIC
                             type: string
                         required:
                         - apiVersion
@@ -1088,9 +1089,10 @@ spec:
                                     type: string
                                   publisher:
                                     description: The type of publisher model to use.
-                                      Currently, only Google is supported.
+                                      Google and Anthropic are supported.
                                     enum:
                                     - GOOGLE
+                                    - ANTHROPIC
                                     type: string
                                 required:
                                 - apiVersion

internal/kgateway/extensions2/plugins/backend/ai/ai_model_cluster.go

@@ -285,6 +285,8 @@ func buildVertexAIEndpoint(data *v1alpha1.VertexAIConfig, host *string, port *gw
 	switch data.Publisher {
 	case v1alpha1.GOOGLE:
 		publisher = "google"
+	case v1alpha1.ANTHROPIC:
+		publisher = "anthropic"
 	default:
 		// TODO(npolshak): add support for other publishers
 		slog.Warn("unsupported Vertex AI publisher, defaulting to Google", "publisher", string(data.Publisher))
@@ -401,6 +403,8 @@ func getTransformation(provider *v1alpha1.LLMProvider) (string, string, string,
 			switch provider.VertexAI.Publisher {
 			case v1alpha1.GOOGLE:
 				modelPath = getVertexAIGeminiModelPath()
+			case v1alpha1.ANTHROPIC:
+				modelPath = getVertexAIAnthropicModelPath()
 			default:
 				// TODO(npolshak): add support for other publishers
 				slog.Warn("unsupported Vertex AI publisher, defaulting to Google", "publisher", string(provider.VertexAI.Publisher))
@@ -439,6 +443,11 @@ func getVertexAIGeminiModelPath() string {
 	return `models/{{host_metadata("model")}}:{% if dynamic_metadata("route_type") == "CHAT_STREAMING" %}streamGenerateContent?alt=sse{% else %}generateContent{% endif %}`
 }
 
+func getVertexAIAnthropicModelPath() string {
+	return `models/{{host_metadata("model")}}:{% if dynamic_metadata("route_type") == "CHAT_STREAMING" %}streamGenerateContent{% else %}generateContent{% endif %}`
+	// return `models/{{host_metadata("model")}}:{% if dynamic_metadata("route_type") == "CHAT_STREAMING" %}streamRawPredict?alt=sse{% else %}rawPredict{% endif %}`
+}
+
 func defaultBodyTransformation() *envoytransformation.TransformationTemplate_MergeJsonKeys {
 	return &envoytransformation.TransformationTemplate_MergeJsonKeys{
 		MergeJsonKeys: &envoytransformation.MergeJsonKeys{

internal/kgateway/extensions2/plugins/backend/ai/ai_model_cluster_test.go

@@ -236,7 +236,7 @@ func TestProcessAIBackend_Gemini(t *testing.T) {
 	assert.Equal(t, "v1", filterMeta.Fields["api_version"].GetStringValue())
 }
 
-func TestProcessAIBackend_VertexAI(t *testing.T) {
+func TestProcessAIBackend_VertexAIGoogle(t *testing.T) {
 	cluster := &envoyclusterv3.Cluster{
 		Name: "vertex-ai-cluster",
 	}
@@ -290,6 +290,60 @@ func TestProcessAIBackend_VertexAI(t *testing.T) {
 	assert.Equal(t, "google", filterMeta.Fields["publisher"].GetStringValue())
 }
 
+func TestProcessAIBackend_VertexAIAnthropic(t *testing.T) {
+	cluster := &envoyclusterv3.Cluster{
+		Name: "vertex-ai-cluster",
+	}
+
+	aiBackend := &v1alpha1.AIBackend{
+		LLM: &v1alpha1.LLMProvider{
+			VertexAI: &v1alpha1.VertexAIConfig{
+				Model:      "claude-3-opus-20240229",
+				ApiVersion: "v1",
+				Location:   "us-central1",
+				ProjectId:  "my-project",
+				Publisher:  v1alpha1.ANTHROPIC,
+				AuthToken: v1alpha1.SingleAuthToken{
+					Kind:   v1alpha1.Inline,
+					Inline: ptr.To("vertex-token"),
+				},
+			},
+		},
+	}
+
+	secrets := &ir.Secret{}
+	multiSecrets := map[string]*ir.Secret{}
+
+	err := ProcessAIBackend(aiBackend, secrets, multiSecrets, cluster)
+
+	require.NoError(t, err)
+
+	// Verify endpoint
+	endpoints := cluster.LoadAssignment.Endpoints[0].LbEndpoints
+	require.Len(t, endpoints, 1)
+
+	endpoint := endpoints[0]
+	hostEndpoint := endpoint.GetEndpoint()
+	require.NotNil(t, hostEndpoint)
+
+	// Verify socket address
+	address := hostEndpoint.Address.GetSocketAddress()
+	require.NotNil(t, address)
+	assert.Equal(t, "us-central1-aiplatform.googleapis.com", address.Address)
+	assert.Equal(t, uint32(443), address.GetPortValue())
+
+	// Verify metadata
+	require.NotNil(t, endpoint.Metadata)
+	filterMeta := endpoint.Metadata.FilterMetadata["io.solo.transformation"]
+	require.NotNil(t, filterMeta)
+	assert.Equal(t, "vertex-token", filterMeta.Fields["auth_token"].GetStringValue())
+	assert.Equal(t, "claude-3-opus-20240229", filterMeta.Fields["model"].GetStringValue())
+	assert.Equal(t, "v1", filterMeta.Fields["api_version"].GetStringValue())
+	assert.Equal(t, "us-central1", filterMeta.Fields["location"].GetStringValue())
+	assert.Equal(t, "my-project", filterMeta.Fields["project"].GetStringValue())
+	assert.Equal(t, "anthropic", filterMeta.Fields["publisher"].GetStringValue())
+}
+
 func TestProcessAIBackend_CustomURL(t *testing.T) {
 	cluster := &envoyclusterv3.Cluster{
 		Name: "custom-host-cluster",