waldiez 0.5.10__py3-none-any.whl → 0.6.1__py3-none-any.whl

waldiez/running/base_runner.py

@@ -2,79 +2,46 @@
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 
 # pyright: reportUnknownMemberType=false, reportAttributeAccessIssue=false
-# pyright: reportUnknownArgumentType=false
+# pyright: reportUnknownArgumentType=false, reportUnusedParameter=false
 """Base runner for Waldiez workflows."""
 
 import importlib.util
-import inspect
+import json
 import shutil
 import sys
 import tempfile
 import threading
 from pathlib import Path
 from types import ModuleType, TracebackType
-from typing import TYPE_CHECKING, Any, Callable, Coroutine, Type, Union
+from typing import Any
 
 from aiofiles.os import wrap
 from anyio.from_thread import start_blocking_portal
-from typing_extensions import Self
+from typing_extensions import Self, override
 
 from waldiez.exporter import WaldiezExporter
 from waldiez.logger import WaldiezLogger, get_logger
 from waldiez.models import Waldiez
 
-from .environment import refresh_environment, reset_env_vars, set_env_vars
+from .dir_utils import a_chdir, chdir
+from .environment import reset_env_vars, set_env_vars
+from .events_mixin import EventsMixin
 from .exceptions import StopRunningException
-from .post_run import after_run
-from .pre_run import RequirementsMixin
 from .protocol import WaldiezRunnerProtocol
-from .utils import (
-    a_chdir,
-    chdir,
-    input_async,
-    input_sync,
-    is_async_callable,
-    syncify,
-)
-
-if TYPE_CHECKING:
-    from autogen.events import BaseEvent  # type: ignore[import-untyped]
-    from autogen.messages import BaseMessage  # type: ignore[import-untyped]
+from .requirements_mixin import RequirementsMixin
+from .results_mixin import ResultsMixin
 
 
 # pylint: disable=too-many-public-methods
-class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
-    """Base runner for Waldiez.
-
-    Initialization parameters:
-        - waldiez: The Waldiez flow to run.
-        - output_path: Path to save the output file.
-        - uploads_root: Root directory for uploads.
-        - structured_io: Whether to use structured I/O.
-        - skip_patch_io: Whether to skip patching I/O functions.
-        - dot_env: Path to a .env file for environment variables.
-
-    Methods to possibly override:
-        - prepare: Prepare the environment and paths for running the flow.
-        - _before_run: Actions to perform before running the flow.
-        - a_prepare: Async version of the prepare method.
-        - _a_before_run: Async actions to perform before running the flow.
-        - _run: Actual implementation of the run logic.
-        - _a_run: Async implementation of the run logic.
-        - _after_run: Actions to perform after running the flow.
-        - _a_after_run: Async actions to perform after running the flow.
-    """
+class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin, ResultsMixin):
+    """Base runner for Waldiez."""
 
     _structured_io: bool
     _output_path: str | Path | None
     _uploads_root: str | Path | None
     _dot_env_path: str | Path | None
-    _skip_patch_io: bool
     _running: bool
-    _is_async: bool
-    _input: Callable[..., str] | Callable[..., Coroutine[Any, Any, str]]
-    _print: Callable[..., None]
-    _send: Callable[[Union["BaseEvent", "BaseMessage"]], None]
+    _waldiez_file: Path
 
     def __init__(
         self,
@@ -82,7 +49,6 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
         output_path: str | Path | None,
         uploads_root: str | Path | None,
         structured_io: bool,
-        skip_patch_io: bool = False,
         dot_env: str | Path | None = None,
         **kwargs: Any,
     ) -> None:
@@ -91,12 +57,12 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
         WaldiezBaseRunner._structured_io = structured_io
         WaldiezBaseRunner._output_path = output_path
         WaldiezBaseRunner._uploads_root = uploads_root
-        WaldiezBaseRunner._skip_patch_io = skip_patch_io
         WaldiezBaseRunner._dot_env_path = dot_env
-        WaldiezBaseRunner._input = input
-        WaldiezBaseRunner._print = print
-        WaldiezBaseRunner._send = print
-        WaldiezBaseRunner._is_async = waldiez.is_async
+        EventsMixin.set_input_function(input)
+        EventsMixin.set_print_function(print)
+        EventsMixin.set_send_function(print)
+        EventsMixin.set_async(waldiez.is_async)
+        RequirementsMixin.__init__(self)
         self._waldiez = waldiez
         self._called_install_requirements = False
         self._exporter = WaldiezExporter(waldiez)
