waldiez 0.5.2__py3-none-any.whl → 0.5.4__py3-none-any.whl

waldiez/running/base_runner.py

@@ -12,7 +12,15 @@ import tempfile
 import threading
 from pathlib import Path
 from types import TracebackType
-from typing import TYPE_CHECKING, Type, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Coroutine,
+    Optional,
+    Type,
+    Union,
+)
 
 from anyio.from_thread import start_blocking_portal
 from typing_extensions import Self
@@ -20,7 +28,6 @@ from typing_extensions import Self
 from waldiez.exporter import WaldiezExporter
 from waldiez.logger import WaldiezLogger, get_logger
 from waldiez.models import Waldiez
-from waldiez.utils import get_waldiez_version
 
 from .environment import refresh_environment, reset_env_vars, set_env_vars
 from .post_run import after_run
@@ -35,7 +42,11 @@ from .utils import (
 )
 
 if TYPE_CHECKING:
-    from autogen import ChatResult  # type: ignore[import-untyped]
+    from autogen.events import BaseEvent  # type: ignore[import-untyped]
+    from autogen.io.run_response import (  # type: ignore[import-untyped]
+        AsyncRunResponseProtocol,
+        RunResponseProtocol,
+    )
 
 
 class WaldiezBaseRunner(WaldiezRunnerProtocol):
@@ -54,9 +65,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         - _a_stop: Async actions to perform when stopping the flow.
     """
 
-    _threaded: bool
     _structured_io: bool
-    _isolated: bool
     _output_path: str | Path | None
     _uploads_root: str | Path | None
     _skip_patch_io: bool
@@ -68,29 +77,30 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         output_path: str | Path | None,
         uploads_root: str | Path | None,
         structured_io: bool,
-        isolated: bool,
-        threaded: bool,
         skip_patch_io: bool = False,
     ) -> None:
         """Initialize the Waldiez manager."""
         self._waldiez = waldiez
         WaldiezBaseRunner._running = False
         WaldiezBaseRunner._structured_io = structured_io
-        WaldiezBaseRunner._isolated = isolated
         WaldiezBaseRunner._output_path = output_path
         WaldiezBaseRunner._uploads_root = uploads_root
-        WaldiezBaseRunner._threaded = threaded
         WaldiezBaseRunner._skip_patch_io = skip_patch_io
         self._called_install_requirements = False
         self._exporter = WaldiezExporter(waldiez)
         self._stop_requested = threading.Event()
         self._logger = get_logger()
+        self._print: Callable[..., None] = print
+        self._send: Callable[["BaseEvent"], None] = self._print
+        self._input: (
+            Callable[..., str] | Callable[..., Coroutine[Any, Any, str]]
+        ) = input
         self._last_results: Union[
-            "ChatResult",
-            list["ChatResult"],
-            dict[int, "ChatResult"],
+            list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]
         ] = []
         self._last_exception: Exception | None = None
+        self._execution_complete_event = threading.Event()
+        self._execution_thread: Optional[threading.Thread] = None
 
     def is_running(self) -> bool:
         """Check if the workflow is currently running.
