chore: switch controllerbuilder to use vertexai gemini endpoint

dev/tools/controllerbuilder/go.mod

@@ -5,52 +5,53 @@ go 1.22.0
 toolchain go1.23.2
 
 require (
+	cloud.google.com/go/vertexai v0.13.2
+	github.com/GoogleCloudPlatform/k8s-config-connector v1.125.0
 	github.com/fatih/color v1.17.0
-	github.com/google/generative-ai-go v0.17.0
 	github.com/spf13/cobra v1.8.0
-	golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d
-	google.golang.org/api v0.186.0
-	google.golang.org/protobuf v1.34.2
+	github.com/spf13/pflag v1.0.5
+	golang.org/x/tools v0.24.0
+	google.golang.org/api v0.203.0
+	google.golang.org/genproto/googleapis/api v0.0.0-20241007155032-5fefd90f89a9
+	google.golang.org/protobuf v1.35.1
 	k8s.io/apimachinery v0.27.11
 	k8s.io/klog/v2 v2.130.1
 )
 
 require (
-	cloud.google.com/go v0.115.0 // indirect
-	cloud.google.com/go/ai v0.8.0 // indirect
-	cloud.google.com/go/auth v0.6.0 // indirect
-	cloud.google.com/go/auth/oauth2adapt v0.2.2 // indirect
-	cloud.google.com/go/compute/metadata v0.3.0 // indirect
-	cloud.google.com/go/longrunning v0.5.7 // indirect
+	cloud.google.com/go v0.116.0 // indirect
+	cloud.google.com/go/aiplatform v1.68.0 // indirect
+	cloud.google.com/go/auth v0.9.9 // indirect
+	cloud.google.com/go/auth/oauth2adapt v0.2.4 // indirect
+	cloud.google.com/go/compute/metadata v0.5.2 // indirect
+	cloud.google.com/go/iam v1.2.1 // indirect
+	cloud.google.com/go/longrunning v0.6.1 // indirect
 	github.com/felixge/httpsnoop v1.0.4 // indirect
 	github.com/go-logr/logr v1.4.2 // indirect
 	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
 	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
-	github.com/golang/protobuf v1.5.4 // indirect
-	github.com/google/s2a-go v0.1.7 // indirect
-	github.com/google/uuid v1.6.0 // indirect
-	github.com/googleapis/enterprise-certificate-proxy v0.3.2 // indirect
-	github.com/googleapis/gax-go/v2 v2.12.5 // indirect
+	github.com/google/s2a-go v0.1.8 // indirect
+	github.com/googleapis/enterprise-certificate-proxy v0.3.4 // indirect
+	github.com/googleapis/gax-go/v2 v2.13.0 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
-	github.com/spf13/pflag v1.0.5 // indirect
 	go.opencensus.io v0.24.0 // indirect
-	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.51.0 // indirect
-	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.51.0 // indirect
-	go.opentelemetry.io/otel v1.26.0 // indirect
-	go.opentelemetry.io/otel/metric v1.26.0 // indirect
-	go.opentelemetry.io/otel/trace v1.26.0 // indirect
-	golang.org/x/crypto v0.24.0 // indirect
-	golang.org/x/mod v0.17.0 // indirect
-	golang.org/x/net v0.26.0 // indirect
-	golang.org/x/oauth2 v0.21.0 // indirect
+	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.54.0 // indirect
+	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.54.0 // indirect
+	go.opentelemetry.io/otel v1.29.0 // indirect
+	go.opentelemetry.io/otel/metric v1.29.0 // indirect
+	go.opentelemetry.io/otel/trace v1.29.0 // indirect
+	golang.org/x/crypto v0.28.0 // indirect
+	golang.org/x/mod v0.20.0 // indirect
+	golang.org/x/net v0.30.0 // indirect
+	golang.org/x/oauth2 v0.23.0 // indirect
 	golang.org/x/sync v0.8.0 // indirect
-	golang.org/x/sys v0.22.0 // indirect
-	golang.org/x/text v0.16.0 // indirect
-	golang.org/x/time v0.5.0 // indirect
-	google.golang.org/genproto/googleapis/api v0.0.0-20240617180043-68d350f18fd4 // indirect
-	google.golang.org/genproto/googleapis/rpc v0.0.0-20240617180043-68d350f18fd4 // indirect
-	google.golang.org/grpc v1.64.1 // indirect
+	golang.org/x/sys v0.26.0 // indirect
+	golang.org/x/text v0.19.0 // indirect
+	golang.org/x/time v0.7.0 // indirect
+	google.golang.org/genproto v0.0.0-20241015192408-796eee8c2d53 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20241015192408-796eee8c2d53 // indirect
+	google.golang.org/grpc v1.67.1 // indirect
 )

dev/tools/controllerbuilder/pkg/llm/gemini.go

@@ -0,0 +1,69 @@
+// 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 llm
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"os"
+	"os/exec"
+	"strings"
+
+	"cloud.google.com/go/vertexai/genai"
+	"google.golang.org/api/option"
+	"k8s.io/klog/v2"
+)
+
+// BuildGeminiClient builds a client for the Gemini API.
+func BuildGeminiClient(ctx context.Context) (*genai.Client, error) {
+	log := klog.FromContext(ctx)
+
+	var opts []option.ClientOption
+
+	if s := os.Getenv("GEMINI_API_KEY"); s != "" {
+		opts = append(opts, option.WithAPIKey(s))
+	}
+	// else {
+	// 	creds, err := google.FindDefaultCredentials(ctx, "https://www.googleapis.com/auth/generative-language", "https://www.googleapis.com/auth/cloud-platform")
+	// 	if err != nil {
+	// 		return nil, fmt.Errorf("finding default credentials: %w", err)
+	// 	}
+	// 	opts = append(opts, option.WithCredentials(creds))
+	// }
+
+	projectID := ""
+	location := ""
+
+	if projectID == "" {
+		cmd := exec.CommandContext(ctx, "gcloud", "config", "get", "project")
+		var stdout bytes.Buffer
+		cmd.Stdout = &stdout
+		if err := cmd.Run(); err != nil {
+			return nil, fmt.Errorf("cannot get project (using gcloud config get project): %w", err)
+		}
+		projectID = strings.TrimSpace(stdout.String())
+		if projectID == "" {
+			return nil, fmt.Errorf("project was not set in gcloud config")
+		}
+		log.Info("got project from gcloud config", "project", projectID)
+	}
+
+	client, err := genai.NewClient(ctx, projectID, location, opts...)
+	if err != nil {
+		return nil, fmt.Errorf("building gemini client: %w", err)
+	}
+	return client, nil
+}

dev/tools/controllerbuilder/pkg/toolbot/csv.go

@@ -23,8 +23,8 @@ import (
 	"os"
 	"path/filepath"
 
-	"github.com/google/generative-ai-go/genai"
-	"google.golang.org/api/option"
+	"cloud.google.com/go/vertexai/genai"
+	"github.com/GoogleCloudPlatform/k8s-config-connector/dev/tools/controllerbuilder/pkg/llm"
 	"k8s.io/apimachinery/pkg/util/sets"
 	"k8s.io/klog/v2"
 )
@@ -194,7 +194,9 @@ func (x *CSVExporter) BuildDataPoints(ctx context.Context, src []byte) ([]*DataP
 
 // RunGemini runs a prompt against Gemini, generating context based on the source code.
 func (x *CSVExporter) RunGemini(ctx context.Context, input *DataPoint, out io.Writer) error {
-	client, err := genai.NewClient(ctx, option.WithAPIKey(os.Getenv("GEMINI_API_KEY")))
+	log := klog.FromContext(ctx)
+
+	client, err := llm.BuildGeminiClient(ctx)
 	if err != nil {
 		return fmt.Errorf("building gemini client: %w", err)
 	}
@@ -219,6 +221,8 @@ func (x *CSVExporter) RunGemini(ctx context.Context, input *DataPoint, out io.Wr
 		parts = append(parts, dataPoint.ToGenAIParts()...)
 	}
 
+	log.Info("context information", "num(parts)", len(parts))
+
 	// We also include the input data point.
 	parts = append(parts, input.ToGenAIParts()...)
 

dev/tools/controllerbuilder/pkg/toolbot/datapoint.go

@@ -19,7 +19,7 @@ import (
 	"fmt"
 	"strings"
 
-	"github.com/google/generative-ai-go/genai"
+	"cloud.google.com/go/vertexai/genai"
 	"k8s.io/apimachinery/pkg/util/sets"
 	"k8s.io/klog/v2"
 )

dev/tools/controllerbuilder/pkg/toolbot/extracttoolmarkers.go

@@ -22,7 +22,7 @@ import (
 	"io"
 	"strings"
 
-	"k8s.io/klog"
+	"k8s.io/klog/v2"
 )
 
 // ExtractToolMarkers extracts tool markers from source code.

dev/tools/controllerbuilder/pkg/typeupdater/insertfield-gemini.go

@@ -22,15 +22,15 @@ import (
 	"path/filepath"
 	"strings"
 
-	"github.com/google/generative-ai-go/genai"
-	"google.golang.org/api/option"
+	"cloud.google.com/go/vertexai/genai"
+	"github.com/GoogleCloudPlatform/k8s-config-connector/dev/tools/controllerbuilder/pkg/llm"
 	"k8s.io/klog/v2"
 )
 
 func (u *TypeUpdater) insertGoFieldGemini() error {
 	klog.Infof("inserting the generated Go code for field %s", u.newField.proto.Name())
 	ctx := context.Background()
-	client, err := genai.NewClient(ctx, option.WithAPIKey(os.Getenv("GEMINI_API_KEY")))
+	client, err := llm.BuildGeminiClient(ctx)
 	if err != nil {
 		log.Fatal(err)
 	}

dev/tools/controllerbuilder/pkg/typeupdater/insertmessage-gemini.go

@@ -18,15 +18,14 @@ import (
 	"context"
 	"fmt"
 	"log"
-	"os"
 
-	"github.com/google/generative-ai-go/genai"
-	"google.golang.org/api/option"
+	"cloud.google.com/go/vertexai/genai"
+	"github.com/GoogleCloudPlatform/k8s-config-connector/dev/tools/controllerbuilder/pkg/llm"
 )
 
 func (u *TypeUpdater) insertGoMessagesGemini() error {
 	ctx := context.Background()
-	client, err := genai.NewClient(ctx, option.WithAPIKey(os.Getenv("GEMINI_API_KEY")))
+	client, err := llm.BuildGeminiClient(ctx)
 	if err != nil {
 		log.Fatal(err)
 	}

dev/tools/controllerbuilder/pkg/typeupdater/typeupdater.go

@@ -27,7 +27,7 @@ import (
 	"google.golang.org/protobuf/reflect/protoreflect"
 	"google.golang.org/protobuf/types/descriptorpb"
 	"k8s.io/apimachinery/pkg/util/sets"
-	"k8s.io/klog"
+	"k8s.io/klog/v2"
 )
 
 const kccProtoPrefix = "+kcc:proto="