@@ -110,100 +76,46 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
             self._logger = logger
         else:
             self._logger = get_logger()
+        waldiez_file = kwargs.get("waldiez_file", "")
+        if isinstance(waldiez_file, str) and waldiez_file:
+            waldiez_file_path = Path(waldiez_file).resolve()
+        elif isinstance(waldiez_file, Path):
+            waldiez_file_path = waldiez_file.resolve()
+        else:
+            waldiez_file_path = waldiez.dump(to=output_path)
+        if not waldiez_file_path or not waldiez_file_path.is_file():
+            raise ValueError("Could not resolve a waldiez file path")
+        WaldiezBaseRunner._waldiez_file = waldiez_file_path
 
-    def is_running(self) -> bool:
-        """Check if the workflow is currently running.
-
-        Returns
-        -------
-        bool
-            True if the workflow is running, False otherwise.
-        """
-        with self._running_lock:
-            return WaldiezBaseRunner._running
-
-    @staticmethod
-    def get_input_function() -> (
-        Callable[..., str] | Callable[..., Coroutine[Any, Any, str]]
-    ):
-        """Get the input function for user interaction.
-
-        Returns
-        -------
-        Callable[[str, bool], str]
-            A function that takes a prompt and a password flag,
-            returning user input.
-        """
-        if hasattr(WaldiezBaseRunner, "_input") and callable(
-            WaldiezBaseRunner._input
-        ):
-            return WaldiezBaseRunner._input
-        if WaldiezBaseRunner._is_async:
-            return input_async
-        return input_sync
-
+    @override
     @staticmethod
-    async def a_get_user_input(prompt: str, *, password: bool = False) -> str:
-        """Get user input with an optional password prompt.
+    def print(*args: Any, **kwargs: Any) -> None:
+        """Print a message to the output stream.
 
         Parameters
         ----------
-        prompt : str
-            The prompt to display to the user.
-        password : bool, optional
-            If True, the input will be hidden (default is False).
-
-        Returns
-        -------
-        str
-            The user input.
+        *args : Any
+            Positional arguments to print.
+        **kwargs : Any
+            Keyword arguments to print.
         """
-        input_function = WaldiezBaseRunner.get_input_function()
-        if is_async_callable(input_function):
-            try:
-                result = await input_function(  # type: ignore
-                    prompt,
-                    password=password,
-                )
-            except TypeError:
-                result = await input_function(prompt)  # type: ignore
+        if len(args) == 1 and isinstance(args[0], dict):
+            arg = json.dumps(args[0], default=str, ensure_ascii=False)
+            EventsMixin.do_print(arg, **kwargs)
         else:
-            try:
-                result = input_function(prompt, password=password)
-            except TypeError:
-                result = input_function(prompt)
-        return result  # pyright: ignore
-
-    @staticmethod
-    def get_user_input(
-        prompt: str,
-        *,
-        password: bool = False,
-    ) -> str:
-        """Get user input with an optional password prompt.
+            EventsMixin.do_print(*args, **kwargs)
 
-        Parameters
-        ----------
-        prompt : str
-            The prompt to display to the user.
-        password : bool, optional
-            If True, the input will be hidden (default is False).
+    @override
+    def is_running(self) -> bool:
+        """Check if the workflow is currently running.
 
         Returns
         -------
-        str
-            The user input.
+        bool
+            True if the workflow is running, False otherwise.
         """
-        input_function = WaldiezBaseRunner.get_input_function()
-        if inspect.iscoroutinefunction(input_function):
-            try:
-                return syncify(input_function)(prompt, password=password)
-            except TypeError:
-                return syncify(input_function)(prompt)
-        try:
-            return str(input_function(prompt, password=password))
-        except TypeError:
-            return str(input_function(prompt))
+        with self._running_lock:
+            return WaldiezBaseRunner._running
 
     def _load_module(self, output_file: Path, temp_dir: Path) -> ModuleType:
         """Load the module from the waldiez file."""
@@ -230,7 +142,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
     ) -> Path:
         """Run before the flow execution."""
         self.log.info("Preparing workflow file: %s", output_file)
