hatchet-sdk 0.42.3__tar.gz → 0.42.4__tar.gz

{hatchet_sdk-0.42.3 → hatchet_sdk-0.42.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: hatchet-sdk
-Version: 0.42.3
+Version: 0.42.4
 Summary: 
 Author: Alexander Belanger
 Author-email: alexander@hatchet.run

{hatchet_sdk-0.42.3 → hatchet_sdk-0.42.4}/hatchet_sdk/clients/admin.py

@@ -41,7 +41,7 @@ def new_admin(config: ClientConfig):
     return AdminClient(config)
 
 
-class ScheduleTriggerWorkflowOptions(TypedDict):
+class ScheduleTriggerWorkflowOptions(TypedDict, total=False):
     parent_id: Optional[str]
     parent_step_run_id: Optional[str]
     child_index: Optional[int]
@@ -49,25 +49,25 @@ class ScheduleTriggerWorkflowOptions(TypedDict):
     namespace: Optional[str]
 
 
-class ChildTriggerWorkflowOptions(TypedDict):
+class ChildTriggerWorkflowOptions(TypedDict, total=False):
     additional_metadata: Dict[str, str] | None = None
     sticky: bool | None = None
 
 
-class ChildWorkflowRunDict(TypedDict):
+class ChildWorkflowRunDict(TypedDict, total=False):
     workflow_name: str
     input: Any
     options: ChildTriggerWorkflowOptions
     key: str | None = None
 
 
-class TriggerWorkflowOptions(ScheduleTriggerWorkflowOptions, TypedDict):
+class TriggerWorkflowOptions(ScheduleTriggerWorkflowOptions, total=False):
     additional_metadata: Dict[str, str] | None = None
     desired_worker_id: str | None = None
     namespace: str | None = None
 
 
-class WorkflowRunDict(TypedDict):
+class WorkflowRunDict(TypedDict, total=False):
     workflow_name: str
     input: Any
     options: TriggerWorkflowOptions | None
@@ -425,7 +425,7 @@ class AdminClient(AdminClientBase):
         self,
         key: str,
         limit: int,
-        duration: RateLimitDuration = RateLimitDuration.SECOND,
+        duration: Union[RateLimitDuration.Value, str] = RateLimitDuration.SECOND,
     ):
         try:
             self.client.PutRateLimit(

{hatchet_sdk-0.42.3 → hatchet_sdk-0.42.4}/hatchet_sdk/clients/dispatcher/dispatcher.py

@@ -1,3 +1,5 @@
+from typing import Any, cast
+
 from google.protobuf.timestamp_pb2 import Timestamp
 
 from hatchet_sdk.clients.dispatcher.action_listener import (
@@ -31,7 +33,7 @@ from ...metadata import get_metadata
 DEFAULT_REGISTER_TIMEOUT = 30
 
 
-def new_dispatcher(config: ClientConfig):
+def new_dispatcher(config: ClientConfig) -> "DispatcherClient":
     return DispatcherClient(config=config)
 
 
@@ -40,10 +42,10 @@ class DispatcherClient:
 
     def __init__(self, config: ClientConfig):
         conn = new_conn(config)
-        self.client = DispatcherStub(conn)
+        self.client = DispatcherStub(conn)  # type: ignore[no-untyped-call]
 
         aio_conn = new_conn(config, True)
-        self.aio_client = DispatcherStub(aio_conn)
+        self.aio_client = DispatcherStub(aio_conn)  # type: ignore[no-untyped-call]
         self.token = config.token
         self.config = config
 
@@ -67,7 +69,7 @@ class DispatcherClient:
 
     async def send_step_action_event(
         self, action: Action, event_type: StepActionEventType, payload: str
-    ):
+    ) -> Any:
         try:
             return await self._try_send_step_action_event(action, event_type, payload)
         except Exception as e:
@@ -87,7 +89,7 @@ class DispatcherClient:
     @tenacity_retry
     async def _try_send_step_action_event(
         self, action: Action, event_type: StepActionEventType, payload: str
-    ):
+    ) -> Any:
         eventTimestamp = Timestamp()
         eventTimestamp.GetCurrentTime()
 
