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

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))

pixeltable/tests/functions/test_together.py

@@ -0,0 +1,111 @@
+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 TestTogether:
+
+    def test_completions(self, test_client: pxt.Client) -> None:
+        skip_test_if_not_installed('together')
+        TestTogether.skip_test_if_no_together_client()
+        cl = test_client
+        t = cl.create_table('test_tbl', {'input': pxt.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_2=completions(
+            prompt=t.input,
+            model='mistralai/Mixtral-8x7B-v0.1',
+            max_tokens=300,
+            stop=['\n'],
+            temperature=0.7,
+            top_p=0.9,
+            top_k=40,
+            repetition_penalty=1.1,
+            logprobs=1,
+            echo=True,
+            n=3,
+            safety_model='Meta-Llama/Llama-Guard-7b'
+        ))
+        validate_update_status(t.insert(input='I am going to the '), 1)
+        result = t.collect()
+        assert len(result['output'][0]['choices'][0]['text']) > 0
+        assert len(result['output_2'][0]['choices'][0]['text']) > 0
+
+    def test_chat_completions(self, test_client: pxt.Client) -> None:
+        skip_test_if_not_installed('together')
+        TestTogether.skip_test_if_no_together_client()
+        cl = test_client
+        t = cl.create_table('test_tbl', {'input': pxt.StringType()})
+        messages = [{'role': 'user', 'content': t.input}]
+        from pixeltable.functions.together import chat_completions
+        t.add_column(output=chat_completions(messages=messages, model='mistralai/Mixtral-8x7B-v0.1', stop=['\n']))
+        t.add_column(output_2=chat_completions(
+            messages=messages,
+            model='mistralai/Mixtral-8x7B-Instruct-v0.1',
+            max_tokens=300,
+            stop=['\n'],
+            temperature=0.7,
+            top_p=0.9,
+            top_k=40,
+            repetition_penalty=1.1,
+            logprobs=1,
+            echo=True,
+            n=3,
+            safety_model='Meta-Llama/Llama-Guard-7b',
+            response_format={'type': 'json_object'}
+        ))
+        validate_update_status(t.insert(input='Give me a typical example of a JSON structure.'), 1)
+        result = t.collect()
+        assert len(result['output'][0]['choices'][0]['message']) > 0
+        assert len(result['output_2'][0]['choices'][0]['message']) > 0
+
+    def test_embeddings(self, test_client: pxt.Client) -> None:
+        skip_test_if_not_installed('together')
+        TestTogether.skip_test_if_no_together_client()
+        cl = test_client
+        t = cl.create_table('test_tbl', {'input': pxt.StringType()})
+        from pixeltable.functions.together import embeddings
+        t.add_column(embed=embeddings(input=t.input, model='togethercomputer/m2-bert-80M-8k-retrieval'))
+        validate_update_status(t.insert(input='Together AI provides a variety of embeddings models.'), 1)
+        assert len(t.collect()['embed'][0]) > 0
+
+    def test_image_generations(self, test_client: pxt.Client) -> None:
+        skip_test_if_not_installed('together')
+        TestTogether.skip_test_if_no_together_client()
+        cl = test_client
+        t = cl.create_table(
+            'test_tbl',
+            {'input': pxt.StringType(), 'negative_prompt': pxt.StringType(nullable=True)}
+        )
+        from pixeltable.functions.together import image_generations
+        t.add_column(img=image_generations(t.input, model='runwayml/stable-diffusion-v1-5'))
+        t.add_column(img_2=image_generations(
+            t.input,
+            model='stabilityai/stable-diffusion-2-1',
+            steps=30,
+            seed=4178780,
+            height=768,
+            width=512,
+            negative_prompt=t.negative_prompt
+        ))
+        validate_update_status(t.insert([
+            {'input': 'A friendly dinosaur playing tennis in a cornfield'},
+            {'input': 'A friendly dinosaur playing tennis in a cornfield',
+             'negative_prompt': 'tennis court'}
+        ]), 2)
+        assert t.collect()['img'][0].size == (512, 512)
+        assert t.collect()['img_2'][0].size == (512, 768)
+        assert t.collect()['img'][1].size == (512, 512)
+        assert t.collect()['img_2'][1].size == (512, 768)
+
+    # 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_together_client() -> None:
+        try:
+            import pixeltable.functions.together
+            _ = pixeltable.functions.together.together_client()
+        except excs.Error as exc:
+            pytest.skip(str(exc))

pixeltable/tests/test_audio.py

@@ -0,0 +1,65 @@
+from typing import Optional
+
+import av
+
+import pixeltable as pxt
+import pixeltable.env as env
+from pixeltable.tests.utils import get_video_files, get_audio_files
+from pixeltable.type_system import VideoType, AudioType
+from pixeltable.utils.media_store import MediaStore
+
+
+class TestAudio:
+    def check_audio_params(self, path: str, format: Optional[str] = None, codec: Optional[str] = None):
+        with av.open(path) as container:
+            audio_stream = container.streams.audio[0]
+            if format is not None:
+                assert format == container.format.name
+            if codec is not None:
+                assert codec == audio_stream.codec_context.codec.name
+
+    def test_basic(self, test_client: pxt.Client) -> None:
+        audio_filepaths = get_audio_files()
+        cl = test_client
+        audio_t = cl.create_table('audio', {'audio_file': AudioType()})
+        status = audio_t.insert({'audio_file': p} for p in audio_filepaths)
+        assert status.num_rows == len(audio_filepaths)
+        assert status.num_excs == 0
+        paths = audio_t.select(output=audio_t.audio_file.localpath).collect()['output']
+        assert set(paths) == set(audio_filepaths)
+
+    def test_extract(self, test_client: pxt.Client) -> None:
+        video_filepaths = get_video_files()
+        cl = test_client
+        video_t = cl.create_table('videos', {'video': VideoType()})
+        from pixeltable.functions.video import extract_audio
+        video_t.add_column(audio=extract_audio(video_t.video))
+
+        # one of the 3 videos doesn't have audio
+        status = video_t.insert({'video': p} for p in video_filepaths)
+        assert status.num_rows == len(video_filepaths)
+        assert status.num_excs == 0
+        assert MediaStore.count(video_t.get_id()) == len(video_filepaths) - 1
+        assert video_t.where(video_t.audio != None).count() == len(video_filepaths) - 1
+        assert env.Env.get().num_tmp_files() == 0
+
+        # make sure everything works with a fresh client
+        cl = pxt.Client()
+        video_t = cl.get_table('videos')
+        assert video_t.where(video_t.audio != None).count() == len(video_filepaths) - 1
+
+        # test generating different formats and codecs
+        paths = video_t.select(output=extract_audio(video_t.video, format='wav', codec='pcm_s16le')).collect()['output']
+        # media files that are created as a part of a query end up in the tmp dir
+        assert env.Env.get().num_tmp_files() == video_t.where(video_t.audio != None).count()
+        for path in [p for p in paths if p is not None]:
+            self.check_audio_params(path, format='wav', codec='pcm_s16le')
+        # higher resolution
+        paths = video_t.select(output=extract_audio(video_t.video, format='wav', codec='pcm_s32le')).collect()['output']
+        for path in [p for p in paths if p is not None]:
+            self.check_audio_params(path, format='wav', codec='pcm_s32le')
+
+        for format in ['mp3', 'flac']:
+            paths = video_t.select(output=extract_audio(video_t.video, format=format)).collect()['output']
+            for path in [p for p in paths if p is not None]:
+                self.check_audio_params(path, format=format)

pixeltable/tests/test_catalog.py

@@ -0,0 +1,27 @@
+from pixeltable.catalog import is_valid_identifier, is_valid_path
+
+class TestCatalog:
+    """Tests for miscellanous catalog functions."""
+    def test_valid_identifier(self) -> None:
+        valid_ids = ['a', 'a1', 'a_1', 'a_']
+        invalid_ids = ['', '_', '__', '_a', '1a', 'a.b', '.a', 'a-b']
+        for valid_id in valid_ids:
+            assert is_valid_identifier(valid_id), valid_ids
+
+        for invalid_id in invalid_ids:
+            assert not is_valid_identifier(invalid_id), invalid_ids
+
+    def test_valid_path(self) -> None:
+        assert is_valid_path('', empty_is_valid=True)
+        assert not is_valid_path('', empty_is_valid=False)
+
+        valid_paths = ['a', 'a_.b_', 'a.b.c', 'a.b.c.d']
+        invalid_paths = ['.', '..', 'a.', '.a', 'a..b']
+        
+        for valid_path in valid_paths:
+            assert is_valid_path(valid_path, empty_is_valid=False), valid_path
+            assert is_valid_path(valid_path, empty_is_valid=True), valid_path
+
+        for invalid_path in invalid_paths:
+            assert not is_valid_path(invalid_path, empty_is_valid=False), invalid_path
+            assert not is_valid_path(invalid_path, empty_is_valid=True), invalid_path

pixeltable/tests/test_client.py

@@ -1,21 +1,21 @@
 import pytest
 
-import pixeltable as pt
-from pixeltable import exceptions as exc
+import pixeltable as pxt
+import pixeltable.exceptions as excs
 
 
 class TestClient:
-    def test_create_db(self, init_db: None) -> None:
-        cl = pt.Client()
-        _ = cl.create_db('test')
-        with pytest.raises(exc.DuplicateNameError):
-            _ = cl.create_db('test')
+    def test_list_functions(self, init_env) -> None:
+        cl = pxt.Client()
+        _ = cl.list_functions()
+        print(_)
 
-        _ = cl.get_db('test')
-        with pytest.raises(exc.UnknownEntityError):
-            _ = cl.get_db('xyz')
-
-        cl.drop_db('test', force=True)
-        with pytest.raises(exc.UnknownEntityError):
-            cl.drop_db('test', force=True)
+    def test_drop_table(self, test_tbl: pxt.Table) -> None:
+        cl = pxt.Client()
+        t = cl.get_table('test_tbl')
+        cl.drop_table('test_tbl')
+        with pytest.raises(excs.Error):
+            _ = cl.get_table('test_tbl')
+        with pytest.raises(excs.Error):
+            _ = t.show(1)