pixeltable 0.0.0__py3-none-any.whl

pixeltable/functions/openai.py

@@ -0,0 +1,289 @@
+import base64
+import io
+import pathlib
+import uuid
+from typing import Optional, TypeVar, Union, Callable
+
+import PIL.Image
+import numpy as np
+import openai
+import tenacity
+from openai._types import NOT_GIVEN, NotGiven
+
+import pixeltable as pxt
+import pixeltable.type_system as ts
+from pixeltable import env
+from pixeltable.func import Batch
+
+
+def openai_client() -> openai.OpenAI:
+    return env.Env.get().get_client('openai', lambda api_key: openai.OpenAI(api_key=api_key))
+
+
+# Exponential backoff decorator using tenacity.
+# TODO(aaron-siegel): Right now this hardwires random exponential backoff with defaults suggested
+# by OpenAI. Should we investigate making this more customizable in the future?
+def _retry(fn: Callable) -> Callable:
+    return tenacity.retry(
+        retry=tenacity.retry_if_exception_type(openai.RateLimitError),
+        wait=tenacity.wait_random_exponential(multiplier=3, max=180),
+        stop=tenacity.stop_after_attempt(20)
+    )(fn)
+
+
+#####################################
+# Audio Endpoints
+
+@pxt.udf(return_type=ts.AudioType())
+@_retry
+def speech(
+        input: str,
+        *,
+        model: str,
+        voice: str,
+        response_format: Optional[str] = None,
+        speed: Optional[float] = None
+) -> str:
+    content = openai_client().audio.speech.create(
+        input=input,
+        model=model,
+        voice=voice,
+        response_format=_opt(response_format),
+        speed=_opt(speed)
+    )
+    ext = response_format or 'mp3'
+    output_filename = str(env.Env.get().tmp_dir / f"{uuid.uuid4()}.{ext}")
+    content.write_to_file(output_filename)
+    return output_filename
+
+
+@pxt.udf(
+    param_types=[ts.AudioType(), ts.StringType(), ts.StringType(nullable=True),
+                 ts.StringType(nullable=True), ts.FloatType(nullable=True)]
+)
+@_retry
+def transcriptions(
+        audio: str,
+        *,
+        model: str,
+        language: Optional[str] = None,
+        prompt: Optional[str] = None,
+        temperature: Optional[float] = None
+) -> dict:
+    file = pathlib.Path(audio)
+    transcription = openai_client().audio.transcriptions.create(
+        file=file,
+        model=model,
+        language=_opt(language),
+        prompt=_opt(prompt),
+        temperature=_opt(temperature)
+    )
+    return transcription.dict()
+
+
+@pxt.udf(
+    param_types=[ts.AudioType(), ts.StringType(), ts.StringType(nullable=True), ts.FloatType(nullable=True)]
+)
+@_retry
+def translations(
+        audio: str,
+        *,
+        model: str,
+        prompt: Optional[str] = None,
+        temperature: Optional[float] = None
+) -> dict:
+    file = pathlib.Path(audio)
+    translation = openai_client().audio.translations.create(
+        file=file,
+        model=model,
+        prompt=_opt(prompt),
+        temperature=_opt(temperature)
+    )
+    return translation.dict()
+
+
+#####################################
+# Chat Endpoints
+
+@pxt.udf
+@_retry
+def chat_completions(
+        messages: list,
+        *,
+        model: str,
+        frequency_penalty: Optional[float] = None,
+        logit_bias: Optional[dict[str, int]] = None,
+        logprobs: Optional[bool] = None,
+        top_logprobs: Optional[int] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        presence_penalty: Optional[float] = None,
+        response_format: Optional[dict] = None,
+        seed: Optional[int] = None,
+        stop: Optional[list[str]] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        tools: Optional[list[dict]] = None,
+        tool_choice: Optional[dict] = None,
+        user: Optional[str] = None
+) -> dict:
+    result = openai_client().chat.completions.create(
+        messages=messages,
+        model=model,
+        frequency_penalty=_opt(frequency_penalty),
+        logit_bias=_opt(logit_bias),
+        logprobs=_opt(logprobs),
+        top_logprobs=_opt(top_logprobs),
+        max_tokens=_opt(max_tokens),
+        n=_opt(n),
+        presence_penalty=_opt(presence_penalty),
+        response_format=_opt(response_format),
+        seed=_opt(seed),
+        stop=_opt(stop),
+        temperature=_opt(temperature),
+        top_p=_opt(top_p),
+        tools=_opt(tools),
+        tool_choice=_opt(tool_choice),
+        user=_opt(user)
+    )
+    return result.dict()
+
+
+@pxt.udf
+@_retry
+def vision(
+        prompt: str,
+        image: PIL.Image.Image,
+        *,
+        model: str = 'gpt-4-vision-preview'
+) -> str:
+    # TODO(aaron-siegel): Decompose CPU/GPU ops into separate functions
+    bytes_arr = io.BytesIO()
+    image.save(bytes_arr, format='png')
+    b64_bytes = base64.b64encode(bytes_arr.getvalue())
+    b64_encoded_image = b64_bytes.decode('utf-8')
+    messages = [
+        {'role': 'user',
+         'content': [
+             {'type': 'text', 'text': prompt},
+             {'type': 'image_url', 'image_url': {
+                 'url': f'data:image/png;base64,{b64_encoded_image}'
+             }}
+         ]}
+    ]
+    result = openai_client().chat.completions.create(
+        messages=messages,
+        model=model
+    )
+    return result.choices[0].message.content
+
+
+#####################################
+# Embeddings Endpoints
+
+_embedding_dimensions_cache: dict[str, int] = {
+    'text-embedding-ada-002': 1536,
+    'text-embedding-3-small': 1536,
+    'text-embedding-3-large': 3072,
+}
+
+
+@pxt.udf(batch_size=32, return_type=ts.ArrayType((None,), dtype=ts.FloatType()))
+@_retry
+def embeddings(
+        input: Batch[str],
+        *,
+        model: str,
+        dimensions: Optional[int] = None,
+        user: Optional[str] = None
+) -> Batch[np.ndarray]:
+    result = openai_client().embeddings.create(
+        input=input,
+        model=model,
+        dimensions=_opt(dimensions),
+        user=_opt(user),
+        encoding_format='float'
+    )
+    return [
+        np.array(data.embedding, dtype=np.float64)
+        for data in result.data
+    ]
+
+
+@embeddings.conditional_return_type
+def _(model: str, dimensions: Optional[int] = None) -> ts.ArrayType:
+    if dimensions is None:
+        if model not in _embedding_dimensions_cache:
+            # TODO: find some other way to retrieve a sample
+            return ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False)
+        dimensions = _embedding_dimensions_cache.get(model, None)
+    return ts.ArrayType((dimensions,), dtype=ts.FloatType(), nullable=False)
+
+
+#####################################
+# Images Endpoints
+
+@pxt.udf
+@_retry
+def image_generations(
+        prompt: str,
+        *,
+        model: Optional[str] = None,
+        quality: Optional[str] = None,
+        size: Optional[str] = None,
+        style: Optional[str] = None,
+        user: Optional[str] = None
+) -> PIL.Image.Image:
+    # TODO(aaron-siegel): Decompose CPU/GPU ops into separate functions
+    result = openai_client().images.generate(
+        prompt=prompt,
+        model=_opt(model),
+        quality=_opt(quality),
+        size=_opt(size),
+        style=_opt(style),
+        user=_opt(user),
+        response_format="b64_json"
+    )
+    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
+
+
+@image_generations.conditional_return_type
+def _(size: Optional[str] = None) -> ts.ImageType:
+    if size is None:
+        return ts.ImageType(size=(1024, 1024))
+    x_pos = size.find('x')
+    if x_pos == -1:
+        return ts.ImageType()
+    try:
+        width, height = int(size[:x_pos]), int(size[x_pos + 1:])
+    except ValueError:
+        return ts.ImageType()
+    return ts.ImageType(size=(width, height))
+
+
+#####################################
+# Moderations Endpoints
+
+@pxt.udf
+@_retry
+def moderations(
+        input: str,
+        *,
+        model: Optional[str] = None
+) -> dict:
+    result = openai_client().moderations.create(
+        input=input,
+        model=_opt(model)
+    )
+    return result.dict()
+
+
+_T = TypeVar('_T')
+
+
+def _opt(arg: _T) -> Union[_T, NotGiven]:
+    return arg if arg is not None else NOT_GIVEN

