Working with Groq in Pixeltable

%pip install -qU pixeltable groq

import getpass
import os

if 'GROQ_API_KEY' not in os.environ:
    os.environ['GROQ_API_KEY'] = getpass.getpass(
        'Enter your Groq API key:'
    )

import pixeltable as pxt

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

from pixeltable.functions import groq

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

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

messages = [{'role': 'user', 'content': t.input}]
t.add_computed_column(
    output=groq.chat_completions(
        messages=messages,
        model='llama-3.3-70b-versatile',
        model_kwargs={
            # Optional dict with parameters for the Groq API
            'max_tokens': 300,
            'top_p': 0.9,
            'temperature': 0.7,
        },
    )
)

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

# Start a conversation
t.insert(
    [{'input': 'How many islands are in the Aleutian island chain?'}]
)
t.select(t.input, t.response).head()