Create gemini client based on project and location

- Keep the api key based one for local execution.
- All the necessary parameters will be provided via environment variables.

Bug: b/455899615
Change-Id: I2f7d9a409951958e340cc1e75a870ecf7dc1e20d
Reviewed-on: https://skia-review.googlesource.com/c/buildbot/+/1096044
Reviewed-by: Farid (Mojtaba) Faridzad <faridzad@google.com>
Commit-Queue: Ashwin Verleker <ashwinpv@google.com>

rag/go/api/server/server.go

diff --git a/rag/go/api/server/server.go b/rag/go/api/server/server.go
index f7e85c7..0ee3301 100644
--- a/rag/go/api/server/server.go
+++ b/rag/go/api/server/server.go
@@ -5,7 +5,6 @@
 	"fmt"
 	"net"
 	"net/http"
-	"os"
 
 	"cloud.google.com/go/spanner"
 	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
@@ -129,18 +128,10 @@
 	// Initialize metrics/
 	metrics2.InitPrometheus(flags.PromPort)
 
-	// Get the api key from the env.
-	// TODO(ashwinpv): Get this from the secrets key store if not in the environment.
-	apiKey := os.Getenv("GEMINI_API_KEY")
-
-	if apiKey == "" {
-		sklog.Fatalf("GEMINI_API_KEY environment variable is not set.")
-	}
-
 	// Define the list of services to be hosted based on the "services" flag.
 	serviceList := []Service{}
 	var serviceMap = map[string]Service{
-		"history": history.NewApiService(ctx, server.dbClient, apiKey, server.queryEmbeddingModel),
+		"history": history.NewApiService(ctx, server.dbClient, server.queryEmbeddingModel),
 	}
 	for _, serviceName := range flags.Services.Value() {
 		service, ok := serviceMap[serviceName]

rag/go/api/services/history/apiService.go

diff --git a/rag/go/api/services/history/apiService.go b/rag/go/api/services/history/apiService.go
index 1174ff9..4f913a2 100644
--- a/rag/go/api/services/history/apiService.go
+++ b/rag/go/api/services/history/apiService.go
@@ -2,6 +2,7 @@
 
 import (
 	"context"
+	"os"
 
 	"cloud.google.com/go/spanner"
 	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
@@ -15,6 +16,12 @@
 	pb "go.skia.org/infra/rag/proto/history/v1"
 )
 
+const (
+	geminiApiKeyEnvVar   = "GEMINI_API_KEY"
+	geminiProjectEnvVar  = "GEMINI_PROJECT"
+	geminiLocationEnvVar = "GEMINI_LOCATION"
+)
+
 // ApiService provides a struct for the HistoryRag api implementation.
 type ApiService struct {
 	pb.UnimplementedHistoryRagApiServiceServer
@@ -33,8 +40,25 @@
 }
 
 // NewApiService returns a new instance of the ApiService struct.
-func NewApiService(ctx context.Context, dbClient *spanner.Client, geminiApiKey string, queryEmbeddingModel string) *ApiService {
-	genAiClient, err := genai.NewGeminiClient(ctx, geminiApiKey)
+func NewApiService(ctx context.Context, dbClient *spanner.Client, queryEmbeddingModel string) *ApiService {
+	var genAiClient *genai.GeminiClient
+	var err error
+	// Get the api key from the env.
+	apiKey := os.Getenv(geminiApiKeyEnvVar)
+
+	if apiKey != "" {
+		sklog.Infof("Gemini api key specified in the environment, creating a local client.")
+		genAiClient, err = genai.NewLocalGeminiClient(ctx, apiKey)
+	} else {
+		projectId := os.Getenv(geminiProjectEnvVar)
+		location := os.Getenv(geminiLocationEnvVar)
+		if projectId == "" || location == "" {
+			sklog.Fatalf("%s and %s environment variables need to be set.", geminiProjectEnvVar, geminiLocationEnvVar)
+		}
+		sklog.Infof("Creating a new Gemini client for project %s and location %s", projectId, location)
+		genAiClient, err = genai.NewGeminiClient(ctx, projectId, location)
+	}
+
 	if err != nil {
 		sklog.Errorf("Error creating new gemini client: %v", err)
 		return nil

rag/go/genai/genai_client.go

diff --git a/rag/go/genai/genai_client.go b/rag/go/genai/genai_client.go
index 0e54a85..123c09e 100644
--- a/rag/go/genai/genai_client.go
+++ b/rag/go/genai/genai_client.go
@@ -20,7 +20,26 @@
 }
 
 // NewGeminiClient returns a new instance of the GeminiClient.
-func NewGeminiClient(ctx context.Context, apiKey string) (*GeminiClient, error) {
+func NewGeminiClient(ctx context.Context, project, location string) (*GeminiClient, error) {
+	genAiClient, err := genai.NewClient(ctx, &genai.ClientConfig{
+		Backend:  genai.BackendGeminiAPI,
+		Project:  project,
+		Location: location,
+	})
+	if err != nil {
+		sklog.Errorf("Error creating new gemini client: %v", err)
+		return nil, err
+	}
+
+	return &GeminiClient{
+		genAiClient: genAiClient,
+	}, nil
+}
+
+// NewLocalGeminiClient returns a new instance of the GeminiClient using an api key.
+//
+// This is not intended for production purposes.
+func NewLocalGeminiClient(ctx context.Context, apiKey string) (*GeminiClient, error) {
 	genAiClient, err := genai.NewClient(ctx, &genai.ClientConfig{
 		APIKey:  apiKey,
 		Backend: genai.BackendGeminiAPI,