pixeltable/functions/pil/image.py

@@ -0,0 +1,147 @@
+from typing import Tuple, Optional
+
+import PIL.Image
+from PIL.Image import Dither
+
+import pixeltable.func as func
+from pixeltable.type_system import FloatType, ImageType, IntType, ArrayType, ColumnType, StringType, JsonType
+
+
+@func.udf(
+    py_fn=PIL.Image.alpha_composite, return_type=ImageType(), param_types=[ImageType(), ImageType()])
+def alpha_composite(im1: PIL.Image.Image, im2: PIL.Image.Image) -> PIL.Image.Image:
+    pass
+@func.udf(
+    py_fn=PIL.Image.blend, return_type=ImageType(), param_types=[ImageType(), ImageType(), FloatType()])
+def blend(im1: PIL.Image.Image, im2: PIL.Image.Image, alpha: float) -> PIL.Image.Image:
+    pass
+@func.udf(
+    py_fn=PIL.Image.composite, return_type=ImageType(), param_types=[ImageType(), ImageType(), ImageType()])
+def composite(image1: PIL.Image.Image, image2: PIL.Image.Image, mask: PIL.Image.Image) -> PIL.Image.Image:
+    pass
+
+
+# PIL.Image.Image methods
+
+# Image.convert()
+@func.udf(param_types=[ImageType(), StringType()])
+def convert(self: PIL.Image.Image, mode: str) -> PIL.Image.Image:
+    return self.convert(mode)
+
+
+@convert.conditional_return_type
+def _(self: PIL.Image.Image, mode: str) -> ColumnType:
+    input_type = self.col_type
+    assert input_type.is_image_type()
+    return ImageType(size=input_type.size, mode=mode, nullable=input_type.nullable)
+
+
+# Image.crop()
+@func.udf(
+    py_fn=PIL.Image.Image.crop,
+    param_types=[ImageType(), ArrayType((4,), dtype=IntType())])
+def crop(self: PIL.Image.Image, box: Tuple[int, int, int, int]) -> PIL.Image.Image:
+    pass
+
+@crop.conditional_return_type
+def _(self: PIL.Image.Image, box: Tuple[int, int, int, int]) -> ColumnType:
+    input_type = self.col_type
+    assert input_type.is_image_type()
+    if isinstance(box, list) and all(isinstance(x, int) for x in box):
+        return ImageType(size=(box[2] - box[0], box[3] - box[1]), mode=input_type.mode, nullable=input_type.nullable)
+    return ImageType(mode=input_type.mode, nullable=input_type.nullable)  # we can't compute the size statically
+
+# Image.getchannel()
+@func.udf(py_fn=PIL.Image.Image.getchannel, param_types=[ImageType(), IntType()])
+def getchannel(self: PIL.Image.Image, channel: int) -> PIL.Image.Image:
+    pass
+
+@getchannel.conditional_return_type
+def _(self: PIL.Image.Image) -> ColumnType:
+    input_type = self.col_type
+    assert input_type.is_image_type()
+    return ImageType(size=input_type.size, mode='L', nullable=input_type.nullable)
+
+
+# Image.resize()
+@func.udf(param_types=[ImageType(), ArrayType((2, ), dtype=IntType())])
+def resize(self: PIL.Image.Image, size: Tuple[int, int]) -> PIL.Image.Image:
+    return self.resize(size)
+
+@resize.conditional_return_type
+def _(self: PIL.Image.Image, size: Tuple[int, int]) -> ColumnType:
+    input_type = self.col_type
+    assert input_type.is_image_type()
+    return ImageType(size=size, mode=input_type.mode, nullable=input_type.nullable)
+
+# Image.rotate()
+@func.udf(param_types=[ImageType(), IntType()])
+def rotate(self: PIL.Image.Image, angle: int) -> PIL.Image.Image:
+    return self.rotate(angle)
+
+@func.udf(py_fn=PIL.Image.Image.effect_spread, param_types=[ImageType(), IntType()])
+def effect_spread(self: PIL.Image.Image, distance: int) -> PIL.Image.Image:
+    pass
+
+@func.udf(py_fn=PIL.Image.Image.transpose, param_types=[ImageType(), IntType()])
+def transpose(self: PIL.Image.Image, method: int) -> PIL.Image.Image:
+    pass
+
+@rotate.conditional_return_type
+@effect_spread.conditional_return_type
+@transpose.conditional_return_type
+def _(self: PIL.Image.Image) -> ColumnType:
+    return self.col_type
+
+@func.udf(
+    py_fn=PIL.Image.Image.entropy, return_type=FloatType(), param_types=[ImageType(), ImageType(), JsonType()])
+def entropy(self: PIL.Image.Image, mask: PIL.Image.Image, extrema: Optional[list] = None) -> float:
+    pass
+
+@func.udf(py_fn=PIL.Image.Image.getbands, return_type=JsonType(), param_types=[ImageType()])
+def getbands(self: PIL.Image.Image) -> Tuple[str]:
+    pass
+
+@func.udf(py_fn=PIL.Image.Image.getbbox, return_type=JsonType(), param_types=[ImageType()])
+def getbbox(self: PIL.Image.Image) -> Tuple[int, int, int, int]:
+    pass
+
+@func.udf(py_fn=PIL.Image.Image.getcolors, return_type=JsonType(), param_types=[ImageType(), IntType()])
+def getcolors(self: PIL.Image.Image, maxcolors: int) -> Tuple[Tuple[int, int, int], int]:
+    pass
+
+@func.udf(py_fn=PIL.Image.Image.getextrema, return_type=JsonType(), param_types=[ImageType()])
+def getextrema(self: PIL.Image.Image) -> Tuple[int, int]:
+    pass
+
+@func.udf(
+    py_fn=PIL.Image.Image.getpalette, return_type=JsonType(), param_types=[ImageType(), StringType()])
+def getpalette(self: PIL.Image.Image, mode: Optional[str] = None) -> Tuple[int]:
+    pass
+
+@func.udf(
+    return_type=JsonType(), param_types=[ImageType(), ArrayType((2,), dtype=IntType())])
+def getpixel(self: PIL.Image.Image, xy: tuple[int, int]) -> Tuple[int]:
+    # `xy` will be a list; `tuple(xy)` is necessary for pillow 9 compatibility
+    return self.getpixel(tuple(xy))
+
+@func.udf(py_fn=PIL.Image.Image.getprojection, return_type=JsonType(), param_types=[ImageType()])
+def getprojection(self: PIL.Image.Image) -> Tuple[int]:
+    pass
+
+@func.udf(py_fn=PIL.Image.Image.histogram, return_type=JsonType(), param_types=[ImageType(), ImageType(), JsonType()])
+def histogram(self: PIL.Image.Image, mask: PIL.Image.Image, extrema: Optional[list] = None) -> Tuple[int]:
+    pass
+
+@func.udf(
+    py_fn=PIL.Image.Image.quantize, return_type=ImageType(),
+    param_types=[ImageType(), IntType(), IntType(nullable=True), IntType(), IntType(nullable=True), IntType()])
+def quantize(
+        self: PIL.Image.Image, colors: int = 256, method: Optional[int] = None, kmeans: int = 0,
+        palette: Optional[int] = None, dither: int = Dither.FLOYDSTEINBERG) -> PIL.Image.Image:
+    pass
+
+@func.udf(
+    py_fn=PIL.Image.Image.reduce, return_type=ImageType(), param_types=[ImageType(), IntType(), JsonType()])
+def reduce(self: PIL.Image.Image, factor: int, box: Optional[Tuple[int]]) -> PIL.Image.Image:
+    pass

