pixeltable 0.2.3__py3-none-any.whl → 0.2.4__py3-none-any.whl

pixeltable/functions/together.py

@@ -1,27 +1,122 @@
+import base64
+import io
 from typing import Optional
 
+import PIL.Image
+import numpy as np
+import together
+
 import pixeltable as pxt
+from pixeltable import env
+from pixeltable.func import Batch
+
+
+def together_client() -> together.Together:
+    return env.Env.get().get_client('together', lambda api_key: together.Together(api_key=api_key))
 
 
 @pxt.udf
 def completions(
         prompt: str,
+        *,
         model: str,
         max_tokens: Optional[int] = None,
-        repetition_penalty: Optional[float] = None,
         stop: Optional[list] = None,
-        top_k: Optional[int] = None,
+        temperature: Optional[float] = None,
         top_p: Optional[float] = None,
-        temperature: Optional[float] = None
+        top_k: Optional[int] = None,
+        repetition_penalty: Optional[float] = None,
+        logprobs: Optional[int] = None,
+        echo: Optional[bool] = None,
+        n: Optional[int] = None,
+        safety_model: Optional[str] = None
 ) -> dict:
-    import together
-    return together.Complete.create(
-        prompt,
-        model,
+    return together_client().completions.create(
+        prompt=prompt,
+        model=model,
         max_tokens=max_tokens,
-        repetition_penalty=repetition_penalty,
         stop=stop,
+        temperature=temperature,
+        top_p=top_p,
         top_k=top_k,
+        repetition_penalty=repetition_penalty,
+        logprobs=logprobs,
+        echo=echo,
+        n=n,
+        safety_model=safety_model
+    ).dict()
+
+
+@pxt.udf
+def chat_completions(
+        messages: list[dict[str, str]],
+        *,
+        model: str,
+        max_tokens: Optional[int] = None,
+        stop: Optional[list[str]] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        top_k: Optional[int] = None,
+        repetition_penalty: Optional[float] = None,
+        logprobs: Optional[int] = None,
+        echo: Optional[bool] = None,
+        n: Optional[int] = None,
+        safety_model: Optional[str] = None,
+        response_format: Optional[dict] = None,
+        tools: Optional[dict] = None,
+        tool_choice: Optional[dict] = None
+) -> dict:
+    return together_client().chat.completions.create(
+        messages=messages,
+        model=model,
+        max_tokens=max_tokens,
+        stop=stop,
+        temperature=temperature,
         top_p=top_p,
-        temperature=temperature
+        top_k=top_k,
+        repetition_penalty=repetition_penalty,
+        logprobs=logprobs,
+        echo=echo,
+        n=n,
+        safety_model=safety_model,
+        response_format=response_format,
+        tools=tools,
+        tool_choice=tool_choice
+    ).dict()
+
+
+@pxt.udf(batch_size=32, return_type=pxt.ArrayType((None,), dtype=pxt.FloatType()))
+def embeddings(input: Batch[str], *, model: str) -> Batch[np.ndarray]:
+    result = together_client().embeddings.create(input=input, model=model)
+    return [
+        np.array(data.embedding, dtype=np.float64)
+        for data in result.data
+    ]
+
+
+@pxt.udf
+def image_generations(
+        prompt: str,
+        *,
+        model: str,
+        steps: Optional[int] = None,
+        seed: Optional[int] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        negative_prompt: Optional[str] = None,
+) -> PIL.Image.Image:
+    # TODO(aaron-siegel): Decompose CPU/GPU ops into separate functions
+    result = together_client().images.generate(
+        prompt=prompt,
+        model=model,
+        steps=steps,
+        seed=seed,
+        height=height,
+        width=width,
+        negative_prompt=negative_prompt
     )
+    b64_str = result.data[0].b64_json
+    b64_bytes = base64.b64decode(b64_str)
+    img = PIL.Image.open(io.BytesIO(b64_bytes))
+    img.load()
+    return img

pixeltable/tests/conftest.py

@@ -80,8 +80,8 @@ def test_tbl(test_client: pxt.Client) -> catalog.Table:
     return create_test_tbl(test_client)
 
 # @pytest.fixture(scope='function')
-# def test_stored_fn(test_client: pt.Client) -> pt.Function:
-#     @pt.udf(return_type=pt.IntType(), param_types=[pt.IntType()])
+# def test_stored_fn(test_client: pxt.Client) -> pxt.Function:
+#     @pxt.udf(return_type=pxt.IntType(), param_types=[pxt.IntType()])
 #     def test_fn(x):
 #         return x + 1
 #     test_client.create_function('test_fn', test_fn)
@@ -89,7 +89,7 @@ def test_tbl(test_client: pxt.Client) -> catalog.Table:
 
 @pytest.fixture(scope='function')
 def test_tbl_exprs(test_tbl: catalog.Table) -> List[exprs.Expr]:
-#def test_tbl_exprs(test_tbl: catalog.Table, test_stored_fn: pt.Function) -> List[exprs.Expr]:
+#def test_tbl_exprs(test_tbl: catalog.Table, test_stored_fn: pxt.Function) -> List[exprs.Expr]:
 
     t = test_tbl
     return [
@@ -156,7 +156,7 @@ def img_tbl_exprs(img_tbl: catalog.Table) -> List[exprs.Expr]:
 # TODO: why does this not work with a session scope? (some user tables don't get created with create_all())
 #@pytest.fixture(scope='session')
 #def indexed_img_tbl(init_env: None) -> catalog.Table:
-#    cl = pt.Client()
+#    cl = pxt.Client()
 #    db = cl.create_db('test_indexed')
 @pytest.fixture(scope='function')
 def indexed_img_tbl(test_client: pxt.Client) -> catalog.Table:

pixeltable/tests/functions/test_fireworks.py

@@ -0,0 +1,42 @@
+import pytest
+
+import pixeltable as pxt
+import pixeltable.exceptions as excs
+from pixeltable.tests.utils import skip_test_if_not_installed, validate_update_status
+
+
+class TestFireworks:
+
+    def test_fireworks(self, test_client: pxt.Client) -> None:
+        skip_test_if_not_installed('fireworks')
+        TestFireworks.skip_test_if_no_fireworks_client()
+        cl = test_client
+        t = cl.create_table('test_tbl', {'input': pxt.StringType()})
+        from pixeltable.functions.fireworks import chat_completions
+        messages = [{'role': 'user', 'content': t.input}]
+        t['output'] = chat_completions(
+            messages=messages,
+            model='accounts/fireworks/models/llama-v2-7b-chat'
+        )
+        t['output_2'] = chat_completions(
+            messages=messages,
+            model='accounts/fireworks/models/llama-v2-7b-chat',
+            max_tokens=300,
+            top_k=40,
+            top_p=0.9,
+            temperature=0.7
+        )
+        validate_update_status(t.insert(input="How's everything going today?"), 1)
+        results = t.collect()
+        assert len(results['output'][0]['choices'][0]['message']['content']) > 0
+        assert len(results['output_2'][0]['choices'][0]['message']['content']) > 0
+
+    # This ensures that the test will be skipped, rather than returning an error, when no API key is
+    # available (for example, when a PR runs in CI).
+    @staticmethod
+    def skip_test_if_no_fireworks_client() -> None:
+        try:
+            import pixeltable.functions.fireworks
+            _ = pixeltable.functions.fireworks.fireworks_client()
+        except excs.Error as exc:
+            pytest.skip(str(exc))

pixeltable/tests/functions/test_functions.py

@@ -0,0 +1,60 @@
+import pixeltable as pxt
+from pixeltable import catalog
+from pixeltable.functions.pil.image import blend
+from pixeltable.iterators import FrameIterator
+from pixeltable.tests.utils import get_video_files, skip_test_if_not_installed
+from pixeltable.type_system import VideoType, StringType
+
+
+class TestFunctions:
+    def test_pil(self, img_tbl: catalog.Table) -> None:
+        t = img_tbl
+        _ = t[t.img, t.img.rotate(90), blend(t.img, t.img.rotate(90), 0.5)].show()
+
+    def test_eval_detections(self, test_client: pxt.Client) -> None:
+        skip_test_if_not_installed('nos')
+        cl = test_client
+        video_t = cl.create_table('video_tbl', {'video': VideoType()})
+        # create frame view
+        args = {'video': video_t.video, 'fps': 1}
+        v = cl.create_view('test_view', video_t, iterator_class=FrameIterator, iterator_args=args)
+
+        files = get_video_files()
+        video_t.insert(video=files[-1])
+        v.add_column(frame_s=v.frame.resize([640, 480]))
+        from pixeltable.functions.nos.object_detection_2d import yolox_nano, yolox_small, yolox_large
+        v.add_column(detections_a=yolox_nano(v.frame_s))
+        v.add_column(detections_b=yolox_small(v.frame_s))
+        v.add_column(gt=yolox_large(v.frame_s))
+        from pixeltable.functions.eval import eval_detections, mean_ap
+        res = v.select(
+            eval_detections(
+                v.detections_a.bboxes, v.detections_a.labels, v.detections_a.scores, v.gt.bboxes, v.gt.labels
+            )).show()
+        v.add_column(
+            eval_a=eval_detections(
+                v.detections_a.bboxes, v.detections_a.labels, v.detections_a.scores, v.gt.bboxes, v.gt.labels))
+        v.add_column(
+            eval_b=eval_detections(
+                v.detections_b.bboxes, v.detections_b.labels, v.detections_b.scores, v.gt.bboxes, v.gt.labels))
+        ap_a = v.select(mean_ap(v.eval_a)).show()[0, 0]
+        ap_b = v.select(mean_ap(v.eval_b)).show()[0, 0]
+        common_classes = set(ap_a.keys()) & set(ap_b.keys())
+
+        ## TODO: following assertion is failing on CI, 
+        # It is not necessarily a bug, as assert codition is not expected to be always true
+        # for k in common_classes:
+        # assert ap_a[k] <= ap_b[k]
+
+    def test_str(self, test_client: pxt.Client) -> None:
+        cl = test_client
+        t = cl.create_table('test_tbl', {'input': StringType()})
+        from pixeltable.functions.string import str_format
+        t.add_column(s1=str_format('ABC {0}', t.input))
+        t.add_column(s2=str_format('DEF {this}', this=t.input))
+        t.add_column(s3=str_format('GHI {0} JKL {this}', t.input, this=t.input))
+        status = t.insert(input='MNO')
+        assert status.num_rows == 1
+        assert status.num_excs == 0
+        row = t.head()[0]
+        assert row == {'input': 'MNO', 's1': 'ABC MNO', 's2': 'DEF MNO', 's3': 'GHI MNO JKL MNO'}

pixeltable/tests/{test_functions.py → functions/test_huggingface.py}

@@ -3,144 +3,12 @@ from typing import Dict, Any
 import pytest
 
 import pixeltable as pxt
-from pixeltable import catalog
-from pixeltable.env import Env
-import pixeltable.exceptions as excs
-from pixeltable.functions.pil.image import blend
-from pixeltable.iterators import FrameIterator
-from pixeltable.tests.utils import get_video_files, skip_test_if_not_installed, get_sentences, get_image_files
-from pixeltable.type_system import VideoType, StringType, JsonType, ImageType, BoolType, FloatType, ArrayType
+from pixeltable.tests.utils import skip_test_if_not_installed, get_sentences, get_image_files, \
+    SAMPLE_IMAGE_URL
+from pixeltable.type_system import StringType, JsonType, ImageType, BoolType, FloatType, ArrayType
 
 
-class TestFunctions:
-    def test_pil(self, img_tbl: catalog.Table) -> None:
-        t = img_tbl
-        _ = t[t.img, t.img.rotate(90), blend(t.img, t.img.rotate(90), 0.5)].show()
-
-    def test_eval_detections(self, test_client: pxt.Client) -> None:
-        skip_test_if_not_installed('nos')
-        cl = test_client
-        video_t = cl.create_table('video_tbl', {'video': VideoType()})
-        # create frame view
-        args = {'video': video_t.video, 'fps': 1}
-        v = cl.create_view('test_view', video_t, iterator_class=FrameIterator, iterator_args=args)
-
-        files = get_video_files()
-        video_t.insert(video=files[-1])
-        v.add_column(frame_s=v.frame.resize([640, 480]))
-        from pixeltable.functions.nos.object_detection_2d import yolox_nano, yolox_small, yolox_large
-        v.add_column(detections_a=yolox_nano(v.frame_s))
-        v.add_column(detections_b=yolox_small(v.frame_s))
-        v.add_column(gt=yolox_large(v.frame_s))
-        from pixeltable.functions.eval import eval_detections, mean_ap
-        res = v.select(
-            eval_detections(
-                v.detections_a.bboxes, v.detections_a.labels, v.detections_a.scores, v.gt.bboxes, v.gt.labels
-            )).show()
-        v.add_column(
-            eval_a=eval_detections(
-                v.detections_a.bboxes, v.detections_a.labels, v.detections_a.scores, v.gt.bboxes, v.gt.labels))
-        v.add_column(
-            eval_b=eval_detections(
-                v.detections_b.bboxes, v.detections_b.labels, v.detections_b.scores, v.gt.bboxes, v.gt.labels))
-        ap_a = v.select(mean_ap(v.eval_a)).show()[0, 0]
-        ap_b = v.select(mean_ap(v.eval_b)).show()[0, 0]
-        common_classes = set(ap_a.keys()) & set(ap_b.keys())
-
-        ## TODO: following assertion is failing on CI, 
-        # It is not necessarily a bug, as assert codition is not expected to be always true
-        # for k in common_classes:
-        # assert ap_a[k] <= ap_b[k]
-
-    def test_str(self, test_client: pxt.Client) -> None:
-        cl = test_client
-        t = cl.create_table('test_tbl', {'input': StringType()})
-        from pixeltable.functions.string import str_format
-        t.add_column(s1=str_format('ABC {0}', t.input))
-        t.add_column(s2=str_format('DEF {this}', this=t.input))
-        t.add_column(s3=str_format('GHI {0} JKL {this}', t.input, this=t.input))
-        status = t.insert(input='MNO')
-        assert status.num_rows == 1
-        assert status.num_excs == 0
-        row = t.head()[0]
-        assert row == {'input': 'MNO', 's1': 'ABC MNO', 's2': 'DEF MNO', 's3': 'GHI MNO JKL MNO'}
-
-    def test_openai(self, test_client: pxt.Client) -> None:
-        skip_test_if_not_installed('openai')
-        TestFunctions.skip_test_if_no_openai_client()
-        cl = test_client
-        t = cl.create_table('test_tbl', {'input': StringType()})
-        from pixeltable.functions.openai import chat_completions, embeddings, moderations
-        msgs = [
-            {"role": "system", "content": "You are a helpful assistant."},
-            {"role": "user", "content": t.input}
-        ]
-        t.add_column(input_msgs=msgs)
-        t.add_column(chat_output=chat_completions(model='gpt-3.5-turbo', messages=t.input_msgs))
-        # with inlined messages
-        t.add_column(chat_output2=chat_completions(model='gpt-3.5-turbo', messages=msgs))
-        t.add_column(ada_embed=embeddings(model='text-embedding-ada-002', input=t.input))
-        t.add_column(text_3=embeddings(model='text-embedding-3-small', input=t.input))
-        t.add_column(moderation=moderations(input=t.input))
-        t.insert(input='I find you really annoying')
-        _ = t.head()
-
-    def test_gpt_4_vision(self, test_client: pxt.Client) -> None:
-        skip_test_if_not_installed('openai')
-        TestFunctions.skip_test_if_no_openai_client()
-        cl = test_client
-        t = cl.create_table('test_tbl', {'prompt': StringType(), 'img': ImageType()})
-        from pixeltable.functions.openai import chat_completions
-        from pixeltable.functions.string import str_format
-        msgs = [
-            {'role': 'user',
-             'content': [
-                 {'type': 'text', 'text': t.prompt},
-                 {'type': 'image_url', 'image_url': {
-                     'url': str_format('data:image/png;base64,{0}', t.img.b64_encode())
-                 }}
-             ]}
-        ]
-        t.add_column(response=chat_completions(model='gpt-4-vision-preview', messages=msgs, max_tokens=300))
-        t.add_column(response_content=t.response.choices[0].message.content)
-        t.insert(prompt="What's in this image?", img=_sample_image_url)
-        result = t.collect()['response_content'][0]
-        assert len(result) > 0
-
-    @staticmethod
-    def skip_test_if_no_openai_client() -> None:
-        try:
-            _ = Env.get().openai_client
-        except excs.Error as exc:
-            pytest.skip(str(exc))
-
-    def test_together(self, test_client: pxt.Client) -> None:
-        skip_test_if_not_installed('together')
-        if not Env.get().has_together_client:
-            pytest.skip(f'Together client does not exist (missing API key?)')
-        cl = test_client
-        t = cl.create_table('test_tbl', {'input': StringType()})
-        from pixeltable.functions.together import completions
-        t.add_column(output=completions(prompt=t.input, model='mistralai/Mixtral-8x7B-v0.1', stop=['\n']))
-        t.add_column(output_text=t.output.output.choices[0].text)
-        t.insert(input='I am going to the ')
-        result = t.select(t.output_text).collect()['output_text'][0]
-        assert len(result) > 0
-
-    def test_fireworks(self, test_client: pxt.Client) -> None:
-        skip_test_if_not_installed('fireworks')
-        try:
-            from pixeltable.functions.fireworks import initialize
-            initialize()
-        except:
-            pytest.skip(f'Fireworks client does not exist (missing API key?)')
-        cl = test_client
-        t = cl.create_table('test_tbl', {'input': StringType()})
-        from pixeltable.functions.fireworks import chat_completions
-        t['output'] = chat_completions(prompt=t.input, model='accounts/fireworks/models/llama-v2-7b-chat', max_tokens=256).choices[0].text
-        t.insert(input='I am going to the ')
-        result = t.select(t.output).collect()['output'][0]
-        assert len(result) > 0
+class TestHuggingface:
 
     def test_hf_function(self, test_client: pxt.Client) -> None:
         skip_test_if_not_installed('sentence_transformers')
@@ -281,14 +149,10 @@ class TestFunctions:
         t = cl.create_table('test_tbl', {'img': ImageType()})
         from pixeltable.functions.huggingface import detr_for_object_detection
         t['detect'] = detr_for_object_detection(t.img, model_id='facebook/detr-resnet-50', threshold=0.8)
-        status = t.insert(img=_sample_image_url)
+        status = t.insert(img=SAMPLE_IMAGE_URL)
         assert status.num_rows == 1
         assert status.num_excs == 0
         result = t.select(t.detect).collect()[0]['detect']
         assert 'orange' in result['label_text']
         assert 'bowl' in result['label_text']
         assert 'broccoli' in result['label_text']
-
-
-_sample_image_url = \
-    'https://raw.githubusercontent.com/pixeltable/pixeltable/master/docs/source/data/images/000000000009.jpg'

pixeltable/tests/functions/test_openai.py

@@ -0,0 +1,152 @@
+import pytest
+
+import pixeltable as pxt
+import pixeltable.exceptions as excs
+from pixeltable.tests.utils import SAMPLE_IMAGE_URL, skip_test_if_not_installed, validate_update_status
+from pixeltable.type_system import StringType, ImageType
+
+
+class TestOpenai:
+
+    def test_audio(self, test_client: pxt.Client) -> None:
+        skip_test_if_not_installed('openai')
+        TestOpenai.skip_test_if_no_openai_client()
+        cl = test_client
+        t = cl.create_table('test_tbl', {'input': StringType()})
+        from pixeltable.functions.openai import speech, transcriptions, translations
+        t.add_column(speech=speech(t.input, model='tts-1', voice='onyx'))
+        t.add_column(speech_2=speech(t.input, model='tts-1', voice='onyx', response_format='flac', speed=1.05))
+        t.add_column(transcription=transcriptions(t.speech, model='whisper-1'))
+        t.add_column(transcription_2=transcriptions(
+            t.speech, model='whisper-1', language='en', prompt='Transcribe the contents of this recording.'
+        ))
+        t.add_column(translation=translations(t.speech, model='whisper-1'))
+        t.add_column(translation_2=translations(
+            t.speech, model='whisper-1', prompt='Translate the recording from Spanish into English.', temperature=0.7
+        ))
+        validate_update_status(t.insert([
+            {'input': 'I am a banana.'},
+            {'input': 'Es fácil traducir del español al inglés.'}
+        ]), expected_rows=2)
+        # The audio generation -> transcription loop on these examples should be simple and clear enough
+        # that the unit test can reliably expect the output closely enough to pass these checks.
+        results = t.collect()
+        assert results[0]['transcription']['text'] in ['I am a banana.', "I'm a banana."]
+        assert results[0]['transcription_2']['text'] in ['I am a banana.', "I'm a banana."]
+        assert 'easy to translate from Spanish' in results[1]['translation']['text']
+        assert 'easy to translate from Spanish' in results[1]['translation_2']['text']
+
+    def test_chat_completions(self, test_client: pxt.Client) -> None:
+        skip_test_if_not_installed('openai')
+        TestOpenai.skip_test_if_no_openai_client()
+        cl = test_client
+        t = cl.create_table('test_tbl', {'input': StringType()})
+        from pixeltable.functions.openai import chat_completions
+        msgs = [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": t.input}
+        ]
+        t.add_column(input_msgs=msgs)
+        t.add_column(chat_output=chat_completions(model='gpt-3.5-turbo', messages=t.input_msgs))
+        # with inlined messages
+        t.add_column(chat_output_2=chat_completions(model='gpt-3.5-turbo', messages=msgs))
+        # test a bunch of the parameters
+        t.add_column(chat_output_3=chat_completions(
+            model='gpt-3.5-turbo', messages=msgs, frequency_penalty=0.1, logprobs=True, top_logprobs=3,
+            max_tokens=500, n=3, presence_penalty=0.1, seed=4171780, stop=['\n'], temperature=0.7, top_p=0.8,
+            user='pixeltable'
+        ))
+        # test with JSON output enforced
+        t.add_column(chat_output_4=chat_completions(
+            model='gpt-3.5-turbo', messages=msgs, response_format={'type': 'json_object'}
+        ))
+        # TODO Also test the `tools` and `tool_choice` parameters.
+        validate_update_status(t.insert(input='Give me an example of a typical JSON structure.'), 1)
+        result = t.collect()
+        assert len(result['chat_output'][0]['choices'][0]['message']['content']) > 0
+        assert len(result['chat_output_2'][0]['choices'][0]['message']['content']) > 0
+        assert len(result['chat_output_3'][0]['choices'][0]['message']['content']) > 0
+        assert len(result['chat_output_4'][0]['choices'][0]['message']['content']) > 0
+
+        # When OpenAI gets a request with `response_format` equal to `json_object`, but the prompt does not
+        # contain the string "json", it refuses the request.
+        # TODO This should probably not be throwing an exception, but rather logging the error in
+        # `t.chat_output_4.errormsg` etc.
+        with pytest.raises(excs.ExprEvalError) as exc_info:
+            t.insert(input='Say something interesting.')
+        assert "\\'messages\\' must contain the word \\'json\\'" in str(exc_info.value)
+
+    def test_gpt_4_vision(self, test_client: pxt.Client) -> None:
+        skip_test_if_not_installed('openai')
+        TestOpenai.skip_test_if_no_openai_client()
+        cl = test_client
+        t = cl.create_table('test_tbl', {'prompt': StringType(), 'img': ImageType()})
+        from pixeltable.functions.openai import chat_completions, vision
+        from pixeltable.functions.string import str_format
+        t.add_column(response=vision(prompt="What's in this image?", image=t.img))
+        # Also get the response the low-level way, by calling chat_completions
+        msgs = [
+            {'role': 'user',
+             'content': [
+                 {'type': 'text', 'text': t.prompt},
+                 {'type': 'image_url', 'image_url': {
+                     'url': str_format('data:image/png;base64,{0}', t.img.b64_encode())
+                 }}
+             ]}
+        ]
+        t.add_column(response_2=chat_completions(model='gpt-4-vision-preview', messages=msgs, max_tokens=300).choices[0].message.content)
+        validate_update_status(t.insert(prompt="What's in this image?", img=SAMPLE_IMAGE_URL), 1)
+        result = t.collect()['response_2'][0]
+        assert len(result) > 0
+
+    def test_embeddings(self, test_client: pxt.Client) -> None:
+        skip_test_if_not_installed('openai')
+        TestOpenai.skip_test_if_no_openai_client()
+        cl = test_client
+        from pixeltable.functions.openai import embeddings
+        t = cl.create_table('test_tbl', {'input': StringType()})
+        t.add_column(ada_embed=embeddings(model='text-embedding-ada-002', input=t.input))
+        t.add_column(text_3=embeddings(model='text-embedding-3-small', input=t.input, user='pixeltable'))
+        validate_update_status(t.insert(input='Say something interesting.'), 1)
+        _ = t.head()
+
+    def test_moderations(self, test_client: pxt.Client) -> None:
+        skip_test_if_not_installed('openai')
+        TestOpenai.skip_test_if_no_openai_client()
+        cl = test_client
+        t = cl.create_table('test_tbl', {'input': StringType()})
+        from pixeltable.functions.openai import moderations
+        t.add_column(moderation=moderations(input=t.input))
+        t.add_column(moderation_2=moderations(input=t.input, model='text-moderation-stable'))
+        validate_update_status(t.insert(input='Say something interesting.'), 1)
+        _ = t.head()
+
+    def test_image_generations(self, test_client: pxt.Client) -> None:
+        skip_test_if_not_installed('openai')
+        TestOpenai.skip_test_if_no_openai_client()
+        cl = test_client
+        t = cl.create_table('test_tbl', {'input': StringType()})
+        from pixeltable.functions.openai import image_generations
+        t.add_column(img=image_generations(t.input))
+        # Test dall-e-2 options
+        t.add_column(img_2=image_generations(
+            t.input, model='dall-e-2', size='512x512', user='pixeltable'
+        ))
+        # Test dall-e-3 options
+        t.add_column(img_3=image_generations(
+            t.input, model='dall-e-3', quality='hd', size='1792x1024', style='natural', user='pixeltable'
+        ))
+        validate_update_status(t.insert(input='A friendly dinosaur playing tennis in a cornfield'), 1)
+        assert t.collect()['img'][0].size == (1024, 1024)
+        assert t.collect()['img_2'][0].size == (512, 512)
+        assert t.collect()['img_3'][0].size == (1792, 1024)
+
+    # This ensures that the test will be skipped, rather than returning an error, when no API key is
+    # available (for example, when a PR runs in CI).
+    @staticmethod
+    def skip_test_if_no_openai_client() -> None:
+        try:
+            import pixeltable.functions.openai
+            _ = pixeltable.functions.openai.openai_client()
+        except excs.Error as exc:
+            pytest.skip(str(exc))