@@ -102,9 +112,64 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         """
         return WaldiezBaseRunner._running
 
-    # ===================================================================
-    # PRIVATE METHODS TO OVERRIDE IN SUBCLASSES
-    # ===================================================================
+    def wait_for_completion(self, timeout: Optional[float] = None) -> bool:
+        """Wait for non-blocking execution to complete.
+
+        This is a base implementation that subclasses can override if needed.
+
+        Parameters
+        ----------
+        timeout: float | None
+            The maximum time to wait for completion, in seconds.
+            If None, wait indefinitely.
+
+        Returns
+        -------
+        bool
+            True if the workflow completed successfully, False if it timed out.
+        """
+        # Default implementation - subclasses can override
+        if not self.is_running():
+            return True
+
+        if self._execution_thread and self._execution_thread.is_alive():
+            self._execution_thread.join(timeout)
+            return not self._execution_thread.is_alive()
+
+        return self._execution_complete_event.wait(timeout or 0)
+
+    def get_execution_stats(self) -> dict[str, Any]:
+        """Get basic execution statistics.
+
+        Returns
+        -------
+        dict[str, Any]
+            A dictionary containing execution statistics.
+        - is_running: Whether the runner is currently running.
+        - has_thread: Whether there is an execution thread.
+        - thread_alive: Whether the execution thread is alive.
+        - has_error: Whether there was an error during execution.
+        """
+        return {
+            "is_running": self.is_running(),
+            "has_thread": self._execution_thread is not None,
+            "thread_alive": (
+                self._execution_thread.is_alive()
+                if self._execution_thread
+                else False
+            ),
+            "has_error": self._last_exception is not None,
+        }
+
+    # Helper for subclasses
+    def _signal_completion(self) -> None:
+        """Signal that execution has completed."""
+        self._execution_complete_event.set()
+
+    def _reset_completion_state(self) -> None:
+        """Reset completion state for new execution."""
+        self._execution_complete_event.clear()
+
     def _before_run(
         self,
         output_file: Path,
@@ -119,10 +184,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
                 path=file_name,
                 force=True,
                 uploads_root=uploads_root,
-                # if not isolated, we use structured IO in a context manager
-                structured_io=WaldiezBaseRunner._structured_io
-                and WaldiezBaseRunner._isolated,
-                skip_patch_io=WaldiezBaseRunner._skip_patch_io,
+                structured_io=WaldiezBaseRunner._structured_io,
             )
         return temp_dir
 
@@ -149,11 +211,8 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         output_file: Path,
         uploads_root: Path | None,
         skip_mmd: bool,
-    ) -> Union[
-        "ChatResult",
-        list["ChatResult"],
-        dict[int, "ChatResult"],
-    ]:  # pyright: ignore
+        skip_timeline: bool,
+    ) -> Union[list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]]:
         """Run the Waldiez flow."""
         raise NotImplementedError(
             "The _run method must be implemented in the subclass."
@@ -165,11 +224,8 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         output_file: Path,
         uploads_root: Path | None,
         skip_mmd: bool,
-    ) -> Union[
-        "ChatResult",
-        list["ChatResult"],
-        dict[int, "ChatResult"],
-    ]:  # pyright: ignore
+        skip_timeline: bool,
+    ) -> Union[list["AsyncRunResponseProtocol"], list["RunResponseProtocol"]]:
         """Run the Waldiez flow asynchronously."""
         raise NotImplementedError(
             "The _a_run method must be implemented in the subclass."
@@ -181,6 +237,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         output_file: Path,
         uploads_root: Path | None,
         skip_mmd: bool,
+        skip_timeline: bool,
     ) -> None:
         """Start running the Waldiez flow in a non-blocking way."""
         raise NotImplementedError(
@@ -193,6 +250,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         output_file: Path,
         uploads_root: Path | None,
         skip_mmd: bool,
+        skip_timeline: bool,
     ) -> None:
         """Start running the Waldiez flow in a non-blocking way asynchronously.
 
@@ -216,14 +274,13 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
     def _after_run(
         self,
         results: Union[
-            "ChatResult",
-            list["ChatResult"],
-            dict[int, "ChatResult"],
+            list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]
         ],
         output_file: Path,
         uploads_root: Path | None,
         temp_dir: Path,
         skip_mmd: bool,
+        skip_timeline: bool,
     ) -> None:
         """Run after the flow execution."""
         # Save results
@@ -237,20 +294,20 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             flow_name=self.waldiez.name,
             uploads_root=uploads_root,
             skip_mmd=skip_mmd,
+            skip_timeline=skip_timeline,
         )
         self.log.info("Cleanup completed")
 
     async def _a_after_run(
         self,
         results: Union[
-            "ChatResult",
-            list["ChatResult"],
-            dict[int, "ChatResult"],
+            list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]
         ],
         output_file: Path,
         uploads_root: Path | None,
         temp_dir: Path,
         skip_mmd: bool,