pixeltable/functions/string.py

@@ -0,0 +1,13 @@
+from typing import Any
+
+from pixeltable.type_system import StringType
+import pixeltable.func as func
+
+
+@func.udf(return_type=StringType(), param_types=[StringType()])
+def str_format(format_str: str, *args: Any, **kwargs: Any) -> str:
+    """ Return a formatted version of format_str, using substitutions from args and kwargs:
+    - {<int>} will be replaced by the corresponding element in args
+    - {<key>} will be replaced by the corresponding value in kwargs
+    """
+    return format_str.format(*args, **kwargs)

pixeltable/functions/together.py

@@ -0,0 +1,143 @@
+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,
+        stop: Optional[list] = 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
+) -> dict:
+    return together_client().completions.create(
+        prompt=prompt,
+        model=model,
+        max_tokens=max_tokens,
+        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,
+        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()
+
+
+_embedding_dimensions_cache = {
+    'togethercomputer/m2-bert-80M-2k-retrieval': 768,
+    'togethercomputer/m2-bert-80M-8k-retrieval': 768,
+    'togethercomputer/m2-bert-80M-32k-retrieval': 768,
+    'WhereIsAI/UAE-Large-V1': 1024,
+    'BAAI/bge-large-en-v1.5': 1024,
+    'BAAI/bge-base-en-v1.5': 768,
+    'sentence-transformers/msmarco-bert-base-dot-v5': 768,
+    'bert-base-uncased': 768,
+}
+
+
+@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
+    ]
+
+
+@embeddings.conditional_return_type
+def _(model: str) -> pxt.ArrayType:
+    if model not in _embedding_dimensions_cache:
+        # TODO: find some other way to retrieve a sample
+        return pxt.ArrayType((None,), dtype=pxt.FloatType())
+    dimensions = _embedding_dimensions_cache[model]
+    return pxt.ArrayType((dimensions,), dtype=pxt.FloatType())
+
+
+@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/functions/util.py

