setups/vertex_ai_setup.py

from typing import Optional
from dotenv import load_dotenv
from google.oauth2 import service_account
import vertexai
import os

load_dotenv()

# TODO: fix it in spaces
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "gcp-default-creds.json"
GOOGLE_APPLICATION_CREDENTIALS = os.environ["GOOGLE_APPLICATION_CREDENTIALS"]
VERTEXAI_PROJECT = os.environ["VERTEXAI_PROJECT"]


def initialize_vertexai_params(location: Optional[str] = "us-central1"):
    service_account.Credentials.from_service_account_file(
        filename=GOOGLE_APPLICATION_CREDENTIALS,
        scopes=["https://www.googleapis.com/auth/cloud-platform"],
    )
    vertexai.init(project=VERTEXAI_PROJECT, location=location)