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

waldiez/running/base_runner.py

@@ -2,11 +2,10 @@
 # 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
@@ -14,72 +13,35 @@ 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, dump_waldiez
 from .protocol import WaldiezRunnerProtocol
-from .utils import (
-    a_chdir,
-    chdir,
-    input_async,
-    input_sync,
-    is_async_callable,
-    syncify,
-)
-
-if TYPE_CHECKING:
-    # pylint: disable=line-too-long
-    from autogen.events import (  # type: ignore[import-untyped,import-not-found,unused-ignore]  # noqa: E501
-        BaseEvent,
-    )
-    from autogen.messages import (  # type: ignore[import-untyped,import-not-found,unused-ignore] # noqa: E501
-        BaseMessage,
-    )
+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.
-        - 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
     _running: bool
-    _is_async: bool
     _waldiez_file: Path
-    _input: Callable[..., str] | Callable[..., Coroutine[Any, Any, str]]
-    _print: Callable[..., None]
-    _send: Callable[[Union["BaseEvent", "BaseMessage"]], None]
 
     def __init__(
         self,
@@ -96,10 +58,11 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
         WaldiezBaseRunner._output_path = output_path
         WaldiezBaseRunner._uploads_root = uploads_root
         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)
@@ -119,14 +82,15 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
         elif isinstance(waldiez_file, Path):
             waldiez_file_path = waldiez_file.resolve()
         else:
-            waldiez_file_path = dump_waldiez(waldiez, output_path=output_path)
+            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
 
+    @override
     @staticmethod
     def print(*args: Any, **kwargs: Any) -> None:
-        """Print a message to the console.
+        """Print a message to the output stream.
 
         Parameters
         ----------
@@ -137,10 +101,11 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
         """
         if len(args) == 1 and isinstance(args[0], dict):
             arg = json.dumps(args[0], default=str, ensure_ascii=False)
-            WaldiezBaseRunner._print(arg, **kwargs)
+            EventsMixin.do_print(arg, **kwargs)
         else:
-            WaldiezBaseRunner._print(*args, **kwargs)
+            EventsMixin.do_print(*args, **kwargs)
 
+    @override
     def is_running(self) -> bool:
         """Check if the workflow is currently running.
 
@@ -152,99 +117,6 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
         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
-
-    @staticmethod
-    async def a_get_user_input(
-        prompt: str, *, password: bool = False, **kwargs: Any
-    ) -> str:
-        """Get user input with an optional password prompt.
-
-        Parameters
-        ----------
-        prompt : str
-            The prompt to display to the user.
-        password : bool, optional
-            If True, the input will be hidden (default is False).
-        **kwargs : Any
-            Additional keyword arguments to pass to the input function.
-
-        Returns
-        -------
-        str
-            The user input.
-        """
-        input_function = WaldiezBaseRunner.get_input_function()
-        if is_async_callable(input_function):
-            try:
-                result = await input_function(  # type: ignore
-                    prompt,
-                    password=password,
-                    **kwargs,
-                )
-            except TypeError:
-                result = await input_function(prompt)  # type: ignore
-        else:
-            try:
-                result = input_function(prompt, password=password, **kwargs)
-            except TypeError:
-                result = input_function(prompt)
-        return result  # pyright: ignore
-
-    @staticmethod
-    def get_user_input(
-        prompt: str,
-        *,
-        password: bool = False,
-        **kwargs: Any,
-    ) -> str:
-        """Get user input with an optional password prompt.
-
-        Parameters
-        ----------
-        prompt : str
-            The prompt to display to the user.
-        password : bool, optional
-            If True, the input will be hidden (default is False).
-        **kwargs : Any
-            Additional keyword arguments to pass to the input function.
-
-        Returns
-        -------
-        str
-            The user input.
-        """
-        input_function = WaldiezBaseRunner.get_input_function()
-        if inspect.iscoroutinefunction(input_function):
-            try:
-                return syncify(input_function)(
-                    prompt, password=password, **kwargs
-                )
-            except TypeError:
-                return syncify(input_function)(prompt)
-        try:
-            return str(input_function(prompt, password=password, **kwargs))
-        except TypeError:
-            return str(input_function(prompt))
-
     def _load_module(self, output_file: Path, temp_dir: Path) -> ModuleType:
         """Load the module from the waldiez file."""
         file_name = output_file.name
@@ -270,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(
@@ -340,13 +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
@@ -355,7 +228,9 @@ 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,
@@ -367,27 +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,
-            waldiez_file=waldiez_file,
-        )
+        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(
@@ -408,68 +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"],
-        skip_send: bool = False,
-    ) -> None:
-        """Process an event or message asynchronously.
-
-        Parameters
-        ----------
-        event : Union[BaseEvent, BaseMessage]
-            The event or message to process.
-        skip_send : bool
-            Skip sending the event.
-        """
-        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)
-            elif not skip_send:
-                WaldiezBaseRunner._send(event)
-
-    @staticmethod
-    def process_event(
-        event: Union["BaseEvent", "BaseMessage"],
-        skip_send: bool = False,
-    ) -> None:
-        """Process an event or message synchronously.
-
-        Parameters
-        ----------
-        event : Union[BaseEvent, BaseMessage]
-            The event or message to process.
-        skip_send : bool
-            Skip sending the event.
-        """
-        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)
-            elif not skip_send:
-                WaldiezBaseRunner._send(event)
-
+    @override
     def before_run(
         self,
         output_file: Path,
@@ -494,6 +320,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
             uploads_root=uploads_root,
         )
 
+    @override
     async def a_before_run(
         self,
         output_file: Path,
@@ -549,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,
@@ -618,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))
@@ -631,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,
@@ -679,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,
@@ -739,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):
@@ -752,40 +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,
-            waldiez_file=WaldiezBaseRunner._waldiez_file,
-            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
@@ -796,9 +641,15 @@ 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,
@@ -807,21 +658,25 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
             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
@@ -832,9 +687,14 @@ 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,
@@ -946,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:
@@ -980,7 +833,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol, RequirementsMixin):
 
     def __exit__(
         self,
-        exc_type: Type[BaseException],
+        exc_type: type[BaseException],
         exc_value: BaseException,
         traceback: TracebackType,
     ) -> None:
@@ -990,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:

waldiez/running/dir_utils.py

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Dir related utilities for the waldiez runner."""
+
+import os
+from collections.abc import AsyncIterator, Iterator
+from contextlib import asynccontextmanager, contextmanager
+from pathlib import Path
+
+
+@contextmanager
+def chdir(to: str | Path) -> Iterator[None]:
+    """Change the current working directory in a context.
+
+    Parameters
+    ----------
+    to : str | Path
+        The directory to change to.
+
+    Yields
+    ------
+    Iterator[None]
+        The context manager.
+    """
+    old_cwd = str(os.getcwd())
+    try:
+        os.chdir(to)
+        yield
+    finally:
+        os.chdir(old_cwd)
+
+
+@asynccontextmanager
+async def a_chdir(to: str | Path) -> AsyncIterator[None]:
+    """Asynchronously change the current working directory in a context.
+
+    Parameters
+    ----------
+    to : str | Path
+        The directory to change to.
+
+    Yields
+    ------
+    AsyncIterator[None]
+        The async context manager.
+    """
+    old_cwd = str(os.getcwd())
+    try:
+        os.chdir(to)
+        yield
+    finally:
+        os.chdir(old_cwd)