-        temp_dir = Path(tempfile.mkdtemp())
+        temp_dir = Path(tempfile.mkdtemp(prefix="wlz-"))
         file_name = output_file.name
         with chdir(to=temp_dir):
             self._exporter.export(
@@ -300,12 +212,14 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
     def _after_run(
         self,
         results: list[dict[str, Any]],
+        error: BaseException | None,
         output_file: Path,
+        waldiez_file: Path,
         uploads_root: Path | None,
         temp_dir: Path,
         skip_mmd: bool,
         skip_timeline: bool,
-    ) -> None:
+    ) -> Path | None:
         """Run after the flow execution."""
         # Save results
         self._last_results = results
@@ -314,10 +228,13 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
         self._stop_requested.clear()
         # pylint: disable=broad-exception-caught
         try:
-            after_run(
+            return ResultsMixin.post_run(
+                results=results,
+                error=error,
                 temp_dir=temp_dir,
                 output_file=output_file,
                 flow_name=self._waldiez.name,
+                waldiez_file=waldiez_file,
                 uploads_root=uploads_root,
                 skip_mmd=skip_mmd,
                 skip_timeline=skip_timeline,
@@ -325,25 +242,39 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
         except BaseException as exc:  # pragma: no cover
             self.log.warning("Error occurred during after_run: %s", exc)
         self.log.info("Cleanup completed")
+        return None
 
     async def _a_after_run(
         self,
         results: list[dict[str, Any]],
+        error: BaseException | None,
         output_file: Path,
+        waldiez_file: Path,
         uploads_root: Path | None,
         temp_dir: Path,
         skip_mmd: bool,
         skip_timeline: bool,
-    ) -> None:
+    ) -> Path | None:
         """Run after the flow execution asynchronously."""
-        self._after_run(
-            results=results,
-            output_file=output_file,
-            uploads_root=uploads_root,
-            temp_dir=temp_dir,
-            skip_mmd=skip_mmd,
-            skip_timeline=skip_timeline,
-        )
+        self._last_results = results
+        self._stop_requested.clear()
+        # pylint: disable=broad-exception-caught
+        try:
+            return await ResultsMixin.a_post_run(
+                results=results,
+                error=error,
+                temp_dir=temp_dir,
+                output_file=output_file,
+                flow_name=self._waldiez.name,
+                waldiez_file=waldiez_file,
+                uploads_root=uploads_root,
+                skip_mmd=skip_mmd,
+                skip_timeline=skip_timeline,
+            )
+        except BaseException as exc:  # pragma: no cover
+            self.log.warning("Error occurred during a_after_run: %s", exc)
+        self.log.info("Cleanup completed")
+        return None
 
     @staticmethod
     def _prepare_paths(
@@ -364,58 +295,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
         output_file: Path = Path(WaldiezBaseRunner._output_path)
         return output_file, uploads_root_path
 
-    @staticmethod
-    async def a_process_event(event: Union["BaseEvent", "BaseMessage"]) -> None:
-        """Process an event or message asynchronously.
-
-        Parameters
-        ----------
-        event : Union[BaseEvent, BaseMessage]
-            The event or message to process.
-        """
-        if hasattr(event, "type"):  # pragma: no branch
-            if event.type == "input_request":
-                prompt = getattr(
-                    event, "prompt", getattr(event.content, "prompt", "> ")
-                )
-                password = getattr(
-                    event,
-                    "password",
-                    getattr(event.content, "password", False),
-                )
-                user_input = await WaldiezBaseRunner.a_get_user_input(
-                    prompt, password=password
-                )
-                await event.content.respond(user_input)
-            else:
-                WaldiezBaseRunner._send(event)
-
-    @staticmethod
-    def process_event(event: Union["BaseEvent", "BaseMessage"]) -> None:
-        """Process an event or message synchronously.
-
-        Parameters
-        ----------
-        event : Union[BaseEvent, BaseMessage]
-            The event or message to process.
-        """
-        if hasattr(event, "type"):  # pragma: no branch
-            if event.type == "input_request":
-                prompt = getattr(
-                    event, "prompt", getattr(event.content, "prompt", "> ")
-                )
-                password = getattr(
-                    event,
-                    "password",
-                    getattr(event.content, "password", False),
-                )
-                user_input = WaldiezBaseRunner.get_user_input(
-                    prompt, password=password
-                )
-                event.content.respond(user_input)
-            else:
-                WaldiezBaseRunner._send(event)
-
+    @override
     def before_run(
         self,
         output_file: Path,
@@ -440,6 +320,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
             uploads_root=uploads_root,
         )
 
+    @override
     async def a_before_run(
         self,
         output_file: Path,
@@ -495,11 +376,11 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
             uploads_root=uploads_root_path,
         )
         self.install_requirements()
-        refresh_environment()
         return temp_dir, output_file, uploads_root_path
 
     # noinspection PyProtocol
     # pylint: disable=too-many-locals,unused-argument
+    @override
     def run(
         self,
         output_path: str | Path | None = None,
@@ -564,8 +445,10 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
             uploads_root=uploads_root,
         )
         WaldiezBaseRunner._running = True
-        results: list[dict[str, Any]]
+        results: list[dict[str, Any]] = []
+        error: BaseException | None = None
         old_env_vars = set_env_vars(self._waldiez.get_flow_env_vars())
+        output_dir = output_file.parent
         try:
             with chdir(to=temp_dir):
                 sys.path.insert(0, str(temp_dir))
@@ -577,25 +460,30 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
                     skip_timeline=skip_timeline,
                 )
         except (SystemExit, StopRunningException, KeyboardInterrupt) as exc:
+            error = exc
+            self.log.warning("Execution stopped: %s", exc)
             raise StopRunningException(StopRunningException.reason) from exc
         except BaseException as exc:  # pylint: disable=broad-exception-caught
             self.log.error("Error occurred while running workflow: %s", exc)
-            results = [{"error": str(exc)}]
+            error = exc
         finally:
             WaldiezBaseRunner._running = False
             reset_env_vars(old_env_vars)
-        self.after_run(
-            results=results,
-            output_file=output_file,
-            uploads_root=uploads_root_path,
-            temp_dir=temp_dir,
-            skip_mmd=skip_mmd,
-            skip_timeline=skip_timeline,
-        )
-        self._print("<Waldiez> - Done running the flow.")
+            output = self.after_run(
+                results=results,
+                error=error,
+                output_file=output_file,
+                uploads_root=uploads_root_path,
+                temp_dir=temp_dir,
+                skip_mmd=skip_mmd,
+                skip_timeline=skip_timeline,
+            )
+            if output:
+                output_dir = output
+        EventsMixin.do_print("<Waldiez> - Done running the flow.")
         if sys.path[0] == str(temp_dir):
             sys.path.pop(0)
-        return results
+        return self.get_results(results, output_dir)
 
     async def a_prepare(
         self,
@@ -625,11 +513,11 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
             uploads_root=uploads_root_path,
         )
         await self.a_install_requirements()
-        refresh_environment()
         return temp_dir, output_file, uploads_root_path
 
     # noinspection DuplicatedCode
     # noinspection PyProtocol
+    @override
     async def a_run(
         self,
         output_path: str | Path | None = None,
@@ -685,7 +573,9 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
             uploads_root=uploads_root,
         )
         WaldiezBaseRunner._running = True
-        results: list[dict[str, Any]]
+        results: list[dict[str, Any]] = []
+        error: BaseException | None = None
+        output_dir = output_file.parent
         old_env_vars = set_env_vars(self._waldiez.get_flow_env_vars())
         try:
             async with a_chdir(to=temp_dir):
@@ -698,39 +588,49 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
                     skip_timeline=skip_timeline,
                 )
         except (SystemExit, StopRunningException, KeyboardInterrupt) as exc:
+            self.log.warning("Execution stopped: %s", exc)
+            error = exc
             raise StopRunningException(StopRunningException.reason) from exc
         except BaseException as exc:  # pylint: disable=broad-exception-caught
-            results = [{"error": str(exc)}]
+            self.log.error("Error occurred while running workflow: %s", exc)
+            error = exc
         finally:
             WaldiezBaseRunner._running = False
             reset_env_vars(old_env_vars)
-        await self._a_after_run(
-            results=results,
-            output_file=output_file,
-            uploads_root=uploads_root_path,
-            temp_dir=temp_dir,
-            skip_mmd=skip_mmd,
-            skip_timeline=skip_timeline,
-        )
+            output = await self.a_after_run(
+                results=results,
+                error=error,
+                output_file=output_file,
+                uploads_root=uploads_root_path,
+                temp_dir=temp_dir,
+                skip_mmd=skip_mmd,
+                skip_timeline=skip_timeline,
+            )
+            if output:
+                output_dir = output
         if sys.path[0] == str(temp_dir):
             sys.path.pop(0)
-        return results
+        return await self.a_get_results(results, output_dir)
 
+    @override
     def after_run(
         self,
         results: list[dict[str, Any]],
+        error: BaseException | None,
         output_file: Path,
         uploads_root: Path | None,
         temp_dir: Path,
         skip_mmd: bool,
         skip_timeline: bool,
-    ) -> None:
+    ) -> Path | None:
         """Actions to perform after running the flow.
 
         Parameters
         ----------
         results : list[dict[str, Any]]
             The results of the flow run.
+        error : BaseException | None
+            Optional error during the run.
         output_file : Path
             The path to the output file.
         uploads_root : Path | None
@@ -741,31 +641,42 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
             Whether to skip generating the mermaid diagram.
         skip_timeline : bool
             Whether to skip generating the timeline JSON.
+
+        Returns
+        -------
+        Path | None
+            The destination directory if output file, else None
         """
