Skip to content

Commit

Permalink
chore: switch controllerbuilder to use vertexai gemini endpoint
Browse files Browse the repository at this point in the history
This means we don't need an API Key (but we do need a GCP account)
  • Loading branch information
justinsb committed Dec 18, 2024
1 parent 95b3225 commit 4d4d2ed
Show file tree
Hide file tree
Showing 9 changed files with 165 additions and 102 deletions.
63 changes: 32 additions & 31 deletions dev/tools/controllerbuilder/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
120 changes: 62 additions & 58 deletions dev/tools/controllerbuilder/go.sum

Large diffs are not rendered by default.

55 changes: 55 additions & 0 deletions dev/tools/controllerbuilder/pkg/llm/gemini.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
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
}
10 changes: 7 additions & 3 deletions dev/tools/controllerbuilder/pkg/toolbot/csv.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)
Expand Down Expand Up @@ -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)
}
Expand All @@ -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()...)

Expand Down
2 changes: 1 addition & 1 deletion dev/tools/controllerbuilder/pkg/toolbot/datapoint.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ import (
"io"
"strings"

"k8s.io/klog"
"k8s.io/klog/v2"
)

// ExtractToolMarkers extracts tool markers from source code.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand Down
2 changes: 1 addition & 1 deletion dev/tools/controllerbuilder/pkg/typeupdater/typeupdater.go
Original file line number Diff line number Diff line change
Expand Up @@ -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="
Expand Down

0 comments on commit 4d4d2ed

Please sign in to comment.