Working with gemini

%pip install -qU pixeltable google-genai

import os

vertex_enabled = (
    os.environ.get('GOOGLE_GENAI_USE_VERTEXAI', '').lower() == 'true'
)

# Option 1: Google AI Studio (API key)
# Set GEMINI_API_KEY or GOOGLE_API_KEY in your environment,
# or add api_key to the [gemini] section of $PIXELTABLE_HOME/config.toml

if (
    not vertex_enabled
    and 'GEMINI_API_KEY' not in os.environ
    and 'GOOGLE_API_KEY' not in os.environ
):
    import getpass

    os.environ['GEMINI_API_KEY'] = getpass.getpass(
        'Google AI Studio API Key:'
    )

# Option 2: Vertex AI (Application Default Credentials)
# Uncomment and set the following environment variables to use Vertex AI instead:
# os.environ['GOOGLE_GENAI_USE_VERTEXAI'] = 'true'
# os.environ['GOOGLE_CLOUD_PROJECT'] = 'your-project-id'
# os.environ['GOOGLE_CLOUD_LOCATION'] = 'us-central1'  # optional, defaults to us-central1
# Then authenticate via: gcloud auth application-default login

import pixeltable as pxt

# Remove the 'gemini_demo' directory and its contents, if it exists
pxt.drop_dir('gemini_demo', force=True)
pxt.create_dir('gemini_demo')

Connected to Pixeltable database at: postgresql+psycopg://postgres:@/pixeltable?host=/Users/asiegel/.pixeltable/pgdata
Created directory 'gemini_demo'.

<pixeltable.catalog.dir.Dir at 0x17c233250>

from google.genai.types import GenerateContentConfigDict
from pixeltable.functions import gemini

# Create a table in Pixeltable and pick a model hosted on Google AI Studio with some parameters

t = pxt.create_table('gemini_demo/text', {'input': pxt.String})

config = GenerateContentConfigDict(
    stop_sequences=['\n'],
    max_output_tokens=300,
    temperature=1.0,
    top_p=0.95,
    top_k=40,
)
t.add_computed_column(
    output=gemini.generate_content(
        t.input, model='gemini-2.5-flash', config=config
    )
)

Created table 'text'.
Added 0 column values with 0 errors in 0.01 s

No rows affected.

# Ask Gemini to generate some content based on the input
t.insert(
    [
        {'input': 'Write a story about a magic backpack.'},
        {'input': 'Tell me a science joke.'},
    ]
)

Inserted 2 rows with 0 errors in 1.43 s (1.39 rows/s)

2 rows inserted.

# Parse the response into a new column
t.add_computed_column(
    response=t.output['candidates'][0]['content']['parts'][0]['text']
)
t.select(t.input, t.response).head()

Added 2 column values with 0 errors in 0.03 s (62.79 rows/s)

from google.genai.types import GenerateImagesConfigDict

images_t = pxt.create_table('gemini_demo/images', {'prompt': pxt.String})

config = GenerateImagesConfigDict(aspect_ratio='16:9')
images_t.add_computed_column(
    generated_image=gemini.generate_images(
        images_t.prompt, model='imagen-4.0-generate-001', config=config
    )
)

Created table 'images'.
Added 0 column values with 0 errors in 0.01 s

No rows affected.

images_t.insert(
    [{'prompt': 'A friendly dinosaur playing tennis in a cornfield'}]
)

Inserted 1 row with 0 errors in 9.41 s (0.11 rows/s)

1 row inserted.

images_t.head()

videos_t = pxt.create_table('gemini_demo/videos', {'prompt': pxt.String})

videos_t.add_computed_column(
    generated_video=gemini.generate_videos(
        videos_t.prompt, model='veo-2.0-generate-001'
    )
)

Created table 'videos'.
Added 0 column values with 0 errors in 0.01 s

No rows affected.

videos_t.insert(
    [
        {
            'prompt': 'A giant pixel floating over the open ocean in a sea of data'
        }
    ]
)

Inserted 1 row with 0 errors in 46.23 s (0.02 rows/s)

1 row inserted.

videos_t.head()

images_t.add_computed_column(
    generated_video=gemini.generate_videos(
        image=images_t.generated_image, model='veo-2.0-generate-001'
    )
)

Added 1 column value with 0 errors in 40.00 s (0.03 rows/s)

1 row updated.

images_t.head()