@@ -0,0 +1,52 @@
+from typing import Tuple, List, Optional
+import types
+import sys
+
+import pixeltable.func as func
+import pixeltable.type_system as ts
+import pixeltable.env as env
+
+
+def create_nos_modules() -> List[types.ModuleType]:
+    """Create module pixeltable.functions.nos with one submodule per task and return the submodules"""
+    models = env.Env.get().nos_client.ListModels()
+    model_info = [env.Env.get().nos_client.GetModelInfo(model) for model in models]
+    model_info.sort(key=lambda info: info.task.value)
+
+    module_name = 'pixeltable.functions.nos'
+    nos_module = types.ModuleType(module_name)
+    nos_module.__package__ = 'pixeltable.functions'
+    sys.modules[module_name] = nos_module
+
+    prev_task = ''
+    new_modules: List[types.ModuleType] = []
+    sub_module: Optional[types.ModuleType] = None
+    for info in model_info:
+        if info.task.value != prev_task:
+            # we construct one submodule per task
+            namespace = info.task.name.lower()
+            submodule_name = f'{module_name}.{namespace}'
+            sub_module = types.ModuleType(submodule_name)
+            sub_module.__package__ = module_name
+            setattr(nos_module, namespace, sub_module)
+            new_modules.append(sub_module)
+            sys.modules[submodule_name] = sub_module
+            prev_task = info.task.value
+
+        # add a Function for this model to the module
+        model_id = info.name.replace("/", "_").replace("-", "_")
+        pt_func = func.NOSFunction(info, f'{submodule_name}.{model_id}')
+        setattr(sub_module, model_id, pt_func)
+
+    return new_modules
+
+
+def resolve_torch_device(device: str) -> str:
+    import torch
+    if device == 'auto':
+        if torch.cuda.is_available():
+            return 'cuda'
+        if torch.backends.mps.is_available():
+            return 'mps'
+        return 'cpu'
+    return device