@@ -103,6 +105,7 @@ class DispatcherClient:
             eventPayload=payload,
         )
 
+        ## TODO: What does this return?
         return await self.aio_client.SendStepActionEvent(
             event,
             metadata=get_metadata(self.token),
@@ -110,7 +113,7 @@ class DispatcherClient:
 
     async def send_group_key_action_event(
         self, action: Action, event_type: GroupKeyActionEventType, payload: str
-    ):
+    ) -> Any:
         eventTimestamp = Timestamp()
         eventTimestamp.GetCurrentTime()
 
@@ -124,19 +127,21 @@ class DispatcherClient:
             eventPayload=payload,
         )
 
+        ## TODO: What does this return?
         return await self.aio_client.SendGroupKeyActionEvent(
             event,
             metadata=get_metadata(self.token),
         )
 
-    def put_overrides_data(self, data: OverridesData):
-        response: ActionEventResponse = self.client.PutOverridesData(
-            data,
-            metadata=get_metadata(self.token),
+    def put_overrides_data(self, data: OverridesData) -> ActionEventResponse:
+        return cast(
+            ActionEventResponse,
+            self.client.PutOverridesData(
+                data,
+                metadata=get_metadata(self.token),
+            ),
         )
 
-        return response
-
     def release_slot(self, step_run_id: str) -> None:
         self.client.ReleaseSlot(
             ReleaseSlotRequest(stepRunId=step_run_id),
@@ -154,7 +159,9 @@ class DispatcherClient:
             metadata=get_metadata(self.token),
         )
 
-    def upsert_worker_labels(self, worker_id: str, labels: dict[str, str | int]):
+    def upsert_worker_labels(
+        self, worker_id: str | None, labels: dict[str, str | int]
+    ) -> None:
         worker_labels = {}
 
         for key, value in labels.items():
@@ -171,9 +178,9 @@ class DispatcherClient:
 
     async def async_upsert_worker_labels(
         self,
-        worker_id: str,
+        worker_id: str | None,
         labels: dict[str, str | int],
-    ):
+    ) -> None:
         worker_labels = {}
 
         for key, value in labels.items():

{hatchet_sdk-0.42.3 → hatchet_sdk-0.42.4}/hatchet_sdk/clients/events.py

@@ -43,16 +43,16 @@ def proto_timestamp_now():
     return timestamp_pb2.Timestamp(seconds=seconds, nanos=nanos)
 
 
-class PushEventOptions(TypedDict):
+class PushEventOptions(TypedDict, total=False):
     additional_metadata: Dict[str, str] | None = None
     namespace: str | None = None
 
 
-class BulkPushEventOptions(TypedDict):
+class BulkPushEventOptions(TypedDict, total=False):
     namespace: str | None = None
 
 
-class BulkPushEventWithMetadata(TypedDict):
+class BulkPushEventWithMetadata(TypedDict, total=False):
     key: str
     payload: Any
     additional_metadata: Optional[Dict[str, Any]]  # Optional metadata

{hatchet_sdk-0.42.3 → hatchet_sdk-0.42.4}/hatchet_sdk/clients/rest/models/workflow_list.py

