1 files changed, 126 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/hatchet_sdk/rate_limit.py b/.venv/lib/python3.12/site-packages/hatchet_sdk/rate_limit.py
new file mode 100644
index 00000000..0d7b9143
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/hatchet_sdk/rate_limit.py
@@ -0,0 +1,126 @@
+from dataclasses import dataclass
+from typing import Union
+
+from celpy import CELEvalError, Environment
+
+from hatchet_sdk.contracts.workflows_pb2 import CreateStepRateLimit
+
+
+def validate_cel_expression(expr: str) -> bool:
+    env = Environment()
+    try:
+        env.compile(expr)
+        return True
+    except CELEvalError:
+        return False
+
+
+class RateLimitDuration:
+    SECOND = "SECOND"
+    MINUTE = "MINUTE"
+    HOUR = "HOUR"
+    DAY = "DAY"
+    WEEK = "WEEK"
+    MONTH = "MONTH"
+    YEAR = "YEAR"
+
+
+@dataclass
+class RateLimit:
+    """
+    Represents a rate limit configuration for a step in a workflow.
+
+    This class allows for both static and dynamic rate limiting based on various parameters.
+    It supports both simple integer values and Common Expression Language (CEL) expressions
+    for dynamic evaluation.
+
+    Attributes:
+        static_key (str, optional): A static key for rate limiting.
+        dynamic_key (str, optional): A CEL expression for dynamic key evaluation.
+        units (int or str, default=1): The number of units or a CEL expression for dynamic unit calculation.
+        limit (int or str, optional): The rate limit value or a CEL expression for dynamic limit calculation.
+        duration (str, default=RateLimitDuration.MINUTE): The window duration of the rate limit.
+        key (str, optional): Deprecated. Use static_key instead.
+
+    Usage:
+        1. Static rate limit:
+           rate_limit = RateLimit(static_key="external-api", units=100)
+           > NOTE: if you want to use a static key, you must first put the rate limit: hatchet.admin.put_rate_limit("external-api", 200, RateLimitDuration.SECOND)
+
+        2. Dynamic rate limit with CEL expressions:
+           rate_limit = RateLimit(
+               dynamic_key="input.user_id",
+               units="input.units",
+               limit="input.limit * input.user_tier"
+           )
+
+    Note:
+        - Either static_key or dynamic_key must be set, but not both.
+        - When using dynamic_key, limit must also be set.
+        - CEL expressions are validated upon instantiation.
+
+    Raises:
+        ValueError: If invalid combinations of attributes are provided or if CEL expressions are invalid.
+        DeprecationWarning: If the deprecated 'key' attribute is used.
+    """
+
+    key: Union[str, None] = None
+    static_key: Union[str, None] = None
+    dynamic_key: Union[str, None] = None
+    units: Union[int, str] = 1
+    limit: Union[int, str, None] = None
+    duration: RateLimitDuration = RateLimitDuration.MINUTE
+
+    _req: CreateStepRateLimit = None
+
+    def __post_init__(self):
+        # juggle the key and key_expr fields
+        key = self.static_key
+        key_expression = self.dynamic_key
+
+        if self.key is not None:
+            DeprecationWarning(
+                "key is deprecated and will be removed in a future release, please use static_key instead"
+            )
+            key = self.key
+
+        if key_expression is not None:
+            if key is not None:
+                raise ValueError("Cannot have both static key and dynamic key set")
+
+            key = key_expression
+            if not validate_cel_expression(key_expression):
+                raise ValueError(f"Invalid CEL expression: {key_expression}")
+
+        # juggle the units and units_expr fields
+        units = None
+        units_expression = None
+        if isinstance(self.units, int):
+            units = self.units
+        else:
+            if not validate_cel_expression(self.units):
+                raise ValueError(f"Invalid CEL expression: {self.units}")
+            units_expression = self.units
+
+        # juggle the limit and limit_expr fields
+        limit_expression = None
+
+        if self.limit:
+            if isinstance(self.limit, int):
+                limit_expression = f"{self.limit}"
+            else:
+                if not validate_cel_expression(self.limit):
+                    raise ValueError(f"Invalid CEL expression: {self.limit}")
+                limit_expression = self.limit
+
+        if key_expression is not None and limit_expression is None:
+            raise ValueError("CEL based keys requires limit to be set")
+
+        self._req = CreateStepRateLimit(
+            key=key,
+            key_expr=key_expression,
+            units=units,
+            units_expr=units_expression,
+            limit_values_expr=limit_expression,
+            duration=self.duration,
+        )