hatchet-sdk 1.0.0__py3-none-any.whl → 1.0.1__py3-none-any.whl

hatchet_sdk/opentelemetry/instrumentor.py

@@ -30,7 +30,7 @@ import hatchet_sdk
 from hatchet_sdk.clients.admin import (
     AdminClient,
     TriggerWorkflowOptions,
-    WorkflowRunDict,
+    WorkflowRunTriggerConfig,
 )
 from hatchet_sdk.clients.dispatcher.action_listener import Action
 from hatchet_sdk.clients.events import (
@@ -352,14 +352,12 @@ class HatchetInstrumentor(BaseInstrumentor):  # type: ignore[misc]
     def _wrap_run_workflows(
         self,
         wrapped: Callable[
-            [list[WorkflowRunDict], TriggerWorkflowOptions | None], list[WorkflowRunRef]
+            [list[WorkflowRunTriggerConfig]],
+            list[WorkflowRunRef],
         ],
         instance: AdminClient,
-        args: tuple[
-            list[WorkflowRunDict],
-            TriggerWorkflowOptions | None,
-        ],
-        kwargs: dict[str, list[WorkflowRunDict] | TriggerWorkflowOptions | None],
+        args: tuple[list[WorkflowRunTriggerConfig],],
+        kwargs: dict[str, list[WorkflowRunTriggerConfig]],
     ) -> list[WorkflowRunRef]:
         with self._tracer.start_as_current_span(
             "hatchet.run_workflows",
@@ -370,15 +368,12 @@ class HatchetInstrumentor(BaseInstrumentor):  # type: ignore[misc]
     async def _wrap_async_run_workflows(
         self,
         wrapped: Callable[
-            [list[WorkflowRunDict], TriggerWorkflowOptions | None],
+            [list[WorkflowRunTriggerConfig]],
             Coroutine[None, None, list[WorkflowRunRef]],
         ],
         instance: AdminClient,
-        args: tuple[
-            list[WorkflowRunDict],
-            TriggerWorkflowOptions | None,
-        ],
-        kwargs: dict[str, list[WorkflowRunDict] | TriggerWorkflowOptions | None],
+        args: tuple[list[WorkflowRunTriggerConfig],],
+        kwargs: dict[str, list[WorkflowRunTriggerConfig]],
     ) -> list[WorkflowRunRef]:
         with self._tracer.start_as_current_span(
             "hatchet.run_workflows",

hatchet_sdk/rate_limit.py

@@ -1,9 +1,9 @@
-from dataclasses import dataclass
 from enum import Enum
 
 from celpy import CELEvalError, Environment  # type: ignore
+from pydantic import BaseModel, model_validator
 
-from hatchet_sdk.contracts.workflows_pb2 import CreateStepRateLimit
+from hatchet_sdk.contracts.v1.workflows_pb2 import CreateTaskRateLimit
 
 
 def validate_cel_expression(expr: str) -> bool:
@@ -25,8 +25,7 @@ class RateLimitDuration(str, Enum):
     YEAR = "YEAR"
 
 
-@dataclass
-class RateLimit:
+class RateLimit(BaseModel):
     """
     Represents a rate limit configuration for a step in a workflow.
 
@@ -68,46 +67,41 @@ class RateLimit:
     limit: int | str | None = None
     duration: RateLimitDuration = RateLimitDuration.MINUTE
 
-    _req: CreateStepRateLimit | None = None
+    @model_validator(mode="after")
+    def validate_rate_limit(self) -> "RateLimit":
+        if self.dynamic_key and self.static_key:
+            raise ValueError("Cannot have both static key and dynamic key set")
 
-    def __post_init__(self) -> None:
-        # juggle the key and key_expr fields
+        if self.dynamic_key and not validate_cel_expression(self.dynamic_key):
+            raise ValueError(f"Invalid CEL expression: {self.dynamic_key}")
+
+        if not isinstance(self.units, int) and not validate_cel_expression(self.units):
+            raise ValueError(f"Invalid CEL expression: {self.units}")
+
+        if (
+            self.limit
+            and not isinstance(self.limit, int)
+            and not validate_cel_expression(self.limit)
+        ):
+            raise ValueError(f"Invalid CEL expression: {self.limit}")
+
+        if self.dynamic_key and not self.limit:
+            raise ValueError("CEL based keys requires limit to be set")
+
+        return self
+
+    def to_proto(self) -> CreateTaskRateLimit:
         key = self.static_key
         key_expression = self.dynamic_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")
+        key = self.static_key or self.dynamic_key
+
+        units = self.units if isinstance(self.units, int) else None
+        units_expression = None if isinstance(self.units, int) else self.units
+
+        limit_expression = None if not self.limit else str(self.limit)
 
-        self._req = CreateStepRateLimit(
+        return CreateTaskRateLimit(
             key=key,
             key_expr=key_expression,
             units=units,

hatchet_sdk/runnables/contextvars.py

@@ -0,0 +1,12 @@
+import asyncio
+from collections import Counter
+from contextvars import ContextVar
+
+ctx_workflow_run_id: ContextVar[str | None] = ContextVar(
+    "ctx_workflow_run_id", default=None
+)
+ctx_step_run_id: ContextVar[str | None] = ContextVar("ctx_step_run_id", default=None)
+ctx_worker_id: ContextVar[str | None] = ContextVar("ctx_worker_id", default=None)
+
+workflow_spawn_indices = Counter[str]()
+spawn_index_lock = asyncio.Lock()

hatchet_sdk/runnables/standalone.py

@@ -0,0 +1,192 @@
+import asyncio
+from datetime import datetime
+from typing import Any, Generic, cast, get_type_hints
+
+from hatchet_sdk.clients.admin import (
+    ScheduleTriggerWorkflowOptions,
+    TriggerWorkflowOptions,
+    WorkflowRunTriggerConfig,
+)
+from hatchet_sdk.clients.rest.models.cron_workflows import CronWorkflows
+from hatchet_sdk.contracts.workflows_pb2 import WorkflowVersion
+from hatchet_sdk.runnables.task import Task
+from hatchet_sdk.runnables.types import EmptyModel, R, TWorkflowInput
+from hatchet_sdk.runnables.workflow import BaseWorkflow, Workflow
+from hatchet_sdk.utils.aio_utils import get_active_event_loop
+from hatchet_sdk.utils.typing import JSONSerializableMapping, is_basemodel_subclass
+from hatchet_sdk.workflow_run import WorkflowRunRef
+
+
+class TaskRunRef(Generic[TWorkflowInput, R]):
+    def __init__(
+        self,
+        standalone: "Standalone[TWorkflowInput, R]",
+        workflow_run_ref: WorkflowRunRef,
+    ):
+        self._s = standalone
+        self._wrr = workflow_run_ref
+
+    async def aio_result(self) -> R:
+        result = await self._wrr.workflow_listener.result(self._wrr.workflow_run_id)
+        return self._s._extract_result(result)
+
+    def result(self) -> R:
+        coro = self._wrr.workflow_listener.result(self._wrr.workflow_run_id)
+
+        loop = get_active_event_loop()
+
+        if loop is None:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            try:
+                result = loop.run_until_complete(coro)
+            finally:
+                asyncio.set_event_loop(None)
+        else:
+            result = loop.run_until_complete(coro)
+
+        return self._s._extract_result(result)
+
+
+class Standalone(BaseWorkflow[TWorkflowInput], Generic[TWorkflowInput, R]):
+    def __init__(
+        self, workflow: Workflow[TWorkflowInput], task: Task[TWorkflowInput, R]
+    ) -> None:
+        super().__init__(config=workflow.config, client=workflow.client)
+
+        ## NOTE: This is a hack to assign the task back to the base workflow,
+        ## since the decorator to mutate the tasks is not being called.
+        self._default_tasks = [task]
+
+        self._workflow = workflow
+        self._task = task
+
+        return_type = get_type_hints(self._task.fn).get("return")
+
+        self._output_validator = (
+            return_type if is_basemodel_subclass(return_type) else None
+        )
+
+        self.config = self._workflow.config
+
+    def _extract_result(self, result: dict[str, Any]) -> R:
+        output = result.get(self._task.name)
+
+        if not self._output_validator:
+            return cast(R, output)
+
+        return cast(R, self._output_validator.model_validate(output))
+
+    def run(
+        self,
+        input: TWorkflowInput = cast(TWorkflowInput, EmptyModel()),
+        options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
+    ) -> R:
+        return self._extract_result(self._workflow.run(input, options))
+
+    async def aio_run(
+        self,
+        input: TWorkflowInput = cast(TWorkflowInput, EmptyModel()),
+        options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
+    ) -> R:
+        result = await self._workflow.aio_run(input, options)
+        return self._extract_result(result)
+
+    def run_no_wait(
+        self,
+        input: TWorkflowInput = cast(TWorkflowInput, EmptyModel()),
+        options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
+    ) -> TaskRunRef[TWorkflowInput, R]:
+        ref = self._workflow.run_no_wait(input, options)
+
+        return TaskRunRef[TWorkflowInput, R](self, ref)
+
+    async def aio_run_no_wait(
+        self,
+        input: TWorkflowInput = cast(TWorkflowInput, EmptyModel()),
+        options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
+    ) -> TaskRunRef[TWorkflowInput, R]:
+        ref = await self._workflow.aio_run_no_wait(input, options)
+
+        return TaskRunRef[TWorkflowInput, R](self, ref)
+
+    def run_many(self, workflows: list[WorkflowRunTriggerConfig]) -> list[R]:
+        return [
+            self._extract_result(result)
+            for result in self._workflow.run_many(workflows)
+        ]
+
+    async def aio_run_many(self, workflows: list[WorkflowRunTriggerConfig]) -> list[R]:
+        return [
+            self._extract_result(result)
+            for result in await self._workflow.aio_run_many(workflows)
+        ]
+
+    def run_many_no_wait(
+        self, workflows: list[WorkflowRunTriggerConfig]
+    ) -> list[TaskRunRef[TWorkflowInput, R]]:
+        refs = self._workflow.run_many_no_wait(workflows)
+
+        return [TaskRunRef[TWorkflowInput, R](self, ref) for ref in refs]
+
+    async def aio_run_many_no_wait(
+        self, workflows: list[WorkflowRunTriggerConfig]
+    ) -> list[TaskRunRef[TWorkflowInput, R]]:
+        refs = await self._workflow.aio_run_many_no_wait(workflows)
+
+        return [TaskRunRef[TWorkflowInput, R](self, ref) for ref in refs]
+
+    def schedule(
+        self,
+        run_at: datetime,
+        input: TWorkflowInput | None = None,
+        options: ScheduleTriggerWorkflowOptions = ScheduleTriggerWorkflowOptions(),
+    ) -> WorkflowVersion:
+        return self._workflow.schedule(
+            run_at=run_at,
+            input=input,
+            options=options,
+        )
+
+    async def aio_schedule(
+        self,
+        run_at: datetime,
+        input: TWorkflowInput,
+        options: ScheduleTriggerWorkflowOptions = ScheduleTriggerWorkflowOptions(),
+    ) -> WorkflowVersion:
+        return await self._workflow.aio_schedule(
+            run_at=run_at,
+            input=input,
+            options=options,
+        )
+
+    def create_cron(
+        self,
+        cron_name: str,
+        expression: str,
+        input: TWorkflowInput,
+        additional_metadata: JSONSerializableMapping,
+    ) -> CronWorkflows:
+        return self._workflow.create_cron(
+            cron_name=cron_name,
+            expression=expression,
+            input=input,
+            additional_metadata=additional_metadata,
+        )
+
+    async def aio_create_cron(
+        self,
+        cron_name: str,
+        expression: str,
+        input: TWorkflowInput,
+        additional_metadata: JSONSerializableMapping,
+    ) -> CronWorkflows:
+        return await self._workflow.aio_create_cron(
+            cron_name=cron_name,
+            expression=expression,
+            input=input,
+            additional_metadata=additional_metadata,
+        )
+
+    def to_task(self) -> Task[TWorkflowInput, R]:
+        return self._task

hatchet_sdk/runnables/task.py

@@ -0,0 +1,144 @@
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Callable,
+    Generic,
+    TypeVar,
+    Union,
+    cast,
+)
+
+from hatchet_sdk.context.context import Context, DurableContext
+from hatchet_sdk.contracts.v1.workflows_pb2 import (
+    CreateTaskRateLimit,
+    DesiredWorkerLabels,
+)
+from hatchet_sdk.runnables.types import (
+    DEFAULT_EXECUTION_TIMEOUT,
+    DEFAULT_SCHEDULE_TIMEOUT,
+    ConcurrencyExpression,
+    R,
+    StepType,
+    TWorkflowInput,
+    is_async_fn,
+    is_durable_sync_fn,
+    is_sync_fn,
+)
+from hatchet_sdk.utils.timedelta_to_expression import Duration
+from hatchet_sdk.waits import Condition, OrGroup
+
+if TYPE_CHECKING:
+    from hatchet_sdk.runnables.workflow import Workflow
+
+
+T = TypeVar("T")
+
+
+def fall_back_to_default(value: T, default: T, fallback_value: T) -> T:
+    ## If the value is not the default, it's set
+    if value != default:
+        return value
+
+    ## Otherwise, it's unset, so return the fallback value
+    return fallback_value
+
+
+class Task(Generic[TWorkflowInput, R]):
+    def __init__(
+        self,
+        _fn: Union[
+            Callable[[TWorkflowInput, Context], R]
+            | Callable[[TWorkflowInput, Context], Awaitable[R]],
+            Callable[[TWorkflowInput, DurableContext], R]
+            | Callable[[TWorkflowInput, DurableContext], Awaitable[R]],
+        ],
+        is_durable: bool,
+        type: StepType,
+        workflow: "Workflow[TWorkflowInput]",
+        name: str,
+        execution_timeout: Duration = DEFAULT_EXECUTION_TIMEOUT,
+        schedule_timeout: Duration = DEFAULT_SCHEDULE_TIMEOUT,
+        parents: "list[Task[TWorkflowInput, Any]]" = [],
+        retries: int = 0,
+        rate_limits: list[CreateTaskRateLimit] = [],
+        desired_worker_labels: dict[str, DesiredWorkerLabels] = {},
+        backoff_factor: float | None = None,
+        backoff_max_seconds: int | None = None,
+        concurrency: list[ConcurrencyExpression] = [],
+        wait_for: list[Condition | OrGroup] = [],
+        skip_if: list[Condition | OrGroup] = [],
+        cancel_if: list[Condition | OrGroup] = [],
+    ) -> None:
+        self.is_durable = is_durable
+
+        self.fn = _fn
+        self.is_async_function = is_async_fn(self.fn)  # type: ignore
+
+        self.workflow = workflow
+
+        self.type = type
+        self.execution_timeout = fall_back_to_default(
+            execution_timeout, DEFAULT_EXECUTION_TIMEOUT, DEFAULT_EXECUTION_TIMEOUT
+        )
+        self.schedule_timeout = fall_back_to_default(
+            schedule_timeout, DEFAULT_SCHEDULE_TIMEOUT, DEFAULT_SCHEDULE_TIMEOUT
+        )
+        self.name = name
+        self.parents = parents
+        self.retries = retries
+        self.rate_limits = rate_limits
+        self.desired_worker_labels = desired_worker_labels
+        self.backoff_factor = backoff_factor
+        self.backoff_max_seconds = backoff_max_seconds
+        self.concurrency = concurrency
+
+        self.wait_for = self._flatten_conditions(wait_for)
+        self.skip_if = self._flatten_conditions(skip_if)
+        self.cancel_if = self._flatten_conditions(cancel_if)
+
+    def _flatten_conditions(
+        self, conditions: list[Condition | OrGroup]
+    ) -> list[Condition]:
+        flattened: list[Condition] = []
+
+        for condition in conditions:
+            if isinstance(condition, OrGroup):
+                for or_condition in condition.conditions:
+                    or_condition.base.or_group_id = condition.or_group_id
+
+                flattened.extend(condition.conditions)
+            else:
+                flattened.append(condition)
+
+        return flattened
+
+    def call(self, ctx: Context | DurableContext) -> R:
+        if self.is_async_function:
+            raise TypeError(f"{self.name} is not a sync function. Use `acall` instead.")
+
+        workflow_input = self.workflow._get_workflow_input(ctx)
+
+        if self.is_durable:
+            fn = cast(Callable[[TWorkflowInput, DurableContext], R], self.fn)
+            if is_durable_sync_fn(fn):
+                return fn(workflow_input, cast(DurableContext, ctx))
+        else:
+            fn = cast(Callable[[TWorkflowInput, Context], R], self.fn)
+            if is_sync_fn(fn):
+                return fn(workflow_input, cast(Context, ctx))
+
+        raise TypeError(f"{self.name} is not a sync function. Use `acall` instead.")
+
+    async def aio_call(self, ctx: Context | DurableContext) -> R:
+        if not self.is_async_function:
+            raise TypeError(
+                f"{self.name} is not an async function. Use `call` instead."
+            )
+
+        workflow_input = self.workflow._get_workflow_input(ctx)
+
+        if is_async_fn(self.fn):  # type: ignore
+            return await self.fn(workflow_input, cast(Context, ctx))  # type: ignore
+
+        raise TypeError(f"{self.name} is not an async function. Use `call` instead.")

hatchet_sdk/runnables/types.py

@@ -0,0 +1,138 @@
+import asyncio
+from datetime import timedelta
+from enum import Enum
+from typing import Any, Awaitable, Callable, ParamSpec, Type, TypeGuard, TypeVar, Union
+
+from pydantic import BaseModel, ConfigDict, Field, StrictInt, model_validator
+
+from hatchet_sdk.context.context import Context, DurableContext
+from hatchet_sdk.utils.timedelta_to_expression import Duration
+from hatchet_sdk.utils.typing import JSONSerializableMapping
+
+ValidTaskReturnType = Union[BaseModel, JSONSerializableMapping, None]
+
+R = TypeVar("R", bound=Union[ValidTaskReturnType, Awaitable[ValidTaskReturnType]])
+P = ParamSpec("P")
+
+
+DEFAULT_EXECUTION_TIMEOUT = timedelta(seconds=60)
+DEFAULT_SCHEDULE_TIMEOUT = timedelta(minutes=5)
+DEFAULT_PRIORITY = 1
+
+
+class EmptyModel(BaseModel):
+    model_config = ConfigDict(extra="allow", frozen=True)
+
+
+class StickyStrategy(str, Enum):
+    SOFT = "SOFT"
+    HARD = "HARD"
+
+
+class ConcurrencyLimitStrategy(str, Enum):
+    CANCEL_IN_PROGRESS = "CANCEL_IN_PROGRESS"
+    DROP_NEWEST = "DROP_NEWEST"
+    QUEUE_NEWEST = "QUEUE_NEWEST"
+    GROUP_ROUND_ROBIN = "GROUP_ROUND_ROBIN"
+    CANCEL_NEWEST = "CANCEL_NEWEST"
+
+
+class ConcurrencyExpression(BaseModel):
+    """
+    Defines concurrency limits for a workflow using a CEL expression.
+    Args:
+        expression (str): CEL expression to determine concurrency grouping. (i.e. "input.user_id")
+        max_runs (int): Maximum number of concurrent workflow runs.
+        limit_strategy (ConcurrencyLimitStrategy): Strategy for handling limit violations.
+    Example:
+        ConcurrencyExpression("input.user_id", 5, ConcurrencyLimitStrategy.CANCEL_IN_PROGRESS)
+    """
+
+    expression: str
+    max_runs: int
+    limit_strategy: ConcurrencyLimitStrategy
+
+
+TWorkflowInput = TypeVar("TWorkflowInput", bound=BaseModel)
+
+
+class TaskDefaults(BaseModel):
+    schedule_timeout: Duration = DEFAULT_SCHEDULE_TIMEOUT
+    execution_timeout: Duration = DEFAULT_EXECUTION_TIMEOUT
+    priority: StrictInt = Field(gt=0, lt=4, default=DEFAULT_PRIORITY)
+
+
+class WorkflowConfig(BaseModel):
+    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
+
+    name: str
+    description: str | None = None
+    version: str | None = None
+    on_events: list[str] = Field(default_factory=list)
+    on_crons: list[str] = Field(default_factory=list)
+    sticky: StickyStrategy | None = None
+    concurrency: ConcurrencyExpression | None = None
+    input_validator: Type[BaseModel] = EmptyModel
+
+    task_defaults: TaskDefaults = TaskDefaults()
+
+    @model_validator(mode="after")
+    def validate_concurrency_expression(self) -> "WorkflowConfig":
+        if not self.concurrency:
+            return self
+
+        expr = self.concurrency.expression
+
+        if not expr.startswith("input."):
+            return self
+
+        _, field = expr.split(".", maxsplit=2)
+
+        if field not in self.input_validator.model_fields.keys():
+            raise ValueError(
+                f"The concurrency expression provided relies on the `{field}` field, which was not present in `{self.input_validator.__name__}`."
+            )
+
+        return self
+
+
+class StepType(str, Enum):
+    DEFAULT = "default"
+    ON_FAILURE = "on_failure"
+    ON_SUCCESS = "on_success"
+
+
+AsyncFunc = Callable[[TWorkflowInput, Context], Awaitable[R]]
+SyncFunc = Callable[[TWorkflowInput, Context], R]
+TaskFunc = Union[AsyncFunc[TWorkflowInput, R], SyncFunc[TWorkflowInput, R]]
+
+
+def is_async_fn(
+    fn: TaskFunc[TWorkflowInput, R]
+) -> TypeGuard[AsyncFunc[TWorkflowInput, R]]:
+    return asyncio.iscoroutinefunction(fn)
+
+
+def is_sync_fn(
+    fn: TaskFunc[TWorkflowInput, R]
+) -> TypeGuard[SyncFunc[TWorkflowInput, R]]:
+    return not asyncio.iscoroutinefunction(fn)
+
+
+DurableAsyncFunc = Callable[[TWorkflowInput, DurableContext], Awaitable[R]]
+DurableSyncFunc = Callable[[TWorkflowInput, DurableContext], R]
+DurableTaskFunc = Union[
+    DurableAsyncFunc[TWorkflowInput, R], DurableSyncFunc[TWorkflowInput, R]
+]
+
+
+def is_durable_async_fn(
+    fn: Callable[..., Any]
+) -> TypeGuard[DurableAsyncFunc[TWorkflowInput, R]]:
+    return asyncio.iscoroutinefunction(fn)
+
+
+def is_durable_sync_fn(
+    fn: DurableTaskFunc[TWorkflowInput, R]
+) -> TypeGuard[DurableSyncFunc[TWorkflowInput, R]]:
+    return not asyncio.iscoroutinefunction(fn)