+        skip_timeline: bool,
     ) -> None:
         """Run after the flow execution asynchronously."""
         self._after_run(
@@ -259,6 +316,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             uploads_root=uploads_root,
             temp_dir=temp_dir,
             skip_mmd=skip_mmd,
+            skip_timeline=skip_timeline,
         )
 
     def _stop(self) -> None:
@@ -309,9 +367,6 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             for req in self.waldiez.requirements
             if req not in sys.modules and "waldiez" not in req
         }
-        waldiez_version = get_waldiez_version()
-        if "waldiez" not in sys.modules:
-            extra_requirements.add(f"waldiez=={waldiez_version}")
         return extra_requirements
 
     def install_requirements(self) -> None:
@@ -387,14 +442,9 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         output_path: str | Path | None = None,
         uploads_root: str | Path | None = None,
         structured_io: bool | None = None,
-        threaded: bool | None = None,
-        skip_patch_io: bool | None = None,
         skip_mmd: bool = False,
-    ) -> Union[
-        "ChatResult",
-        list["ChatResult"],
-        dict[int, "ChatResult"],
-    ]:  # pyright: ignore
+        skip_timeline: bool = False,
+    ) -> Union[list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]]:
         """Run the Waldiez flow in blocking mode.
 
         Parameters
@@ -406,31 +456,24 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         structured_io : bool
             Whether to use structured IO instead of the default 'input/print',
             by default False.
-        threaded : bool | None
-            Whether to run the flow in a threaded environment, by default None.
         skip_mmd : bool
             Whether to skip generating the mermaid diagram, by default False.
-        skip_patch_io : bool | None
-            Whether to skip patching the IO streams, by default None.
+        skip_timeline : bool
+            Whether to skip generating the timeline JSON.
 
         Returns
         -------
-        Union[ChatResult, list[ChatResult], dict[int, ChatResult]]
-            The result of the run, which can be a single ChatResult,
-            a list of ChatResults,
-            or a dictionary mapping indices to ChatResults.
+        Union[list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]]
+            The result of the run, which can be a list of RunResponseProtocol
+            or a list of AsyncRunResponseProtocol.
 
         Raises
         ------
         RuntimeError
             If the runner is already running.
         """
-        if skip_patch_io is not None:
-            WaldiezBaseRunner._skip_patch_io = skip_patch_io
         if structured_io is not None:
             WaldiezBaseRunner._structured_io = structured_io
-        if threaded is not None:
-            WaldiezBaseRunner._threaded = threaded
         if self.is_running():
             raise RuntimeError("Workflow already running")
         if self.waldiez.is_async:
@@ -440,7 +483,6 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
                     output_path,
                     uploads_root,
                     structured_io,
-                    skip_patch_io,
                     skip_mmd,
                 )
         output_file, uploads_root_path = self._prepare_paths(
@@ -455,9 +497,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         refresh_environment()
         WaldiezBaseRunner._running = True
         results: Union[
-            "ChatResult",
-            list["ChatResult"],
-            dict[int, "ChatResult"],
+            list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]
         ] = []
         old_env_vars = set_env_vars(self.waldiez.get_flow_env_vars())
         try:
@@ -468,6 +508,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
                     output_file=output_file,
                     uploads_root=uploads_root_path,
                     skip_mmd=skip_mmd,
+                    skip_timeline=skip_timeline,
                 )
         finally:
             WaldiezBaseRunner._running = False
@@ -478,7 +519,9 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             uploads_root=uploads_root_path,
             temp_dir=temp_dir,
             skip_mmd=skip_mmd,
+            skip_timeline=skip_timeline,
         )
+        self._print("<Waldiez> - Done running the flow.")
         if sys.path[0] == str(temp_dir):
             sys.path.pop(0)
         return results
@@ -488,13 +531,9 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         output_path: str | Path | None = None,
         uploads_root: str | Path | None = None,
         structured_io: bool | None = None,
-        skip_patch_io: bool | None = None,
         skip_mmd: bool = False,
