pixeltable 0.3.2__py3-none-any.whl → 0.3.4__py3-none-any.whl

pixeltable/functions/openai.py

@@ -10,15 +10,15 @@ import datetime
 import io
 import json
 import logging
+import math
 import pathlib
 import re
 import uuid
-from typing import TYPE_CHECKING, Callable, Optional, TypeVar, Union, cast, Any, Type
+from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Type, TypeVar, Union, cast
 
-import PIL.Image
 import httpx
 import numpy as np
-import tenacity
+import PIL
 
 import pixeltable as pxt
 from pixeltable import env, exprs
@@ -32,36 +32,18 @@ _logger = logging.getLogger('pixeltable')
 
 
 @env.register_client('openai')
-def _(api_key: str) -> tuple['openai.OpenAI', 'openai.AsyncOpenAI']:
+def _(api_key: str) -> 'openai.AsyncOpenAI':
     import openai
-    return (
-        openai.OpenAI(api_key=api_key),
-        openai.AsyncOpenAI(
-            api_key=api_key,
-            # recommended to increase limits for async client to avoid connection errors
-            http_client=httpx.AsyncClient(limits=httpx.Limits(max_keepalive_connections=100, max_connections=500)),
-        )
-    )
-
-
-def _openai_client() -> 'openai.OpenAI':
-    return env.Env.get().get_client('openai')[0]
-
 
-def _async_openai_client() -> 'openai.AsyncOpenAI':
-    return env.Env.get().get_client('openai')[1]
+    return openai.AsyncOpenAI(
+        api_key=api_key,
+        # recommended to increase limits for async client to avoid connection errors
+        http_client=httpx.AsyncClient(limits=httpx.Limits(max_keepalive_connections=100, max_connections=500)),
+    )
 
 
-# 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:
-    import openai
-    return tenacity.retry(
-        retry=tenacity.retry_if_exception_type(openai.RateLimitError),
-        wait=tenacity.wait_random_exponential(multiplier=1, max=60),
-        stop=tenacity.stop_after_attempt(20),
-    )(fn)
+def _openai_client() -> 'openai.AsyncOpenAI':
+    return env.Env.get().get_client('openai')
 
 
 # models that share rate limits; see https://platform.openai.com/settings/organization/limits for details
@@ -72,7 +54,7 @@ _shared_rate_limits = {
         'gpt-4-turbo-2024-04-09',
         'gpt-4-turbo-preview',
         'gpt-4-0125-preview',
-        'gpt-4-1106-preview'
+        'gpt-4-1106-preview',
     ],
     'gpt-4o': [
         'gpt-4o',
@@ -82,24 +64,24 @@ _shared_rate_limits = {
         'gpt-4o-2024-11-20',
         'gpt-4o-audio-preview',
         'gpt-4o-audio-preview-2024-10-01',
-        'gpt-4o-audio-preview-2024-12-17'
+        'gpt-4o-audio-preview-2024-12-17',
     ],
     'gpt-4o-mini': [
         'gpt-4o-mini',
         'gpt-4o-mini-latest',
         'gpt-4o-mini-2024-07-18',
         'gpt-4o-mini-audio-preview',
-        'gpt-4o-mini-audio-preview-2024-12-17'
+        'gpt-4o-mini-audio-preview-2024-12-17',
     ],
     'gpt-4o-mini-realtime-preview': [
         'gpt-4o-mini-realtime-preview',
         'gpt-4o-mini-realtime-preview-latest',
-        'gpt-4o-mini-realtime-preview-2024-12-17'
-    ]
+        'gpt-4o-mini-realtime-preview-2024-12-17',
+    ],
 }
 
 
-def _resource_pool(model: str) -> str:
+def _rate_limits_pool(model: str) -> str:
     for model_family, models in _shared_rate_limits.items():
         if model in models:
             return f'rate-limits:openai:{model_family}'
@@ -112,13 +94,13 @@ class OpenAIRateLimitsInfo(env.RateLimitsInfo):
     def __init__(self, get_request_resources: Callable[..., dict[str, int]]):
         super().__init__(get_request_resources)
         import openai
