diff --git a/Taskfile.yml b/Taskfile.yml index 77c6906d9..13fa084a1 100644 --- a/Taskfile.yml +++ b/Taskfile.yml @@ -63,7 +63,7 @@ tasks: test: aliases: [t] desc: Run tests - cmd: CAGENT_MODELS_GATEWAY= OPENAI_API_KEY= ANTHROPIC_API_KEY= GOOGLE_API_KEY= MISTRAL_API_KEY= GITHUB_TOKEN= go test {{.CLI_ARGS}} ./... + cmd: CAGENT_MODELS_GATEWAY= OPENAI_API_KEY= ANTHROPIC_API_KEY= GOOGLE_API_KEY= GOOGLE_GENAI_USE_VERTEXAI= MISTRAL_API_KEY= GITHUB_TOKEN= go test {{.CLI_ARGS}} ./... build-local: desc: Build binaries for local host platform diff --git a/pkg/config/auto.go b/pkg/config/auto.go index 73ff67062..ef395e419 100644 --- a/pkg/config/auto.go +++ b/pkg/config/auto.go @@ -22,7 +22,11 @@ type providerConfig struct { var cloudProviders = []providerConfig{ {"anthropic", []string{"ANTHROPIC_API_KEY"}, "ANTHROPIC_API_KEY"}, {"openai", []string{"OPENAI_API_KEY"}, "OPENAI_API_KEY"}, - {"google", []string{"GOOGLE_API_KEY"}, "GOOGLE_API_KEY"}, + {"google", []string{ + "GOOGLE_API_KEY", + "GEMINI_API_KEY", + "GOOGLE_GENAI_USE_VERTEXAI", + }, "GOOGLE_API_KEY (or GEMINI_API_KEY, GOOGLE_GENAI_USE_VERTEXAI)"}, {"mistral", []string{"MISTRAL_API_KEY"}, "MISTRAL_API_KEY"}, {"amazon-bedrock", []string{ "AWS_BEARER_TOKEN_BEDROCK", diff --git a/pkg/config/gather.go b/pkg/config/gather.go index 60d990db5..435a9348f 100644 --- a/pkg/config/gather.go +++ b/pkg/config/gather.go @@ -111,7 +111,12 @@ func addEnvVarsForModelConfig(model *latest.ModelConfig, customProviders map[str requiredEnv["ANTHROPIC_API_KEY"] = true case "google": if model.ProviderOpts["project"] == nil && model.ProviderOpts["location"] == nil { - requiredEnv["GOOGLE_API_KEY"] = true + if os.Getenv("GOOGLE_GENAI_USE_VERTEXAI") != "" { + requiredEnv["GOOGLE_CLOUD_PROJECT"] = true + requiredEnv["GOOGLE_CLOUD_LOCATION"] = true + } else if _, exist := os.LookupEnv("GEMINI_API_KEY"); !exist { + requiredEnv["GOOGLE_API_KEY"] = true + } } case "mistral": requiredEnv["MISTRAL_API_KEY"] = true diff --git a/pkg/model/provider/gemini/client.go b/pkg/model/provider/gemini/client.go index 02e6b0d55..fac8b9f3a 100644 --- a/pkg/model/provider/gemini/client.go +++ b/pkg/model/provider/gemini/client.go @@ -76,12 +76,22 @@ func NewClient(ctx context.Context, cfg *latest.ModelConfig, env environment.Pro return nil, errors.New("location must be set") } + backend = genai.BackendVertexAI + httpClient = nil // Use default client + } else if _, exist := env.Get(ctx, "GOOGLE_GENAI_USE_VERTEXAI"); exist { + project, _ = env.Get(ctx, "GOOGLE_CLOUD_PROJECT") + location, _ = env.Get(ctx, "GOOGLE_CLOUD_LOCATION") backend = genai.BackendVertexAI httpClient = nil // Use default client } else { - apiKey, _ = env.Get(ctx, "GOOGLE_API_KEY") + if value, exist := env.Get(ctx, "GEMINI_API_KEY"); exist { + apiKey = value + } + if value, exist := env.Get(ctx, "GOOGLE_API_KEY"); exist { + apiKey = value + } if apiKey == "" { - return nil, errors.New("GOOGLE_API_KEY environment variable is required") + return nil, errors.New("GOOGLE_API_KEY or GEMINI_API_KEY environment variable is required") } backend = genai.BackendGeminiAPI