fix: fix vertex provider

Kujtim Hoxha created 4 months ago

Change summary

internal/config/load.go           | 11 +++++------
internal/llm/provider/vertexai.go |  4 ++--
2 files changed, 7 insertions(+), 8 deletions(-)

Detailed changes

internal/config/load.go

@@ -168,8 +168,8 @@ func (c *Config) configureProviders(env env.Env, resolver VariableResolver, know
 				}
 				continue
 			}
-			prepared.ExtraParams["project"] = env.Get("GOOGLE_CLOUD_PROJECT")
-			prepared.ExtraParams["location"] = env.Get("GOOGLE_CLOUD_LOCATION")
+			prepared.ExtraParams["project"] = env.Get("VERTEXAI_PROJECT")
+			prepared.ExtraParams["location"] = env.Get("VERTEXAI_LOCATION")
 		case catwalk.InferenceProviderAzure:
 			endpoint, err := resolver.ResolveValue(p.APIEndpoint)
 			if err != nil || endpoint == "" {
@@ -470,10 +470,9 @@ func loadFromReaders(readers []io.Reader) (*Config, error) {
 }
 
 func hasVertexCredentials(env env.Env) bool {
-	useVertex := env.Get("GOOGLE_GENAI_USE_VERTEXAI") == "true"
-	hasProject := env.Get("GOOGLE_CLOUD_PROJECT") != ""
-	hasLocation := env.Get("GOOGLE_CLOUD_LOCATION") != ""
-	return useVertex && hasProject && hasLocation
+	hasProject := env.Get("VERTEXAI_PROJECT") != ""
+	hasLocation := env.Get("VERTEXAI_LOCATION") != ""
+	return hasProject && hasLocation
 }
 
 func hasAWSCredentials(env env.Env) bool {

internal/llm/provider/vertexai.go

@@ -10,8 +10,8 @@ import (
 type VertexAIClient ProviderClient
 
 func newVertexAIClient(opts providerClientOptions) VertexAIClient {
-	project := opts.extraHeaders["project"]
-	location := opts.extraHeaders["location"]
+	project := opts.extraParams["project"]
+	location := opts.extraParams["location"]
 	client, err := genai.NewClient(context.Background(), &genai.ClientConfig{
 		Project:  project,
 		Location: location,