-        self._after_run(
+        return self._after_run(
             results=results,
+            error=error,
             output_file=output_file,
+            waldiez_file=WaldiezBaseRunner._waldiez_file,
             uploads_root=uploads_root,
             temp_dir=temp_dir,
             skip_mmd=skip_mmd,
             skip_timeline=skip_timeline,
         )
 
+    @override
     async def a_after_run(
         self,
         results: list[dict[str, Any]],
+        error: BaseException | None,
         output_file: Path,
         uploads_root: Path | None,
         temp_dir: Path,
         skip_mmd: bool,
         skip_timeline: bool,
-    ) -> None:
+    ) -> Path | None:
         """Asynchronously perform actions after running the flow.
 
         Parameters
         ----------
         results : list[dict[str, Any]]
             The results of the flow run.
+        error : BaseException | None
+            Optional error during the run.
         output_file : Path
             The path to the output file.
         uploads_root : Path | None
@@ -776,14 +687,20 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
             Whether to skip generating the mermaid diagram.
         skip_timeline : bool
 
+        Returns
+        -------
+        Path | None
+            The destination directory if output file, else None
         """
-        await self._a_after_run(
+        return await self._a_after_run(
             results=results,
+            error=error,
             output_file=output_file,
             uploads_root=uploads_root,
             temp_dir=temp_dir,
             skip_mmd=skip_mmd,
             skip_timeline=skip_timeline,
+            waldiez_file=WaldiezBaseRunner._waldiez_file,
         )
 
     @property
@@ -791,6 +708,11 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
         """Get the Waldiez instance."""
         return self._waldiez
 
+    @property
+    def waldiez_file(self) -> Path:
+        """Get the path to the waldiez file."""
+        return self._waldiez_file
+
     @property
     def is_async(self) -> bool:
         """Check if the workflow is async."""
@@ -826,11 +748,6 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
         """Get the uploads root path for the runner."""
         return WaldiezBaseRunner._uploads_root
 
-    @property
-    def skip_patch_io(self) -> bool:
-        """Check if the runner is skipping patching IO."""
-        return WaldiezBaseRunner._skip_patch_io
-
     @classmethod
     def load(
         cls,
@@ -889,14 +806,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
         )
 
     def stop(self) -> None:
-        """Stop the workflow execution.
-
-        This method sets the stop flag that will be checked by the event
-        handlers and other parts of the workflow execution to gracefully
-        terminate the workflow execution.
-        Note: Stopping will occur at the "next" AutoGen event, not immediately.
-        """
-        self.log.info("Stop requested - setting stop flag")
+        """Stop the workflow execution."""
         self._stop_requested.set()
 
     def is_stop_requested(self) -> bool:
@@ -909,6 +819,10 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
         """
         return self._stop_requested.is_set()
 
+    def set_stop_requested(self) -> None:
+        """Set the stop requested flag."""
+        self._stop_requested.set()
+
     def __enter__(self) -> Self:
         """Enter the context manager."""
         return self
@@ -919,7 +833,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
 
     def __exit__(
         self,
-        exc_type: Type[BaseException],
+        exc_type: type[BaseException],
         exc_value: BaseException,
         traceback: TracebackType,
     ) -> None:
@@ -929,7 +843,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
 
     async def __aexit__(
         self,
-        exc_type: Type[BaseException],
+        exc_type: type[BaseException],
         exc_value: BaseException,
         traceback: TracebackType,
     ) -> None: