From f5d56eb03558fce093a5b9947ae041fba4d844b2 Mon Sep 17 00:00:00 2001
From: Eli Bendersky <eliben@gmail.com>
Date: Thu, 8 Feb 2024 12:36:16 -0800
Subject: [PATCH] feat(vertexai): add WithREST option to vertexai client
 (#9389)

Initially configure whether the client uses REST or gRPC
---
 vertexai/genai/client.go      | 16 ++++++++--
 vertexai/genai/client_test.go | 21 +++++++++++++
 vertexai/genai/option.go      | 57 +++++++++++++++++++++++++++++++++++
 3 files changed, 91 insertions(+), 3 deletions(-)
 create mode 100644 vertexai/genai/option.go

diff --git a/vertexai/genai/client.go b/vertexai/genai/client.go
index cfbdef56c9f..5be3fc96ed9 100644
--- a/vertexai/genai/client.go
+++ b/vertexai/genai/client.go
@@ -46,16 +46,26 @@ type Client struct {
 // Clients should be reused instead of created as needed. The methods of Client
 // are safe for concurrent use by multiple goroutines.
 //
-// You may configure the client by passing in options from the [google.golang.org/api/option]
-// package.
+// You may configure the client by passing in options from the
+// [google.golang.org/api/option] package. You may also use options defined in
+// this package, such as [WithREST].
 func NewClient(ctx context.Context, projectID, location string, opts ...option.ClientOption) (*Client, error) {
 	opts = append([]option.ClientOption{
 		option.WithEndpoint(fmt.Sprintf("%s-aiplatform.googleapis.com:443", location)),
 	}, opts...)
-	c, err := aiplatform.NewPredictionClient(ctx, opts...)
+	conf := newConfig(opts...)
+
+	var c *aiplatform.PredictionClient
+	var err error
+	if conf.withREST {
+		c, err = aiplatform.NewPredictionRESTClient(ctx, opts...)
+	} else {
+		c, err = aiplatform.NewPredictionClient(ctx, opts...)
+	}
 	if err != nil {
 		return nil, err
 	}
+
 	c.SetGoogleClientInfo("gccl", internal.Version)
 	return &Client{
 		c:         c,
diff --git a/vertexai/genai/client_test.go b/vertexai/genai/client_test.go
index 319e92d4c98..e2e2ccded55 100644
--- a/vertexai/genai/client_test.go
+++ b/vertexai/genai/client_test.go
@@ -263,6 +263,27 @@ func TestLive(t *testing.T) {
 	})
 }
 
+func TestLiveREST(t *testing.T) {
+	if *projectID == "" || *modelName == "" {
+		t.Skip("need -project and -model")
+	}
+	ctx := context.Background()
+	client, err := NewClient(ctx, *projectID, "us-central1", WithREST())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer client.Close()
+	model := client.GenerativeModel(*modelName)
+	model.SetTemperature(0.0)
+
+	resp, err := model.GenerateContent(ctx, Text("What is the average size of a swallow?"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	got := responseString(resp)
+	checkMatch(t, got, `15.* cm|[1-9].* inches`)
+}
+
 func TestJoinResponses(t *testing.T) {
 	r1 := &GenerateContentResponse{
 		Candidates: []*Candidate{
diff --git a/vertexai/genai/option.go b/vertexai/genai/option.go
new file mode 100644
index 00000000000..2f98fcf310c
--- /dev/null
+++ b/vertexai/genai/option.go
@@ -0,0 +1,57 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package genai
+
+import (
+	"google.golang.org/api/option"
+	"google.golang.org/api/option/internaloption"
+)
+
+// WithREST is an option that enables REST transport for the client.
+// The default transport (if this option isn't provided) is gRPC.
+func WithREST() option.ClientOption {
+	return &withREST{}
+}
+
+func (w *withREST) applyVertexaiOpt(c *config) {
+	c.withREST = true
+}
+
+type config struct {
+	// withREST tells the client to use REST as the underlying transport.
+	withREST bool
+}
+
+// newConfig generates a new config with all the given
+// vertexaiClientOptions applied.
+func newConfig(opts ...option.ClientOption) config {
+	var conf config
+	for _, opt := range opts {
+		if vOpt, ok := opt.(vertexaiClientOption); ok {
+			vOpt.applyVertexaiOpt(&conf)
+		}
+	}
+	return conf
+}
+
+// A vertexaiClientOption is an option for a vertexai client.
+type vertexaiClientOption interface {
+	option.ClientOption
+	applyVertexaiOpt(*config)
+}
+
+type withREST struct {
+	internaloption.EmbeddableAdapter
+}