@@ -32,10 +32,10 @@ class WorkflowList(BaseModel):
     WorkflowList
     """  # noqa: E501
 
-    metadata: Optional[APIResourceMeta] = None
-    rows: Optional[List[Workflow]] = None
-    pagination: Optional[PaginationResponse] = None
-    __properties: ClassVar[List[str]] = ["metadata", "rows", "pagination"]
+    metadata: APIResourceMeta | None = None
+    rows: list[Workflow] | None = None
+    pagination: PaginationResponse | None = None
+    __properties: ClassVar[list[str]] = ["metadata", "rows", "pagination"]
 
     model_config = ConfigDict(
         populate_by_name=True,

{hatchet_sdk-0.42.3 → hatchet_sdk-0.42.4}/hatchet_sdk/clients/rest/models/workflow_run.py

@@ -39,28 +39,28 @@ class WorkflowRun(BaseModel):
     metadata: APIResourceMeta
     tenant_id: StrictStr = Field(alias="tenantId")
     workflow_version_id: StrictStr = Field(alias="workflowVersionId")
-    workflow_version: Optional[WorkflowVersion] = Field(
+    workflow_version: WorkflowVersion | None = Field(
         default=None, alias="workflowVersion"
     )
     status: WorkflowRunStatus
-    display_name: Optional[StrictStr] = Field(default=None, alias="displayName")
-    job_runs: Optional[List[JobRun]] = Field(default=None, alias="jobRuns")
+    display_name: StrictStr | None = Field(default=None, alias="displayName")
+    job_runs: list[JobRun] | None = Field(default=None, alias="jobRuns")
     triggered_by: WorkflowRunTriggeredBy = Field(alias="triggeredBy")
-    input: Optional[Dict[str, Any]] = None
-    error: Optional[StrictStr] = None
-    started_at: Optional[datetime] = Field(default=None, alias="startedAt")
-    finished_at: Optional[datetime] = Field(default=None, alias="finishedAt")
-    duration: Optional[StrictInt] = None
+    input: dict[str, Any] | None = None
+    error: StrictStr | None = None
+    started_at: datetime | None = Field(default=None, alias="startedAt")
+    finished_at: datetime | None = Field(default=None, alias="finishedAt")
+    duration: StrictInt | None = None
     parent_id: Optional[
         Annotated[str, Field(min_length=36, strict=True, max_length=36)]
     ] = Field(default=None, alias="parentId")
     parent_step_run_id: Optional[
         Annotated[str, Field(min_length=36, strict=True, max_length=36)]
     ] = Field(default=None, alias="parentStepRunId")
-    additional_metadata: Optional[Dict[str, Any]] = Field(
+    additional_metadata: dict[str, Any] | None = Field(
         default=None, alias="additionalMetadata"
     )
-    __properties: ClassVar[List[str]] = [
+    __properties: ClassVar[list[str]] = [
         "metadata",
         "tenantId",
         "workflowVersionId",

{hatchet_sdk-0.42.3 → hatchet_sdk-0.42.4}/hatchet_sdk/context/context.py

@@ -13,8 +13,8 @@ from hatchet_sdk.clients.rest_client import RestApi
 from hatchet_sdk.clients.run_event_listener import RunEventListenerClient
 from hatchet_sdk.clients.workflow_listener import PooledWorkflowRunListener
 from hatchet_sdk.context.worker_context import WorkerContext
-from hatchet_sdk.contracts.dispatcher_pb2 import OverridesData  # type: ignore
-from hatchet_sdk.contracts.workflows_pb2 import (  # type: ignore[attr-defined]
+from hatchet_sdk.contracts.dispatcher_pb2 import OverridesData
+from hatchet_sdk.contracts.workflows_pb2 import (
     BulkTriggerWorkflowRequest,
     TriggerWorkflowRequest,
 )
@@ -69,7 +69,7 @@ class BaseContext:
             meta = options["additional_metadata"]
 
         ## TODO: Pydantic here to simplify this
-        trigger_options: TriggerWorkflowOptions = {  # type: ignore[typeddict-item]
+        trigger_options: TriggerWorkflowOptions = {
             "parent_id": workflow_run_id,
             "parent_step_run_id": step_run_id,
             "child_key": key,
@@ -149,8 +149,7 @@ class ContextAioImpl(BaseContext):
             key = child_workflow_run.get("key")
             options = child_workflow_run.get("options", {})
 
-            ## TODO: figure out why this is failing
-            trigger_options = self._prepare_workflow_options(key, options, worker_id)  # type: ignore[arg-type]
+            trigger_options = self._prepare_workflow_options(key, options, worker_id)
 
             bulk_trigger_workflow_runs.append(
                 WorkflowRunDict(
@@ -238,14 +237,17 @@ class Context(BaseContext):
             self.input = self.data.get("input", {})
 
     def step_output(self, step: str) -> dict[str, Any] | BaseModel:
-        validators = self.validator_registry.get(step)
+        workflow_validator = next(
+            (v for k, v in self.validator_registry.items() if k.split(":")[-1] == step),
+            None,
+        )
 
         try:
             parent_step_data = cast(dict[str, Any], self.data["parents"][step])
         except KeyError:
             raise ValueError(f"Step output for '{step}' not found")
 
-        if validators and (v := validators.step_output):
+        if workflow_validator and (v := workflow_validator.step_output):
             return v.model_validate(parent_step_data)
 
         return parent_step_data

{hatchet_sdk-0.42.3 → hatchet_sdk-0.42.4}/hatchet_sdk/context/worker_context.py

@@ -2,7 +2,7 @@ from hatchet_sdk.clients.dispatcher.dispatcher import DispatcherClient
 
 
 class WorkerContext:
-    _worker_id: str = None
+    _worker_id: str | None = None
     _registered_workflow_names: list[str] = []
     _labels: dict[str, str | int] = {}
 
@@ -10,18 +10,18 @@ class WorkerContext:
         self._labels = labels
         self.client = client
 
-    def labels(self):
+    def labels(self) -> dict[str, str | int]:
         return self._labels
 
-    def upsert_labels(self, labels: dict[str, str | int]):
+    def upsert_labels(self, labels: dict[str, str | int]) -> None:
         self.client.upsert_worker_labels(self._worker_id, labels)
         self._labels.update(labels)
 
-    async def async_upsert_labels(self, labels: dict[str, str | int]):
+    async def async_upsert_labels(self, labels: dict[str, str | int]) -> None:
         await self.client.async_upsert_worker_labels(self._worker_id, labels)
         self._labels.update(labels)
 
-    def id(self) -> str:
+    def id(self) -> str | None:
         return self._worker_id
 
     # def has_workflow(self, workflow_name: str):

{hatchet_sdk-0.42.3 → hatchet_sdk-0.42.4}/hatchet_sdk/hatchet.py

@@ -1,15 +1,13 @@
 import asyncio
 import logging
-from typing import Any, Callable, Optional, Type, TypeVar, cast, get_type_hints
+from typing import Any, Callable, Optional, Type, TypeVar, Union
 
 from pydantic import BaseModel
 from typing_extensions import deprecated
 
 from hatchet_sdk.clients.rest_client import RestApi
-
-## TODO: These type stubs need to be updated to mass MyPy, and then we can remove this ignore
-## There are file-level type ignore lines in the corresponding .pyi files.
-from hatchet_sdk.contracts.workflows_pb2 import (  # type: ignore[attr-defined]
+from hatchet_sdk.context.context import Context
+from hatchet_sdk.contracts.workflows_pb2 import (
     ConcurrencyLimitStrategy,
     CreateStepRateLimit,
     DesiredWorkerLabels,
@@ -20,6 +18,7 @@ from hatchet_sdk.features.scheduled import ScheduledClient
 from hatchet_sdk.labels import DesiredWorkerLabel
 from hatchet_sdk.loader import ClientConfig, ConfigLoader
 from hatchet_sdk.rate_limit import RateLimit
+from hatchet_sdk.v2.callable import HatchetCallable
 
 from .client import Client, new_client, new_client_raw
 from .clients.admin import AdminClient
@@ -36,6 +35,7 @@ from .workflow import (
 )
 
 T = TypeVar("T", bound=BaseModel)
+TWorkflow = TypeVar("TWorkflow", bound=object)
 
 
 def workflow(
@@ -45,30 +45,35 @@ def workflow(
     version: str = "",
     timeout: str = "60m",
     schedule_timeout: str = "5m",
-    sticky: StickyStrategy = None,
+    sticky: Union[StickyStrategy.Value, None] = None,  # type: ignore[name-defined]
     default_priority: int | None = None,
     concurrency: ConcurrencyExpression | None = None,
     input_validator: Type[T] | None = None,
-) -> Callable[[Type[WorkflowInterface]], WorkflowMeta]:
+) -> Callable[[Type[TWorkflow]], WorkflowMeta]:
     on_events = on_events or []
     on_crons = on_crons or []
 
-    def inner(cls: Type[WorkflowInterface]) -> WorkflowMeta:
-        cls.on_events = on_events
-        cls.on_crons = on_crons
-        cls.name = name or str(cls.__name__)
-        cls.version = version
-        cls.timeout = timeout
-        cls.schedule_timeout = schedule_timeout
-        cls.sticky = sticky
-        cls.default_priority = default_priority
-        cls.concurrency_expression = concurrency
+    def inner(cls: Type[TWorkflow]) -> WorkflowMeta:
+        nonlocal name
+        name = name or str(cls.__name__)
+
+        setattr(cls, "on_events", on_events)
+        setattr(cls, "on_crons", on_crons)
+        setattr(cls, "name", name)
+        setattr(cls, "version", version)
+        setattr(cls, "timeout", timeout)
+        setattr(cls, "schedule_timeout", schedule_timeout)
+        setattr(cls, "sticky", sticky)
+        setattr(cls, "default_priority", default_priority)
+        setattr(cls, "concurrency_expression", concurrency)
+
         # Define a new class with the same name and bases as the original, but
         # with WorkflowMeta as its metaclass
 
         ## TODO: Figure out how to type this metaclass correctly
-        cls.input_validator = input_validator
-        return WorkflowMeta(cls.name, cls.__bases__, dict(cls.__dict__))
+        setattr(cls, "input_validator", input_validator)
+
+        return WorkflowMeta(name, cls.__bases__, dict(cls.__dict__))
 
     return inner
 
@@ -82,34 +87,39 @@ def step(
     desired_worker_labels: dict[str, DesiredWorkerLabel] = {},
     backoff_factor: float | None = None,
     backoff_max_seconds: int | None = None,
-) -> Callable[[WorkflowStepProtocol], WorkflowStepProtocol]:
+) -> Callable[..., Any]:
     parents = parents or []
 
-    def inner(func: WorkflowStepProtocol) -> WorkflowStepProtocol:
+    def inner(func: Callable[[Context], Any]) -> Callable[[Context], Any]:
         limits = None
         if rate_limits:
             limits = [rate_limit._req for rate_limit in rate_limits or []]
 
-        func._step_name = name.lower() or str(func.__name__).lower()
-        func._step_parents = parents
-        func._step_timeout = timeout
-        func._step_retries = retries
-        func._step_rate_limits = limits
-        func._step_backoff_factor = backoff_factor
-        func._step_backoff_max_seconds = backoff_max_seconds
-
-        func._step_desired_worker_labels = {}
+        setattr(func, "_step_name", name.lower() or str(func.__name__).lower())
+        setattr(func, "_step_parents", parents)
+        setattr(func, "_step_timeout", timeout)
+        setattr(func, "_step_retries", retries)
+        setattr(func, "_step_rate_limits", retries)
+        setattr(func, "_step_rate_limits", limits)
+        setattr(func, "_step_backoff_factor", backoff_factor)
+        setattr(func, "_step_backoff_max_seconds", backoff_max_seconds)
 
-        for key, d in desired_worker_labels.items():
+        def create_label(d: DesiredWorkerLabel) -> DesiredWorkerLabels:
             value = d["value"] if "value" in d else None
-            func._step_desired_worker_labels[key] = DesiredWorkerLabels(
+            return DesiredWorkerLabels(
                 strValue=str(value) if not isinstance(value, int) else None,
                 intValue=value if isinstance(value, int) else None,
-                required=d["required"] if "required" in d else None,
+                required=d["required"] if "required" in d else None,  # type: ignore[arg-type]
                 weight=d["weight"] if "weight" in d else None,
-                comparator=d["comparator"] if "comparator" in d else None,
+                comparator=d["comparator"] if "comparator" in d else None,  # type: ignore[arg-type]
             )
 
+        setattr(
+            func,
+            "_step_desired_worker_labels",
+            {key: create_label(d) for key, d in desired_worker_labels.items()},
+        )
+
         return func
 
     return inner
@@ -122,21 +132,23 @@ def on_failure_step(
     rate_limits: list[RateLimit] | None = None,
     backoff_factor: float | None = None,
     backoff_max_seconds: int | None = None,
-) -> Callable[[WorkflowStepProtocol], WorkflowStepProtocol]:
-    def inner(func: WorkflowStepProtocol) -> WorkflowStepProtocol:
+) -> Callable[..., Any]:
+    def inner(func: Callable[[Context], Any]) -> Callable[[Context], Any]:
         limits = None
         if rate_limits:
             limits = [
-                CreateStepRateLimit(key=rate_limit.static_key, units=rate_limit.units)
+                CreateStepRateLimit(key=rate_limit.static_key, units=rate_limit.units)  # type: ignore[arg-type]
                 for rate_limit in rate_limits or []
             ]
 
-        func._on_failure_step_name = name.lower() or str(func.__name__).lower()
-        func._on_failure_step_timeout = timeout
-        func._on_failure_step_retries = retries
-        func._on_failure_step_rate_limits = limits
-        func._on_failure_step_backoff_factor = backoff_factor
-        func._on_failure_step_backoff_max_seconds = backoff_max_seconds
+        setattr(
+            func, "_on_failure_step_name", name.lower() or str(func.__name__).lower()
+        )
+        setattr(func, "_on_failure_step_timeout", timeout)
+        setattr(func, "_on_failure_step_retries", retries)
+        setattr(func, "_on_failure_step_rate_limits", limits)
+        setattr(func, "_on_failure_step_backoff_factor", backoff_factor)
+        setattr(func, "_on_failure_step_backoff_max_seconds", backoff_max_seconds)
 
         return func
 
@@ -147,11 +159,15 @@ def concurrency(
     name: str = "",
     max_runs: int = 1,
     limit_strategy: ConcurrencyLimitStrategy = ConcurrencyLimitStrategy.CANCEL_IN_PROGRESS,
-) -> Callable[[WorkflowStepProtocol], WorkflowStepProtocol]:
-    def inner(func: WorkflowStepProtocol) -> WorkflowStepProtocol:
-        func._concurrency_fn_name = name.lower() or str(func.__name__).lower()
-        func._concurrency_max_runs = max_runs
-        func._concurrency_limit_strategy = limit_strategy
+) -> Callable[..., Any]:
+    def inner(func: Callable[[Context], Any]) -> Callable[[Context], Any]:
+        setattr(
+            func,
+            "_concurrency_fn_name",
+            name.lower() or str(func.__name__).lower(),
+        )
+        setattr(func, "_concurrency_max_runs", max_runs)
+        setattr(func, "_concurrency_limit_strategy", limit_strategy)
 
         return func
 

{hatchet_sdk-0.42.3 → hatchet_sdk-0.42.4}/hatchet_sdk/labels.py

@@ -1,7 +1,7 @@
 from typing import TypedDict
 
 
-class DesiredWorkerLabel(TypedDict):
+class DesiredWorkerLabel(TypedDict, total=False):
     value: str | int
     required: bool | None = None
     weight: int | None = None

{hatchet_sdk-0.42.3 → hatchet_sdk-0.42.4}/hatchet_sdk/utils/aio_utils.py

@@ -1,10 +1,13 @@
 import asyncio
 import inspect
+from concurrent.futures import Executor
 from functools import partial, wraps
 from threading import Thread
+from typing import Any
 
 
-def sync_to_async(func):
+## TODO: Stricter typing here
+def sync_to_async(func: Any) -> Any:
     """
     A decorator to run a synchronous function or coroutine in an asynchronous context with added
     asyncio loop safety.
