Working with Fireworks AI in Pixeltable

%pip install -qU pixeltable fireworks-ai

import getpass
import os

if 'FIREWORKS_API_KEY' not in os.environ:
    os.environ['FIREWORKS_API_KEY'] = getpass.getpass(
        'Fireworks API Key:'
    )

import pixeltable as pxt

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

from pixeltable.functions.fireworks import chat_completions

# Create a table in Pixeltable and pick a model hosted on Fireworks with some parameters

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

messages = [{'role': 'user', 'content': t.input}]
t.add_computed_column(
    output=chat_completions(
        messages=messages,
        model='accounts/fireworks/models/gpt-oss-20b',
        model_kwargs={
            # Optional dict with parameters for the Fireworks API
            'max_tokens': 300,
            'top_k': 40,
            'top_p': 0.9,
            'temperature': 0.7,
        },
    )
)

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

# Start a conversation
t.insert(
    [{'input': 'Can you tell me who was President of the US in 1961?'}]
)
t.select(t.input, t.response).show()