Working with OpenAI in Pixeltable

%pip install -qU pixeltable openai

import getpass
import os

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

import os

os.environ['OPENAI_API_KEY'] = '...'  # Azure OpenAI key
os.environ['OPENAI_BASE_URL'] = 'https://YOUR_RESOURCE.openai.azure.com/openai/v1/'
os.environ['OPENAI_API_VERSION'] = 'preview'

import pixeltable as pxt

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

from pixeltable.functions import openai

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

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

messages = [{'role': 'user', 'content': t.input}]
t.add_computed_column(
    output=openai.chat_completions(
        messages=messages,
        model='gpt-4o-mini',
        model_kwargs={
            # Optional dict with parameters for the OpenAI 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()

# Create a table and add a computed column using the Responses API

resp_t = pxt.create_table('openai_demo/responses', {'input': pxt.String})

messages = [{'role': 'user', 'content': resp_t.input}]
resp_t.add_computed_column(
    output=openai.responses(
        input=messages,
        model='gpt-4o-mini',
        model_kwargs={
            'instructions': 'You are a helpful assistant. Be concise.',
            'temperature': 0.7,
            'max_output_tokens': 300,
        },
    )
)

# The response text is the first content part of the first output item
resp_t.add_computed_column(
    response=resp_t.output.output[0].content[0].text
)

# Insert data and view results
resp_t.insert([{'input': 'What are the three largest moons of Jupiter?'}])
resp_t.select(resp_t.input, resp_t.response).head()

@pxt.udf
def reply_prompt(complaint: str, supported: float | None) -> list | None:
    if supported is None or supported < 0.5:
        return None
    return [
        {
            'role': 'user',
            'content': f'Draft a two-sentence apology with a refund for: {complaint}',
        }
    ]


dec_t = pxt.create_table(
    'openai_demo/complaints',
    {'complaint': pxt.String, 'photo': pxt.Image},
)
messages = [
    {
        'role': 'user',
        'content': [
            {'type': 'input_text', 'text': dec_t.complaint},
            {'type': 'input_image', 'image_url': dec_t.photo},
        ],
    }
]
questions = [
    {
        'type': 'predicate',
        'name': 'supported',
        'instructions': "Does the photo support the customer's complaint?",
    },
    {
        'type': 'choice',
        'name': 'issue',
        'instructions': 'What is the complaint about?',
        'choices': [
            {'value': 'missing item'},
            {'value': 'unwanted item'},
            {'value': 'food quality'},
            {'value': 'other'},
        ],
    },
    {
        'type': 'score',
        'name': 'urgency',
        'instructions': 'How upset is the customer?',
        'levels': [
            {'label': 'calm'},
            {'label': 'annoyed'},
            {'label': 'angry'},
        ],
    },
]
dec_t.add_computed_column(
    triage=openai.decisions(
        messages, questions=questions, model='gpt-6-luna'
    )
)
dec_t.add_computed_column(supported=dec_t.triage.answers[0].probability)
dec_t.add_computed_column(issue=dec_t.triage.answers[1].choice)
dec_t.add_computed_column(urgency=dec_t.triage.answers[2].score)
dec_t.add_computed_column(
    reply=openai.responses(
        reply_prompt(dec_t.complaint, dec_t.supported),
        model='gpt-4o-mini',
    )
)
dec_t.add_computed_column(
    reply_text=dec_t.reply.output[0].content[0].text
)

photo_url = 'https://raw.githubusercontent.com/pixeltable/pixeltable/release/docs/resources/images/000000000009.jpg'
dec_t.insert(
    [
        {
            'complaint': 'I asked for no fruit, but the box came with oranges and pineapple. Not happy.',
            'photo': photo_url,
        },
        {
            'complaint': 'My order is missing the broccoli I paid for!',
            'photo': photo_url,
        },
    ]
)

dec_t.order_by(dec_t.urgency, asc=False).select(
    dec_t.complaint,
    dec_t.issue,
    dec_t.supported,
    dec_t.urgency,
    dec_t.reply_text,
).collect()

emb_t = pxt.create_table('openai_demo/embeddings', {'input': pxt.String})
emb_t.add_computed_column(
    embedding=openai.embeddings(
        input=emb_t.input, model='text-embedding-3-small'
    )
)

emb_t.insert(
    [{'input': 'OpenAI provides a variety of embeddings models.'}]
)

emb_t.head()

image_t = pxt.create_table('openai_demo/images', {'input': pxt.String})
image_t.add_computed_column(
    img=openai.image_generations(image_t.input, model='gpt-image-2')
)

image_t.insert(
    [
        {
            'input': 'A giant Pixel floating in the open ocean in a sea of data'
        }
    ]
)

image_t

image_t.head()

audio_t = pxt.create_table('openai_demo/audio', {'input': pxt.Audio})
audio_t.add_computed_column(
    result=openai.transcriptions(
        audio_t.input,
        model='whisper-1',
        model_kwargs={
            'language': 'en',
            'prompt': 'Transcribe the contents of this recording.',
        },
    )
)

url = (
    'https://github.com/pixeltable/pixeltable/raw/release/tests/data/audio/'
    'jfk_1961_0109_cityuponahill-excerpt.flac'
)
audio_t.insert([{'input': url}])

audio_t.head()

audio_t.head()[0]['result']['text']