-    ) -> Union[
-        "ChatResult",
-        list["ChatResult"],
-        dict[int, "ChatResult"],
-    ]:  # pyright: ignore
+        skip_timeline: bool = False,
+    ) -> Union[list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]]:
         """Run the Waldiez flow asynchronously.
 
         Parameters
@@ -506,25 +545,22 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         structured_io : bool
             Whether to use structured IO instead of the default 'input/print',
             by default False.
-        skip_patch_io : bool | None
-            Whether to skip patching I/O, by default None.
         skip_mmd : bool
             Whether to skip generating the mermaid diagram, by default False.
+        skip_timeline : bool
+            Whether to skip generating the timeline JSON, by default False.
 
         Returns
         -------
-        Union[ChatResult, list[ChatResult], dict[int, ChatResult]]
-            The result of the run, which can be a single ChatResult,
-            a list of ChatResults,
-            or a dictionary mapping indices to ChatResults.
+        Union[list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]]
+            The result of the run, which can be a list of RunResponseProtocol
+            or a list of AsyncRunResponseProtocol.
 
         Raises
         ------
         RuntimeError
             If the runner is already running.
         """
-        if skip_patch_io is not None:
-            WaldiezBaseRunner._skip_patch_io = skip_patch_io
         if structured_io is not None:
             WaldiezBaseRunner._structured_io = structured_io
         if self.is_running():
@@ -541,10 +577,9 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         refresh_environment()
         WaldiezBaseRunner._running = True
         results: Union[
-            "ChatResult",
-            list["ChatResult"],
-            dict[int, "ChatResult"],
+            list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]
         ] = []
+        old_env_vars = set_env_vars(self.waldiez.get_flow_env_vars())
         try:
             async with a_chdir(to=temp_dir):
                 sys.path.insert(0, str(temp_dir))
@@ -553,15 +588,18 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
                     output_file=output_file,
                     uploads_root=uploads_root_path,
                     skip_mmd=skip_mmd,
+                    skip_timeline=skip_timeline,
                 )
         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,
         )
         if sys.path[0] == str(temp_dir):
             sys.path.pop(0)
@@ -572,8 +610,8 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         output_path: str | Path | None,
         uploads_root: str | Path | None,
         structured_io: bool | None = None,
-        skip_patch_io: bool | None = None,
         skip_mmd: bool = False,
+        skip_timeline: bool = False,
     ) -> None:
         """Start running the Waldiez flow in a non-blocking way.
 
@@ -585,18 +623,16 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             The runtime uploads root.
         structured_io : bool | None
             Whether to use structured IO instead of the default 'input/print'.
-        skip_patch_io : bool | None
-            Whether to skip patching I/O, by default None.
         skip_mmd : bool
             Whether to skip generating the mermaid diagram, by default False.
+        skip_timeline : bool
+            Whether to skip generating the timeline JSON, by default False.
 
         Raises
         ------
         RuntimeError
             If the runner is already running.
         """
-        if skip_patch_io is not None:
-            WaldiezBaseRunner._skip_patch_io = skip_patch_io
         if structured_io is not None:
             WaldiezBaseRunner._structured_io = structured_io
         if self.is_running():
@@ -617,6 +653,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             output_file=output_file,
             uploads_root=uploads_root_path,
             skip_mmd=skip_mmd,
+            skip_timeline=skip_timeline,
         )
 
     async def a_start(
@@ -624,8 +661,8 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         output_path: str | Path | None,
         uploads_root: str | Path | None,
         structured_io: bool | None = None,
-        skip_patch_io: bool | None = None,
         skip_mmd: bool = False,
+        skip_timeline: bool = False,
     ) -> None:
         """Asynchronously start running the Waldiez flow in a non-blocking way.
 
@@ -637,18 +674,16 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             The runtime uploads root.
         structured_io : bool | None = None
             Whether to use structured IO instead of the default 'input/print'.
-        skip_patch_io : bool | None = None
-            Whether to skip patching I/O, by default None.
-        skip_mmd : bool | None = None
-            Whether to skip generating the mermaid diagram, by default None.
+        skip_mmd : bool = False
+            Whether to skip generating the mermaid diagram, by default False.
+        skip_timeline : bool = False
+            Whether to skip generating the timeline JSON, by default False.
 
         Raises
         ------
         RuntimeError
             If the runner is already running.
         """