@@ -40,8 +43,14 @@ def sync_to_async(func):
         asyncio.run(main())
     """
 
+    ## TODO: Stricter typing here
     @wraps(func)
-    async def run(*args, loop=None, executor=None, **kwargs):
+    async def run(
+        *args: Any,
+        loop: asyncio.AbstractEventLoop | None = None,
+        executor: Executor | None = None,
+        **kwargs: Any
+    ) -> Any:
         """
         The asynchronous wrapper function that runs the given function in an executor.
 
@@ -59,7 +68,7 @@ def sync_to_async(func):
 
         if inspect.iscoroutinefunction(func):
             # Wrap the coroutine to run it in an executor
-            async def wrapper():
+            async def wrapper() -> Any:
                 return await func(*args, **kwargs)
 
             pfunc = partial(asyncio.run, wrapper())
@@ -75,7 +84,7 @@ def sync_to_async(func):
 class EventLoopThread:
     """A class that manages an asyncio event loop running in a separate thread."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         """
         Initializes the EventLoopThread by creating an event loop
         and setting up a thread to run the loop.

{hatchet_sdk-0.42.3 → hatchet_sdk-0.42.4}/hatchet_sdk/utils/typing.py

@@ -6,4 +6,7 @@ T = TypeVar("T", bound=BaseModel)
 
 
 def is_basemodel_subclass(model: Any) -> bool:
