chore(vertexai): clean up chat example

vertexai/chat/chat.go

@@ -24,11 +24,8 @@ import (
 	"cloud.google.com/go/vertexai/genai"
 )
 
-func makeChatRequests(w io.Writer, projectID string, location string, modelName string) error {
-	// location := "us-central1"
-	// modelName := "gemini-1.0-pro-002"
-	ctx := context.Background()
-	client, err := genai.NewClient(ctx, projectID, location)
+func makeChatRequests(ctx context.Context, w io.Writer, projectID, region, modelName string) error {
+	client, err := genai.NewClient(ctx, projectID, region)
 	if err != nil {
 		return fmt.Errorf("error creating client: %w", err)
 	}
@@ -37,43 +34,26 @@ func makeChatRequests(w io.Writer, projectID string, location string, modelName
 	gemini := client.GenerativeModel(modelName)
 	chat := gemini.StartChat()
 
-	r, err := chat.SendMessage(
-		ctx,
-		genai.Text("Hello"))
-	if err != nil {
-		return err
-	}
-	rb, err := json.MarshalIndent(r, "", "  ")
-	if err != nil {
-		return fmt.Errorf("json.MarshalIndent: %w", err)
+	send := func(message string) error {
+		r, err := chat.SendMessage(ctx, genai.Text(message))
+		if err != nil {
+			return err
+		}
+		rb, err := json.MarshalIndent(r, "", "  ")
+		if err != nil {
+			return err
+		}
+		fmt.Fprintln(w, string(rb))
+		return nil
 	}
-	fmt.Fprintln(w, string(rb))
 
-	r, err = chat.SendMessage(
-		ctx,
-		genai.Text("What are all the colors in a rainbow?"))
-	if err != nil {
+	if err := send("Hello"); err != nil {
 		return err
 	}
-	rb, err = json.MarshalIndent(r, "", "  ")
-	if err != nil {
-		return fmt.Errorf("json.MarshalIndent: %w", err)
-	}
-	fmt.Fprintln(w, string(rb))
-
-	r, err = chat.SendMessage(
-		ctx,
-		genai.Text("Why does it appear when it rains?"))
-	if err != nil {
-		return fmt.Errorf("chat.SendMessage: %w", err)
-	}
-	rb, err = json.MarshalIndent(r, "", "  ")
-	if err != nil {
-		return fmt.Errorf("json.MarshalIndent: %w", err)
+	if err := send("What are all the colors in a rainbow?"); err != nil {
+		return err
 	}
-	fmt.Fprintln(w, string(rb))
-
-	return nil
+	return send("Why does it appear when it rains?")
 }
 
 // [END aiplatform_gemini_multiturn_chat]

vertexai/chat/chat_test.go

@@ -15,16 +15,17 @@
 package chat
 
 import (
-	"io"
+	"bytes"
+	"context"
 	"testing"
 
 	"github.com/GoogleCloudPlatform/golang-samples/internal/testutil"
 )
 
 func Test_makeChatRequests(t *testing.T) {
-	w := io.Discard
 	tc := testutil.SystemTest(t)
-	err := makeChatRequests(w, tc.ProjectID, "us-central1", "gemini-1.0-pro-002")
+	w := &bytes.Buffer{}
+	err := makeChatRequests(context.Background(), w, tc.ProjectID, "us-central1", "gemini-pro-vision")
 	if err != nil {
 		t.Errorf("unexpected error: %v", err.Error())
 	}