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

pixeltable/functions/openrouter.py

@@ -0,0 +1,143 @@
+"""
+Pixeltable UDFs that wrap the OpenRouter API.
+
+OpenRouter provides a unified interface to multiple LLM providers. In order to use it,
+you must first sign up at https://openrouter.ai, create an API key, and configure it
+as described in the Working with OpenRouter tutorial.
+"""
+
+from typing import TYPE_CHECKING, Any
+
+import pixeltable as pxt
+from pixeltable.env import Env, register_client
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    import openai
+
+
+@register_client('openrouter')
+def _(api_key: str, site_url: str | None = None, app_name: str | None = None) -> 'openai.AsyncOpenAI':
+    import openai
+
+    # Create default headers for OpenRouter
+    default_headers: dict[str, Any] = {}
+    if site_url:
+        default_headers['HTTP-Referer'] = site_url
+    if app_name:
+        default_headers['X-Title'] = app_name
+
+    return openai.AsyncOpenAI(base_url='https://openrouter.ai/api/v1', api_key=api_key, default_headers=default_headers)
+
+
+def _openrouter_client() -> 'openai.AsyncOpenAI':
+    return Env.get().get_client('openrouter')
+
+
+@pxt.udf(resource_pool='request-rate:openrouter')
+async def chat_completions(
+    messages: list,
+    *,
+    model: str,
+    model_kwargs: dict[str, Any] | None = None,
+    tools: list[dict[str, Any]] | None = None,
+    tool_choice: dict[str, Any] | None = None,
+    provider: dict[str, Any] | None = None,
+    transforms: list[str] | None = None,
+) -> dict:
+    """
+    Chat Completion API via OpenRouter.
+
+    OpenRouter provides access to multiple LLM providers through a unified API.
+    For additional details, see: <https://openrouter.ai/docs>
+
+    Supported models can be found at: <https://openrouter.ai/models>
+
+    Request throttling:
+    Applies the rate limit set in the config (section `openrouter`, key `rate_limit`). If no rate
+    limit is configured, uses a default of 600 RPM.
+
+    __Requirements:__
+
+    - `pip install openai`
+
+    Args:
+        messages: A list of messages comprising the conversation so far.
+        model: ID of the model to use (e.g., 'anthropic/claude-3.5-sonnet', 'openai/gpt-4').
+        model_kwargs: Additional OpenAI-compatible parameters.
+        tools: List of tools available to the model.
+        tool_choice: Controls which (if any) tool is called by the model.
+        provider: OpenRouter-specific provider preferences (e.g., {'order': ['Anthropic', 'OpenAI']}).
+        transforms: List of message transforms to apply (e.g., ['middle-out']).
+
+    Returns:
+        A dictionary containing the response in OpenAI format.
+
+    Examples:
+        Basic chat completion:
+
+        >>> messages = [{'role': 'user', 'content': tbl.prompt}]
+        ... tbl.add_computed_column(
+        ...     response=chat_completions(
+        ...         messages,
+        ...         model='anthropic/claude-3.5-sonnet'
+        ...     )
+        ... )
+
+        With provider routing:
+
+        >>> tbl.add_computed_column(
+        ...     response=chat_completions(
+        ...         messages,
+        ...         model='anthropic/claude-3.5-sonnet',
+        ...         provider={'require_parameters': True, 'order': ['Anthropic']}
+        ...     )
+        ... )
+
+        With transforms:
+
+        >>> tbl.add_computed_column(
+        ...     response=chat_completions(
+        ...         messages,
+        ...         model='openai/gpt-4',
+        ...         transforms=['middle-out']  # Optimize for long contexts
+        ...     )
+        ... )
+    """
+    if model_kwargs is None:
+        model_kwargs = {}
+
+    Env.get().require_package('openai')
+
+    # Handle tools if provided
+    if tools is not None:
+        model_kwargs['tools'] = [{'type': 'function', 'function': tool} for tool in tools]
+
+    if tool_choice is not None:
+        if tool_choice['auto']:
+            model_kwargs['tool_choice'] = 'auto'
+        elif tool_choice['required']:
+            model_kwargs['tool_choice'] = 'required'
+        else:
+            assert tool_choice['tool'] is not None
+            model_kwargs['tool_choice'] = {'type': 'function', 'function': {'name': tool_choice['tool']}}
+
+    # Prepare OpenRouter-specific parameters for extra_body
+    extra_body: dict[str, Any] = {}
+    if provider is not None:
+        extra_body['provider'] = provider
+    if transforms is not None:
+        extra_body['transforms'] = transforms
+
+    # Make the API call
+    result = await _openrouter_client().chat.completions.create(
+        messages=messages, model=model, extra_body=extra_body if extra_body else None, **model_kwargs
+    )
+    return result.model_dump()
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__

