fal 1.3.4__tar.gz → 1.4.1__tar.gz

{fal-1.3.4 → fal-1.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fal
-Version: 1.3.4
+Version: 1.4.1
 Summary: fal is an easy-to-use Serverless Python Framework
 Author: Features & Labels <support@fal.ai>
 Requires-Python: >=3.8

{fal-1.3.4 → fal-1.4.1}/fal.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fal
-Version: 1.3.4
+Version: 1.4.1
 Summary: fal is an easy-to-use Serverless Python Framework
 Author: Features & Labels <support@fal.ai>
 Requires-Python: >=3.8

{fal-1.3.4 → fal-1.4.1}/fal.egg-info/SOURCES.txt

@@ -137,6 +137,7 @@ src/fal/toolkit/image/nsfw_filter/model.py
 src/fal/toolkit/image/nsfw_filter/requirements.txt
 src/fal/toolkit/utils/__init__.py
 src/fal/toolkit/utils/download_utils.py
+src/fal/toolkit/utils/retry.py
 tests/__init__.py
 tests/conftest.py
 tests/integration_test.py
@@ -155,4 +156,5 @@ tests/mainify_package/impl.py
 tests/mainify_package/utils.py
 tests/toolkit/file_test.py
 tests/toolkit/image_test.py
+tests/toolkit/utils/retry.py
 tools/demo_script.py

{fal-1.3.4 → fal-1.4.1}/src/fal/_fal_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '1.3.4'
-__version_tuple__ = version_tuple = (1, 3, 4)
+__version__ = version = '1.4.1'
+__version_tuple__ = version_tuple = (1, 4, 1)

{fal-1.3.4 → fal-1.4.1}/src/fal/api.py

@@ -389,6 +389,8 @@ class FalServerlessHost(Host):
     _SUPPORTED_KEYS = frozenset(
         {
             "machine_type",
+            "machine_types",
+            "num_gpus",
             "keep_alive",
             "max_concurrency",
             "min_concurrency",
@@ -431,7 +433,7 @@ class FalServerlessHost(Host):
         environment_options.setdefault("python_version", active_python())
         environments = [self._connection.define_environment(**environment_options)]
 
-        machine_type = options.host.get(
+        machine_type: list[str] | str = options.host.get(
             "machine_type", FAL_SERVERLESS_DEFAULT_MACHINE_TYPE
         )
         keep_alive = options.host.get("keep_alive", FAL_SERVERLESS_DEFAULT_KEEP_ALIVE)
@@ -444,7 +446,8 @@ class FalServerlessHost(Host):
         exposed_port = options.get_exposed_port()
 
         machine_requirements = MachineRequirements(
-            machine_type=machine_type,
+            machine_types=machine_type,  # type: ignore
+            num_gpus=options.host.get("num_gpus"),
             keep_alive=keep_alive,
             base_image=base_image,
             exposed_port=exposed_port,
@@ -501,7 +504,7 @@ class FalServerlessHost(Host):
         environment_options.setdefault("python_version", active_python())
         environments = [self._connection.define_environment(**environment_options)]
 
-        machine_type = options.host.get(
+        machine_type: list[str] | str = options.host.get(
             "machine_type", FAL_SERVERLESS_DEFAULT_MACHINE_TYPE
         )
         keep_alive = options.host.get("keep_alive", FAL_SERVERLESS_DEFAULT_KEEP_ALIVE)
@@ -515,7 +518,8 @@ class FalServerlessHost(Host):
         setup_function = options.host.get("setup_function", None)
 
         machine_requirements = MachineRequirements(
-            machine_type=machine_type,
+            machine_types=machine_type,  # type: ignore
+            num_gpus=options.host.get("num_gpus"),
             keep_alive=keep_alive,
             base_image=base_image,
             exposed_port=exposed_port,
@@ -684,7 +688,8 @@ def function(
     max_concurrency: int | None = None,
     # FalServerlessHost options
     metadata: dict[str, Any] | None = None,
-    machine_type: str = FAL_SERVERLESS_DEFAULT_MACHINE_TYPE,
+    machine_type: str | list[str] = FAL_SERVERLESS_DEFAULT_MACHINE_TYPE,
+    num_gpus: int | None = None,
     keep_alive: int = FAL_SERVERLESS_DEFAULT_KEEP_ALIVE,
     max_multiplexing: int = FAL_SERVERLESS_DEFAULT_MAX_MULTIPLEXING,
     min_concurrency: int = FAL_SERVERLESS_DEFAULT_MIN_CONCURRENCY,
@@ -709,7 +714,8 @@ def function(
     max_concurrency: int | None = None,
     # FalServerlessHost options
     metadata: dict[str, Any] | None = None,
-    machine_type: str = FAL_SERVERLESS_DEFAULT_MACHINE_TYPE,
+    machine_type: str | list[str] = FAL_SERVERLESS_DEFAULT_MACHINE_TYPE,
+    num_gpus: int | None = None,
     keep_alive: int = FAL_SERVERLESS_DEFAULT_KEEP_ALIVE,
     max_multiplexing: int = FAL_SERVERLESS_DEFAULT_MAX_MULTIPLEXING,
     min_concurrency: int = FAL_SERVERLESS_DEFAULT_MIN_CONCURRENCY,
@@ -784,7 +790,8 @@ def function(
     max_concurrency: int | None = None,
     # FalServerlessHost options
     metadata: dict[str, Any] | None = None,
-    machine_type: str = FAL_SERVERLESS_DEFAULT_MACHINE_TYPE,
+    machine_type: str | list[str] = FAL_SERVERLESS_DEFAULT_MACHINE_TYPE,
+    num_gpus: int | None = None,
     keep_alive: int = FAL_SERVERLESS_DEFAULT_KEEP_ALIVE,
     max_multiplexing: int = FAL_SERVERLESS_DEFAULT_MAX_MULTIPLEXING,
     min_concurrency: int = FAL_SERVERLESS_DEFAULT_MIN_CONCURRENCY,
@@ -814,7 +821,8 @@ def function(
     max_concurrency: int | None = None,
     # FalServerlessHost options
     metadata: dict[str, Any] | None = None,
-    machine_type: str = FAL_SERVERLESS_DEFAULT_MACHINE_TYPE,
+    machine_type: str | list[str] = FAL_SERVERLESS_DEFAULT_MACHINE_TYPE,
+    num_gpus: int | None = None,
     keep_alive: int = FAL_SERVERLESS_DEFAULT_KEEP_ALIVE,
     max_multiplexing: int = FAL_SERVERLESS_DEFAULT_MAX_MULTIPLEXING,
     min_concurrency: int = FAL_SERVERLESS_DEFAULT_MIN_CONCURRENCY,
@@ -838,7 +846,8 @@ def function(
     max_concurrency: int | None = None,
     # FalServerlessHost options
     metadata: dict[str, Any] | None = None,
-    machine_type: str = FAL_SERVERLESS_DEFAULT_MACHINE_TYPE,
+    machine_type: str | list[str] = FAL_SERVERLESS_DEFAULT_MACHINE_TYPE,
+    num_gpus: int | None = None,
     keep_alive: int = FAL_SERVERLESS_DEFAULT_KEEP_ALIVE,
     max_multiplexing: int = FAL_SERVERLESS_DEFAULT_MAX_MULTIPLEXING,
     min_concurrency: int = FAL_SERVERLESS_DEFAULT_MIN_CONCURRENCY,
@@ -862,7 +871,8 @@ def function(
     max_concurrency: int | None = None,
     # FalServerlessHost options
     metadata: dict[str, Any] | None = None,
-    machine_type: str = FAL_SERVERLESS_DEFAULT_MACHINE_TYPE,
+    machine_type: str | list[str] = FAL_SERVERLESS_DEFAULT_MACHINE_TYPE,
+    num_gpus: int | None = None,
     keep_alive: int = FAL_SERVERLESS_DEFAULT_KEEP_ALIVE,
     max_multiplexing: int = FAL_SERVERLESS_DEFAULT_MAX_MULTIPLEXING,
     min_concurrency: int = FAL_SERVERLESS_DEFAULT_MIN_CONCURRENCY,

{fal-1.3.4 → fal-1.4.1}/src/fal/app.py

@@ -60,6 +60,7 @@ def wrap_app(cls: type[App], **kwargs) -> fal.api.IsolatedFunction:
         kind,
         requirements=cls.requirements,
         machine_type=cls.machine_type,
+        num_gpus=cls.num_gpus,
         **cls.host_kwargs,
         **kwargs,
         metadata=metadata,
@@ -177,6 +178,7 @@ def _to_fal_app_name(name: str) -> str:
 class App(fal.api.BaseServable):
     requirements: ClassVar[list[str]] = []
     machine_type: ClassVar[str] = "S"
+    num_gpus: ClassVar[int | None] = None
     host_kwargs: ClassVar[dict[str, Any]] = {
         "_scheduler": "nomad",
         "_scheduler_options": {

{fal-1.3.4 → fal-1.4.1}/src/fal/cli/_utils.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from fal.files import find_pyproject_toml, parse_pyproject_toml
+from fal.files import find_project_root, find_pyproject_toml, parse_pyproject_toml
 
 
 def is_app_name(app_ref: tuple[str, str | None]) -> bool:
@@ -29,6 +29,10 @@ def get_app_data_from_toml(app_name):
     except KeyError:
         raise ValueError(f"App {app_name} does not have a ref key in pyproject.toml")
 
+    # Convert the app_ref to a path relative to the project root
+    project_root, _ = find_project_root(None)
+    app_ref = str(project_root / app_ref)
+
     app_auth = app_data.get("auth", "private")
 
     return app_ref, app_auth

{fal-1.3.4 → fal-1.4.1}/src/fal/sdk.py

@@ -389,7 +389,8 @@ def _from_grpc_hosted_run_result(
 
 @dataclass
 class MachineRequirements:
-    machine_type: str
+    machine_types: list[str]
+    num_gpus: int | None = field(default=None)
     keep_alive: int = FAL_SERVERLESS_DEFAULT_KEEP_ALIVE
     base_image: str | None = None
     exposed_port: int | None = None
@@ -399,6 +400,16 @@ class MachineRequirements:
     max_multiplexing: int | None = None
     min_concurrency: int | None = None
 
+    def __post_init__(self):
+        if isinstance(self.machine_types, str):
+            self.machine_types = [self.machine_types]
+
+        if not isinstance(self.machine_types, list):
+            raise ValueError("machine_types must be a list of strings.")
+
+        if not self.machine_types:
+            raise ValueError("No machine type provided.")
+
 
 @dataclass
 class FalServerlessConnection:
@@ -489,7 +500,10 @@ class FalServerlessConnection:
         wrapped_function = to_serialized_object(function, serialization_method)
         if machine_requirements:
             wrapped_requirements = isolate_proto.MachineRequirements(
-                machine_type=machine_requirements.machine_type,
+                # NOTE: backwards compatibility with old API
+                machine_type=machine_requirements.machine_types[0],
+                machine_types=machine_requirements.machine_types,
+                num_gpus=machine_requirements.num_gpus,
                 keep_alive=machine_requirements.keep_alive,
                 base_image=machine_requirements.base_image,
                 exposed_port=machine_requirements.exposed_port,
@@ -579,7 +593,10 @@ class FalServerlessConnection:
         wrapped_function = to_serialized_object(function, serialization_method)
         if machine_requirements:
             wrapped_requirements = isolate_proto.MachineRequirements(
-                machine_type=machine_requirements.machine_type,
+                # NOTE: backwards compatibility with old API
+                machine_type=machine_requirements.machine_types[0],
+                machine_types=machine_requirements.machine_types,
+                num_gpus=machine_requirements.num_gpus,
                 keep_alive=machine_requirements.keep_alive,
                 base_image=machine_requirements.base_image,
                 exposed_port=machine_requirements.exposed_port,

{fal-1.3.4 → fal-1.4.1}/src/fal/toolkit/file/file.py

@@ -51,7 +51,8 @@ def get_builtin_repository(id: RepositoryId) -> FileRepository:
 
 get_builtin_repository.__module__ = "__main__"
 
-DEFAULT_REPOSITORY: FileRepository | RepositoryId = "fal"
+DEFAULT_REPOSITORY: FileRepository | RepositoryId = "fal_v2"
+FALLBACK_REPOSITORY: FileRepository | RepositoryId = "cdn"
 
 
 class File(BaseModel):
@@ -126,6 +127,9 @@ class File(BaseModel):
         content_type: Optional[str] = None,
         file_name: Optional[str] = None,
         repository: FileRepository | RepositoryId = DEFAULT_REPOSITORY,
+        fallback_repository: Optional[
+            FileRepository | RepositoryId
+        ] = FALLBACK_REPOSITORY,
     ) -> File:
         repo = (
             repository
@@ -135,8 +139,22 @@ class File(BaseModel):
 
         fdata = FileData(data, content_type, file_name)
 
+        try:
+            url = repo.save(fdata)
+        except Exception:
+            if not fallback_repository:
+                raise
+
+            fallback_repo = (
+                fallback_repository
+                if isinstance(fallback_repository, FileRepository)
+                else get_builtin_repository(fallback_repository)
+            )
+
+            url = fallback_repo.save(fdata)
+
         return cls(
-            url=repo.save(fdata),
+            url=url,
             content_type=fdata.content_type,
             file_name=fdata.file_name,
             file_size=len(data),
@@ -150,6 +168,9 @@ class File(BaseModel):
         content_type: Optional[str] = None,
         repository: FileRepository | RepositoryId = DEFAULT_REPOSITORY,
         multipart: bool | None = None,
+        fallback_repository: Optional[
+            FileRepository | RepositoryId
+        ] = FALLBACK_REPOSITORY,
     ) -> File:
         file_path = Path(path)
         if not file_path.exists():
@@ -163,11 +184,28 @@ class File(BaseModel):
 
         content_type = content_type or "application/octet-stream"
 
-        url, data = repo.save_file(
-            file_path,
-            content_type=content_type,
-            multipart=multipart,
-        )
+        try:
+            url, data = repo.save_file(
+                file_path,
+                content_type=content_type,
+                multipart=multipart,
+            )
+        except Exception:
+            if not fallback_repository:
+                raise
+
+            fallback_repo = (
+                fallback_repository
+                if isinstance(fallback_repository, FileRepository)
+                else get_builtin_repository(fallback_repository)
+            )
+
+            url, data = fallback_repo.save_file(
+                file_path,
+                content_type=content_type,
+                multipart=multipart,
+            )
+
         return cls(
             url=url,
             file_data=data.data if data else None,

{fal-1.3.4 → fal-1.4.1}/src/fal/toolkit/file/providers/fal.py

@@ -4,19 +4,91 @@ import dataclasses
 import json
 import math
 import os
+import threading
 from base64 import b64encode
 from dataclasses import dataclass
+from datetime import datetime, timezone
 from pathlib import Path
 from urllib.error import HTTPError
+from urllib.parse import urlparse, urlunparse
 from urllib.request import Request, urlopen
 
 from fal.auth import key_credentials
 from fal.toolkit.exceptions import FileUploadException
 from fal.toolkit.file.types import FileData, FileRepository
+from fal.toolkit.utils.retry import retry
 
 _FAL_CDN = "https://fal.media"
 
 
+@dataclass
+class FalV2Token:
+    token: str
+    token_type: str
+    base_upload_url: str
+    expires_at: datetime
+
+    def is_expired(self) -> bool:
+        return datetime.now(timezone.utc) >= self.expires_at
+
+
+class FalV2TokenManager:
+    def __init__(self):
+        self._token: FalV2Token = FalV2Token(
+            token="",
+            token_type="",
+            base_upload_url="",
+            expires_at=datetime.min.replace(tzinfo=timezone.utc),
+        )
+        self._lock: threading.Lock = threading.Lock()
+
+    def get_token(self) -> FalV2Token:
+        with self._lock:
+            if self._token.is_expired():
+                self._refresh_token()
+            return self._token
+
+    def _refresh_token(self) -> None:
+        key_creds = key_credentials()
+        if not key_creds:
+            raise FileUploadException("FAL_KEY must be set")
+
+        key_id, key_secret = key_creds
+        headers = {
+            "Authorization": f"Key {key_id}:{key_secret}",
+            "Accept": "application/json",
+            "Content-Type": "application/json",
+        }
+
+        grpc_host = os.environ.get("FAL_HOST", "api.alpha.fal.ai")
+        rest_host = grpc_host.replace("api", "rest", 1)
+        url = f"https://{rest_host}/storage/auth/token"
+
+        req = Request(
+            url,
+            headers=headers,
+            data=b"{}",
+            method="POST",
+        )
+        with urlopen(req) as response:
+            result = json.load(response)
+
+        parsed_base_url = urlparse(result["base_url"])
+        base_upload_url = urlunparse(
+            parsed_base_url._replace(netloc="upload." + parsed_base_url.netloc)
+        )
+
+        self._token = FalV2Token(
+            token=result["token"],
+            token_type=result["token_type"],
+            base_upload_url=base_upload_url,
+            expires_at=datetime.fromisoformat(result["expires_at"]),
+        )
+
+
+fal_v2_token_manager = FalV2TokenManager()
+
+
 @dataclass
 class ObjectLifecyclePreference:
     expriation_duration_seconds: int
@@ -29,6 +101,7 @@ GLOBAL_LIFECYCLE_PREFERENCE = ObjectLifecyclePreference(
 
 @dataclass
 class FalFileRepositoryBase(FileRepository):
+    @retry(max_retries=3, base_delay=1, backoff_type="exponential", jitter=True)
     def _save(self, file: FileData, storage_type: str) -> str:
         key_creds = key_credentials()
         if not key_creds:
@@ -108,26 +181,14 @@ class MultipartUpload:
 
         self._parts: list[dict] = []
 
-        key_creds = key_credentials()
-        if not key_creds:
-            raise FileUploadException("FAL_KEY must be set")
-
-        key_id, key_secret = key_creds
-
-        self._auth_headers = {
-            "Authorization": f"Key {key_id}:{key_secret}",
-        }
-        grpc_host = os.environ.get("FAL_HOST", "api.alpha.fal.ai")
-        rest_host = grpc_host.replace("api", "rest", 1)
-        self._storage_upload_url = f"https://{rest_host}/storage/upload"
-
     def create(self):
+        token = fal_v2_token_manager.get_token()
         try:
             req = Request(
-                f"{self._storage_upload_url}/initiate-multipart",
+                f"{token.base_upload_url}/upload/initiate-multipart",
                 method="POST",
                 headers={
-                    **self._auth_headers,
+                    "Authorization": f"{token.token_type} {token.token}",
                     "Accept": "application/json",
                     "Content-Type": "application/json",
                 },
@@ -140,7 +201,7 @@ class MultipartUpload:
             )
             with urlopen(req) as response:
                 result = json.load(response)
-                self._upload_id = result["upload_id"]
+                self._upload_url = result["upload_url"]
                 self._file_url = result["file_url"]
         except HTTPError as exc:
             raise FileUploadException(
@@ -180,10 +241,7 @@ class MultipartUpload:
         ) as executor:
             futures = []
             for part_number in range(1, parts + 1):
-                upload_url = (
-                    f"{self._file_url}?upload_id={self._upload_id}"
-                    f"&part_number={part_number}"
-                )
+                upload_url = f"{self._upload_url}&part_number={part_number}"
                 futures.append(
                     executor.submit(self._upload_part, upload_url, part_number)
                 )
@@ -193,7 +251,7 @@ class MultipartUpload:
                 self._parts.append(entry)
 
     def complete(self):
-        url = f"{self._file_url}?upload_id={self._upload_id}"
+        url = self._upload_url
         try:
             req = Request(
                 url,
@@ -216,8 +274,33 @@ class MultipartUpload:
 
 @dataclass
 class FalFileRepositoryV2(FalFileRepositoryBase):
+    @retry(max_retries=3, base_delay=1, backoff_type="exponential", jitter=True)
     def save(self, file: FileData) -> str:
-        return self._save(file, "fal-cdn")
+        token = fal_v2_token_manager.get_token()
+        headers = {
+            "Authorization": f"{token.token_type} {token.token}",
+            "Accept": "application/json",
+            "X-Fal-File-Name": file.file_name,
+            "Content-Type": file.content_type,
+        }
+
+        storage_url = f"{token.base_upload_url}/upload"
+
+        try:
+            req = Request(
+                storage_url,
+                data=file.data,
+                headers=headers,
+                method="PUT",
+            )
+            with urlopen(req) as response:
+                result = json.load(response)
+
+            return result["file_url"]
+        except HTTPError as e:
+            raise FileUploadException(
+                f"Error initiating upload. Status {e.status}: {e.reason}"
+            )
 
     def _save_multipart(
         self,
@@ -280,6 +363,7 @@ class InMemoryRepository(FileRepository):
 
 @dataclass
 class FalCDNFileRepository(FileRepository):
+    @retry(max_retries=3, base_delay=1, backoff_type="exponential", jitter=True)
     def save(
         self,
         file: FileData,

{fal-1.3.4 → fal-1.4.1}/src/fal/toolkit/file/providers/gcp.py

@@ -8,6 +8,7 @@ import uuid
 from dataclasses import dataclass
 
 from fal.toolkit.file.types import FileData, FileRepository
+from fal.toolkit.utils.retry import retry
 
 DEFAULT_URL_TIMEOUT = 60 * 15  # 15 minutes
 
@@ -50,6 +51,7 @@ class GoogleStorageRepository(FileRepository):
 
         return self._bucket
 
+    @retry(max_retries=3, base_delay=1, backoff_type="exponential", jitter=True)
     def save(self, data: FileData) -> str:
         destination_path = posixpath.join(
             self.folder,

{fal-1.3.4 → fal-1.4.1}/src/fal/toolkit/file/providers/r2.py

@@ -8,6 +8,7 @@ from dataclasses import dataclass
 from io import BytesIO
 
 from fal.toolkit.file.types import FileData, FileRepository
+from fal.toolkit.utils.retry import retry
 
 DEFAULT_URL_TIMEOUT = 60 * 15  # 15 minutes
 
@@ -67,6 +68,7 @@ class R2Repository(FileRepository):
 
         return self._bucket
 
+    @retry(max_retries=3, base_delay=1, backoff_type="exponential", jitter=True)
     def save(self, data: FileData) -> str:
         destination_path = posixpath.join(
             self.key,

{fal-1.3.4 → fal-1.4.1}/src/fal/toolkit/image/image.py

@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Literal, Optional, Union
 
 from pydantic import BaseModel, Field
 
-from fal.toolkit.file.file import DEFAULT_REPOSITORY, File
+from fal.toolkit.file.file import DEFAULT_REPOSITORY, FALLBACK_REPOSITORY, File
 from fal.toolkit.file.types import FileRepository, RepositoryId
 from fal.toolkit.utils.download_utils import _download_file_python
 
@@ -79,12 +79,16 @@ class Image(File):
         size: ImageSize | None = None,
         file_name: str | None = None,
         repository: FileRepository | RepositoryId = DEFAULT_REPOSITORY,
+        fallback_repository: Optional[
+            FileRepository | RepositoryId
+        ] = FALLBACK_REPOSITORY,
     ) -> Image:
         obj = super().from_bytes(
             data,
             content_type=f"image/{format}",
             file_name=file_name,
             repository=repository,
+            fallback_repository=fallback_repository,
         )
         obj.width = size.width if size else None
         obj.height = size.height if size else None
@@ -97,6 +101,9 @@ class Image(File):
         format: ImageFormat | None = None,
         file_name: str | None = None,
         repository: FileRepository | RepositoryId = DEFAULT_REPOSITORY,
+        fallback_repository: Optional[
+            FileRepository | RepositoryId
+        ] = FALLBACK_REPOSITORY,
     ) -> Image:
         size = ImageSize(width=pil_image.width, height=pil_image.height)
         if format is None:
@@ -119,7 +126,14 @@ class Image(File):
             pil_image.save(f, format=format, **saving_options)
             raw_image = f.getvalue()
 
-        return cls.from_bytes(raw_image, format, size, file_name, repository)
+        return cls.from_bytes(
+            raw_image,
+            format,
+            size,
+            file_name,
+            repository,
+            fallback_repository=fallback_repository,
+        )
 
     def to_pil(self, mode: str = "RGB") -> PILImage.Image:
         try:

fal-1.4.1/src/fal/toolkit/utils/retry.py

@@ -0,0 +1,42 @@
+import functools
+import random
+import time
+from typing import Any, Callable, Literal
+
+BackoffType = Literal["exponential", "fixed"]
+
+
+def retry(
+    max_retries: int = 3,
+    base_delay: float = 1.0,
+    max_delay: float = 60.0,
+    backoff_type: BackoffType = "exponential",
+    jitter: bool = False,
+) -> Callable:
+    def decorator(func: Callable) -> Callable:
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            retries = 0
+            while retries < max_retries:
+                try:
+                    return func(*args, **kwargs)
+                except Exception as e:
+                    retries += 1
+                    print(f"Retrying {retries} of {max_retries}...")
+                    if retries == max_retries:
+                        print(f"Max retries reached. Raising exception: {e}")
+                        raise e
+
+                    if backoff_type == "exponential":
+                        delay = min(base_delay * (2 ** (retries - 1)), max_delay)
+                    else:  # fixed
+                        delay = min(base_delay, max_delay)
+
+                    if jitter:
+                        delay *= random.uniform(0.5, 1.5)
+
+                    time.sleep(delay)
+
+        return wrapper
+
+    return decorator