+
         self.retryable_errors = (
             # ConnectionError: we occasionally see this error when the AsyncConnectionPool is trying to close
             # expired connections
             # (AsyncConnectionPool._close_expired_connections() fails with ConnectionError when executing
-            # 'await connection.aclose()', which is potentially a bug in AsyncConnectionPool)
+            # 'await connection.aclose()', which is very likely a bug in AsyncConnectionPool)
             openai.APIConnectionError,
-
             # the following errors are retryable according to OpenAI's API documentation
             openai.RateLimitError,
             openai.APITimeoutError,
@@ -143,7 +125,7 @@ _header_duration_pattern = re.compile(r'(?:(\d+)d)?(?:(\d+)h)?(?:(\d+)ms)|(?:(\d
 def _parse_header_duration(duration_str):
     match = _header_duration_pattern.match(duration_str)
     if not match:
-        raise ValueError("Invalid duration format")
+        raise ValueError('Invalid duration format')
 
     days = int(match.group(1) or 0)
     hours = int(match.group(2) or 0)
@@ -151,17 +133,11 @@ def _parse_header_duration(duration_str):
     minutes = int(match.group(4) or 0)
     seconds = float(match.group(5) or 0)
 
-    return datetime.timedelta(
-        days=days,
-        hours=hours,
-        minutes=minutes,
-        seconds=seconds,
-        milliseconds=milliseconds
-    )
+    return datetime.timedelta(days=days, hours=hours, minutes=minutes, seconds=seconds, milliseconds=milliseconds)
 
 
 def _get_header_info(
-        headers: httpx.Headers, *, requests: bool = True, tokens: bool = True
+    headers: httpx.Headers, *, requests: bool = True, tokens: bool = True
 ) -> tuple[Optional[tuple[int, int, datetime.datetime]], Optional[tuple[int, int, datetime.datetime]]]:
     assert requests or tokens
     now = datetime.datetime.now(tz=datetime.timezone.utc)
@@ -194,8 +170,14 @@ def _get_header_info(
 
 
 @pxt.udf
-def speech(
-        input: str, *, model: str, voice: str, response_format: Optional[str] = None, speed: Optional[float] = None
+async def speech(
+    input: str,
+    *,
+    model: str,
+    voice: str,
+    response_format: Optional[str] = None,
+    speed: Optional[float] = None,
+    timeout: Optional[float] = None,
 ) -> pxt.Audio:
     """
     Generates audio from the input text.
@@ -203,6 +185,10 @@ def speech(
     Equivalent to the OpenAI `audio/speech` API endpoint.
     For additional details, see: [https://platform.openai.com/docs/guides/text-to-speech](https://platform.openai.com/docs/guides/text-to-speech)
 
+    Request throttling:
+    Applies the rate limit set in the config (section `openai.rate_limits`; use the model id as the key). If no rate
+    limit is configured, uses a default of 600 RPM.
+
     __Requirements:__
 
     - `pip install openai`
@@ -222,10 +208,15 @@ def speech(
         Add a computed column that applies the model `tts-1` to an existing Pixeltable column `tbl.text`
         of the table `tbl`:
 
-        >>> tbl['audio'] = speech(tbl.text, model='tts-1', voice='nova')
+        >>> tbl.add_computed_column(audio=speech(tbl.text, model='tts-1', voice='nova'))
     """
-    content = _retry(_openai_client().audio.speech.create)(
-        input=input, model=model, voice=voice, response_format=_opt(response_format), speed=_opt(speed)
+    content = await _openai_client().audio.speech.create(
+        input=input,
+        model=model,
+        voice=voice,  # type: ignore
+        response_format=_opt(response_format),  # type: ignore
+        speed=_opt(speed),
+        timeout=_opt(timeout),
     )
     ext = response_format or 'mp3'
     output_filename = str(env.Env.get().tmp_dir / f'{uuid.uuid4()}.{ext}')
@@ -234,13 +225,14 @@ def speech(
 
 
 @pxt.udf
-def transcriptions(
+async def transcriptions(
     audio: pxt.Audio,
     *,
     model: str,
     language: Optional[str] = None,
     prompt: Optional[str] = None,
     temperature: Optional[float] = None,
+    timeout: Optional[float] = None,
 ) -> dict:
     """
     Transcribes audio into the input language.
@@ -248,6 +240,10 @@ def transcriptions(
     Equivalent to the OpenAI `audio/transcriptions` API endpoint.
     For additional details, see: [https://platform.openai.com/docs/guides/speech-to-text](https://platform.openai.com/docs/guides/speech-to-text)
 
+    Request throttling:
+    Applies the rate limit set in the config (section `openai.rate_limits`; use the model id as the key). If no rate
+    limit is configured, uses a default of 600 RPM.
+
     __Requirements:__
 
     - `pip install openai`
@@ -265,22 +261,28 @@ def transcriptions(
         Add a computed column that applies the model `whisper-1` to an existing Pixeltable column `tbl.audio`
         of the table `tbl`:
 
-        >>> tbl['transcription'] = transcriptions(tbl.audio, model='whisper-1', language='en')
+        >>> tbl.add_computed_column(transcription=transcriptions(tbl.audio, model='whisper-1', language='en'))
     """
     file = pathlib.Path(audio)
-    transcription = _retry(_openai_client().audio.transcriptions.create)(
-        file=file, model=model, language=_opt(language), prompt=_opt(prompt), temperature=_opt(temperature)
+    transcription = await _openai_client().audio.transcriptions.create(
+        file=file,
+        model=model,
+        language=_opt(language),
+        prompt=_opt(prompt),
+        temperature=_opt(temperature),
+        timeout=_opt(timeout),
     )
     return transcription.dict()
 
 
 @pxt.udf
-def translations(
+async def translations(
     audio: pxt.Audio,
     *,
     model: str,
     prompt: Optional[str] = None,
-    temperature: Optional[float] = None
+    temperature: Optional[float] = None,
+    timeout: Optional[float] = None,
 ) -> dict:
     """
     Translates audio into English.
@@ -288,6 +290,10 @@ def translations(
     Equivalent to the OpenAI `audio/translations` API endpoint.
     For additional details, see: [https://platform.openai.com/docs/guides/speech-to-text](https://platform.openai.com/docs/guides/speech-to-text)
 
+    Request throttling:
+    Applies the rate limit set in the config (section `openai.rate_limits`; use the model id as the key). If no rate
+    limit is configured, uses a default of 600 RPM.
+
     __Requirements:__
 
     - `pip install openai`
@@ -305,11 +311,11 @@ def translations(
         Add a computed column that applies the model `whisper-1` to an existing Pixeltable column `tbl.audio`
         of the table `tbl`:
 
-        >>> tbl['translation'] = translations(tbl.audio, model='whisper-1', language='en')
+        >>> tbl.add_computed_column(translation=translations(tbl.audio, model='whisper-1', language='en'))
     """
     file = pathlib.Path(audio)
-    translation = _retry(_openai_client().audio.translations.create)(
-        file=file, model=model, prompt=_opt(prompt), temperature=_opt(temperature)
+    translation = await _openai_client().audio.translations.create(
+        file=file, model=model, prompt=_opt(prompt), temperature=_opt(temperature), timeout=_opt(timeout)
     )
     return translation.dict()
 
@@ -318,17 +324,24 @@ def translations(
 # Chat Endpoints
 
 
+def _default_max_tokens(model: str) -> int:
+    if model in ('o1', 'o3-mini'):
+        return 65536
+    else:
+        return 1024
+
+
 def _chat_completions_get_request_resources(
-        messages: list, max_tokens: Optional[int], n: Optional[int]
+    messages: list, model: str, max_completion_tokens: Optional[int], max_tokens: Optional[int], n: Optional[int]
 ) -> dict[str, int]:
-    completion_tokens = n * max_tokens
+    completion_tokens = (n or 1) * (max_completion_tokens or max_tokens or _default_max_tokens(model))
 
     num_tokens = 0.0
     for message in messages:
         num_tokens += 4  # every message follows <im_start>{role/name}\n{content}<im_end>\n
         for key, value in message.items():
             num_tokens += len(value) / 4
-            if key == "name":  # if there's a name, the role is omitted
+            if key == 'name':  # if there's a name, the role is omitted
                 num_tokens -= 1  # role is always required and always 1 token
     num_tokens += 2  # every reply is primed with <im_start>assistant
     return {'requests': 1, 'tokens': int(num_tokens) + completion_tokens}
@@ -343,17 +356,20 @@ async def chat_completions(
     logit_bias: Optional[dict[str, int]] = None,
     logprobs: Optional[bool] = None,
     top_logprobs: Optional[int] = None,
-    max_tokens: Optional[int] = 1024,
-    n: Optional[int] = 1,
+    max_completion_tokens: Optional[int] = None,
+    max_tokens: Optional[int] = None,
+    n: Optional[int] = None,
     presence_penalty: Optional[float] = None,
+    reasoning_effort: Optional[Literal['low', 'medium', 'high']] = 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,
+    top_p: Optional[float] = None,
     user: Optional[str] = None,
+    timeout: Optional[float] = None,
 ) -> dict:
     """
     Creates a model response for the given chat conversation.
@@ -361,6 +377,10 @@ async def chat_completions(
     Equivalent to the OpenAI `chat/completions` API endpoint.
     For additional details, see: [https://platform.openai.com/docs/guides/chat-completions](https://platform.openai.com/docs/guides/chat-completions)
 
+    Request throttling:
+    Uses the rate limit-related headers returned by the API to throttle requests adaptively, based on available
+    request and token capacity. No configuration is necessary.
+
     __Requirements:__
 
     - `pip install openai`
@@ -382,16 +402,10 @@ async def chat_completions(
                 {'role': 'system', 'content': 'You are a helpful assistant.'},
                 {'role': 'user', 'content': tbl.prompt}
             ]
-            tbl['response'] = chat_completions(messages, model='gpt-4o-mini')
+            tbl.add_computed_column(response=chat_completions(messages, model='gpt-4o-mini'))
     """
     if tools is not None:
-        tools = [
-            {
-                'type': 'function',
-                'function': tool
-            }
-            for tool in tools
-        ]
+        tools = [{'type': 'function', 'function': tool} for tool in tools]
 
     tool_choice_: Union[str, dict, None] = None
     if tool_choice is not None:
@@ -401,40 +415,43 @@ async def chat_completions(
             tool_choice_ = 'required'
         else:
             assert tool_choice['tool'] is not None
-            tool_choice_ = {
-                'type': 'function',
-                'function': {'name': tool_choice['tool']}
-            }
+            tool_choice_ = {'type': 'function', 'function': {'name': tool_choice['tool']}}
 
     extra_body: Optional[dict[str, Any]] = None
     if tool_choice is not None and not tool_choice['parallel_tool_calls']:
         extra_body = {'parallel_tool_calls': False}
 
     # make sure the pool info exists prior to making the request
-    resource_pool = _resource_pool(model)
+    resource_pool = _rate_limits_pool(model)
     rate_limits_info = env.Env.get().get_resource_pool_info(
-        resource_pool, lambda: OpenAIRateLimitsInfo(_chat_completions_get_request_resources))
+        resource_pool, lambda: OpenAIRateLimitsInfo(_chat_completions_get_request_resources)
+    )
+
+    if max_completion_tokens is None and max_tokens is None:
+        max_completion_tokens = _default_max_tokens(model)
 
     # cast(Any, ...): avoid mypy errors
-    result = await _async_openai_client().chat.completions.with_raw_response.create(
+    result = await _openai_client().chat.completions.with_raw_response.create(
         messages=messages,
         model=model,
         frequency_penalty=_opt(frequency_penalty),
         logit_bias=_opt(logit_bias),
         logprobs=_opt(logprobs),
         top_logprobs=_opt(top_logprobs),
+        max_completion_tokens=_opt(max_completion_tokens),
         max_tokens=_opt(max_tokens),
         n=_opt(n),
         presence_penalty=_opt(presence_penalty),
+        reasoning_effort=_opt(reasoning_effort),
         response_format=_opt(cast(Any, response_format)),
         seed=_opt(seed),
         stop=_opt(stop),
         temperature=_opt(temperature),
-        top_p=_opt(top_p),
         tools=_opt(cast(Any, tools)),
         tool_choice=_opt(cast(Any, tool_choice_)),
+        top_p=_opt(top_p),
         user=_opt(user),
-        timeout=10,
+        timeout=_opt(timeout),
         extra_body=extra_body,
     )
 
@@ -444,14 +461,61 @@ async def chat_completions(
     return json.loads(result.text)
 
 
+def _vision_get_request_resources(
+    prompt: str,
+    image: PIL.Image.Image,
+    model: str,
+    max_completion_tokens: Optional[int],
+    max_tokens: Optional[int],
+    n: Optional[int],
+) -> dict[str, int]:
+    completion_tokens = (n or 1) * (max_completion_tokens or max_tokens or _default_max_tokens(model))
+    prompt_tokens = len(prompt) / 4
+
+    # calculate image tokens based on
+    # https://platform.openai.com/docs/guides/vision/calculating-costs#calculating-costs
+    # assuming detail='high' (which appears to be the default, according to community forum posts)
+
+    # number of 512x512 crops; ceil(): partial crops still count as full crops
+    crops_width = math.ceil(image.width / 512)
+    crops_height = math.ceil(image.height / 512)
+    total_crops = crops_width * crops_height
+
+    BASE_TOKENS = 85  # base cost for the initial 512x512 overview
+    CROP_TOKENS = 170  # cost per additional 512x512 crop
+    img_tokens = BASE_TOKENS + (CROP_TOKENS * total_crops)
+
+    total_tokens = (
+        prompt_tokens
+        + img_tokens
+        + completion_tokens
+        + 4  # for <im_start>{role/name}\n{content}<im_end>\n
+        + 2  # for reply's <im_start>assistant
+    )
+    return {'requests': 1, 'tokens': int(total_tokens)}
+
+
 @pxt.udf
-def vision(prompt: str, image: PIL.Image.Image, *, model: str) -> str:
+async def vision(
+    prompt: str,
+    image: PIL.Image.Image,
+    *,
+    model: str,
+    max_completion_tokens: Optional[int] = None,
+    max_tokens: Optional[int] = None,
+    n: Optional[int] = 1,
+    timeout: Optional[float] = None,
+) -> str:
     """
     Analyzes an image with the OpenAI vision capability. This is a convenience function that takes an image and
     prompt, and constructs a chat completion request that utilizes OpenAI vision.
 
     For additional details, see: [https://platform.openai.com/docs/guides/vision](https://platform.openai.com/docs/guides/vision)
 
+    Request throttling:
+    Uses the rate limit-related headers returned by the API to throttle requests adaptively, based on available
+    request and token capacity. No configuration is necessary.
+
     __Requirements:__
 
     - `pip install openai`
@@ -468,7 +532,7 @@ def vision(prompt: str, image: PIL.Image.Image, *, model: str) -> str:
         Add a computed column that applies the model `gpt-4o-mini` to an existing Pixeltable column `tbl.image`
         of the table `tbl`:
 
-        >>> tbl['response'] = vision("What's in this image?", tbl.image, model='gpt-4o-mini')
+        >>> tbl.add_computed_column(response=vision("What's in this image?", tbl.image, model='gpt-4o-mini'))
     """
     # TODO(aaron-siegel): Decompose CPU/GPU ops into separate functions
     bytes_arr = io.BytesIO()
@@ -484,8 +548,30 @@ def vision(prompt: str, image: PIL.Image.Image, *, model: str) -> str:
             ],
         }
     ]
-    result = _retry(_openai_client().chat.completions.create)(messages=messages, model=model)
-    return result.choices[0].message.content
+
+    # make sure the pool info exists prior to making the request
+    resource_pool = _rate_limits_pool(model)
+    rate_limits_info = env.Env.get().get_resource_pool_info(
+        resource_pool, lambda: OpenAIRateLimitsInfo(_vision_get_request_resources)
+    )
+
+    if max_completion_tokens is None and max_tokens is None:
+        max_completion_tokens = _default_max_tokens(model)
+
+    result = await _openai_client().chat.completions.with_raw_response.create(
+        messages=messages,  # type: ignore
+        model=model,
+        max_completion_tokens=_opt(max_completion_tokens),
+        max_tokens=_opt(max_tokens),
+        n=_opt(n),
+        timeout=_opt(timeout),
+    )
+
+    requests_info, tokens_info = _get_header_info(result.headers)
+    rate_limits_info.record(requests=requests_info, tokens=tokens_info)
+
+    result = json.loads(result.text)
+    return result['choices'][0]['message']['content']
 
 
 #####################################
@@ -505,7 +591,12 @@ def _embeddings_get_request_resources(input: list[str]) -> dict[str, int]:
 
 @pxt.udf(batch_size=32)
 async def embeddings(
-    input: Batch[str], *, model: str, dimensions: Optional[int] = None, user: Optional[str] = None
+    input: Batch[str],
+    *,
+    model: str,
+    dimensions: Optional[int] = None,
+    user: Optional[str] = None,
+    timeout: Optional[float] = None,
 ) -> Batch[pxt.Array[(None,), pxt.Float]]:
     """
     Creates an embedding vector representing the input text.
@@ -513,6 +604,10 @@ async def embeddings(
     Equivalent to the OpenAI `embeddings` API endpoint.
     For additional details, see: [https://platform.openai.com/docs/guides/embeddings](https://platform.openai.com/docs/guides/embeddings)
 
+    Request throttling:
+    Uses the rate limit-related headers returned by the API to throttle requests adaptively, based on available
+    request and token capacity. No configuration is necessary.
+
     __Requirements:__
 
     - `pip install openai`
@@ -532,14 +627,24 @@ async def embeddings(
         Add a computed column that applies the model `text-embedding-3-small` to an existing
         Pixeltable column `tbl.text` of the table `tbl`:
 
-        >>> tbl['embed'] = embeddings(tbl.text, model='text-embedding-3-small')
+        >>> tbl.add_computed_column(embed=embeddings(tbl.text, model='text-embedding-3-small'))
+
+        Add an embedding index to an existing column `text`, using the model `text-embedding-3-small`:
+
+        >>> tbl.add_embedding_index(embedding=embeddings.using(model='text-embedding-3-small'))
     """
     _logger.debug(f'embeddings: batch_size={len(input)}')
-    resource_pool = _resource_pool(model)
+    resource_pool = _rate_limits_pool(model)
     rate_limits_info = env.Env.get().get_resource_pool_info(
-        resource_pool, lambda: OpenAIRateLimitsInfo(_embeddings_get_request_resources))
-    result = await _async_openai_client().embeddings.with_raw_response.create(
-        input=input, model=model, dimensions=_opt(dimensions), user=_opt(user), encoding_format='float'
+        resource_pool, lambda: OpenAIRateLimitsInfo(_embeddings_get_request_resources)
+    )
+    result = await _openai_client().embeddings.with_raw_response.create(
+        input=input,
+        model=model,
+        dimensions=_opt(dimensions),
+        user=_opt(user),
+        encoding_format='float',
+        timeout=_opt(timeout),
     )
     requests_info, tokens_info = _get_header_info(result.headers)
     rate_limits_info.record(requests=requests_info, tokens=tokens_info)
@@ -561,7 +666,7 @@ def _(model: str, dimensions: Optional[int] = None) -> pxt.ArrayType:
 
 
 @pxt.udf
-def image_generations(
+async def image_generations(
     prompt: str,
     *,
     model: str = 'dall-e-2',
@@ -569,6 +674,7 @@ def image_generations(
     size: Optional[str] = None,
     style: Optional[str] = None,
     user: Optional[str] = None,
+    timeout: Optional[float] = None,
 ) -> PIL.Image.Image:
     """
     Creates an image given a prompt.
@@ -576,6 +682,10 @@ def image_generations(
     Equivalent to the OpenAI `images/generations` API endpoint.
     For additional details, see: [https://platform.openai.com/docs/guides/images](https://platform.openai.com/docs/guides/images)
 
+    Request throttling:
+    Applies the rate limit set in the config (section `openai.rate_limits`; use the model id as the key). If no rate
+    limit is configured, uses a default of 600 RPM.
+
     __Requirements:__
 
     - `pip install openai`
@@ -593,17 +703,18 @@ def image_generations(
         Add a computed column that applies the model `dall-e-2` to an existing
         Pixeltable column `tbl.text` of the table `tbl`:
 
-        >>> tbl['gen_image'] = image_generations(tbl.text, model='dall-e-2')
+        >>> tbl.add_computed_column(gen_image=image_generations(tbl.text, model='dall-e-2'))
     """
     # TODO(aaron-siegel): Decompose CPU/GPU ops into separate functions
-    result = _retry(_openai_client().images.generate)(
+    result = await _openai_client().images.generate(
         prompt=prompt,
         model=_opt(model),
-        quality=_opt(quality),
-        size=_opt(size),
-        style=_opt(style),
+        quality=_opt(quality),  # type: ignore
+        size=_opt(size),  # type: ignore
+        style=_opt(style),  # type: ignore
         user=_opt(user),
         response_format='b64_json',
+        timeout=_opt(timeout),
     )
     b64_str = result.data[0].b64_json
     b64_bytes = base64.b64decode(b64_str)
@@ -620,7 +731,7 @@ def _(size: Optional[str] = None) -> pxt.ImageType:
     if x_pos == -1:
         return pxt.ImageType()
     try:
-        width, height = int(size[:x_pos]), int(size[x_pos + 1:])
+        width, height = int(size[:x_pos]), int(size[x_pos + 1 :])
     except ValueError:
         return pxt.ImageType()
     return pxt.ImageType(size=(width, height))
@@ -631,13 +742,17 @@ def _(size: Optional[str] = None) -> pxt.ImageType:
 
 
 @pxt.udf
-def moderations(input: str, *, model: str = 'omni-moderation-latest') -> dict:
+async def moderations(input: str, *, model: str = 'omni-moderation-latest') -> dict:
     """
     Classifies if text is potentially harmful.
 
     Equivalent to the OpenAI `moderations` API endpoint.
     For additional details, see: [https://platform.openai.com/docs/guides/moderation](https://platform.openai.com/docs/guides/moderation)
 
+    Request throttling:
+    Applies the rate limit set in the config (section `openai.rate_limits`; use the model id as the key). If no rate
+    limit is configured, uses a default of 600 RPM.
+
     __Requirements:__
 
     - `pip install openai`
@@ -655,22 +770,26 @@ def moderations(input: str, *, model: str = 'omni-moderation-latest') -> dict:
         Add a computed column that applies the model `text-moderation-stable` to an existing
         Pixeltable column `tbl.input` of the table `tbl`:
 
-        >>> tbl['moderations'] = moderations(tbl.text, model='text-moderation-stable')
+        >>> tbl.add_computed_column(moderations=moderations(tbl.text, model='text-moderation-stable'))
     """
-    result = _retry(_openai_client().moderations.create)(input=input, model=_opt(model))
+    result = await _openai_client().moderations.create(input=input, model=_opt(model))
     return result.dict()
 
 
-# @speech.resource_pool
-# @transcriptions.resource_pool
-# @translations.resource_pool
+@speech.resource_pool
+@transcriptions.resource_pool
+@translations.resource_pool
+@image_generations.resource_pool
+@moderations.resource_pool
+def _(model: str) -> str:
+    return f'request-rate:openai:{model}'
+
+
 @chat_completions.resource_pool
-# @vision.resource_pool
+@vision.resource_pool
 @embeddings.resource_pool
-# @image_generations.resource_pool
-# @moderations.resource_pool
 def _(model: str) -> str:
-    return _resource_pool(model)
+    return _rate_limits_pool(model)
 
 
 def invoke_tools(tools: Tools, response: exprs.Expr) -> exprs.InlineDict:
@@ -684,9 +803,7 @@ def _openai_response_to_pxt_tool_calls(response: dict) -> Optional[dict]:
         return None
     openai_tool_calls = response['choices'][0]['message']['tool_calls']
     return {
-        tool_call['function']['name']: {
-            'args': json.loads(tool_call['function']['arguments'])
-        }
+        tool_call['function']['name']: {'args': json.loads(tool_call['function']['arguments'])}
         for tool_call in openai_tool_calls
     }
 
@@ -696,6 +813,7 @@ _T = TypeVar('_T')
 
 def _opt(arg: _T) -> Union[_T, 'openai.NotGiven']:
     import openai
+
     return arg if arg is not None else openai.NOT_GIVEN