Working with OpenAI in Pixeltable

Pixeltable unifies data and computation into a table interface. In the Pixeltable Basics tutorial, we saw how OpenAI API calls can be incorporated into Pixeltable workflows. In this tutorial, we'll go into more depth on OpenAI integration. You'll need an OpenAI API key to run this demo.

Prerequisites

• An OpenAI account with an API key (https://openai.com/index/openai-api/)

Important Notes

• OpenAI usage may incur costs based on your OpenAI plan.
• Be mindful of sensitive data and consider security measures when integrating with external services.

First you'll need to install required libraries and enter your OpenAI API key.

%pip install -qU pixeltable openai

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

Enter your OpenAI API key: ········

Now let's create a Pixeltable directory to hold the tables for our demo.

import pixeltable as pxt

pxt.drop_dir('demo', force=True)  # Ensure a clean slate for the demo
pxt.create_dir('demo')

Connected to Pixeltable database at:
postgresql://postgres:@/pixeltable?host=/Users/asiegel/.pixeltable/pgdata
Created directory `demo`.
<pixeltable.catalog.dir.Dir at 0x314cef550>

Creating the Table

First, we'll create a table and populate it with some sample data.

t = pxt.create_table('demo.openai', {'id': pxt.Int, 'input': pxt.String})

Created table `openai`.

# text from https://en.wikipedia.org/wiki/Global_financial_crisis_in_September_2008

wikipedia_text = '''On Sunday, September 14, it was announced that Lehman Brothers would file for bankruptcy after the Federal Reserve Bank declined to participate in creating a financial support facility for Lehman Brothers.
The significance of the Lehman Brothers bankruptcy is disputed with some assigning it a pivotal role in the unfolding of subsequent events.
The principals involved, Ben Bernanke and Henry Paulson, dispute this view, citing a volume of toxic assets at Lehman which made a rescue impossible.[16][17] Immediately following the bankruptcy, JPMorgan Chase provided the broker dealer unit of Lehman Brothers with $138 billion to "settle securities transactions with customers of Lehman and its clearance parties" according to a statement made in a New York City Bankruptcy court filing.[18]
The same day, the sale of Merrill Lynch to Bank of America was announced.[19] The beginning of the week was marked by extreme instability in global stock markets, with dramatic drops in market values on Monday, September 15, and Wednesday, September 17.
On September 16, the large insurer American International Group (AIG), a significant participant in the credit default swaps markets, suffered a liquidity crisis following the downgrade of its credit rating.
The Federal Reserve, at AIG's request, and after AIG had shown that it could not find lenders willing to save it from insolvency, created a credit facility for up to US$85 billion in exchange for a 79.9% equity interest, and the right to suspend dividends to previously issued common and preferred stock.[20]'''

sample_inputs = wikipedia_text.split('\n')

# Insert a single sample row into the table
t.insert(id=0, input=sample_inputs[0])
t.show()

Inserting rows into `openai`: 1 rows [00:00, 333.54 rows/s]
Inserted 1 row with 0 errors.

Making OpenAI API calls

Calling OpenAI API endpoints involves constructing a message object, which we can express in Pixeltable by adding a new computed column.

prompt = "For the following sentence, extract all company names from the text."

msgs = [
    { "role": "system", "content": prompt },
    { "role": "user", "content": t.input }
]

t.add_column(input_msgs=msgs)

Computing cells: 100%|████████████████████████████████████████████| 1/1 [00:00<00:00, 80.67 cells/s]
Added 1 column value with 0 errors.

UpdateStatus(num_rows=1, num_computed_values=1, num_excs=0, updated_cols=[], cols_with_excs=[])

Unlike the values of theinput column, which users provide, the t.input_msgs column is computed automatically from the t.input column values:

t.show()

In Pixeltable, OpenAI API calls are exposed as Pixeltable functions, which can be used to create computed columns. We can run the prompt against our input data using OpenAI's chat_completions API.

from pixeltable.functions import openai

t['chat_output'] = openai.chat_completions(model='gpt-4o-mini', messages=t.input_msgs)
t.show()

Computing cells: 100%|████████████████████████████████████████████| 1/1 [00:00<00:00,  1.28 cells/s]
Added 1 column value with 0 errors.

The output of the OpenAI API calls are generally complex JSON structures, which require some navigation to extract the response. We can express this as JSON path expressions and create another computed column:

t['response'] = t.chat_output.choices[0].message.content
t.show()

Computing cells: 100%|████████████████████████████████████████████| 1/1 [00:00<00:00, 81.02 cells/s]
Added 1 column value with 0 errors.

Let's run a query to look only at the input and output:

t.select(t.input, t.response).show()

Once we have defined these computed columns, much like with a spreadsheet, newly inserted t.input values trigger computation of all derived columns, such as t.response. Let's insert the rest of our data and see how this works.

t.insert({'id': i, 'input': sample_inputs[i]} for i in range(1, len(sample_inputs)))

Computing cells: 100%|██████████████████████████████████████████| 20/20 [00:02<00:00,  9.19 cells/s]
Inserting rows into `openai`: 5 rows [00:00, 710.97 rows/s]
Computing cells: 100%|██████████████████████████████████████████| 20/20 [00:02<00:00,  9.13 cells/s]
Inserted 5 rows with 0 errors.

UpdateStatus(num_rows=5, num_computed_values=20, num_excs=0, updated_cols=[], cols_with_excs=[])

t.select(t.input, t.response).show()

Adding Ground Truth Data

Now let's see how Pixeltable can be used to evaluate a model against ground truth data. We'll start by manually populating a ground_truth column in our table.

t['ground_truth'] = pxt.StringType(nullable=True)

ground_truth = [
    'Lehman Brothers',
    'Lehman Brothers',
    'JP Morgan Chase, Lehman Brothers',
    'Merill Lynch, Bank of America',
    'American International Group',
    'American International Group',
]

for i, gt in enumerate(ground_truth):
    t.update({'ground_truth': gt}, where=(t.id == i))

Added 6 column values with 0 errors.
Inserting rows into `openai`: 1 rows [00:00, 726.16 rows/s]
Inserting rows into `openai`: 1 rows [00:00, 1356.50 rows/s]
Inserting rows into `openai`: 1 rows [00:00, 1471.17 rows/s]
Inserting rows into `openai`: 1 rows [00:00, 1718.27 rows/s]
Inserting rows into `openai`: 1 rows [00:00, 2376.38 rows/s]
Inserting rows into `openai`: 1 rows [00:00, 1886.78 rows/s]

Inserting rows into openai: 1 rows [00:00, 2308.37 rows/s]

Inserting rows into openai: 0 rows [00:00, ? rows/s]

Inserting rows into openai: 1 rows [00:00, 2432.89 rows/s]

Inserting rows into openai: 0 rows [00:00, ? rows/s]

Inserting rows into openai: 1 rows [00:00, 2777.68 rows/s]

Inserting rows into openai: 0 rows [00:00, ? rows/s]

Inserting rows into openai: 1 rows [00:00, 2855.21 rows/s]

And this is what we have so far:

t.select(t.input, t.response, t.ground_truth).show()

Evaluation

Now that we have some ground truth available, we can carry out basic evaluations of the GPT outputs, in this case by asking ChatGPT to decide whether the two are equivalent.

To start with, we'll create an evaluation prompt. In this case, the prompt requires some bespoke string substitution, so it's easiest to do using a UDF. (See the Pixeltable Basics tutorial and the UDFs in Pixeltable guide for more details on UDFs.)

system_prompt = '''
Compare the following listA and listB of entities, and check if they contain the same entities.
Return a JSON object with the following format:
{"reasoning": explaining your reasoning, "decision": 1 if the lists matched, 0 otherwise}
'''

@pxt.udf
def eval_prompt(listA: str, listB: str) -> list[dict]:
    return [
        {'role': 'system', 'content': system_prompt},
        {'role': 'user', 'content': f'listA: "{listA}" \n listB: "{listB}"'}
    ]

t['eval_prompt'] = eval_prompt(t.response, t.ground_truth)

Computing cells: 100%|███████████████████████████████████████████| 6/6 [00:00<00:00, 219.90 cells/s]
Added 6 column values with 0 errors.

t.select(t.eval_prompt).show()

The actual evaluation happens in another computed column. We can use OpenAI's handy response_format parameter to enforce that the output is properly formed JSON.

t['eval'] = openai.chat_completions(model='gpt-4o-mini', messages=t.eval_prompt, response_format={'type': 'json_object'})
t['eval_output'] = t.eval.choices[0].message.content

Computing cells: 100%|████████████████████████████████████████████| 6/6 [00:06<00:00,  1.03s/ cells]
Added 6 column values with 0 errors.
Computing cells: 100%|███████████████████████████████████████████| 6/6 [00:00<00:00, 610.26 cells/s]
Added 6 column values with 0 errors.

Let's take a look:

t.select(t.response, t.ground_truth, t.eval_output).show()

Finally, it's time to pull the decision out of the JSON structs returned by OpenAI. There's just one complication: the chat_completions responses are strings, not JSON structs. We can resolve this with Pixeltable's handy apply method, which turns any Python function into a Pixeltable function. In this case, we'll apply the Python function json.loads to parse the string into a JSON struct.

import json

t['eval_decision'] = t.eval_output.apply(json.loads).decision

Computing cells: 100%|███████████████████████████████████████████| 6/6 [00:00<00:00, 987.94 cells/s]
Added 6 column values with 0 errors.

t.select(t.response, t.ground_truth, t.eval_output, t.eval_decision).show()