-        if skip_patch_io is not None:
-            WaldiezBaseRunner._skip_patch_io = skip_patch_io
         if structured_io is not None:
             WaldiezBaseRunner._structured_io = structured_io
         if self.is_running():
@@ -669,19 +704,19 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             output_file=output_file,
             uploads_root=uploads_root_path,
             skip_mmd=skip_mmd,
+            skip_timeline=skip_timeline,
         )
 
     def after_run(
         self,
         results: Union[
-            "ChatResult",
-            list["ChatResult"],
-            dict[int, "ChatResult"],
+            list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]
         ],
         output_file: Path,
         uploads_root: Path | None,
         temp_dir: Path,
         skip_mmd: bool,
+        skip_timeline: bool,
     ) -> None:
         """Actions to perform after running the flow.
 
@@ -697,6 +732,8 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             The path to the temporary directory used during the run.
         skip_mmd : bool
             Whether to skip generating the mermaid diagram.
+        skip_timeline : bool
+            Whether to skip generating the timeline JSON.
         """
         self._after_run(
             results=results,
@@ -704,19 +741,19 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             uploads_root=uploads_root,
             temp_dir=temp_dir,
             skip_mmd=skip_mmd,
+            skip_timeline=skip_timeline,
         )
 
     async def a_after_run(
         self,
         results: Union[
-            "ChatResult",
-            list["ChatResult"],
-            dict[int, "ChatResult"],
+            list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]
         ],
         output_file: Path,
         uploads_root: Path | None,
         temp_dir: Path,
         skip_mmd: bool,
+        skip_timeline: bool,
     ) -> None:
         """Asynchronously perform actions after running the flow.
 
@@ -732,6 +769,8 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             The path to the temporary directory used during the run.
         skip_mmd : bool
             Whether to skip generating the mermaid diagram.
+        skip_timeline : bool
+
         """
         await self._a_after_run(
             results=results,
@@ -739,6 +778,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             uploads_root=uploads_root,
             temp_dir=temp_dir,
             skip_mmd=skip_mmd,
+            skip_timeline=skip_timeline,
         )
 
     def stop(self) -> None:
@@ -783,21 +823,11 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         """Get the logger for the runner."""
         return self._logger
 
-    @property
-    def threaded(self) -> bool:
-        """Check if the runner is running in a threaded environment."""
-        return WaldiezBaseRunner._threaded
-
     @property
     def structured_io(self) -> bool:
         """Check if the runner is using structured IO."""
         return WaldiezBaseRunner._structured_io
 
-    @property
-    def isolated(self) -> bool:
-        """Check if the runner is running in an isolated environment."""
-        return WaldiezBaseRunner._isolated
-
     @property
     def output_path(self) -> str | Path | None:
         """Get the output path for the runner."""
@@ -824,9 +854,6 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         output_path: str | Path | None = None,
         uploads_root: str | Path | None = None,
         structured_io: bool = False,
-        isolated: bool = False,
-        threaded: bool = False,
-        skip_patch_io: bool = True,
     ) -> "WaldiezBaseRunner":
         """Load a waldiez flow from a file and create a runner.
 
@@ -849,12 +876,6 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         structured_io : bool, optional
             Whether to use structured IO instead of the default 'input/print',
             by default False.
-        isolated : bool, optional
-            Whether to run the flow in an isolated environment, default False.
-        threaded : bool, optional
-            Whether to run the flow in a threaded environment, default False.
-        skip_patch_io : bool, optional
-            Whether to skip patching IO, by default True.
 
         Returns
         -------
@@ -873,9 +894,6 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             output_path=output_path,
             uploads_root=uploads_root,
             structured_io=structured_io,
-            isolated=isolated,
-            threaded=threaded,
-            skip_patch_io=skip_patch_io,
         )
 
     def __enter__(self) -> Self: