Repository: GoogleCloudPlatform/kubectl-ai
Pull request #88: replace VertexAI implimentation to use google.golang.org/genai
Author: login: sakajunquality
Created: 2025-04-04T06:54:09Z | Merged: 2025-04-10T04:28:46Z



Diff:
diff --git a/gollm/factory.go b/gollm/factory.go
index ee6178b1..50fc426c 100644
--- a/gollm/factory.go
+++ b/gollm/factory.go
@@ -44,9 +44,9 @@ func NewClient(ctx context.Context, providerID string) (Client, error) {
 
 	switch providerID {
 	case "gemini":
-		return NewGeminiClient(ctx)
+		return NewGeminiClient(ctx, geminiBackend)
 	case "vertexai":
-		return NewVertexAIClient(ctx)
+		return NewGeminiClient(ctx, vertexaiBackend)
 	case "ollama":
 		return NewOllamaClient(ctx)
 	case "llamacpp":
diff --git a/gollm/gemini.go b/gollm/gemini.go
index 785cc9c8..cbefaa31 100644
--- a/gollm/gemini.go
+++ b/gollm/gemini.go
@@ -31,20 +31,35 @@ import (
 
 const (
 	geminiDefaultModel = "gemini-2.0-pro-exp-02-05"
+	geminiBackend      = "gemini"
+	vertexaiBackend    = "vertexai"
 )
 
 // NewGeminiClient builds a client for the Gemini API.
-func NewGeminiClient(ctx context.Context) (*GeminiClient, error) {
-	apiKey := os.Getenv("GEMINI_API_KEY")
-	if apiKey == "" {
-		return nil, fmt.Errorf("GEMINI_API_KEY environment variable not set")
+func NewGeminiClient(ctx context.Context, backend string) (*GeminiClient, error) {
+	var cc *genai.ClientConfig
+
+	switch backend {
+	case geminiBackend:
+		apiKey := os.Getenv("GEMINI_API_KEY")
+		if apiKey == "" {
+			return nil, fmt.Errorf("GEMINI_API_KEY environment variable not set")
+		}
+		cc = &genai.ClientConfig{
+			APIKey:  apiKey,
+			Backend: genai.BackendGeminiAPI,
+		}
+	case vertexaiBackend:
+		cc = &genai.ClientConfig{
+			// Project ID is loaded from the GOOGLE_CLOUD_PROJECT environment variable
+			// Location/Region is loaded from either GOOGLE_CLOUD_LOCATION or GOOGLE_CLOUD_REGION environment variable
+			Backend: genai.BackendVertexAI,
+		}
+	default:
+		return nil, fmt.Errorf("unknown backend %q", backend)
 	}
 
-	client, err := genai.NewClient(ctx, &genai.ClientConfig{
-		APIKey:  apiKey,
-		Backend: genai.BackendGeminiAPI,
-	})
-
+	client, err := genai.NewClient(ctx, cc)
 	if err != nil {
 		return nil, fmt.Errorf("building gemini client: %w", err)
 	}
@@ -178,21 +193,23 @@ type GeminiChat struct {
 // SetFunctionDefinitions sets the function definitions for the chat.
 // This allows the LLM to call user-defined functions.
 func (c *GeminiChat) SetFunctionDefinitions(functionDefinitions []*FunctionDefinition) error {
+	var genaiFunctionDeclarations []*genai.FunctionDeclaration
 	for _, functionDefinition := range functionDefinitions {
 		parameters, err := toGeminiSchema(functionDefinition.Parameters)
 		if err != nil {
 			return err
 		}
-		c.genConfig.Tools = append(c.genConfig.Tools, &genai.Tool{
-			FunctionDeclarations: []*genai.FunctionDeclaration{
-				{
-					Name:        functionDefinition.Name,
-					Description: functionDefinition.Description,
-					Parameters:  parameters,
-				},
-			},
+		genaiFunctionDeclarations = append(genaiFunctionDeclarations, &genai.FunctionDeclaration{
+			Name:        functionDefinition.Name,
+			Description: functionDefinition.Description,
+			Parameters:  parameters,
 		})
 	}
+	c.genConfig.Tools = []*genai.Tool{
+		{
+			FunctionDeclarations: genaiFunctionDeclarations,
+		},
+	}
 	return nil
 }
 
diff --git a/gollm/go.mod b/gollm/go.mod
index 6ae7f9d4..a311e93b 100644
--- a/gollm/go.mod
+++ b/gollm/go.mod
@@ -5,45 +5,18 @@ go 1.24.0
 toolchain go1.24.1
 
 require (
-	cloud.google.com/go/vertexai v0.13.3
 	github.com/ollama/ollama v0.5.13
-	golang.org/x/oauth2 v0.26.0
-	google.golang.org/api v0.222.0
 	google.golang.org/genai v0.6.0
 	k8s.io/klog/v2 v2.130.1
 )
 
 require (
 	cloud.google.com/go v0.116.0 // indirect
-	cloud.google.com/go/aiplatform v1.69.0 // indirect
-	cloud.google.com/go/auth v0.14.1 // indirect
-	cloud.google.com/go/auth/oauth2adapt v0.2.7 // indirect
 	cloud.google.com/go/compute/metadata v0.6.0 // indirect
-	cloud.google.com/go/iam v1.2.2 // indirect
-	cloud.google.com/go/longrunning v0.6.2 // 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/google/go-cmp v0.6.0 // indirect
-	github.com/google/s2a-go v0.1.9 // indirect
-	github.com/googleapis/enterprise-certificate-proxy v0.3.4 // indirect
-	github.com/googleapis/gax-go/v2 v2.14.1 // indirect
 	github.com/gorilla/websocket v1.5.3 // indirect
-	go.opentelemetry.io/auto/sdk v1.1.0 // indirect
-	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.58.0 // indirect
-	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.58.0 // indirect
-	go.opentelemetry.io/otel v1.34.0 // indirect
-	go.opentelemetry.io/otel/metric v1.34.0 // indirect
-	go.opentelemetry.io/otel/trace v1.34.0 // indirect
-	golang.org/x/crypto v0.35.0 // indirect
-	golang.org/x/net v0.36.0 // indirect
-	golang.org/x/sync v0.11.0 // indirect
+	github.com/stretchr/testify v1.10.0 // indirect
+	golang.org/x/oauth2 v0.26.0 // indirect
 	golang.org/x/sys v0.30.0 // indirect
-	golang.org/x/text v0.22.0 // indirect
-	golang.org/x/time v0.10.0 // indirect
-	google.golang.org/genproto v0.0.0-20241118233622-e639e219e697 // indirect
-	google.golang.org/genproto/googleapis/api v0.0.0-20241209162323-e6fa225c2576 // indirect
-	google.golang.org/genproto/googleapis/rpc v0.0.0-20250212204824-5a70512c5d8b // indirect
-	google.golang.org/grpc v1.70.0 // indirect
-	google.golang.org/protobuf v1.36.5 // indirect
 )
diff --git a/gollm/go.sum b/gollm/go.sum
index 5824cb23..0dd179eb 100644
--- a/gollm/go.sum
+++ b/gollm/go.sum
@@ -1,40 +1,13 @@
 cloud.google.com/go v0.116.0 h1:B3fRrSDkLRt5qSHWe40ERJvhvnQwdZiHu0bJOpldweE=
 cloud.google.com/go v0.116.0/go.mod h1:cEPSRWPzZEswwdr9BxE6ChEn01dWlTaF05LiC2Xs70U=
-cloud.google.com/go/aiplatform v1.69.0 h1:XvBzK8e6/6ufbi/i129Vmn/gVqFwbNPmRQ89K+MGlgc=
-cloud.google.com/go/aiplatform v1.69.0/go.mod h1:nUsIqzS3khlnWvpjfJbP+2+h+VrFyYsTm7RNCAViiY8=
-cloud.google.com/go/auth v0.14.1 h1:AwoJbzUdxA/whv1qj3TLKwh3XX5sikny2fc40wUl+h0=
-cloud.google.com/go/auth v0.14.1/go.mod h1:4JHUxlGXisL0AW8kXPtUF6ztuOksyfUQNFjfsOCXkPM=
-cloud.google.com/go/auth/oauth2adapt v0.2.7 h1:/Lc7xODdqcEw8IrZ9SvwnlLX6j9FHQM74z6cBk9Rw6M=
-cloud.google.com/go/auth/oauth2adapt v0.2.7/go.mod h1:NTbTTzfvPl1Y3V1nPpOgl2w6d/FjO7NNUQaWSox6ZMc=
 cloud.google.com/go/compute/metadata v0.6.0 h1:A6hENjEsCDtC1k8byVsgwvVcioamEHvZ4j01OwKxG9I=
 cloud.google.com/go/compute/metadata v0.6.0/go.mod h1:FjyFAW1MW0C203CEOMDTu3Dk1FlqW3Rga40jzHL4hfg=
-cloud.google.com/go/iam v1.2.2 h1:ozUSofHUGf/F4tCNy/mu9tHLTaxZFLOUiKzjcgWHGIA=
-cloud.google.com/go/iam v1.2.2/go.mod h1:0Ys8ccaZHdI1dEUilwzqng/6ps2YB6vRsjIe00/+6JY=
-cloud.google.com/go/longrunning v0.6.2 h1:xjDfh1pQcWPEvnfjZmwjKQEcHnpz6lHjfy7Fo0MK+hc=
-cloud.google.com/go/longrunning v0.6.2/go.mod h1:k/vIs83RN4bE3YCswdXC5PFfWVILjm3hpEUlSko4PiI=
-cloud.google.com/go/vertexai v0.13.3 h1:pbw1KfpdE8ZDrXxBKcIsS/j+EixyQRsyu6gxRkXq8/k=
-cloud.google.com/go/vertexai v0.13.3/go.mod h1:AxzUNrd36yhfOZedO+Y1v0ajVgGKOdv1njeQChL8IFY=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg=
-github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
-github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
 github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY=
 github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
-github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
-github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
-github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
-github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
 github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
 github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
-github.com/google/s2a-go v0.1.9 h1:LGD7gtMgezd8a/Xak7mEWL0PjoTQFvpRudN895yqKW0=
-github.com/google/s2a-go v0.1.9/go.mod h1:YA0Ei2ZQL3acow2O62kdp9UlnvMmU7kA6Eutn0dXayM=
-github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
-github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
-github.com/googleapis/enterprise-certificate-proxy v0.3.4 h1:XYIDZApgAnrN1c855gTgghdIA6Stxb52D5RnLI1SLyw=
-github.com/googleapis/enterprise-certificate-proxy v0.3.4/go.mod h1:YKe7cfqYXjKGpGvmSg28/fFvhNzinZQm8DGnaburhGA=
-github.com/googleapis/gax-go/v2 v2.14.1 h1:hb0FFeiPaQskmvakKu5EbCbpntQn48jyHuvrkurSS/Q=
-github.com/googleapis/gax-go/v2 v2.14.1/go.mod h1:Hb/NubMaVM88SrNkvl8X/o8XWwDJEPqouaLeN2IUxoA=
 github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
 github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/ollama/ollama v0.5.13 h1:URBx4e6nyAaVhEGXH6AWVqORhebcSQcJ7hLTS0xkAPg=
@@ -43,50 +16,12 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
 github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
-go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA=
-go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A=
-go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.58.0 h1:PS8wXpbyaDJQ2VDHHncMe9Vct0Zn1fEjpsjrLxGJoSc=
-go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.58.0/go.mod h1:HDBUsEjOuRC0EzKZ1bSaRGZWUBAzo+MhAcUUORSr4D0=
-go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.58.0 h1:yd02MEjBdJkG3uabWP9apV+OuWRIXGDuJEUJbOHmCFU=
-go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.58.0/go.mod h1:umTcuxiv1n/s/S6/c2AT/g2CQ7u5C59sHDNmfSwgz7Q=
-go.opentelemetry.io/otel v1.34.0 h1:zRLXxLCgL1WyKsPVrgbSdMN4c0FMkDAskSTQP+0hdUY=
-go.opentelemetry.io/otel v1.34.0/go.mod h1:OWFPOQ+h4G8xpyjgqo4SxJYdDQ/qmRH+wivy7zzx9oI=
-go.opentelemetry.io/otel/metric v1.34.0 h1:+eTR3U0MyfWjRDhmFMxe2SsW64QrZ84AOhvqS7Y+PoQ=
-go.opentelemetry.io/otel/metric v1.34.0/go.mod h1:CEDrp0fy2D0MvkXE+dPV7cMi8tWZwX3dmaIhwPOaqHE=
-go.opentelemetry.io/otel/sdk v1.34.0 h1:95zS4k/2GOy069d321O8jWgYsW3MzVV+KuSPKp7Wr1A=
-go.opentelemetry.io/otel/sdk v1.34.0/go.mod h1:0e/pNiaMAqaykJGKbi+tSjWfNNHMTxoC9qANsCzbyxU=
-go.opentelemetry.io/otel/sdk/metric v1.32.0 h1:rZvFnvmvawYb0alrYkjraqJq0Z4ZUJAiyYCU9snn1CU=
-go.opentelemetry.io/otel/sdk/metric v1.32.0/go.mod h1:PWeZlq0zt9YkYAp3gjKZ0eicRYvOh1Gd+X99x6GHpCQ=
-go.opentelemetry.io/otel/trace v1.34.0 h1:+ouXS2V8Rd4hp4580a8q23bg0azF2nI8cqLYnC8mh/k=
-go.opentelemetry.io/otel/trace v1.34.0/go.mod h1:Svm7lSjQD7kG7KJ/MUHPVXSDGz2OX4h0M2jHBhmSfRE=
-golang.org/x/crypto v0.35.0 h1:b15kiHdrGCHrP6LvwaQ3c03kgNhhiMgvlhxHQhmg2Xs=
-golang.org/x/crypto v0.35.0/go.mod h1:dy7dXNW32cAb/6/PRuTNsix8T+vJAqvuIy5Bli/x0YQ=
-golang.org/x/net v0.36.0 h1:vWF2fRbw4qslQsQzgFqZff+BItCvGFQqKzKIzx1rmoA=
-golang.org/x/net v0.36.0/go.mod h1:bFmbeoIPfrw4sMHNhb4J9f6+tPziuGjq7Jk/38fxi1I=
 golang.org/x/oauth2 v0.26.0 h1:afQXWNNaeC4nvZ0Ed9XvCCzXM6UHJG7iCg0W4fPqSBE=
 golang.org/x/oauth2 v0.26.0/go.mod h1:XYTD2NtWslqkgxebSiOHnXEap4TF09sJSc7H1sXbhtI=
-golang.org/x/sync v0.11.0 h1:GGz8+XQP4FvTTrjZPzNKTMFtSXH80RAzG+5ghFPgK9w=
-golang.org/x/sync v0.11.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
 golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc=
 golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
-golang.org/x/text v0.22.0 h1:bofq7m3/HAFvbF51jz3Q9wLg3jkvSPuiZu/pD1XwgtM=
-golang.org/x/text v0.22.0/go.mod h1:YRoo4H8PVmsu+E3Ou7cqLVH8oXWIHVoX0jqUWALQhfY=
-golang.org/x/time v0.10.0 h1:3usCWA8tQn0L8+hFJQNgzpWbd89begxN66o1Ojdn5L4=
-golang.org/x/time v0.10.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
-google.golang.org/api v0.222.0 h1:Aiewy7BKLCuq6cUCeOUrsAlzjXPqBkEeQ/iwGHVQa/4=
-google.golang.org/api v0.222.0/go.mod h1:efZia3nXpWELrwMlN5vyQrD4GmJN1Vw0x68Et3r+a9c=
 google.golang.org/genai v0.6.0 h1:S9eDmXHPPqiWrKO2G7ydTNQ70fG1y1+ttR6zsFCPJd0=
 google.golang.org/genai v0.6.0/go.mod h1:yPyKKBezIg2rqZziLhHQ5CD62HWr7sLDLc2PDzdrNVs=
-google.golang.org/genproto v0.0.0-20241118233622-e639e219e697 h1:ToEetK57OidYuqD4Q5w+vfEnPvPpuTwedCNVohYJfNk=
-google.golang.org/genproto v0.0.0-20241118233622-e639e219e697/go.mod h1:JJrvXBWRZaFMxBufik1a4RpFw4HhgVtBBWQeQgUj2cc=
-google.golang.org/genproto/googleapis/api v0.0.0-20241209162323-e6fa225c2576 h1:CkkIfIt50+lT6NHAVoRYEyAvQGFM7xEwXUUywFvEb3Q=
-google.golang.org/genproto/googleapis/api v0.0.0-20241209162323-e6fa225c2576/go.mod h1:1R3kvZ1dtP3+4p4d3G8uJ8rFk/fWlScl38vanWACI08=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20250212204824-5a70512c5d8b h1:FQtJ1MxbXoIIrZHZ33M+w5+dAP9o86rgpjoKr/ZmT7k=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20250212204824-5a70512c5d8b/go.mod h1:8BS3B93F/U1juMFq9+EDk+qOT5CO1R9IzXxG3PTqiRk=
-google.golang.org/grpc v1.70.0 h1:pWFv03aZoHzlRKHWicjsZytKAiYCtNS0dHbXnIdq7jQ=
-google.golang.org/grpc v1.70.0/go.mod h1:ofIJqVKDXx/JiXrwr2IG4/zwdH9txy3IlF40RmcJSQw=
-google.golang.org/protobuf v1.36.5 h1:tPhr+woSbjfYvY6/GPufUoYizxw1cF/yFoxJ2fmpwlM=
-google.golang.org/protobuf v1.36.5/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 k8s.io/klog/v2 v2.130.1 h1:n9Xl7H1Xvksem4KFG4PYbdQCQxqc/tTUyrgXaOhHSzk=
diff --git a/gollm/vertexai.go b/gollm/vertexai.go
deleted file mode 100644
index 6668ffaa..00000000
--- a/gollm/vertexai.go
+++ /dev/null
@@ -1,427 +0,0 @@
-// 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 gollm
-
-import (
-	"bytes"
-	"context"
-	"encoding/json"
-	"fmt"
-	"os/exec"
-	"strings"
-
-	"cloud.google.com/go/vertexai/genai"
-	"golang.org/x/oauth2/google"
-	"google.golang.org/api/option"
-	"k8s.io/klog/v2"
-)
-
-const (
-	vertexaiDefaultModel = "gemini-2.0-pro-exp-02-05"
-)
-
-// NewVertexAIClient builds a client for the VertexAI API.
-func NewVertexAIClient(ctx context.Context) (*VertexAIClient, error) {
-	log := klog.FromContext(ctx)
-
-	var opts []option.ClientOption
-
-	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 := "us-central1"
-
-	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)
-	}
-
-	// TODO: Detect and/or auto-enable "gcloud services enable aiplatform.googleapis.com" if not enabled?
-
-	client, err := genai.NewClient(ctx, projectID, location, opts...)
-	if err != nil {
-		return nil, fmt.Errorf("building vertexai client: %w", err)
-	}
-	return &VertexAIClient{client: client}, nil
-}
-
-// VertexAIClient is a client for the VertexAI API.
-// It implements the Client interface.
-type VertexAIClient struct {
-	client *genai.Client
-
-	// responseSchema will constrain the output to match the given schema
-	responseSchema *genai.Schema
-}
-
-var _ Client = &VertexAIClient{}
-
-// Close frees the resources used by the client.
-func (c *VertexAIClient) Close() error {
-	return c.client.Close()
-}
-
-// ListModels lists the models available in the VertexAI API.
-func (c *VertexAIClient) ListModels(ctx context.Context) ([]string, error) {
-	return nil, fmt.Errorf("listing models not supported yet for vertexai")
-}
-
-// SetResponseSchema constrains LLM responses to match the provided schema.
-// Calling with nil will clear the current schema.
-func (c *VertexAIClient) SetResponseSchema(responseSchema *Schema) error {
-	if responseSchema == nil {
-		c.responseSchema = nil
-		return nil
-	}
-
-	vertexAISchema, err := toVertexAISchema(responseSchema)
-	if err != nil {
-		return err
-	}
-
-	c.responseSchema = vertexAISchema
-	return nil
-}
-
-func (c *VertexAIClient) GenerateCompletion(ctx context.Context, request *CompletionRequest) (CompletionResponse, error) {
-	log := klog.FromContext(ctx)
-
-	model := c.client.GenerativeModel(request.Model)
-
-	if c.responseSchema != nil {
-		model.ResponseSchema = c.responseSchema
-		model.ResponseMIMEType = "application/json"
-	}
-
-	var vertexaiParts []genai.Part
-
-	vertexaiParts = append(vertexaiParts, genai.Text(request.Prompt))
-
-	log.Info("sending GenerateContent request to vertexai", "parts", vertexaiParts)
-	vertexaiResponse, err := model.GenerateContent(ctx, vertexaiParts...)
-	if err != nil {
-		return nil, err
-	}
-
-	if len(vertexaiResponse.Candidates) > 1 {
-		klog.Infof("only considering first candidate")
-		for i := 1; i < len(vertexaiResponse.Candidates); i++ {
-			candidate := vertexaiResponse.Candidates[i]
-			klog.Infof("ignoring candidate: %q", candidate.Content)
-		}
-	}
-	var response strings.Builder
-	candidate := vertexaiResponse.Candidates[0]
-	for _, part := range candidate.Content.Parts {
-		switch part := part.(type) {
-		case genai.Text:
-			if response.Len() != 0 {
-				response.WriteString("\n")
-			}
-			response.WriteString(string(part))
-		default:
-			return nil, fmt.Errorf("unexpected type of content part: %T", part)
-		}
-	}
-
-	return &VertexAICompletionResponse{vertexaiResponse: vertexaiResponse, text: response.String()}, nil
-}
-
-// StartChat starts a new chat with the model.
-func (c *VertexAIClient) StartChat(systemPrompt, aiModel string) Chat {
-	model := c.client.GenerativeModel(aiModel)
-
-	// Some values that are recommended by aistudio
-	model.SetTemperature(1)
-	model.SetTopK(40)
-	model.SetTopP(0.95)
-	model.SetMaxOutputTokens(8192)
-	model.ResponseMIMEType = "text/plain"
-
-	if c.responseSchema != nil {
-		model.ResponseSchema = c.responseSchema
-		model.ResponseMIMEType = "application/json"
-	}
-
-	if systemPrompt != "" {
-		model.SystemInstruction = &genai.Content{
-			Parts: []genai.Part{
-				genai.Text(systemPrompt),
-			},
-		}
-	} else {
-		klog.Warningf("systemPrompt not provided")
-	}
-
-	chat := model.StartChat()
-
-	return &VertexAIChat{
-		model: model,
-		chat:  chat,
-	}
-}
-
-// VertexAIChat is a chat with the model.
-// It implements the Chat interface.
-type VertexAIChat struct {
-	model *genai.GenerativeModel
-	chat  *genai.ChatSession
-}
-
-func (c *VertexAIChat) IsRetryableError(err error) bool {
-	// TODO(droot): Implement this
-	return false
-}
-
-// SetFunctionDefinitions sets the function definitions for the chat.
-// This allows the LLM to call user-defined functions.
-func (c *VertexAIChat) SetFunctionDefinitions(functionDefinitions []*FunctionDefinition) error {
-	var vertexaiFunctionDefinitions []*genai.FunctionDeclaration
-	for _, functionDefinition := range functionDefinitions {
-		parameters, err := toVertexAISchema(functionDefinition.Parameters)
-		if err != nil {
-			return err
-		}
-		vertexaiFunctionDefinitions = append(vertexaiFunctionDefinitions, &genai.FunctionDeclaration{
-			Name:        functionDefinition.Name,
-			Description: functionDefinition.Description,
-			Parameters:  parameters,
-		})
-	}
-
-	c.model.Tools = append(c.model.Tools, &genai.Tool{
-		FunctionDeclarations: vertexaiFunctionDefinitions,
-	})
-	return nil
-}
-
-// toVertexAISchema converts our generic Schema to a genai.Schema
-func toVertexAISchema(schema *Schema) (*genai.Schema, error) {
-	ret := &genai.Schema{
-		Description: schema.Description,
-		Required:    schema.Required,
-	}
-
-	switch schema.Type {
-	case TypeObject:
-		ret.Type = genai.TypeObject
-	case TypeString:
-		ret.Type = genai.TypeString
-	case TypeBoolean:
-		ret.Type = genai.TypeBoolean
-	case TypeInteger:
-		ret.Type = genai.TypeInteger
-	case TypeArray:
-		ret.Type = genai.TypeArray
-	default:
-		return nil, fmt.Errorf("type %q not handled by genai.Schema", schema.Type)
-	}
-	if schema.Properties != nil {
-		ret.Properties = make(map[string]*genai.Schema)
-		for k, v := range schema.Properties {
-			vertexaiValue, err := toVertexAISchema(v)
-			if err != nil {
-				return nil, err
-			}
-			ret.Properties[k] = vertexaiValue
-		}
-	}
-	if schema.Items != nil {
-		vertexItems, err := toVertexAISchema(schema.Items)
-		if err != nil {
-			return nil, err
-		}
-		ret.Items = vertexItems
-	}
-	return ret, nil
-}
-
-func (c *VertexAIChat) Send(ctx context.Context, contents ...any) (ChatResponse, error) {
-	log := klog.FromContext(ctx)
-	log.Info("sending LLM request", "user", contents)
-
-	var vertexaiParts []genai.Part
-	for _, content := range contents {
-		switch v := content.(type) {
-		case string:
-			vertexaiParts = append(vertexaiParts, genai.Text(v))
-		case FunctionCallResult:
-			vertexaiParts = append(vertexaiParts, genai.FunctionResponse{
-				Name:     v.Name,
-				Response: v.Result,
-			})
-		default:
-			return nil, fmt.Errorf("unexpected type of content: %T", content)
-		}
-	}
-	vertexaiResponse, err := c.chat.SendMessage(ctx, vertexaiParts...)
-	if err != nil {
-		return nil, err
-	}
-	return &VertexAIChatResponse{vertexaiResponse: vertexaiResponse}, nil
-}
-
-// VertexAIChatResponse is a response from the VertexAI API.
-// It implements the ChatResponse interface.
-type VertexAIChatResponse struct {
-	vertexaiResponse *genai.GenerateContentResponse
-}
-
-var _ ChatResponse = &VertexAIChatResponse{}
-
-func (r *VertexAIChatResponse) MarshalJSON() ([]byte, error) {
-	formatted := RecordChatResponse{
-		Raw: r.vertexaiResponse,
-	}
-	return json.Marshal(&formatted)
-}
-
-// String returns a string representation of the response.
-func (r *VertexAIChatResponse) String() string {
-	var response strings.Builder
-	response.WriteString("{candidates=[")
-	for i, candidate := range r.Candidates() {
-		if i > 0 {
-			response.WriteString(", ")
-		}
-		response.WriteString(candidate.String())
-	}
-	response.WriteString("]}")
-	return response.String()
-}
-
-// UsageMetadata returns the usage metadata for the response.
-func (r *VertexAIChatResponse) UsageMetadata() any {
-	return r.vertexaiResponse.UsageMetadata
-}
-
-// Candidates returns the candidates for the response.
-func (r *VertexAIChatResponse) Candidates() []Candidate {
-	var candidates []Candidate
-	for _, candidate := range r.vertexaiResponse.Candidates {
-		// klog.Infof("candidate: %+v", candidate)
-		candidates = append(candidates, &VertexAICandidate{candidate: candidate})
-	}
-	return candidates
-}
-
-// VertexAICandidate is a candidate for the response.
-// It implements the Candidate interface.
-type VertexAICandidate struct {
-	candidate *genai.Candidate
-}
-
-// String returns a string representation of the response.
-func (r *VertexAICandidate) String() string {
-	var response strings.Builder
-	response.WriteString("[")
-	for i, parts := range r.Parts() {
-		if i > 0 {
-			response.WriteString(", ")
-		}
-		text, ok := parts.AsText()
-		if ok {
-			response.WriteString(text)
-		}
-		functionCalls, ok := parts.AsFunctionCalls()
-		if ok {
-			response.WriteString("functionCalls=[")
-			for _, functionCall := range functionCalls {
-				response.WriteString(fmt.Sprintf("%q(args=%v)", functionCall.Name, functionCall.Arguments))
-			}
-			response.WriteString("]}")
-		}
-	}
-	response.WriteString("]}")
-	return response.String()
-}
-
-// Parts returns the parts of the candidate.
-func (r *VertexAICandidate) Parts() []Part {
-	var parts []Part
-	if r.candidate.Content != nil {
-		for _, part := range r.candidate.Content.Parts {
-			parts = append(parts, &VertexAIPart{part: part})
-		}
-	}
-	return parts
-}
-
-// VertexAIPart is a part of a candidate.
-// It implements the Part interface.
-type VertexAIPart struct {
-	part genai.Part
-}
-
-// AsText returns the text of the part.
-func (p *VertexAIPart) AsText() (string, bool) {
-	if text, ok := p.part.(genai.Text); ok {
-		return string(text), true
-	}
-	return "", false
-}
-
-// AsFunctionCalls returns the function calls of the part.
-func (p *VertexAIPart) AsFunctionCalls() ([]FunctionCall, bool) {
-	if functionCall, ok := p.part.(genai.FunctionCall); ok {
-		var ret []FunctionCall
-		ret = append(ret, FunctionCall{
-			Name:      functionCall.Name,
-			Arguments: functionCall.Args,
-		})
-		return ret, true
-	}
-	return nil, false
-}
-
-type VertexAICompletionResponse struct {
-	vertexaiResponse *genai.GenerateContentResponse
-	text             string
-}
-
-var _ CompletionResponse = &VertexAICompletionResponse{}
-
-func (r *VertexAICompletionResponse) MarshalJSON() ([]byte, error) {
-	formatted := RecordCompletionResponse{
-		Text: r.text,
-		Raw:  r.vertexaiResponse,
-	}
-	return json.Marshal(&formatted)
-}
-
-func (r *VertexAICompletionResponse) Response() string {
-	return r.text
-}
-
-func (r *VertexAICompletionResponse) UsageMetadata() any {
-	return r.vertexaiResponse.UsageMetadata
-}
-
-func (r *VertexAICompletionResponse) String() string {
-	return fmt.Sprintf("{text=%q}", r.text)
-}


Reviews:
user: login: copilot-pull-request-reviewer[bot]
body: 

Copilot reviewed 4 out of 5 changed files in this pull request and generated no comments.

<details>
<summary>Files not reviewed (1)</summary>

* **gollm/go.mod**: Language not supported
</details>

<details>
<summary>Comments suppressed due to low confidence (1)</summary>

**gollm/gemini.go:39**
* [nitpick] The function name 'NewGeminiClient' may be misleading now that it handles both Gemini and VertexAI backends; consider renaming it (e.g., to 'NewGenaiClient') to better reflect its purpose.
```
func NewGeminiClient(ctx context.Context, backend string) (*GeminiClient, error) {
```
</details>


state: COMMENTED
author_association: NONE
submitted_at: 2025-04-04T16:50:14Z
user: login: sakajunquality
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-08T14:05:51Z
user: login: droot
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-09T19:00:40Z
user: login: sakajunquality
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-10T00:03:54Z

Review comments:
diff_hunk: @@ -178,21 +193,23 @@ type GeminiChat struct {
 // SetFunctionDefinitions sets the function definitions for the chat.
 // This allows the LLM to call user-defined functions.
 func (c *GeminiChat) SetFunctionDefinitions(functionDefinitions []*FunctionDefinition) error {
+	var genaiFunctionDeclarations []*genai.FunctionDeclaration
 	for _, functionDefinition := range functionDefinitions {
 		parameters, err := toGeminiSchema(functionDefinition.Parameters)
 		if err != nil {
 			return err
 		}
-		c.genConfig.Tools = append(c.genConfig.Tools, &genai.Tool{
-			FunctionDeclarations: []*genai.FunctionDeclaration{
-				{
-					Name:        functionDefinition.Name,
-					Description: functionDefinition.Description,
-					Parameters:  parameters,
-				},
-			},
+		genaiFunctionDeclarations = append(genaiFunctionDeclarations, &genai.FunctionDeclaration{
+			Name:        functionDefinition.Name,
+			Description: functionDefinition.Description,
+			Parameters:  parameters,
 		})
 	}
+	c.genConfig.Tools = []*genai.Tool{
path: gollm/gemini.go
user: login: sakajunquality
body: > Declare a [Tool](https://cloud.google.com/vertex-ai/generative-ai/docs/reference/python/latest/vertexai.generative_models.Tool) that contains up to 128 [FunctionDeclarations](https://cloud.google.com/vertex-ai/generative-ai/docs/reference/python/latest/vertexai.generative_models.FunctionDeclaration). You will later pass this tool to the model when submitting the prompt. The model can use the functions in the tool to process the prompt. At most one tool can be provided with the request.

https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/function-calling#function-declarations
created_at: 2025-04-08T14:05:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -31,20 +31,35 @@ import (
 
 const (
 	geminiDefaultModel = "gemini-2.0-pro-exp-02-05"
+	geminiBackend      = "gemini"
+	vertexaiBackend    = "vertexai"
 )
 
 // NewGeminiClient builds a client for the Gemini API.
-func NewGeminiClient(ctx context.Context) (*GeminiClient, error) {
-	apiKey := os.Getenv("GEMINI_API_KEY")
-	if apiKey == "" {
-		return nil, fmt.Errorf("GEMINI_API_KEY environment variable not set")
+func NewGeminiClient(ctx context.Context, backend string) (*GeminiClient, error) {
+	var cc *genai.ClientConfig
+
+	switch backend {
+	case geminiBackend:
+		apiKey := os.Getenv("GEMINI_API_KEY")
+		if apiKey == "" {
+			return nil, fmt.Errorf("GEMINI_API_KEY environment variable not set")
+		}
+		cc = &genai.ClientConfig{
+			APIKey: apiKey,
+			// Project ID is loaded from the GOOGLE_CLOUD_PROJECT environment variable
+			// Location/Region is loaded from either GOOGLE_CLOUD_LOCATION or GOOGLE_CLOUD_REGION environment variable
path: gollm/gemini.go
user: login: droot
body: These comments seems relevant for vertex backend only, right ? Can we move them the vertex section.

It will be nice if we can check presence of these environment variables and return error like we are doing for gemini backend. 
created_at: 2025-04-09T18:59:11Z
author_association: MEMBER
diff_hunk: @@ -31,20 +31,35 @@ import (
 
 const (
 	geminiDefaultModel = "gemini-2.0-pro-exp-02-05"
+	geminiBackend      = "gemini"
+	vertexaiBackend    = "vertexai"
 )
 
 // NewGeminiClient builds a client for the Gemini API.
-func NewGeminiClient(ctx context.Context) (*GeminiClient, error) {
-	apiKey := os.Getenv("GEMINI_API_KEY")
-	if apiKey == "" {
-		return nil, fmt.Errorf("GEMINI_API_KEY environment variable not set")
+func NewGeminiClient(ctx context.Context, backend string) (*GeminiClient, error) {
+	var cc *genai.ClientConfig
+
+	switch backend {
+	case geminiBackend:
+		apiKey := os.Getenv("GEMINI_API_KEY")
+		if apiKey == "" {
+			return nil, fmt.Errorf("GEMINI_API_KEY environment variable not set")
+		}
+		cc = &genai.ClientConfig{
+			APIKey: apiKey,
+			// Project ID is loaded from the GOOGLE_CLOUD_PROJECT environment variable
+			// Location/Region is loaded from either GOOGLE_CLOUD_LOCATION or GOOGLE_CLOUD_REGION environment variable
path: gollm/gemini.go
user: login: sakajunquality
body: oops,
Fixed.
https://github.com/GoogleCloudPlatform/kubectl-ai/pull/88/commits/41248569c450b5441b89ce4ce8a06117afa1c167
created_at: 2025-04-10T00:03:54Z
author_association: CONTRIBUTOR

PR comments:
user: login: droot
created_at: 2025-04-04T16:33:37Z
body: Thanks for re-working this one, were you able to test out vertex backend ?
author_association: MEMBER
user: login: sakajunquality
created_at: 2025-04-07T04:08:09Z
body: @droot I got error with this change. I'm going to investigate.
author_association: CONTRIBUTOR
user: login: sakajunquality
created_at: 2025-04-08T14:08:37Z
body: @droot Resolved an error by implementing changes based on the documentation. Verified functionality with both Gemini and Vertex AI.
author_association: CONTRIBUTOR
user: login: droot
created_at: 2025-04-09T19:02:06Z
body: > @droot Resolved an error by implementing changes based on the documentation. Verified functionality with both Gemini and Vertex AI.

Thank you so much. I have a minor comment (variable lookup is optional and its okay to defer). I am OOO so expect some delay in the feedback.
author_association: MEMBER
user: login: sakajunquality
created_at: 2025-04-10T00:17:51Z
body: @droot Let me leave the variable lookup for now.
I think it's better to enhance when upgrading the library 
Looks like [v1.0.0](https://github.com/googleapis/go-genai/releases/tag/v1.0.0) has been just released.
author_association: CONTRIBUTOR
user: login: droot
created_at: 2025-04-10T04:28:18Z
body: > @droot Let me leave the variable lookup for now. I think it's better to enhance when upgrading the library Looks like [v1.0.0](https://github.com/googleapis/go-genai/releases/tag/v1.0.0) has been just released.

sounds good. Feel free to send the follow up PR. Thanks once again.
author_association: MEMBER