pixeltable/functions/video.py

@@ -0,0 +1,62 @@
+from typing import Optional
+import uuid
+import av
+import sys
+
+import pixeltable.env as env
+import pixeltable.func as func
+import pixeltable.type_system as ts
+
+
+_format_defaults = { # format -> (codec, ext)
+    'wav': ('pcm_s16le', 'wav'),
+    'mp3': ('libmp3lame', 'mp3'),
+    'flac': ('flac', 'flac'),
+    #'mp4': ('aac', 'm4a'),
+}
+
+# for mp4:
+# - extract_audio() fails with "Application provided invalid, non monotonically increasing dts to muxer in stream 0: 1146 >= 290"
+# - chatgpt suggests this can be fixed in the following manner
+#     for packet in container.demux(audio_stream):
+#         packet.pts = None  # Reset the PTS and DTS to allow FFmpeg to set them automatically
+#         packet.dts = None
+#         for frame in packet.decode():
+#             frame.pts = None
+#             for packet in output_stream.encode(frame):
+#                 output_container.mux(packet)
+#
+#     # Flush remaining packets
+#     for packet in output_stream.encode():
+#         output_container.mux(packet)
+
+
+_extract_audio_param_types = [
+    ts.VideoType(nullable=False),
+    ts.IntType(nullable=False),
+    ts.StringType(nullable=False),
+    ts.StringType(nullable=False)
+]
+@func.udf(return_type=ts.AudioType(nullable=True), param_types=_extract_audio_param_types)
+def extract_audio(
+        video_path: str, stream_idx: int = 0, format: str = 'wav', codec: Optional[str] = None
+) -> Optional[str]:
+    """Extract an audio stream from a video file, save it as a media file and return its path"""
+    if format not in _format_defaults:
+        raise ValueError(f'extract_audio(): unsupported audio format: {format}')
+    default_codec, ext = _format_defaults[format]
+
+    with av.open(video_path) as container:
+        if len(container.streams.audio) <= stream_idx:
+            return None
+        audio_stream = container.streams.audio[stream_idx]
+        # create this in our tmp directory, so it'll get cleaned up if it's being generated as part of a query
+        output_filename = str(env.Env.get().tmp_dir / f"{uuid.uuid4()}.{ext}")
+
+        with av.open(output_filename, "w", format=format) as output_container:
+            output_stream = output_container.add_stream(codec or default_codec)
+            for packet in container.demux(audio_stream):
+                for frame in packet.decode():
+                    output_container.mux(output_stream.encode(frame))
+
+        return output_filename