stepfunction 0.0.5__tar.gz → 0.0.7__tar.gz

{stepfunction-0.0.5/src/StepFunction.egg-info → stepfunction-0.0.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stepfunction
-Version: 0.0.5
+Version: 0.0.7
 Summary: Step Function Workflow Orchestration Library
 Author: Vineeth Penugonda
 License-Expression: MIT

{stepfunction-0.0.5 → stepfunction-0.0.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "stepfunction"
-version = "0.0.5"
+version = "0.0.7"
 authors = [{ name = "Vineeth Penugonda" }]
 description = "Step Function Workflow Orchestration Library"
 readme = "README.md"

{stepfunction-0.0.5 → stepfunction-0.0.7/src/StepFunction.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stepfunction
-Version: 0.0.5
+Version: 0.0.7
 Summary: Step Function Workflow Orchestration Library
 Author: Vineeth Penugonda
 License-Expression: MIT

{stepfunction-0.0.5 → stepfunction-0.0.7}/src/stepfunction/core/step_function/step_function.py

@@ -3,10 +3,9 @@
 Author: Vineeth Penugonda
 """
 
-from asyncio import run as asyncio_run
-from concurrent.futures import ThreadPoolExecutor, as_completed
+from asyncio import gather, get_running_loop
 from inspect import iscoroutinefunction
-from typing import Any, Callable, Dict, Optional, Union
+from typing import Any, Callable, Dict, Optional, Union, cast
 
 from stepfunction.constants.enums import StepFunctionStatus
 from stepfunction.exceptions.step_errors import (
@@ -34,7 +33,7 @@ class StepFunction:
 
     Properties:
         name (str): The name of the step function.
-        steps (StepParams): A dictionary containing the steps of the workflow.
+        steps (Dict[str, StepParams]): A dictionary containing the steps of the workflow.
         last_result (Any): The result of the last step.
         context (Dict[str, Any]): Stores step names and results, including exceptions if any occur.
         status (StepFunctionStatus): The current status of the workflow. Possible values are:
@@ -53,6 +52,10 @@ class StepFunction:
         add_sub_step_function(name, sub_step_function, next_step=None, on_failure=None):
             Add a sub-step function to be executed as a step.
 
+        validate():
+            Validate the workflow configuration. Raises ValueError if the start step is not set
+            or if any next_step or on_failure reference an unknown step.
+
         execute(initial_input=None):
             Execute the workflow starting from the specified start step. This will run each step in sequence or in parallel,
             depending on the step configuration.
@@ -91,11 +94,11 @@ class StepFunction:
         await step_function.execute()
 
     Parallel Example:
-        step_function.add_step("Step1", func1, parallel=True)
+        step_function.add_step("Step1", func1, next_step="ParallelStep")
         step_function.add_step("ParallelStep", {
             "task1": func2,
             "task2": func3
-        }, next_step="Step2", parallel=True)
+        }, parallel=True)
 
     Sub-step function example:
         sub_step_function = StepFunction("SubStepFunction")
@@ -110,11 +113,13 @@ class StepFunction:
     def __init__(self, name: str):
         self.__name = name  # Name of the step function
 
-        self.__steps: StepParams = {}  # Steps of the workflow
+        self.__steps: Dict[str, StepParams] = {}  # Steps of the workflow
         self.__current_step = None  # The current step being executed
 
         self.__last_result = None  # To hold the result of the last step
-        self.__context = {}  # To hold the step names and results of those steps
+        self.__context: Dict[
+            str, Any
+        ] = {}  # To hold the step names and results of those steps
 
         # Status of the step function
         self.__status: StepFunctionStatus = StepFunctionStatus.INITIALIZED
@@ -131,7 +136,7 @@ class StepFunction:
         func: Union[Callable[[Any], Any], Dict[str, Callable[[Any], Any]], BaseStep],
         next_step: Optional[str] = None,
         on_failure: Optional[str] = None,
-        branch: Optional[Dict[Any, str]] = None,
+        branch: Optional[Union[Dict[Any, str], Callable[[Any], str]]] = None,
         parallel: bool = False,
         stop_on_failure: bool = False,
     ):
@@ -140,6 +145,12 @@ class StepFunction:
         if name in self.__steps:
             raise ValueError(f"Step '{name}' already exists in steps")
 
+        if branch is not None and next_step is not None:
+            raise ValueError(
+                f"Step '{name}' cannot have both 'branch' and 'next_step' set. "
+                "Use 'branch' to control routing or 'next_step' for a fixed transition, not both."
+            )
+
         step_type = None
         if isinstance(func, BaseStep):
             step_type = func.step_type
@@ -152,6 +163,7 @@ class StepFunction:
             "branch": branch,
             "parallel": parallel,
             "stop_on_failure": stop_on_failure,
+            "is_sub_step_function": False,
             "step_type": step_type,
         }
 
@@ -184,6 +196,7 @@ class StepFunction:
             "parallel": False,
             "stop_on_failure": False,
             "is_sub_step_function": True,
+            "step_type": None,
         }
 
     def set_start_step(self, name: str):
@@ -193,9 +206,32 @@ class StepFunction:
 
         self.__current_step = name
 
+    def validate(self):
+        """Validate the workflow configuration before execution."""
+
+        if self.__current_step is None:
+            raise ValueError(
+                "No start step set. Call set_start_step() before executing the workflow."
+            )
+
+        for step_name, step in self.__steps.items():
+            if step["next_step"] is not None and step["next_step"] not in self.__steps:
+                raise ValueError(
+                    f"Step '{step_name}' has unknown next_step '{step['next_step']}'."
+                )
+            if (
+                step["on_failure"] is not None
+                and step["on_failure"] not in self.__steps
+            ):
+                raise ValueError(
+                    f"Step '{step_name}' has unknown on_failure '{step['on_failure']}'."
+                )
+
     async def execute(self, initial_input: Any = None):
         """Execute the workflow."""
 
+        self.validate()
+
         self.__status = StepFunctionStatus.RUNNING
 
         self.__logger.debug(
@@ -208,8 +244,9 @@ class StepFunction:
             step = self.__steps[self.__current_step]
             try:
                 if step["parallel"]:
-                    results = self._execute_parallel(
-                        step["func"], step["stop_on_failure"]
+                    results = await self._execute_parallel(
+                        cast(Dict[str, Callable[[Any], Any]], step["func"]),
+                        step["stop_on_failure"],
                     )
 
                     self.__last_result = results
@@ -219,7 +256,9 @@ class StepFunction:
                         f"Parallel step '{self.__current_step}' succeeded with results: {results}"
                     )
                 else:
-                    result = await self._execute_step(step["func"], self.__last_result)
+                    result = await self._execute_step(
+                        cast(Callable[[Any], Any], step["func"]), self.__last_result
+                    )
 
                     self.__last_result = result
                     self.__context[self.__current_step] = result
@@ -234,6 +273,17 @@ class StepFunction:
                     else:
                         next_step = step["branch"].get(self.__last_result)
 
+                    if next_step is None:
+                        self.__logger.warning(
+                            f"Step '{self.__current_step}': branch did not resolve to a next step "
+                            f"for result '{self.__last_result}'. Workflow will end."
+                        )
+                    elif next_step not in self.__steps:
+                        self.__logger.warning(
+                            f"Step '{self.__current_step}': branch resolved to '{next_step}' "
+                            "which does not exist in steps. This will cause a failure."
+                        )
+
                 self.__current_step = next_step or step["next_step"]
 
             except Exception as exc:
@@ -241,7 +291,9 @@ class StepFunction:
                     f"Step '{self.__current_step}' failed. Exception: {exc}"
                 )
 
-                self.__context[self.__current_step] = exc.args[0]
+                exc_value = exc.args[0] if exc.args else exc
+
+                self.__context[cast(str, self.__current_step)] = exc_value
 
                 if step["on_failure"]:
                     self.__logger.exception(
@@ -249,7 +301,7 @@ class StepFunction:
                     )
 
                     self.__current_step = step["on_failure"]
-                    self.__last_result = exc.args[0]
+                    self.__last_result = exc_value
 
                     self.__status = StepFunctionStatus.FAILED
 
@@ -285,47 +337,31 @@ class StepFunction:
         else:
             return func(input_value)
 
-    def _execute_parallel(
+    async def _execute_parallel(
         self, func_dict: Dict[str, Callable[[Any], Any]], stop_on_failure: bool = False
     ):
-        """Execute the steps in parallel."""
+        """Execute the steps in parallel without blocking the event loop."""
+        loop = get_running_loop()
         results = {}
         errors = []
-        should_stop_execution = False
 
-        def _run(func, arg):
+        async def _run_one(func: Callable[[Any], Any]) -> Any:
             if iscoroutinefunction(func):
-                return asyncio_run(func(arg))
-            return func(arg)
-
-        with ThreadPoolExecutor() as executor:
-            futures = {
-                executor.submit(_run, func, self.__last_result): step_name
-                for step_name, func in func_dict.items()
-            }
-
-            for future in as_completed(futures):
-                step_name = futures[future]
+                return await func(self.__last_result)
+            return await loop.run_in_executor(None, func, self.__last_result)
 
-                if should_stop_execution:
-                    break
-
-                try:
-                    result = future.result()
-                    results[step_name] = result
-                except Exception as exc:
-                    self.__logger.exception(
-                        f"Parallel task '{step_name}' failed: {exc}"
-                    )
-
-                    results[step_name] = exc.args[0]
-
-                    errors.append((step_name, exc))
+        task_results = await gather(
+            *[_run_one(func) for func in func_dict.values()],
+            return_exceptions=True,
+        )
 
-                    if stop_on_failure:
-                        should_stop_execution = True
-                        for f in futures:
-                            f.cancel()
+        for step_name, result in zip(func_dict.keys(), task_results):
+            if isinstance(result, Exception):
+                self.__logger.exception(f"Parallel task '{step_name}' failed: {result}")
+                results[step_name] = result.args[0] if result.args else result
+                errors.append((step_name, result))
+            else:
+                results[step_name] = result
 
         if errors:
             self.__logger.error(f"Some parallel tasks failed: {errors}")
@@ -374,7 +410,7 @@ class StepFunction:
     @property
     def steps(self):
         """Returns the steps of the step function."""
-        return self.__steps
+        return self.__steps.copy()
 
     @property
     def last_result(self):
@@ -384,7 +420,7 @@ class StepFunction:
     @property
     def context(self):
         """Returns the context of the step function."""
-        return self.__context
+        return self.__context.copy()
 
     @property
     def status(self):

{stepfunction-0.0.5 → stepfunction-0.0.7}/src/stepfunction/core/visualizer/visualizer.py

@@ -4,6 +4,7 @@ Author: Vineeth Penugonda
 """
 
 from os import getcwd
+from typing import Dict, Optional
 
 from graphviz import Digraph
 
@@ -29,7 +30,7 @@ from stepfunction.types.visualizer_types import RenderStepFunctionParams
 class Visualizer:
     """This class is responsible for visualizing the graph model."""
 
-    def __init__(self, graph_name: str, steps: StepParams = None):
+    def __init__(self, graph_name: str, steps: Optional[Dict[str, StepParams]] = None):
         """Initializes the visualizer."""
 
         self.graph_name = graph_name
@@ -47,10 +48,7 @@ class Visualizer:
             raise ValueError("No steps found to visualize.")
 
         for step_name, step_info in self.__steps.items():
-            if (
-                "is_sub_step_function" in step_info
-                and step_info["is_sub_step_function"]
-            ):
+            if step_info["is_sub_step_function"]:
                 self.__dot.node(
                     step_name,
                     step_name,
@@ -100,7 +98,7 @@ class Visualizer:
                         if step_info["next_step"]:
                             self.__dot.edge(parallel_step_name, step_info["next_step"])
 
-            if step_info.get("branch"):
+            if isinstance(step_info["branch"], dict):
                 for result, next_step in step_info["branch"].items():
                     self.__dot.edge(step_name, next_step, label=f"Branch: {result}")
 

{stepfunction-0.0.5 → stepfunction-0.0.7}/src/stepfunction/exceptions/step_errors.py

@@ -9,6 +9,6 @@ class StepExecutionError(Exception):
 class ParallelStepExecutionError(Exception):
     """Error raised when a parallel step fails."""
 
-    def __init__(self, exc: Exception):
+    def __init__(self, exc: list[tuple[str, Exception]]):
         self.message = f"Parallel step generated an exception: {exc}"
         super().__init__(self.message)

{stepfunction-0.0.5 → stepfunction-0.0.7}/src/stepfunction/steps/retry_step.py

@@ -5,10 +5,11 @@ Author: Vineeth Penugonda
 
 from asyncio import sleep
 from inspect import iscoroutinefunction
-from typing import Any, Callable
+from typing import Any, Callable, Optional
 
 from stepfunction.constants.enums import StepType
 from stepfunction.steps.base import BaseStep
+from stepfunction.utils.logger import setup_logger
 
 
 class RetryStep(BaseStep):
@@ -43,6 +44,7 @@ class RetryStep(BaseStep):
         self.func = func
         self.max_retries = max_retries
         self.delay = delay
+        self.__logger = setup_logger(__name__)
 
     def build(self) -> Callable[[Any], Any]:
         """Return an async function that retries ``func`` on failure."""
@@ -51,7 +53,10 @@ class RetryStep(BaseStep):
         delay = self.delay
 
         async def run(input_value: Any) -> Any:
-            last_exc: Exception = None
+            self.__logger.debug(
+                f"RetryStep - executing with max_retries={max_retries}, delay={delay}s"
+            )
+            last_exc: Optional[Exception] = None
 
             for attempt in range(max_retries + 1):
                 try:
@@ -60,9 +65,14 @@ class RetryStep(BaseStep):
                     return func(input_value)
                 except Exception as exc:
                     last_exc = exc
+                    self.__logger.warning(
+                        f"RetryStep - attempt {attempt + 1}/{max_retries + 1} failed: {exc}"
+                    )
                     if attempt < max_retries:
+                        self.__logger.debug(f"RetryStep - retrying in {delay}s")
                         await sleep(delay)
 
+            self.__logger.error(f"RetryStep - all {max_retries + 1} attempts exhausted")
             raise last_exc
 
         return run

{stepfunction-0.0.5 → stepfunction-0.0.7}/src/stepfunction/steps/timeout_step.py

@@ -11,6 +11,7 @@ from typing import Any, Callable
 from stepfunction.constants.enums import StepType
 from stepfunction.steps.base import BaseStep
 from stepfunction.steps.exceptions import StepTimeoutError
+from stepfunction.utils.logger import setup_logger
 
 
 class TimeoutStep(BaseStep):
@@ -39,6 +40,7 @@ class TimeoutStep(BaseStep):
 
         self.func = func
         self.timeout = timeout
+        self.__logger = setup_logger(__name__)
 
     def build(self) -> Callable[[Any], Any]:
         """Return an async function that enforces a timeout on ``func``."""
@@ -46,6 +48,7 @@ class TimeoutStep(BaseStep):
         timeout = self.timeout
 
         async def run(input_value: Any) -> Any:
+            self.__logger.debug(f"TimeoutStep - executing with timeout of {timeout}s")
             try:
                 if iscoroutinefunction(func):
                     return await wait_for(func(input_value), timeout=timeout)
@@ -56,6 +59,7 @@ class TimeoutStep(BaseStep):
                         timeout=timeout,
                     )
             except AsyncTimeoutError:
+                self.__logger.error(f"TimeoutStep - exceeded timeout of {timeout}s")
                 raise StepTimeoutError(timeout)
 
         return run

{stepfunction-0.0.5 → stepfunction-0.0.7}/src/stepfunction/steps/wait_step.py

@@ -8,6 +8,7 @@ from typing import Any, Callable
 
 from stepfunction.constants.enums import StepType
 from stepfunction.steps.base import BaseStep
+from stepfunction.utils.logger import setup_logger
 
 
 class WaitStep(BaseStep):
@@ -29,13 +30,16 @@ class WaitStep(BaseStep):
         if duration < 0:
             raise ValueError("duration must be a non-negative number")
         self.duration = duration
+        self.__logger = setup_logger(__name__)
 
     def build(self) -> Callable[[Any], Any]:
         """Return an async function that sleeps for ``duration`` seconds."""
         duration = self.duration
 
         async def wait(input_value: Any) -> Any:
+            self.__logger.debug(f"WaitStep - sleeping for {duration}s")
             await sleep(duration)
+            self.__logger.debug(f"WaitStep - resumed after {duration}s")
             return input_value
 
         return wait

{stepfunction-0.0.5 → stepfunction-0.0.7}/src/stepfunction/types/step_types.py

@@ -3,8 +3,8 @@ from typing import Any, Callable, Dict, Optional, TypedDict, Union
 from stepfunction.constants.enums import StepType
 
 
-class StepParams(TypedDict, total=False):
-    func: Callable[[Any], Any]
+class StepParams(TypedDict):
+    func: Union[Callable[[Any], Any], Dict[str, Callable[[Any], Any]]]
     next_step: Optional[str]
     on_failure: Optional[str]
     branch: Optional[Union[Dict[Any, str], Callable[[Any], Optional[str]]]]