pixeltable/functions/replicate.py

@@ -12,7 +12,7 @@ from pixeltable.env import Env, register_client
 from pixeltable.utils.code import local_public_names
 
 if TYPE_CHECKING:
-    import replicate  # type: ignore[import-untyped]
+    import replicate
 
 
 @register_client('replicate')
@@ -27,7 +27,7 @@ def _replicate_client() -> 'replicate.Client':
 
 
 @pxt.udf(resource_pool='request-rate:replicate')
-async def run(input: dict[str, Any], *, ref: str) -> dict[str, Any]:
+async def run(input: dict[str, Any], *, ref: str) -> pxt.Json:
     """
     Run a model on Replicate.
 

pixeltable/functions/reve.py

@@ -0,0 +1,250 @@
+"""
+Pixeltable [UDFs](https://docs.pixeltable.com/datastore/custom-functions) that wrap [Reve](https://app.reve.com/) image
+generation API. In order to use them, the API key must be specified either with `REVE_API_KEY` environment variable,
+or as `api_key` in the `reve` section of the Pixeltable config file.
+"""
+
+import asyncio
+import atexit
+import logging
+import re
+from io import BytesIO
+
+import aiohttp
+import PIL.Image
+
+import pixeltable as pxt
+from pixeltable.env import Env, register_client
+from pixeltable.utils.code import local_public_names
+from pixeltable.utils.image import to_base64
+
+_logger = logging.getLogger('pixeltable')
+
+
+class ReveRateLimitedError(Exception):
+    pass
+
+
+class ReveContentViolationError(Exception):
+    pass
+
+
+class ReveUnexpectedError(Exception):
+    pass
+
+
+class _ReveClient:
+    """
+    Client for interacting with the Reve API. Maintains a long-lived HTTP session to the service.
+    """
+
+    api_key: str
+    session: aiohttp.ClientSession
+
+    def __init__(self, api_key: str):
+        self.api_key = api_key
+        self.session = Env.get().event_loop.run_until_complete(self._start_session())
+        atexit.register(lambda: asyncio.run(self.session.close()))
+
+    async def _start_session(self) -> aiohttp.ClientSession:
+        # Maintains a long-lived TPC connection. The default keepalive timeout is 15 seconds.
+        return aiohttp.ClientSession(base_url='https://api.reve.com')
+
+    async def _post(self, endpoint: str, *, payload: dict) -> PIL.Image.Image:
+        # Reve supports other formats as well, but we only use PNG for now
+        requested_content_type = 'image/png'
+        request_headers = {
+            'Authorization': f'Bearer {self.api_key}',
+            'Content-Type': 'application/json',
+            'Accept': requested_content_type,
+        }
+
+        async with self.session.post(endpoint, json=payload, headers=request_headers) as resp:
+            request_id = resp.headers.get('X-Reve-Request-Id')
+            error_code = resp.headers.get('X-Reve-Error-Code')
+            match resp.status:
+                case 200:
+                    if error_code is not None:
+                        raise ReveUnexpectedError(
+                            f'Reve request {request_id} returned an unexpected error {error_code}'
+                        )
+                    content_violation = resp.headers.get('X-Reve-Content-Violation', 'false')
+                    if content_violation.lower() != 'false':
+                        raise ReveContentViolationError(
+                            f'Reve request {request_id} resulted in a content violation error'
+                        )
+                    if resp.content_type != requested_content_type:
+                        raise ReveUnexpectedError(
+                            f'Reve request {request_id} expected content type {requested_content_type}, '
+                            f'got {resp.content_type}'
+                        )
+
+                    img_data = await resp.read()
+                    if len(img_data) == 0:
+                        raise ReveUnexpectedError(f'Reve request {request_id} resulted in an empty image')
+                    img = PIL.Image.open(BytesIO(img_data))
+                    img.load()
+                    _logger.debug(
+                        f'Reve request {request_id} successful. Image bytes: {len(img_data)}, size: {img.size}'
+                        f', format: {img.format}, mode: {img.mode}'
+                    )
+                    return img
+                case 429:
+                    # Try to honor the server-provided Retry-After value if present
+                    # Note: Retry-After value can also be given in the form of HTTP Date, which we don't currently
+                    # handle
+                    retry_after_seconds = None
+                    retry_after_header = resp.headers.get('Retry-After')
+                    if retry_after_header is not None and re.fullmatch(r'\d{1,2}', retry_after_header):
+                        retry_after_seconds = int(retry_after_header)
+                    _logger.info(
+                        f'Reve request {request_id} failed due to rate limiting, retry after header value: '
+                        f'{retry_after_header}'
+                    )
+                    # This error message is formatted specifically so that RequestRateScheduler can extract the retry
+                    # delay from it
+                    raise ReveRateLimitedError(
+                        f'Reve request {request_id} failed due to rate limiting (429). retry-after:'
+                        f'{retry_after_seconds}'
+                    )
+                case _:
+                    _logger.info(
+                        f'Reve request {request_id} failed with status code {resp.status} and error code {error_code}'
+                    )
+                    raise ReveUnexpectedError(
+                        f'Reve request failed with status code {resp.status} and error code {error_code}'
+                    )
+
+
+@register_client('reve')
+def _(api_key: str) -> _ReveClient:
+    return _ReveClient(api_key=api_key)
+
+
+def _client() -> _ReveClient:
+    return Env.get().get_client('reve')
+
+
+# TODO Regarding rate limiting: Reve appears to be going for a credits per minute rate limiting model, but does not
+# currently communicate rate limit information in responses. Therefore neither of the currently implemented limiting
+# strategies is a perfect match, but "request-rate" is the closest. Reve does not currently enforce the rate limits,
+# but when it does, we can revisit this choice.
+@pxt.udf(resource_pool='request-rate:reve')
+async def create(prompt: str, *, aspect_ratio: str | None = None, version: str | None = None) -> PIL.Image.Image:
+    """
+    Creates an image from a text prompt.
+
+    This UDF wraps the `https://api.reve.com/v1/image/create` endpoint. For more information, refer to the official
+    [API documentation](https://api.reve.com/console/docs/create).
+
+    Args:
+        prompt: prompt describing the desired image
+        aspect_ratio: desired image aspect ratio, e.g. '3:2', '16:9', '1:1', etc.
+        version: specific model version to use. Latest if not specified.
+
+    Returns:
+        A generated image
+
+    Examples:
+        Add a computed column with generated square images to a table with text prompts:
+
+        >>> t.add_computed_column(
+        ...     img=reve.create(t.prompt, aspect_ratio='1:1')
+        ... )
+    """
+    payload = {'prompt': prompt}
+    if aspect_ratio is not None:
+        payload['aspect_ratio'] = aspect_ratio
+    if version is not None:
+        payload['version'] = version
+
+    result = await _client()._post('/v1/image/create', payload=payload)
+    return result
+
+
+@pxt.udf(resource_pool='request-rate:reve')
+async def edit(image: PIL.Image.Image, edit_instruction: str, *, version: str | None = None) -> PIL.Image.Image:
+    """
+    Edits images based on a text prompt.
+
+    This UDF wraps the `https://api.reve.com/v1/image/edit` endpoint. For more information, refer to the official
+    [API documentation](https://api.reve.com/console/docs/edit)
+
+    Args:
+        image: image to edit
+        edit_instruction: text prompt describing the desired edit
+        version: specific model version to use. Latest if not specified.
+
+    Returns:
+        A generated image
+
+    Examples:
+        Add a computed column with catalog-ready images to the table with product pictures:
+
+        >>> t.add_computed_column(
+        ...     catalog_img=reve.edit(
+        ...         t.product_img,
+        ...         'Remove background and distractions from the product picture, improve lighting.'
+        ...     )
+        ... )
+    """
+    payload = {'edit_instruction': edit_instruction, 'reference_image': to_base64(image)}
+    if version is not None:
+        payload['version'] = version
+
+    result = await _client()._post('/v1/image/edit', payload=payload)
+    return result
+
+
+@pxt.udf(resource_pool='request-rate:reve')
+async def remix(
+    prompt: str, images: list[PIL.Image.Image], *, aspect_ratio: str | None = None, version: str | None = None
+) -> PIL.Image.Image:
+    """
+    Creates images based on a text prompt and reference images.
+
+    The prompt may include `<img>0</img>`, `<img>1</img>`, etc. tags to refer to the images in the `images` argument.
+
+    This UDF wraps the `https://api.reve.com/v1/image/remix` endpoint. For more information, refer to the official
+    [API documentation](https://api.reve.com/console/docs/remix)
+
+    Args:
+        prompt: prompt describing the desired image
+        images: list of reference images
+        aspect_ratio: desired image aspect ratio, e.g. '3:2', '16:9', '1:1', etc.
+        version: specific model version to use. Latest by default.
+
+    Returns:
+        A generated image
+
+    Examples:
+        Add a computed column with promotional collages to a table with original images:
+
+        >>> t.add_computed_column(
+        ...     promo_img=(
+        ...         reve.remix(
+        ...             'Generate a product promotional image by combining the image of the product'
+        ...             ' from <img>0</img> with the landmark scene from <img>1</img>',
+        ...             images=[t.product_img, t.local_landmark_img],
+        ...             aspect_ratio='16:9',
+        ...         )
+        ...     )
+        ... )
+    """
+    if len(images) == 0:
+        raise pxt.Error('Must include at least 1 reference image')
+
+    payload = {'prompt': prompt, 'reference_images': [to_base64(img) for img in images]}
+    if version is not None:
+        payload['version'] = version
+    if aspect_ratio is not None:
+        payload['aspect_ratio'] = aspect_ratio
+    result = await _client()._post('/v1/image/remix', payload=payload)
+    return result
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__

pixeltable/functions/string.py

@@ -15,7 +15,7 @@ import builtins
 import re
 import textwrap
 from string import whitespace
-from typing import Any, Optional
+from typing import Any
 
 import sqlalchemy as sql
 
@@ -78,9 +78,7 @@ def contains(self: str, substr: str, case: bool = True) -> bool:
 
 
 @contains.to_sql
-def _(
-    self: sql.ColumnElement, substr: sql.ColumnElement, case: Optional[sql.ColumnElement] = None
-) -> sql.ColumnElement:
+def _(self: sql.ColumnElement, substr: sql.ColumnElement, case: sql.ColumnElement | None = None) -> sql.ColumnElement:
     # Replace all occurrences of `%`, `_`, and `\` with escaped versions
     escaped_substr = sql.func.regexp_replace(substr, r'(%|_|\\)', r'\\\1', 'g')
     if case is None:
@@ -153,7 +151,7 @@ def fill(self: str, width: int, **kwargs: Any) -> str:
 
 
 @pxt.udf(is_method=True)
-def find(self: str, substr: str, start: int = 0, end: Optional[int] = None) -> int:
+def find(self: str, substr: str, start: int = 0, end: int | None = None) -> int:
     """
     Return the lowest index in string where `substr` is found within the slice `s[start:end]`.
 
@@ -169,10 +167,7 @@ def find(self: str, substr: str, start: int = 0, end: Optional[int] = None) -> i
 
 @find.to_sql
 def _(
-    self: sql.ColumnElement,
-    substr: sql.ColumnElement,
-    start: sql.ColumnElement,
-    end: Optional[sql.ColumnElement] = None,
+    self: sql.ColumnElement, substr: sql.ColumnElement, start: sql.ColumnElement, end: sql.ColumnElement | None = None
 ) -> sql.ColumnElement:
     sl = pxt.functions.string.slice._to_sql(self, start, end)
     if sl is None:
@@ -227,7 +222,7 @@ def fullmatch(self: str, pattern: str, case: bool = True, flags: int = 0) -> boo
 
 
 @pxt.udf(is_method=True)
-def index(self: str, substr: str, start: int = 0, end: Optional[int] = None) -> int:
+def index(self: str, substr: str, start: int = 0, end: int | None = None) -> int:
     """
     Return the lowest index in string where `substr` is found within the slice `[start:end]`.
     Raises ValueError if `substr` is not found.
@@ -413,7 +408,7 @@ def _(self: sql.ColumnElement) -> sql.ColumnElement:
 
 
 @pxt.udf(is_method=True)
-def lstrip(self: str, chars: Optional[str] = None) -> str:
+def lstrip(self: str, chars: str | None = None) -> str:
     """
     Return a copy of the string with leading characters removed. The `chars` argument is a string specifying the set of
     characters to be removed. If omitted or `None`, whitespace characters are removed.
@@ -427,7 +422,7 @@ def lstrip(self: str, chars: Optional[str] = None) -> str:
 
 
 @lstrip.to_sql
-def _(self: sql.ColumnElement, chars: Optional[sql.ColumnElement] = None) -> sql.ColumnElement:
+def _(self: sql.ColumnElement, chars: sql.ColumnElement | None = None) -> sql.ColumnElement:
     return sql.func.ltrim(self, chars if chars is not None else whitespace)
 
 
@@ -535,7 +530,7 @@ def _(self: sql.ColumnElement, n: sql.ColumnElement) -> sql.ColumnElement:
 
 
 @pxt.udf(is_method=True)
-def replace(self: str, substr: str, repl: str, n: Optional[int] = None) -> str:
+def replace(self: str, substr: str, repl: str, n: int | None = None) -> str:
     """
     Replace occurrences of `substr` with `repl`.
 
@@ -551,7 +546,7 @@ def replace(self: str, substr: str, repl: str, n: Optional[int] = None) -> str:
 
 @replace.to_sql
 def _(
-    self: sql.ColumnElement, substr: sql.ColumnElement, repl: sql.ColumnElement, n: Optional[sql.ColumnElement] = None
+    self: sql.ColumnElement, substr: sql.ColumnElement, repl: sql.ColumnElement, n: sql.ColumnElement | None = None
 ) -> sql.ColumnElement:
     if n is not None:
         return None  # SQL does not support bounding the number of replacements
@@ -560,7 +555,7 @@ def _(
 
 
 @pxt.udf(is_method=True)
-def replace_re(self: str, pattern: str, repl: str, n: Optional[int] = None, flags: int = 0) -> str:
+def replace_re(self: str, pattern: str, repl: str, n: int | None = None, flags: int = 0) -> str:
     """
     Replace occurrences of a regular expression pattern with `repl`.
 
@@ -591,7 +586,7 @@ def _(self: sql.ColumnElement) -> sql.ColumnElement:
 
 
 @pxt.udf(is_method=True)
-def rfind(self: str, substr: str, start: Optional[int] = 0, end: Optional[int] = None) -> int:
+def rfind(self: str, substr: str, start: int | None = 0, end: int | None = None) -> int:
     """
     Return the highest index where `substr` is found, such that `substr` is contained within `[start:end]`.
 
@@ -606,7 +601,7 @@ def rfind(self: str, substr: str, start: Optional[int] = 0, end: Optional[int] =
 
 
 @pxt.udf(is_method=True)
-def rindex(self: str, substr: str, start: Optional[int] = 0, end: Optional[int] = None) -> int:
+def rindex(self: str, substr: str, start: int | None = 0, end: int | None = None) -> int:
     """
     Return the highest index where `substr` is found, such that `substr` is contained within `[start:end]`.
     Raises ValueError if `substr` is not found.
@@ -643,7 +638,7 @@ def rpartition(self: str, sep: str = ' ') -> list:
 
 
 @pxt.udf(is_method=True)
-def rstrip(self: str, chars: Optional[str] = None) -> str:
+def rstrip(self: str, chars: str | None = None) -> str:
     """
     Return a copy of string with trailing characters removed.
 
@@ -656,12 +651,12 @@ def rstrip(self: str, chars: Optional[str] = None) -> str:
 
 
 @rstrip.to_sql
-def _(self: sql.ColumnElement, chars: Optional[sql.ColumnElement] = None) -> sql.ColumnElement:
+def _(self: sql.ColumnElement, chars: sql.ColumnElement | None = None) -> sql.ColumnElement:
     return sql.func.rtrim(self, chars if chars is not None else whitespace)
 
 
 @pxt.udf(is_method=True)
-def slice(self: str, start: Optional[int] = None, stop: Optional[int] = None, step: Optional[int] = None) -> str:
+def slice(self: str, start: int | None = None, stop: int | None = None, step: int | None = None) -> str:
     """
     Return a slice.
 
@@ -676,9 +671,9 @@ def slice(self: str, start: Optional[int] = None, stop: Optional[int] = None, st
 @slice.to_sql
 def _(
     self: sql.ColumnElement,
-    start: Optional[sql.ColumnElement] = None,
-    stop: Optional[sql.ColumnElement] = None,
-    step: Optional[sql.ColumnElement] = None,
+    start: sql.ColumnElement | None = None,
+    stop: sql.ColumnElement | None = None,
+    step: sql.ColumnElement | None = None,
 ) -> sql.ColumnElement:
     if step is not None:
         return None
@@ -709,9 +704,7 @@ def _(
 
 
 @pxt.udf(is_method=True)
-def slice_replace(
-    self: str, start: Optional[int] = None, stop: Optional[int] = None, repl: Optional[str] = None
-) -> str:
+def slice_replace(self: str, start: int | None = None, stop: int | None = None, repl: str | None = None) -> str:
     """
     Replace a positional slice with another value.
 
@@ -744,7 +737,7 @@ def _(self: sql.ColumnElement, substr: sql.ColumnElement) -> sql.ColumnElement:
 
 
 @pxt.udf(is_method=True)
-def strip(self: str, chars: Optional[str] = None) -> str:
+def strip(self: str, chars: str | None = None) -> str:
     """
     Return a copy of string with leading and trailing characters removed.
 
@@ -757,7 +750,7 @@ def strip(self: str, chars: Optional[str] = None) -> str:
 
 
 @strip.to_sql
-def _(self: sql.ColumnElement, chars: Optional[sql.ColumnElement] = None) -> sql.ColumnElement:
+def _(self: sql.ColumnElement, chars: sql.ColumnElement | None = None) -> sql.ColumnElement:
     return sql.func.trim(self, chars if chars is not None else whitespace)
 
 

pixeltable/functions/timestamp.py

@@ -11,7 +11,6 @@ t.select(t.timestamp_col.year, t.timestamp_col.weekday()).collect()
 """
 
 from datetime import datetime
-from typing import Optional
 
 import sqlalchemy as sql
 
@@ -237,12 +236,12 @@ def _(
     microsecond: sql.ColumnElement = _SQL_ZERO,
 ) -> sql.ColumnElement:
     return sql.func.make_timestamptz(
-        sql.cast(year, sql.Integer),
-        sql.cast(month, sql.Integer),
-        sql.cast(day, sql.Integer),
-        sql.cast(hour, sql.Integer),
-        sql.cast(minute, sql.Integer),
-        sql.cast(second + microsecond / 1000000.0, sql.Float),
+        year.cast(sql.Integer),
+        month.cast(sql.Integer),
+        day.cast(sql.Integer),
+        hour.cast(sql.Integer),
+        minute.cast(sql.Integer),
+        (second + microsecond / 1000000.0).cast(sql.Float),
     )
 
 
@@ -271,13 +270,13 @@ def _(
 @pxt.udf(is_method=True)
 def replace(
     self: datetime,
-    year: Optional[int] = None,
-    month: Optional[int] = None,
-    day: Optional[int] = None,
-    hour: Optional[int] = None,
-    minute: Optional[int] = None,
-    second: Optional[int] = None,
-    microsecond: Optional[int] = None,
+    year: int | None = None,
+    month: int | None = None,
+    day: int | None = None,
+    hour: int | None = None,
+    minute: int | None = None,
+    second: int | None = None,
+    microsecond: int | None = None,
 ) -> datetime:
     """
     Return a datetime with the same attributes, except for those attributes given new values by whichever keyword

pixeltable/functions/together.py

@@ -7,7 +7,7 @@ the [Working with Together AI](https://pixeltable.readme.io/docs/together-ai) tu
 
 import base64
 import io
-from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar
+from typing import TYPE_CHECKING, Any, Callable, TypeVar
 
 import numpy as np
 import PIL.Image
@@ -50,7 +50,7 @@ def _retry(fn: Callable[..., T]) -> Callable[..., T]:
 
 
 @pxt.udf(resource_pool='request-rate:together:chat')
-async def completions(prompt: str, *, model: str, model_kwargs: Optional[dict[str, Any]] = None) -> dict:
+async def completions(prompt: str, *, model: str, model_kwargs: dict[str, Any] | None = None) -> dict:
     """
     Generate completions based on a given prompt using a specified model.
 
@@ -89,7 +89,7 @@ async def completions(prompt: str, *, model: str, model_kwargs: Optional[dict[st
 
 @pxt.udf(resource_pool='request-rate:together:chat')
 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:
     """
     Generate chat completions based on a given prompt using a specified model.
@@ -183,9 +183,7 @@ def _(model: str) -> ts.ArrayType:
 
 
 @pxt.udf(resource_pool='request-rate:together:images')
-async def image_generations(
-    prompt: str, *, model: str, model_kwargs: Optional[dict[str, Any]] = None
-) -> PIL.Image.Image:
+async def image_generations(prompt: str, *, model: str, model_kwargs: dict[str, Any] | None = None) -> PIL.Image.Image:
     """
     Generate images based on a given prompt using a specified model.