hatchet-sdk 0.42.0__py3-none-any.whl → 0.42.2__py3-none-any.whl

hatchet_sdk/clients/admin.py

@@ -54,16 +54,10 @@ class ChildTriggerWorkflowOptions(TypedDict):
     sticky: bool | None = None
 
 
-class WorkflowRunDict(TypedDict):
-    workflow_name: str
-    input: Any
-    options: Optional[dict]
-
-
 class ChildWorkflowRunDict(TypedDict):
     workflow_name: str
     input: Any
-    options: ChildTriggerWorkflowOptions[dict]
+    options: ChildTriggerWorkflowOptions
     key: str | None = None
 
 
@@ -73,6 +67,12 @@ class TriggerWorkflowOptions(ScheduleTriggerWorkflowOptions, TypedDict):
     namespace: str | None = None
 
 
+class WorkflowRunDict(TypedDict):
+    workflow_name: str
+    input: Any
+    options: TriggerWorkflowOptions | None
+
+
 class DedupeViolationErr(Exception):
     """Raised by the Hatchet library to indicate that a workflow has already been run with this deduplication value."""
 
@@ -260,7 +260,9 @@ class AdminClientAioImpl(AdminClientBase):
 
     @tenacity_retry
     async def run_workflows(
-        self, workflows: List[WorkflowRunDict], options: TriggerWorkflowOptions = None
+        self,
+        workflows: list[WorkflowRunDict],
+        options: TriggerWorkflowOptions | None = None,
     ) -> List[WorkflowRunRef]:
         if len(workflows) == 0:
             raise ValueError("No workflows to run")

hatchet_sdk/clients/dispatcher/action_listener.py

@@ -61,7 +61,7 @@ class Action:
     worker_id: str
     tenant_id: str
     workflow_run_id: str
-    get_group_key_run_id: Optional[str]
+    get_group_key_run_id: str
     job_id: str
     job_name: str
     job_run_id: str

hatchet_sdk/clients/dispatcher/dispatcher.py

@@ -137,14 +137,14 @@ class DispatcherClient:
 
         return response
 
-    def release_slot(self, step_run_id: str):
+    def release_slot(self, step_run_id: str) -> None:
         self.client.ReleaseSlot(
             ReleaseSlotRequest(stepRunId=step_run_id),
             timeout=DEFAULT_REGISTER_TIMEOUT,
             metadata=get_metadata(self.token),
         )
 
