Working with Bedrock in Pixeltable

%pip install -qU pixeltable boto3

import getpass
import os

if 'BEDROCK_API_KEY' not in os.environ:
    os.environ['BEDROCK_API_KEY'] = getpass.getpass(
        'Enter your Bedrock API Key: '
    )

# Optional: set the region if your Bedrock endpoint is not in us-east-1
# os.environ['BEDROCK_REGION_NAME'] = 'us-west-2'

import pixeltable as pxt

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

from pixeltable.functions import bedrock

# Create a table in Pixeltable and add a computed column that calls Bedrock

t = pxt.create_table('bedrock_demo/chat', {'input': pxt.String})

t.add_computed_column(
    output=bedrock.converse(
        model_id='amazon.nova-pro-v1:0',
        messages=[{'role': 'user', 'content': [{'text': t.input}]}],
    )
)

# Parse the response into a new column
t.add_computed_column(response=t.output.output.message.content[0].text)

# Start a conversation
t.insert(
    [
        {
            'input': 'What was the outcome of the 1904 US Presidential election?'
        }
    ]
)
t.select(t.input, t.response).show()