pixeltable 0.4.0rc3__py3-none-any.whl → 0.4.20__py3-none-any.whl

pixeltable/functions/image.py

@@ -10,15 +10,13 @@ t.select(t.img_col.convert('L')).collect()
 ```
 """
 
-import base64
-from typing import Optional
-
 import PIL.Image
 
 import pixeltable as pxt
 import pixeltable.type_system as ts
 from pixeltable.exprs import Expr
 from pixeltable.utils.code import local_public_names
+from pixeltable.utils.image import to_base64
 
 
 @pxt.udf(is_method=True)
@@ -30,12 +28,7 @@ def b64_encode(img: PIL.Image.Image, image_format: str = 'png') -> str:
         img: image
         image_format: image format [supported by PIL](https://pillow.readthedocs.io/en/stable/handbook/image-file-formats.html#fully-supported-formats)
     """
-    import io
-
-    bytes_arr = io.BytesIO()
-    img.save(bytes_arr, format=image_format)
-    b64_bytes = base64.b64encode(bytes_arr.getvalue())
-    return b64_bytes.decode('utf-8')
+    return to_base64(img, format=image_format)
 
 
 @pxt.udf(substitute_fn=PIL.Image.alpha_composite, is_method=True)
@@ -156,7 +149,7 @@ def get_metadata(self: PIL.Image.Image) -> dict:
 
 # Image.point()
 @pxt.udf(is_method=True)
-def point(self: PIL.Image.Image, lut: list[int], mode: Optional[str] = None) -> PIL.Image.Image:
+def point(self: PIL.Image.Image, lut: list[int], mode: str | None = None) -> PIL.Image.Image:
     """
     Map image pixels through a lookup table.
 
@@ -241,7 +234,7 @@ def _(self: Expr) -> ts.ColumnType:
 
 
 @pxt.udf(substitute_fn=PIL.Image.Image.entropy, is_method=True)
-def entropy(self: PIL.Image.Image, mask: Optional[PIL.Image.Image] = None, extrema: Optional[list] = None) -> float:
+def entropy(self: PIL.Image.Image, mask: PIL.Image.Image | None = None, extrema: list | None = None) -> float:
     """
     Returns the entropy of the image, optionally using a mask and extrema.
 
@@ -306,7 +299,7 @@ def getextrema(self: PIL.Image.Image) -> tuple[int, int]:
 
 
 @pxt.udf(substitute_fn=PIL.Image.Image.getpalette, is_method=True)
-def getpalette(self: PIL.Image.Image, mode: Optional[str] = None) -> tuple[int]:
+def getpalette(self: PIL.Image.Image, mode: str | None = None) -> tuple[int]:
     """
     Return the palette of the image, optionally converting it to a different mode.
 
@@ -346,9 +339,7 @@ def getprojection(self: PIL.Image.Image) -> tuple[int]:
 
 
 @pxt.udf(substitute_fn=PIL.Image.Image.histogram, is_method=True)
-def histogram(
-    self: PIL.Image.Image, mask: Optional[PIL.Image.Image] = None, extrema: Optional[list] = None
-) -> list[int]:
+def histogram(self: PIL.Image.Image, mask: PIL.Image.Image | None = None, extrema: list | None = None) -> list[int]:
     """
     Return a histogram for the image.
 
@@ -366,9 +357,9 @@ def histogram(
 def quantize(
     self: PIL.Image.Image,
     colors: int = 256,
-    method: Optional[int] = None,
+    method: int | None = None,
     kmeans: int = 0,
-    palette: Optional[int] = None,
+    palette: int | None = None,
     dither: int = PIL.Image.Dither.FLOYDSTEINBERG,
 ) -> PIL.Image.Image:
     """
@@ -392,7 +383,7 @@ def quantize(
 
 
 @pxt.udf(substitute_fn=PIL.Image.Image.reduce, is_method=True)
-def reduce(self: PIL.Image.Image, factor: int, box: Optional[tuple[int, int, int, int]] = None) -> PIL.Image.Image:
+def reduce(self: PIL.Image.Image, factor: int, box: tuple[int, int, int, int] | None = None) -> PIL.Image.Image:
     """
     Reduce the image by the given factor.
 

pixeltable/functions/llama_cpp.py

@@ -1,5 +1,12 @@
+"""
+Pixeltable UDFs for llama.cpp models.
+
+Provides integration with llama.cpp for running quantized language models locally,
+supporting chat completions and embeddings with GGUF format models.
+"""
+
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 
 import pixeltable as pxt
 import pixeltable.exceptions as excs
@@ -14,10 +21,10 @@ if TYPE_CHECKING:
 def create_chat_completion(
     messages: list[dict],
     *,
-    model_path: Optional[str] = None,
-    repo_id: Optional[str] = None,
-    repo_filename: Optional[str] = None,
-    model_kwargs: Optional[dict[str, Any]] = None,
+    model_path: str | None = None,
+    repo_id: str | None = None,
+    repo_filename: str | None = None,
+    model_kwargs: dict[str, Any] | None = None,
 ) -> dict:
     """
     Generate a chat completion from a list of messages.
@@ -81,7 +88,7 @@ def _lookup_local_model(model_path: str, n_gpu_layers: int) -> 'llama_cpp.Llama'
     return _model_cache[key]
 
 
-def _lookup_pretrained_model(repo_id: str, filename: Optional[str], n_gpu_layers: int) -> 'llama_cpp.Llama':
+def _lookup_pretrained_model(repo_id: str, filename: str | None, n_gpu_layers: int) -> 'llama_cpp.Llama':
     import llama_cpp
 
     key = (repo_id, filename, n_gpu_layers)
@@ -93,8 +100,16 @@ def _lookup_pretrained_model(repo_id: str, filename: Optional[str], n_gpu_layers
     return _model_cache[key]
 
 
-_model_cache: dict[tuple[str, str, int], Any] = {}
-_IS_GPU_AVAILABLE: Optional[bool] = None
+_model_cache: dict[tuple[str, str, int], 'llama_cpp.Llama'] = {}
+_IS_GPU_AVAILABLE: bool | None = None
+
+
+def cleanup() -> None:
+    for model in _model_cache.values():
+        if model._sampler is not None:
+            model._sampler.close()
+        model.close()
+    _model_cache.clear()
 
 
 __all__ = local_public_names(__name__)

pixeltable/functions/math.py

@@ -12,7 +12,6 @@ t.select(t.float_col.floor()).collect()
 
 import builtins
 import math
-from typing import Optional
 
 import sqlalchemy as sql
 
@@ -80,7 +79,7 @@ def _(self: sql.ColumnElement) -> sql.ColumnElement:
 
 
 @pxt.udf(is_method=True)
-def round(self: float, digits: Optional[int] = None) -> float:
+def round(self: float, digits: int | None = None) -> float:
     """
     Round a number to a given precision in decimal digits.
 
@@ -93,11 +92,11 @@ def round(self: float, digits: Optional[int] = None) -> float:
 
 
 @round.to_sql
-def _(self: sql.ColumnElement, digits: Optional[sql.ColumnElement] = None) -> sql.ColumnElement:
+def _(self: sql.ColumnElement, digits: sql.ColumnElement | None = None) -> sql.ColumnElement:
     if digits is None:
         return sql.func.round(self)
     else:
-        return sql.func.round(sql.cast(self, sql.Numeric), sql.cast(digits, sql.Integer))
+        return sql.cast(sql.func.round(sql.cast(self, sql.Numeric), sql.cast(digits, sql.Integer)), sql.Float)
 
 
 @pxt.udf(is_method=True)

pixeltable/functions/mistralai.py

@@ -5,7 +5,7 @@ first `pip install mistralai` and configure your Mistral AI credentials, as desc
 the [Working with Mistral AI](https://pixeltable.readme.io/docs/working-with-mistralai) tutorial.
 """
 
-from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
@@ -16,7 +16,7 @@ from pixeltable.func.signature import Batch
 from pixeltable.utils.code import local_public_names
 
 if TYPE_CHECKING:
-    import mistralai.types.basemodel
+    import mistralai
 
 
 @register_client('mistral')
@@ -32,7 +32,7 @@ def _mistralai_client() -> 'mistralai.Mistral':
 
 @pxt.udf(resource_pool='request-rate:mistral')
 async def chat_completions(
-    messages: list[dict[str, str]], *, model: str, model_kwargs: Optional[dict[str, Any]] = None
+    messages: list[dict[str, str]], *, model: str, model_kwargs: dict[str, Any] | None = None
 ) -> dict:
     """
     Chat Completion API.
@@ -54,8 +54,6 @@ async def chat_completions(
         model_kwargs: Additional keyword args for the Mistral `chat/completions` API.
             For details on the available parameters, see: <https://docs.mistral.ai/api/#tag/chat>
 
-    For details on the other parameters, see: <https://docs.mistral.ai/api/#tag/chat>
-
     Returns:
         A dictionary containing the response and other metadata.
 
@@ -79,7 +77,7 @@ async def chat_completions(
 
 
 @pxt.udf(resource_pool='request-rate:mistral')
-async def fim_completions(prompt: str, *, model: str, model_kwargs: Optional[dict[str, Any]] = None) -> dict:
+async def fim_completions(prompt: str, *, model: str, model_kwargs: dict[str, Any] | None = None) -> dict:
     """
     Fill-in-the-middle Completion API.
 
@@ -156,15 +154,6 @@ def _(model: str) -> ts.ArrayType:
     return ts.ArrayType((dimensions,), dtype=ts.FloatType())
 
 
-_T = TypeVar('_T')
-
-
-def _opt(arg: Optional[_T]) -> Union[_T, 'mistralai.types.basemodel.Unset']:
-    from mistralai.types import UNSET
-
-    return arg if arg is not None else UNSET
-
-
 __all__ = local_public_names(__name__)
 
 

pixeltable/functions/ollama.py

@@ -1,4 +1,11 @@
-from typing import TYPE_CHECKING, Optional
+"""
+Pixeltable UDFs for Ollama local models.
+
+Provides integration with Ollama for running large language models locally,
+including chat completions and embeddings.
+"""
+
+from typing import TYPE_CHECKING
 
 import numpy as np
 
@@ -18,7 +25,7 @@ def _(host: str) -> 'ollama.Client':
     return ollama.Client(host=host)
 
 
-def _ollama_client() -> Optional['ollama.Client']:
+def _ollama_client() -> 'ollama.Client | None':
     try:
         return env.Env.get().get_client('ollama')
     except Exception:
@@ -33,10 +40,10 @@ def generate(
     suffix: str = '',
     system: str = '',
     template: str = '',
-    context: Optional[list[int]] = None,
+    context: list[int] | None = None,
     raw: bool = False,
-    format: Optional[str] = None,
-    options: Optional[dict] = None,
+    format: str | None = None,
+    options: dict | None = None,
 ) -> dict:
     """
     Generate a response for a given prompt with a provided model.
@@ -77,9 +84,9 @@ def chat(
     messages: list[dict],
     *,
     model: str,
-    tools: Optional[list[dict]] = None,
-    format: Optional[str] = None,
-    options: Optional[dict] = None,
+    tools: list[dict] | None = None,
+    format: str | None = None,
+    options: dict | None = None,
 ) -> dict:
     """
     Generate the next message in a chat with a provided model.
@@ -103,7 +110,7 @@ def chat(
 
 @pxt.udf(batch_size=16)
 def embed(
-    input: Batch[str], *, model: str, truncate: bool = True, options: Optional[dict] = None
+    input: Batch[str], *, model: str, truncate: bool = True, options: dict | None = None
 ) -> Batch[pxt.Array[(None,), pxt.Float]]:
     """
     Generate embeddings from a model.

pixeltable/functions/openai.py

@@ -13,8 +13,7 @@ import logging
 import math
 import pathlib
 import re
-import uuid
-from typing import TYPE_CHECKING, Any, Callable, Optional, Type
+from typing import TYPE_CHECKING, Any, Callable, Type
 
 import httpx
 import numpy as np
@@ -24,6 +23,7 @@ import pixeltable as pxt
 from pixeltable import env, exprs, type_system as ts
 from pixeltable.func import Batch, Tools
 from pixeltable.utils.code import local_public_names
+from pixeltable.utils.local_store import TempStore
 
 if TYPE_CHECKING:
     import openai
@@ -32,11 +32,15 @@ _logger = logging.getLogger('pixeltable')
 
 
 @env.register_client('openai')
-def _(api_key: str) -> 'openai.AsyncOpenAI':
+def _(api_key: str, base_url: str | None = None, api_version: str | None = None) -> 'openai.AsyncOpenAI':
     import openai
 
+    default_query = None if api_version is None else {'api-version': api_version}
+
     return openai.AsyncOpenAI(
         api_key=api_key,
+        base_url=base_url,
+        default_query=default_query,
         # 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)),
     )
@@ -88,6 +92,52 @@ def _rate_limits_pool(model: str) -> str:
     return f'rate-limits:openai:{model}'
 
 
+# RE pattern for duration in '*-reset' headers;
+# examples: 1d2h3ms, 4m5.6s; # fractional seconds can be reported as 0.5s or 500ms
+_header_duration_pattern = re.compile(r'(?:(\d+)d)?(?:(\d+)h)?(?:(\d+)ms)|(?:(\d+)m)?(?:([\d.]+)s)?')
+
+
+def _parse_header_duration(duration_str: str) -> datetime.timedelta:
+    match = _header_duration_pattern.match(duration_str)
+    if not match:
+        raise ValueError(f'Invalid duration format: {duration_str}')
+
+    days = int(match.group(1) or 0)
+    hours = int(match.group(2) or 0)
+    milliseconds = int(match.group(3) or 0)
+    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)
+
+
+def _get_header_info(
+    headers: httpx.Headers,
+) -> tuple[tuple[int, int, datetime.datetime] | None, tuple[int, int, datetime.datetime] | None]:
+    now = datetime.datetime.now(tz=datetime.timezone.utc)
+
+    requests_limit_str = headers.get('x-ratelimit-limit-requests')
+    requests_limit = int(requests_limit_str) if requests_limit_str is not None else None
+    requests_remaining_str = headers.get('x-ratelimit-remaining-requests')
+    requests_remaining = int(requests_remaining_str) if requests_remaining_str is not None else None
+    requests_reset_str = headers.get('x-ratelimit-reset-requests', '5s')  # Default to 5 seconds
+    requests_reset_ts = now + _parse_header_duration(requests_reset_str)
+    requests_info = (requests_limit, requests_remaining, requests_reset_ts) if requests_remaining is not None else None
+
+    tokens_limit_str = headers.get('x-ratelimit-limit-tokens')
+    tokens_limit = int(tokens_limit_str) if tokens_limit_str is not None else None
+    tokens_remaining_str = headers.get('x-ratelimit-remaining-tokens')
+    tokens_remaining = int(tokens_remaining_str) if tokens_remaining_str is not None else None
+    tokens_reset_str = headers.get('x-ratelimit-reset-tokens', '5s')  # Default to 5 seconds
+    tokens_reset_ts = now + _parse_header_duration(tokens_reset_str)
+    tokens_info = (tokens_limit, tokens_remaining, tokens_reset_ts) if tokens_remaining is not None else None
+
+    if requests_info is None or tokens_info is None:
+        _logger.debug(f'get_header_info(): incomplete rate limit info: {headers}')
+
+    return requests_info, tokens_info
+
+
 class OpenAIRateLimitsInfo(env.RateLimitsInfo):
     retryable_errors: tuple[Type[Exception], ...]
 
@@ -108,61 +158,24 @@ class OpenAIRateLimitsInfo(env.RateLimitsInfo):
             openai.InternalServerError,
         )
 
-    def get_retry_delay(self, exc: Exception) -> Optional[float]:
+    def record_exc(self, exc: Exception) -> None:
+        import openai
+
+        _ = isinstance(exc, openai.APIError)
+        if not isinstance(exc, openai.APIError) or not hasattr(exc, 'response') or not hasattr(exc.response, 'headers'):
+            return
+        requests_info, tokens_info = _get_header_info(exc.response.headers)
+        _logger.debug(f'record_exc(): requests_info={requests_info} tokens_info={tokens_info}')
+        self.record(requests=requests_info, tokens=tokens_info)
+        self.has_exc = True
+
+    def get_retry_delay(self, exc: Exception) -> float | None:
         import openai
 
         if not isinstance(exc, self.retryable_errors):
             return None
         assert isinstance(exc, openai.APIError)
-        return 1.0
-
-
-# RE pattern for duration in '*-reset' headers;
-# examples: 1d2h3ms, 4m5.6s; # fractional seconds can be reported as 0.5s or 500ms
-_header_duration_pattern = re.compile(r'(?:(\d+)d)?(?:(\d+)h)?(?:(\d+)ms)|(?:(\d+)m)?(?:([\d.]+)s)?')
-
-
-def _parse_header_duration(duration_str: str) -> datetime.timedelta:
-    match = _header_duration_pattern.match(duration_str)
-    if not match:
-        raise ValueError('Invalid duration format')
-
-    days = int(match.group(1) or 0)
-    hours = int(match.group(2) or 0)
-    milliseconds = int(match.group(3) or 0)
-    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)
-
-
-def _get_header_info(
-    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)
-
-    requests_info: Optional[tuple[int, int, datetime.datetime]] = None
-    if requests:
-        requests_limit_str = headers.get('x-ratelimit-limit-requests')
-        requests_limit = int(requests_limit_str) if requests_limit_str is not None else None
-        requests_remaining_str = headers.get('x-ratelimit-remaining-requests')
-        requests_remaining = int(requests_remaining_str) if requests_remaining_str is not None else None
-        requests_reset_str = headers.get('x-ratelimit-reset-requests')
-        requests_reset_ts = now + _parse_header_duration(requests_reset_str)
-        requests_info = (requests_limit, requests_remaining, requests_reset_ts)
-
-    tokens_info: Optional[tuple[int, int, datetime.datetime]] = None
-    if tokens:
-        tokens_limit_str = headers.get('x-ratelimit-limit-tokens')
-        tokens_limit = int(tokens_limit_str) if tokens_limit_str is not None else None
-        tokens_remaining_str = headers.get('x-ratelimit-remaining-tokens')
-        tokens_remaining = int(tokens_remaining_str) if tokens_remaining_str is not None else None
-        tokens_reset_str = headers.get('x-ratelimit-reset-tokens')
-        tokens_reset_ts = now + _parse_header_duration(tokens_reset_str)
-        tokens_info = (tokens_limit, tokens_remaining, tokens_reset_ts)
-
-    return requests_info, tokens_info
+        return super().get_retry_delay(exc)
 
 
 #####################################
@@ -170,7 +183,7 @@ def _get_header_info(
 
 
 @pxt.udf
-async def speech(input: str, *, model: str, voice: str, model_kwargs: Optional[dict[str, Any]] = None) -> pxt.Audio:
+async def speech(input: str, *, model: str, voice: str, model_kwargs: dict[str, Any] | None = None) -> pxt.Audio:
     """
     Generates audio from the input text.
 
@@ -205,20 +218,15 @@ async def speech(input: str, *, model: str, voice: str, model_kwargs: Optional[d
     if model_kwargs is None:
         model_kwargs = {}
 
-    content = await _openai_client().audio.speech.create(
-        input=input,
-        model=model,
-        voice=voice,  # type: ignore
-        **model_kwargs,
-    )
+    content = await _openai_client().audio.speech.create(input=input, model=model, voice=voice, **model_kwargs)
     ext = model_kwargs.get('response_format', 'mp3')
-    output_filename = str(env.Env.get().tmp_dir / f'{uuid.uuid4()}.{ext}')
+    output_filename = str(TempStore.create_path(extension=f'.{ext}'))
     content.write_to_file(output_filename)
     return output_filename
 
 
 @pxt.udf
-async def transcriptions(audio: pxt.Audio, *, model: str, model_kwargs: Optional[dict[str, Any]] = None) -> dict:
+async def transcriptions(audio: pxt.Audio, *, model: str, model_kwargs: dict[str, Any] | None = None) -> dict:
     """
     Transcribes audio into the input language.
 
@@ -257,7 +265,7 @@ async def transcriptions(audio: pxt.Audio, *, model: str, model_kwargs: Optional
 
 
 @pxt.udf
-async def translations(audio: pxt.Audio, *, model: str, model_kwargs: Optional[dict[str, Any]] = None) -> dict:
+async def translations(audio: pxt.Audio, *, model: str, model_kwargs: dict[str, Any] | None = None) -> dict:
     """
     Translates audio into English.
 
@@ -327,7 +335,7 @@ def _is_model_family(model: str, family: str) -> bool:
 
 
 def _chat_completions_get_request_resources(
-    messages: list, model: str, model_kwargs: Optional[dict[str, Any]]
+    messages: list, model: str, model_kwargs: dict[str, Any] | None
 ) -> dict[str, int]:
     if model_kwargs is None:
         model_kwargs = {}
@@ -354,9 +362,10 @@ async def chat_completions(
     messages: list,
     *,
     model: str,
-    model_kwargs: Optional[dict[str, Any]] = None,
-    tools: Optional[list[dict[str, Any]]] = None,
-    tool_choice: Optional[dict[str, Any]] = None,
+    model_kwargs: dict[str, Any] | None = None,
+    tools: list[dict[str, Any]] | None = None,
+    tool_choice: dict[str, Any] | None = None,
+    _runtime_ctx: env.RuntimeCtx | None = None,
 ) -> dict:
     """
     Creates a model response for the given chat conversation.
@@ -386,10 +395,10 @@ async def chat_completions(
         of the table `tbl`:
 
         >>> messages = [
-                {'role': 'system', 'content': 'You are a helpful assistant.'},
-                {'role': 'user', 'content': tbl.prompt}
-            ]
-            tbl.add_computed_column(response=chat_completions(messages, model='gpt-4o-mini'))
+        ...     {'role': 'system', 'content': 'You are a helpful assistant.'},
+        ...     {'role': 'user', 'content': tbl.prompt}
+        ... ]
+        >>> tbl.add_computed_column(response=chat_completions(messages, model='gpt-4o-mini'))
     """
     if model_kwargs is None:
         model_kwargs = {}
@@ -420,13 +429,14 @@ async def chat_completions(
     )
 
     requests_info, tokens_info = _get_header_info(result.headers)
-    rate_limits_info.record(requests=requests_info, tokens=tokens_info)
+    is_retry = _runtime_ctx is not None and _runtime_ctx.is_retry
+    rate_limits_info.record(requests=requests_info, tokens=tokens_info, reset_exc=is_retry)
 
     return json.loads(result.text)
 
 
 def _vision_get_request_resources(
-    prompt: str, image: PIL.Image.Image, model: str, model_kwargs: Optional[dict[str, Any]] = None
+    prompt: str, image: PIL.Image.Image, model: str, model_kwargs: dict[str, Any] | None = None
 ) -> dict[str, int]:
     if model_kwargs is None:
         model_kwargs = {}
@@ -463,7 +473,12 @@ def _vision_get_request_resources(
 
 @pxt.udf
 async def vision(
-    prompt: str, image: PIL.Image.Image, *, model: str, model_kwargs: Optional[dict[str, Any]] = None
+    prompt: str,
+    image: PIL.Image.Image,
+    *,
+    model: str,
+    model_kwargs: dict[str, Any] | None = None,
+    _runtime_ctx: env.RuntimeCtx | None = None,
 ) -> str:
     """
     Analyzes an image with the OpenAI vision capability. This is a convenience function that takes an image and
@@ -523,8 +538,10 @@ async def vision(
         **model_kwargs,
     )
 
+    # _logger.debug(f'vision(): headers={result.headers}')
     requests_info, tokens_info = _get_header_info(result.headers)
-    rate_limits_info.record(requests=requests_info, tokens=tokens_info)
+    is_retry = _runtime_ctx is not None and _runtime_ctx.is_retry
+    rate_limits_info.record(requests=requests_info, tokens=tokens_info, reset_exc=is_retry)
 
     result = json.loads(result.text)
     return result['choices'][0]['message']['content']
@@ -547,7 +564,11 @@ def _embeddings_get_request_resources(input: list[str]) -> dict[str, int]:
 
 @pxt.udf(batch_size=32)
 async def embeddings(
-    input: Batch[str], *, model: str, model_kwargs: Optional[dict[str, Any]] = None
+    input: Batch[str],
+    *,
+    model: str,
+    model_kwargs: dict[str, Any] | None = None,
+    _runtime_ctx: env.RuntimeCtx | None = None,
 ) -> Batch[pxt.Array[(None,), pxt.Float]]:
     """
     Creates an embedding vector representing the input text.
@@ -594,13 +615,14 @@ async def embeddings(
         input=input, model=model, encoding_format='float', **model_kwargs
     )
     requests_info, tokens_info = _get_header_info(result.headers)
-    rate_limits_info.record(requests=requests_info, tokens=tokens_info)
+    is_retry = _runtime_ctx is not None and _runtime_ctx.is_retry
+    rate_limits_info.record(requests=requests_info, tokens=tokens_info, reset_exc=is_retry)
     return [np.array(data['embedding'], dtype=np.float64) for data in json.loads(result.content)['data']]
 
 
 @embeddings.conditional_return_type
-def _(model: str, model_kwargs: Optional[dict[str, Any]] = None) -> ts.ArrayType:
-    dimensions: Optional[int] = None
+def _(model: str, model_kwargs: dict[str, Any] | None = None) -> ts.ArrayType:
+    dimensions: int | None = None
     if model_kwargs is not None:
         dimensions = model_kwargs.get('dimensions')
     if dimensions is None:
@@ -617,7 +639,7 @@ def _(model: str, model_kwargs: Optional[dict[str, Any]] = None) -> ts.ArrayType
 
 @pxt.udf
 async def image_generations(
-    prompt: str, *, model: str = 'dall-e-2', model_kwargs: Optional[dict[str, Any]] = None
+    prompt: str, *, model: str = 'dall-e-2', model_kwargs: dict[str, Any] | None = None
 ) -> PIL.Image.Image:
     """
     Creates an image given a prompt.
@@ -663,7 +685,7 @@ async def image_generations(
 
 
 @image_generations.conditional_return_type
-def _(model_kwargs: Optional[dict[str, Any]] = None) -> ts.ImageType:
+def _(model_kwargs: dict[str, Any] | None = None) -> ts.ImageType:
     if model_kwargs is None or 'size' not in model_kwargs:
         # default size is 1024x1024
         return ts.ImageType(size=(1024, 1024))
@@ -739,7 +761,7 @@ def invoke_tools(tools: Tools, response: exprs.Expr) -> exprs.InlineDict:
 
 
 @pxt.udf
-def _openai_response_to_pxt_tool_calls(response: dict) -> Optional[dict]:
+def _openai_response_to_pxt_tool_calls(response: dict) -> dict | None:
     if 'tool_calls' not in response['choices'][0]['message'] or response['choices'][0]['message']['tool_calls'] is None:
         return None
     openai_tool_calls = response['choices'][0]['message']['tool_calls']