-    return isinstance(model, type) and issubclass(model, BaseModel)
+    try:
+        return issubclass(model, BaseModel)
+    except TypeError:
+        return False

{hatchet_sdk-0.42.3 → hatchet_sdk-0.42.4}/hatchet_sdk/v2/callable.py

@@ -1,8 +1,19 @@
 import asyncio
-from typing import Callable, Dict, Generic, List, Optional, TypedDict, TypeVar, Union
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Generic,
+    List,
+    Optional,
+    TypedDict,
+    TypeVar,
+    Union,
+)
 
-from hatchet_sdk.context import Context
-from hatchet_sdk.contracts.workflows_pb2 import (
+from hatchet_sdk.clients.admin import ChildTriggerWorkflowOptions
+from hatchet_sdk.context.context import Context
+from hatchet_sdk.contracts.workflows_pb2 import (  # type: ignore[attr-defined]
     CreateStepRateLimit,
     CreateWorkflowJobOpts,
     CreateWorkflowStepOpts,
@@ -28,8 +39,8 @@ class HatchetCallable(Generic[T]):
         durable: bool = False,
         name: str = "",
         auto_register: bool = True,
-        on_events: list | None = None,
-        on_crons: list | None = None,
+        on_events: list[str] | None = None,
+        on_crons: list[str] | None = None,
         version: str = "",
         timeout: str = "60m",
         schedule_timeout: str = "5m",
@@ -37,8 +48,8 @@ class HatchetCallable(Generic[T]):
         retries: int = 0,
         rate_limits: List[RateLimit] | None = None,
         concurrency: ConcurrencyFunction | None = None,
-        on_failure: Optional["HatchetCallable"] = None,
-        desired_worker_labels: dict[str:DesiredWorkerLabel] = {},
+        on_failure: Union["HatchetCallable[T]", None] = None,
+        desired_worker_labels: dict[str, DesiredWorkerLabel] = {},
         default_priority: int | None = None,
     ):
         self.func = func
@@ -85,7 +96,7 @@ class HatchetCallable(Generic[T]):
     def __call__(self, context: Context) -> T:
         return self.func(context)
 
-    def with_namespace(self, namespace: str):
+    def with_namespace(self, namespace: str) -> None:
         if namespace is not None and namespace != "":
             self.function_namespace = namespace
             self.function_name = namespace + self.function_name
@@ -161,21 +172,13 @@ class HatchetCallable(Generic[T]):
         return self.function_namespace + ":" + self.function_name
 
 
-T = TypeVar("T")
-
-
-class TriggerOptions(TypedDict):
-    additional_metadata: Dict[str, str] | None = None
-    sticky: bool | None = None
-
-
 class DurableContext(Context):
     def run(
         self,
-        function: Union[str, HatchetCallable[T]],
-        input: dict = {},
-        key: str = None,
-        options: TriggerOptions = None,
+        function: str | Callable[[Context], Any],
+        input: dict[Any, Any] = {},
+        key: str | None = None,
+        options: ChildTriggerWorkflowOptions | None = None,
     ) -> "RunRef[T]":
         worker_id = self.worker.id()