feat: add configurable Gemini provider with VertexAI and GeminiAPI support (#240)

- Add support for Gemini provider configuration, including project, location, backend, and API key options
- Ensure gemini.api_key is treated as sensitive and hidden in config listing
- Allow Gemini client to fall back to openai.api_key if gemini.api_key is not set
- Pass Gemini backend, project, and location settings to the Gemini client
- Refactor Gemini client initialization to support both VertexAI and GeminiAPI backends with appropriate configuration
- Add validation for required project and location fields when using the VertexAI backend
- Introduce new option functions: WithProject, WithLocation, and WithBackend for flexible Gemini configuration

Signed-off-by: Bo-Yi Wu <[email protected]>

Author: appleboy

cmd/config_list.go

@@ -37,6 +37,10 @@ var availableKeys = map[string]string{
 	"openai.frequency_penalty": "Parameter to reduce repetition by penalizing tokens based on their frequency",
 	"openai.presence_penalty":  "Parameter to encourage topic diversity by penalizing previously used tokens",
 	"prompt.folder":            "Directory path for custom prompt templates",
+	"gemini.project":           "VertexAI project for Gemini provider",
+	"gemini.location":          "VertexAI location for Gemini provider",
+	"gemini.backend":           "Gemini backend (BackendGeminiAPI or BackendVertexAI)",
+	"gemini.api_key":           "API key for Gemini provider",
 }
 
 // configListCmd represents the command to list the configuration values.
@@ -65,7 +69,7 @@ var configListCmd = &cobra.Command{
 		// Add the key and value to the table
 		for _, v := range keys {
 			// Hide the api key
-			if v == "openai.api_key" {
+			if v == "openai.api_key" || v == "gemini.api_key" {
 				tbl.AddRow(v, "****************")
 				continue
 			}

cmd/config_set.go

@@ -33,6 +33,12 @@ func init() {
 	configSetCmd.Flags().StringP("headers", "", "", availableKeys["openai.headers"])
 	configSetCmd.Flags().StringP("api_version", "", "", availableKeys["openai.api_version"])
 	configSetCmd.Flags().StringP("prompt_folder", "", "", availableKeys["prompt.folder"])
+	// Gemini flags
+	configSetCmd.Flags().String("gemini.project", "", availableKeys["gemini.project"])
+	configSetCmd.Flags().String("gemini.location", "", availableKeys["gemini.location"])
+	configSetCmd.Flags().String("gemini.backend", "BackendGeminiAPI", availableKeys["gemini.backend"])
+	configSetCmd.Flags().String("gemini.api_key", "", availableKeys["gemini.api_key"])
+
 	_ = viper.BindPFlag("openai.base_url", configSetCmd.Flags().Lookup("base_url"))
 	_ = viper.BindPFlag("openai.org_id", configSetCmd.Flags().Lookup("org_id"))
 	_ = viper.BindPFlag("openai.api_key", configSetCmd.Flags().Lookup("api_key"))
@@ -52,6 +58,10 @@ func init() {
 	_ = viper.BindPFlag("openai.headers", configSetCmd.Flags().Lookup("headers"))
 	_ = viper.BindPFlag("openai.api_version", configSetCmd.Flags().Lookup("api_version"))
 	_ = viper.BindPFlag("prompt.folder", configSetCmd.Flags().Lookup("prompt_folder"))
+	_ = viper.BindPFlag("gemini.project", configSetCmd.Flags().Lookup("gemini.project"))
+	_ = viper.BindPFlag("gemini.location", configSetCmd.Flags().Lookup("gemini.location"))
+	_ = viper.BindPFlag("gemini.backend", configSetCmd.Flags().Lookup("gemini.backend"))
+	_ = viper.BindPFlag("gemini.api_key", configSetCmd.Flags().Lookup("gemini.api_key"))
 }
 
 // configSetCmd updates the config value.

cmd/provider.go

@@ -35,13 +35,20 @@ func NewOpenAI() (*openai.Client, error) {
 
 // NewGemini returns a new Gemini client
 func NewGemini(ctx context.Context) (*gemini.Client, error) {
+	apiKey := viper.GetString("gemini.api_key")
+	if apiKey == "" {
+		apiKey = viper.GetString("openai.api_key")
+	}
 	return gemini.New(
 		ctx,
-		gemini.WithToken(viper.GetString("openai.api_key")),
+		gemini.WithToken(apiKey),
 		gemini.WithModel(viper.GetString("openai.model")),
 		gemini.WithMaxTokens(viper.GetInt32("openai.max_tokens")),
 		gemini.WithTemperature(float32(viper.GetFloat64("openai.temperature"))),
 		gemini.WithTopP(float32(viper.GetFloat64("openai.top_p"))),
+		gemini.WithBackend(viper.GetString("gemini.backend")),
+		gemini.WithProject(viper.GetString("gemini.project")),
+		gemini.WithLocation(viper.GetString("gemini.location")),
 	)
 }
 

provider/gemini/gemini.go

@@ -159,11 +159,26 @@ func New(ctx context.Context, opts ...Option) (c *Client, err error) {
 		},
 	}
 
-	client, err := genai.NewClient(ctx, &genai.ClientConfig{
-		APIKey:     cfg.token,
-		HTTPClient: httpClient,
-		Backend:    genai.BackendGeminiAPI,
-	})
+	var clientConfig *genai.ClientConfig
+	switch cfg.backend {
+	case genai.BackendVertexAI:
+		clientConfig = &genai.ClientConfig{
+			HTTPClient: httpClient,
+			Backend:    cfg.backend,
+			Project:    cfg.project,
+			Location:   cfg.location,
+		}
+	case genai.BackendGeminiAPI, genai.BackendUnspecified:
+		fallthrough
+	default:
+		cfg.backend = genai.BackendGeminiAPI
+		clientConfig = &genai.ClientConfig{
+			APIKey:     cfg.token,
+			HTTPClient: httpClient,
+			Backend:    cfg.backend,
+		}
+	}
+	client, err := genai.NewClient(ctx, clientConfig)
 	if err != nil {
 		return nil, err
 	}

provider/gemini/options.go

@@ -2,11 +2,14 @@ package gemini
 
 import (
 	"errors"
+
+	"google.golang.org/genai"
 )
 
 var (
-	errorsMissingToken = errors.New("missing gemini api key")
-	errorsMissingModel = errors.New("missing model")
+	errorsMissingToken          = errors.New("missing gemini api key")
+	errorsMissingTokenOrProject = errors.New("missing token or project")
+	errorsMissingModel          = errors.New("missing model")
 )
 
 const (
@@ -79,45 +82,71 @@ func WithTopP(val float32) Option {
 	})
 }
 
-// config is a struct that stores configuration options for the instrumentation.
 type config struct {
 	token       string
 	model       string
 	maxTokens   int32
 	temperature float32
 	topP        float32
+	project     string
+	location    string
+	backend     genai.Backend
 }
 
-// valid checks whether a config object is valid, returning an error if it is not.
 func (cfg *config) valid() error {
-	// Check that the token is not empty.
-	if cfg.token == "" {
-		return errorsMissingToken
+	if cfg.backend == genai.BackendVertexAI {
+		if cfg.project == "" || cfg.location == "" {
+			return errorsMissingTokenOrProject
+		}
+	} else {
+		if cfg.token == "" {
+			return errorsMissingToken
+		}
 	}
-
 	if cfg.model == "" {
 		return errorsMissingModel
 	}
-
-	// If all checks pass, return nil (no error).
 	return nil
 }
 
+func WithLocation(val string) Option {
+	return optionFunc(func(c *config) {
+		c.location = val
+	})
+}
+
+func WithBackend(val string) Option {
+	return optionFunc(func(c *config) {
+		switch val {
+		case "BackendVertexAI":
+			c.backend = genai.BackendVertexAI
+		case "BackendGeminiAPI":
+			c.backend = genai.BackendGeminiAPI
+		case "BackendUnspecified":
+			fallthrough
+		default:
+			c.backend = genai.BackendGeminiAPI
+		}
+	})
+}
+
+func WithProject(val string) Option {
+	return optionFunc(func(c *config) {
+		c.project = val
+	})
+}
+
 // newConfig creates a new config object with default values, and applies the given options.
 func newConfig(opts ...Option) *config {
-	// Create a new config object with default values.
 	c := &config{
 		model:       defaultModel,
 		maxTokens:   defaultMaxTokens,
 		temperature: defaultTemperature,
 		topP:        defaultTopP,
+		backend:     genai.BackendGeminiAPI,
 	}
-
-	// Apply each of the given options to the config object.
 	for _, opt := range opts {
 		opt.apply(c)
 	}
-
-	// Return the resulting config object.
 	return c
 }