Working with Mistral AI in Pixeltable

%pip install -qU pixeltable mistralai

import getpass
import os

if 'MISTRAL_API_KEY' not in os.environ:
    os.environ['MISTRAL_API_KEY'] = getpass.getpass('Mistral AI API Key:')

import pixeltable as pxt

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

from pixeltable.functions.mistralai import chat_completions

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

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

messages = [{'role': 'user', 'content': t.input}]
t.add_computed_column(
    output=chat_completions(
        messages=messages,
        model='mistral-small-latest',
        model_kwargs={
            # Optional dict with parameters for the Mistral 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': 'What three species of fish have the highest mercury content?'
        }
    ]
)
t.select(t.input, t.response).show()