-    def refresh_timeout(self, step_run_id: str, increment_by: str):
+    def refresh_timeout(self, step_run_id: str, increment_by: str) -> None:
         self.client.RefreshTimeout(
             RefreshTimeoutRequest(
                 stepRunId=step_run_id,

hatchet_sdk/clients/rest/tenacity_utils.py

@@ -1,10 +1,15 @@
+from typing import Callable, ParamSpec, TypeVar
+
 import grpc
 import tenacity
 
 from hatchet_sdk.logger import logger
 
+P = ParamSpec("P")
+R = TypeVar("R")
+
 
-def tenacity_retry(func):
+def tenacity_retry(func: Callable[P, R]) -> Callable[P, R]:
     return tenacity.retry(
         reraise=True,
         wait=tenacity.wait_exponential_jitter(),

hatchet_sdk/context/context.py

@@ -2,7 +2,10 @@ import inspect
 import json
 import traceback
 from concurrent.futures import Future, ThreadPoolExecutor
-from typing import List
+from typing import Any, Generic, Type, TypeVar, cast, overload
+from warnings import warn
+
+from pydantic import BaseModel, StrictStr
 
 from hatchet_sdk.clients.events import EventClient
 from hatchet_sdk.clients.rest.tenacity_utils import tenacity_retry
@@ -10,11 +13,13 @@ 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
-from hatchet_sdk.contracts.workflows_pb2 import (
+from hatchet_sdk.contracts.dispatcher_pb2 import OverridesData  # type: ignore
+from hatchet_sdk.contracts.workflows_pb2 import (  # type: ignore[attr-defined]
     BulkTriggerWorkflowRequest,
     TriggerWorkflowRequest,
 )
+from hatchet_sdk.utils.types import WorkflowValidator
+from hatchet_sdk.utils.typing import is_basemodel_subclass
 from hatchet_sdk.workflow_run import WorkflowRunRef
 
 from ..clients.admin import (
@@ -24,25 +29,34 @@ from ..clients.admin import (
     TriggerWorkflowOptions,
     WorkflowRunDict,
 )
-from ..clients.dispatcher.dispatcher import Action, DispatcherClient
+from ..clients.dispatcher.dispatcher import (  # type: ignore[attr-defined]
+    Action,
+    DispatcherClient,
+)
 from ..logger import logger
 
 DEFAULT_WORKFLOW_POLLING_INTERVAL = 5  # Seconds
 
+T = TypeVar("T", bound=BaseModel)
 
-def get_caller_file_path():
+
+def get_caller_file_path() -> str:
     caller_frame = inspect.stack()[2]
 
     return caller_frame.filename
 
 
 class BaseContext:
+
+    action: Action
+    spawn_index: int
+
     def _prepare_workflow_options(
         self,
-        key: str = None,
+        key: str | None = None,
         options: ChildTriggerWorkflowOptions | None = None,
-        worker_id: str = None,
-    ):
+        worker_id: str | None = None,
+    ) -> TriggerWorkflowOptions:
         workflow_run_id = self.action.workflow_run_id
         step_run_id = self.action.step_run_id
 
@@ -54,7 +68,8 @@ class BaseContext:
         if options is not None and "additional_metadata" in options:
             meta = options["additional_metadata"]
 
-        trigger_options: TriggerWorkflowOptions = {
+        ## TODO: Pydantic here to simplify this
+        trigger_options: TriggerWorkflowOptions = {  # type: ignore[typeddict-item]
             "parent_id": workflow_run_id,
             "parent_step_run_id": step_run_id,
             "child_key": key,
@@ -95,9 +110,9 @@ class ContextAioImpl(BaseContext):
     async def spawn_workflow(
         self,
         workflow_name: str,
-        input: dict = {},
-        key: str = None,
-        options: ChildTriggerWorkflowOptions = None,
+        input: dict[str, Any] = {},
+        key: str | None = None,
+        options: ChildTriggerWorkflowOptions | None = None,
     ) -> WorkflowRunRef:
         worker_id = self.worker.id()
         # if (
@@ -118,15 +133,15 @@ class ContextAioImpl(BaseContext):
 
     @tenacity_retry
     async def spawn_workflows(
-        self, child_workflow_runs: List[ChildWorkflowRunDict]
-    ) -> List[WorkflowRunRef]:
+        self, child_workflow_runs: list[ChildWorkflowRunDict]
+    ) -> list[WorkflowRunRef]:
 
         if len(child_workflow_runs) == 0:
             raise Exception("no child workflows to spawn")
 
         worker_id = self.worker.id()
 
-        bulk_trigger_workflow_runs: WorkflowRunDict = []
+        bulk_trigger_workflow_runs: list[WorkflowRunDict] = []
         for child_workflow_run in child_workflow_runs:
             workflow_name = child_workflow_run["workflow_name"]
             input = child_workflow_run["input"]
@@ -134,7 +149,8 @@ class ContextAioImpl(BaseContext):
             key = child_workflow_run.get("key")
             options = child_workflow_run.get("options", {})
 
-            trigger_options = self._prepare_workflow_options(key, options, worker_id)
+            ## TODO: figure out why this is failing
+            trigger_options = self._prepare_workflow_options(key, options, worker_id)  # type: ignore[arg-type]
 
             bulk_trigger_workflow_runs.append(
                 WorkflowRunDict(
@@ -161,8 +177,10 @@ class Context(BaseContext):
         workflow_run_event_listener: RunEventListenerClient,
         worker: WorkerContext,
         namespace: str = "",
+        validator_registry: dict[str, WorkflowValidator] = {},
     ):
         self.worker = worker
+        self.validator_registry = validator_registry
 
         self.aio = ContextAioImpl(
             action,
@@ -179,11 +197,11 @@ class Context(BaseContext):
         # Check the type of action.action_payload before attempting to load it as JSON
         if isinstance(action.action_payload, (str, bytes, bytearray)):
             try:
-                self.data = json.loads(action.action_payload)
+                self.data = cast(dict[str, Any], json.loads(action.action_payload))
             except Exception as e:
                 logger.error(f"Error parsing action payload: {e}")
                 # Assign an empty dictionary if parsing fails
-                self.data = {}
+                self.data: dict[str, Any] = {}  # type: ignore[no-redef]
         else:
             # Directly assign the payload to self.data if it's already a dict
             self.data = (
@@ -191,6 +209,7 @@ class Context(BaseContext):
             )
 
         self.action = action
+
         # FIXME: stepRunId is a legacy field, we should remove it
         self.stepRunId = action.step_run_id
 
@@ -218,33 +237,53 @@ class Context(BaseContext):
         else:
             self.input = self.data.get("input", {})
 
-    def step_output(self, step: str):
+    def step_output(self, step: str) -> dict[str, Any] | BaseModel:
+        validators = self.validator_registry.get(step)
+
         try:
-            return self.data["parents"][step]
+            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):
+            return v.model_validate(parent_step_data)
+
+        return parent_step_data
+
     def triggered_by_event(self) -> bool:
-        return self.data.get("triggered_by", "") == "event"
+        return cast(str, self.data.get("triggered_by", "")) == "event"
+
+    def workflow_input(self) -> dict[str, Any] | T:
+        if (r := self.validator_registry.get(self.action.action_id)) and (
+            i := r.workflow_input
+        ):
+            return cast(
+                T,
+                i.model_validate(self.input),
+            )
 
-    def workflow_input(self):
         return self.input
 
-    def workflow_run_id(self):
+    def workflow_run_id(self) -> str:
         return self.action.workflow_run_id
 
-    def cancel(self):
+    def cancel(self) -> None:
         logger.debug("cancelling step...")
         self.exit_flag = True
 
     # done returns true if the context has been cancelled
-    def done(self):
+    def done(self) -> bool:
         return self.exit_flag
 
-    def playground(self, name: str, default: str = None):
+    def playground(self, name: str, default: str | None = None) -> str | None:
         # if the key exists in the overrides_data field, return the value
         if name in self.overrides_data:
-            return self.overrides_data[name]
+            warn(
+                "Use of `overrides_data` is deprecated.",
+                DeprecationWarning,
+                stacklevel=1,
+            )
+            return str(self.overrides_data[name])
 
         caller_file = get_caller_file_path()
 
@@ -259,7 +298,7 @@ class Context(BaseContext):
 
         return default
 
-    def _log(self, line: str) -> (bool, Exception):  # type: ignore
+    def _log(self, line: str) -> tuple[bool, Exception | None]:
         try:
             self.event_client.log(message=line, step_run_id=self.stepRunId)
             return True, None
@@ -267,7 +306,7 @@ class Context(BaseContext):
             # we don't want to raise an exception here, as it will kill the log thread
             return False, e
 
-    def log(self, line, raise_on_error: bool = False):
+    def log(self, line: Any, raise_on_error: bool = False) -> None:
         if self.stepRunId == "":
             return
 
@@ -277,9 +316,9 @@ class Context(BaseContext):
             except Exception:
                 line = str(line)
 
-        future: Future = self.logger_thread_pool.submit(self._log, line)
+        future = self.logger_thread_pool.submit(self._log, line)
 
-        def handle_result(future: Future):
+        def handle_result(future: Future[tuple[bool, Exception | None]]) -> None:
             success, exception = future.result()
             if not success and exception:
                 if raise_on_error:
@@ -297,22 +336,22 @@ class Context(BaseContext):
 
         future.add_done_callback(handle_result)
 
-    def release_slot(self):
+    def release_slot(self) -> None:
         return self.dispatcher_client.release_slot(self.stepRunId)
 
-    def _put_stream(self, data: str | bytes):
+    def _put_stream(self, data: str | bytes) -> None:
         try:
             self.event_client.stream(data=data, step_run_id=self.stepRunId)
         except Exception as e:
             logger.error(f"Error putting stream event: {e}")
 
-    def put_stream(self, data: str | bytes):
+    def put_stream(self, data: str | bytes) -> None:
         if self.stepRunId == "":
             return
 
         self.stream_event_thread_pool.submit(self._put_stream, data)
 
-    def refresh_timeout(self, increment_by: str):
+    def refresh_timeout(self, increment_by: str) -> None:
         try:
             return self.dispatcher_client.refresh_timeout(
                 step_run_id=self.stepRunId, increment_by=increment_by
@@ -320,28 +359,28 @@ class Context(BaseContext):
         except Exception as e:
             logger.error(f"Error refreshing timeout: {e}")
 
-    def retry_count(self):
+    def retry_count(self) -> int:
         return self.action.retry_count
 
-    def additional_metadata(self):
+    def additional_metadata(self) -> dict[str, Any] | None:
         return self.action.additional_metadata
 
-    def child_index(self):
+    def child_index(self) -> int | None:
         return self.action.child_workflow_index
 
-    def child_key(self):
+    def child_key(self) -> str | None:
         return self.action.child_workflow_key
 
-    def parent_workflow_run_id(self):
+    def parent_workflow_run_id(self) -> str | None:
         return self.action.parent_workflow_run_id
 
-    def fetch_run_failures(self):
+    def fetch_run_failures(self) -> list[dict[str, StrictStr]]:
         data = self.rest_client.workflow_run_get(self.action.workflow_run_id)
         other_job_runs = [
-            run for run in data.job_runs if run.job_id != self.action.job_id
+            run for run in (data.job_runs or []) if run.job_id != self.action.job_id
         ]
         # TODO: Parse Step Runs using a Pydantic Model rather than a hand crafted dictionary
-        failed_step_runs = [
+        return [
             {
                 "step_id": step_run.step_id,
                 "step_run_action_name": step_run.step.action,
@@ -350,7 +389,5 @@ class Context(BaseContext):
             for job_run in other_job_runs
             if job_run.step_runs
             for step_run in job_run.step_runs
-            if step_run.error
+            if step_run.error and step_run.step
         ]
-
-        return failed_step_runs

hatchet_sdk/context/worker_context.py

@@ -21,7 +21,7 @@ class WorkerContext:
         await self.client.async_upsert_worker_labels(self._worker_id, labels)
         self._labels.update(labels)
 
-    def id(self):
+    def id(self) -> str:
         return self._worker_id
 
     # def has_workflow(self, workflow_name: str):

hatchet_sdk/hatchet.py

@@ -1,7 +1,8 @@
 import asyncio
 import logging
-from typing import Any, Callable, Optional, ParamSpec, TypeVar
+from typing import Any, Callable, Optional, Type, TypeVar, cast, get_type_hints
 
+from pydantic import BaseModel
 from typing_extensions import deprecated
 
 from hatchet_sdk.clients.rest_client import RestApi
@@ -27,14 +28,17 @@ from .clients.events import EventClient
 from .clients.run_event_listener import RunEventListenerClient
 from .logger import logger
 from .worker.worker import Worker
-from .workflow import ConcurrencyExpression, WorkflowMeta
+from .workflow import (
+    ConcurrencyExpression,
+    WorkflowInterface,
+    WorkflowMeta,
+    WorkflowStepProtocol,
+)
 
-P = ParamSpec("P")
-R = TypeVar("R")
+T = TypeVar("T", bound=BaseModel)
 
 
-## TODO: Fix return type here to properly type hint the metaclass
-def workflow(  # type: ignore[no-untyped-def]
+def workflow(
     name: str = "",
     on_events: list[str] | None = None,
     on_crons: list[str] | None = None,
@@ -44,11 +48,12 @@ def workflow(  # type: ignore[no-untyped-def]
     sticky: StickyStrategy = None,
     default_priority: int | None = None,
     concurrency: ConcurrencyExpression | None = None,
-):
+    input_validator: Type[T] | None = None,
+) -> Callable[[Type[WorkflowInterface]], WorkflowMeta]:
     on_events = on_events or []
     on_crons = on_crons or []
 
-    def inner(cls: Any) -> WorkflowMeta:
+    def inner(cls: Type[WorkflowInterface]) -> WorkflowMeta:
         cls.on_events = on_events
         cls.on_crons = on_crons
         cls.name = name or str(cls.__name__)
@@ -62,7 +67,8 @@ def workflow(  # type: ignore[no-untyped-def]
         # with WorkflowMeta as its metaclass
 
         ## TODO: Figure out how to type this metaclass correctly
-        return WorkflowMeta(cls.name, cls.__bases__, dict(cls.__dict__))  # type: ignore[no-untyped-call]
+        cls.input_validator = input_validator
+        return WorkflowMeta(cls.name, cls.__bases__, dict(cls.__dict__))
 
     return inner
 
@@ -76,10 +82,10 @@ def step(
     desired_worker_labels: dict[str, DesiredWorkerLabel] = {},
     backoff_factor: float | None = None,
     backoff_max_seconds: int | None = None,
-) -> Callable[[Callable[P, R]], Callable[P, R]]:
+) -> Callable[[WorkflowStepProtocol], WorkflowStepProtocol]:
     parents = parents or []
 
-    def inner(func: Callable[P, R]) -> Callable[P, R]:
+    def inner(func: WorkflowStepProtocol) -> WorkflowStepProtocol:
         limits = None
         if rate_limits:
             limits = [
@@ -87,20 +93,19 @@ def step(
                 for rate_limit in rate_limits or []
             ]
 
-        ## TODO: Use Protocol here to help with MyPy errors
-        func._step_name = name.lower() or str(func.__name__).lower()  # type: ignore[attr-defined]
-        func._step_parents = parents  # type: ignore[attr-defined]
-        func._step_timeout = timeout  # type: ignore[attr-defined]
-        func._step_retries = retries  # type: ignore[attr-defined]
-        func._step_rate_limits = limits  # type: ignore[attr-defined]
-        func._step_backoff_factor = backoff_factor  # type: ignore[attr-defined]
-        func._step_backoff_max_seconds = backoff_max_seconds  # type: ignore[attr-defined]
+        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 = {}  # type: ignore[attr-defined]
+        func._step_desired_worker_labels = {}
 
         for key, d in desired_worker_labels.items():
             value = d["value"] if "value" in d else None
-            func._step_desired_worker_labels[key] = DesiredWorkerLabels(  # type: ignore[attr-defined]
+            func._step_desired_worker_labels[key] = 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,
@@ -120,8 +125,8 @@ def on_failure_step(
     rate_limits: list[RateLimit] | None = None,
     backoff_factor: float | None = None,
     backoff_max_seconds: int | None = None,
-) -> Callable[[Callable[P, R]], Callable[P, R]]:
-    def inner(func: Callable[P, R]) -> Callable[P, R]:
+) -> Callable[[WorkflowStepProtocol], WorkflowStepProtocol]:
+    def inner(func: WorkflowStepProtocol) -> WorkflowStepProtocol:
         limits = None
         if rate_limits:
             limits = [
@@ -129,13 +134,12 @@ def on_failure_step(
                 for rate_limit in rate_limits or []
             ]
 
-        ## TODO: Use Protocol here to help with MyPy errors
-        func._on_failure_step_name = name.lower() or str(func.__name__).lower()  # type: ignore[attr-defined]
-        func._on_failure_step_timeout = timeout  # type: ignore[attr-defined]
-        func._on_failure_step_retries = retries  # type: ignore[attr-defined]
-        func._on_failure_step_rate_limits = limits  # type: ignore[attr-defined]
-        func._on_failure_step_backoff_factor = backoff_factor  # type: ignore[attr-defined]
-        func._on_failure_step_backoff_max_seconds = backoff_max_seconds  # type: ignore[attr-defined]
+        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
 
         return func
 
@@ -146,12 +150,11 @@ def concurrency(
     name: str = "",
     max_runs: int = 1,
     limit_strategy: ConcurrencyLimitStrategy = ConcurrencyLimitStrategy.CANCEL_IN_PROGRESS,
-) -> Callable[[Callable[P, R]], Callable[P, R]]:
-    def inner(func: Callable[P, R]) -> Callable[P, R]:
-        ## TODO: Use Protocol here to help with MyPy errors
-        func._concurrency_fn_name = name.lower() or str(func.__name__).lower()  # type: ignore[attr-defined]
-        func._concurrency_max_runs = max_runs  # type: ignore[attr-defined]
-        func._concurrency_limit_strategy = limit_strategy  # type: ignore[attr-defined]
+) -> 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
 
         return func
 

hatchet_sdk/utils/backoff.py

@@ -2,7 +2,7 @@ import asyncio
 import random
 
 
-async def exp_backoff_sleep(attempt: int, max_sleep_time: float = 5):
+async def exp_backoff_sleep(attempt: int, max_sleep_time: float = 5) -> None:
     base_time = 0.1  # starting sleep time in seconds (100 milliseconds)
     jitter = random.uniform(0, base_time)  # add random jitter
     sleep_time = min(base_time * (2**attempt) + jitter, max_sleep_time)

hatchet_sdk/utils/serialization.py

@@ -2,7 +2,10 @@ from typing import Any
 
 
 def flatten(xs: dict[str, Any], parent_key: str, separator: str) -> dict[str, Any]:
-    items = []
+    if not xs:
+        return {}
+
+    items: list[tuple[str, Any]] = []
 
     for k, v in xs.items():
         new_key = parent_key + separator + k if parent_key else k

hatchet_sdk/utils/tracing.py

@@ -6,9 +6,9 @@ from opentelemetry import trace
 from opentelemetry.context import Context
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 from opentelemetry.sdk.resources import SERVICE_NAME, Resource
-from opentelemetry.sdk.trace import Tracer, TracerProvider
+from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
-from opentelemetry.trace import NoOpTracerProvider
+from opentelemetry.trace import NoOpTracerProvider, Tracer
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 
 from hatchet_sdk.loader import ClientConfig
@@ -44,7 +44,7 @@ def create_tracer(config: ClientConfig) -> Tracer:
 
 
 def create_carrier() -> dict[str, str]:
-    carrier = {}
+    carrier: dict[str, str] = {}
     TraceContextTextMapPropagator().inject(carrier)
 
     return carrier
@@ -59,7 +59,10 @@ def inject_carrier_into_metadata(
     return metadata
 
 
-def parse_carrier_from_metadata(metadata: dict[str, Any]) -> Context:
+def parse_carrier_from_metadata(metadata: dict[str, Any] | None) -> Context | None:
+    if not metadata:
+        return None
+
     return (
         TraceContextTextMapPropagator().extract(_ctx)
         if (_ctx := metadata.get(OTEL_CARRIER_KEY))

hatchet_sdk/utils/types.py

@@ -0,0 +1,8 @@
+from typing import Type
+
+from pydantic import BaseModel
+
+
+class WorkflowValidator(BaseModel):
+    workflow_input: Type[BaseModel] | None = None
+    step_output: Type[BaseModel] | None = None

hatchet_sdk/utils/typing.py

@@ -0,0 +1,9 @@
+from typing import Any, Type, TypeGuard, TypeVar
+
+from pydantic import BaseModel
+
+T = TypeVar("T", bound=BaseModel)
+
+
+def is_basemodel_subclass(model: Any) -> bool:
+    return isinstance(model, type) and issubclass(model, BaseModel)

hatchet_sdk/worker/action_listener_process.py

@@ -87,15 +87,13 @@ class WorkerActionListenerProcess:
         try:
             self.dispatcher_client = new_dispatcher(self.config)
 
-            self.listener: ActionListener = (
-                await self.dispatcher_client.get_action_listener(
-                    GetActionListenerRequest(
-                        worker_name=self.name,
-                        services=["default"],
-                        actions=self.actions,
-                        max_runs=self.max_runs,
-                        _labels=self.labels,
-                    )
+            self.listener = await self.dispatcher_client.get_action_listener(
+                GetActionListenerRequest(
+                    worker_name=self.name,
+                    services=["default"],
+                    actions=self.actions,
+                    max_runs=self.max_runs,
+                    _labels=self.labels,
                 )
             )