Google Vertex AI

Google Vertex AI has an OpenAI-compatible endpoint, but it takes a short-lived OAuth access token rather than a static API key. The token expires after 1 hour by default, so it doesn’t fit a static secret.

Instead, store a service account key as a secret, and have the task request a token when it starts.

Setting	Value
Base URL	`https://aiplatform.googleapis.com/v1/projects/<project>/locations/<location>/endpoints/openapi`
Credentials	Google Cloud service account key
Environment variable	None

Store your service account key

Create a file secret from the key file:

        
    
flyte create secret vertex-service-account --from-file <path-to-key.json> --project <your-project> --domain development

Call Vertex AI from a task

call_vertex_ai.py
                
                    
                
            
                
            

                
                    
                
            
import flyte

BASE_URL = "https://aiplatform.googleapis.com/v1/projects/<project>/locations/<location>/endpoints/openapi"

env = flyte.TaskEnvironment(
    name="vertex-ai-inference",
    image=flyte.Image.from_debian_base().with_pip_packages("google-auth[requests]", "openai"),
    secrets=flyte.Secret(key="vertex-service-account", mount="/etc/flyte/secrets"),
)

@env.task
async def summarize(text: str) -> str:
    from google.auth.transport.requests import Request
    from google.oauth2 import service_account
    from openai import AsyncOpenAI

    credentials = service_account.Credentials.from_service_account_file(
        "/etc/flyte/secrets/vertex-service-account",
        scopes=["https://www.googleapis.com/auth/cloud-platform"],
    )
    credentials.refresh(Request())

    client = AsyncOpenAI(base_url=BASE_URL, api_key=credentials.token)
    response = await client.chat.completions.create(
        model="<model-id>",
        messages=[{"role": "user", "content": f"Summarize in two sentences:\n\n{text}"}],
    )
    return response.choices[0].message.content

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(summarize, text="First document...")
    print(run.url)

Replace <project>, <location>, and <model-id> with your own values.

Call Vertex AI from an app

An AppEnvironment takes the same secrets argument, and the app reads the same environment variable. See Call the provider from an app.