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

waldiez/running/environment.py

@@ -63,6 +63,7 @@ def refresh_environment() -> None:
 
 
 # pylint: disable=too-complex,too-many-try-statements,unused-import
+# noinspection TryExceptPass
 def reload_autogen() -> None:  # noqa: C901  # pragma: no cover
     """Reload the autogen package.
 

waldiez/running/exceptions.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Custom exceptions for Waldiez runners."""
+
+
+class StopRunningException(Exception):
+    """Exception to stop the running process."""
+
+    reason: str = "Execution stopped by user"

waldiez/running/post_run.py

@@ -11,8 +11,8 @@ from pathlib import Path
 from typing import Optional, Union
 
 from .gen_seq_diagram import generate_sequence_diagram
-from .patch_io_stream import get_printer
 from .timeline_processor import TimelineProcessor
+from .utils import get_printer
 
 
 # noinspection PyUnusedLocal
@@ -20,6 +20,7 @@ def after_run(
     temp_dir: Path,
     output_file: Optional[Union[str, Path]],
     flow_name: str,
+    waldiez_file: Path,
     uploads_root: Optional[Path] = None,
     skip_mmd: bool = False,
     skip_timeline: bool = False,
@@ -34,6 +35,8 @@ def after_run(
         The output file.
     flow_name : str
         The flow name.
+    waldiez_file : Path
+        The path of the waldiez file used (or dumped) for the run.
     uploads_root : Optional[Path], optional
         The runtime uploads root, by default None
     skip_mmd : bool, optional
@@ -52,7 +55,7 @@ def after_run(
             flow_name=flow_name,
             mmd_dir=mmd_dir,
         )
-    if skip_timeline is False:
+    if skip_timeline is False:  # pragma: no branch
         _make_timeline_json(temp_dir=temp_dir)
     if output_file:
         destination_dir = output_file.parent
@@ -69,6 +72,9 @@ def after_run(
             output_file=output_file,
             destination_dir=destination_dir,
         )
+        dst_waldiez = destination_dir / waldiez_file.name
+        if not dst_waldiez.exists() and waldiez_file.is_file():
+            shutil.copyfile(waldiez_file, dst_waldiez)
     shutil.rmtree(temp_dir)
 
 
@@ -103,7 +109,7 @@ def _make_timeline_json(
     events_csv_path = temp_dir / "logs" / "events.csv"
     if events_csv_path.exists():
         log_files = TimelineProcessor.get_files(temp_dir / "logs")
-        if any(log_files.values()):
+        if any(log_files.values()):  # pragma: no branch
             output_file = temp_dir / "timeline.json"
             # pylint: disable=too-many-try-statements
             try:
@@ -171,7 +177,7 @@ def copy_results(
     if output_file.is_file():
         if output_file.suffix == ".waldiez":
             output_file = output_file.with_suffix(".py")
-        if output_file.suffix == ".py":
+        if output_file.suffix == ".py":  # pragma: no branch
             src = temp_dir / output_file.name
             if src.exists():
                 dst = destination_dir / output_file.name

waldiez/running/pre_run.py

@@ -7,10 +7,61 @@ import io
 import os
 import subprocess
 import sys
+import tempfile
+from pathlib import Path
 from typing import Callable
 
+from waldiez.models import Waldiez
+
 from .environment import in_virtualenv, is_root
-from .utils import strip_ansi
+from .utils import (
+    ensure_pip,
+    get_pip_install_location,
+    get_python_executable,
+    safe_filename,
+    strip_ansi,
+)
+
+
+class RequirementsMixin:
+    """Mixin class to handle requirements installation."""
+
+    _waldiez: Waldiez
+    _called_install_requirements: bool
+
+    def gather_requirements(self) -> set[str]:
+        """Gather extra requirements to install before running the flow.
+
+        Returns
+        -------
+        set[str]
+            A set of requirements that are not already installed and do not
+            include 'waldiez' in their name.
+        """
+        extra_requirements = {
+            req
+            for req in self._waldiez.requirements
+            if req not in sys.modules and "waldiez" not in req
+        }
+        if "python-dotenv" not in extra_requirements:  # pragma: no branch
+            extra_requirements.add("python-dotenv")
+        return extra_requirements
+
+    def install_requirements(self) -> None:
+        """Install the requirements for the flow."""
+        if not self._called_install_requirements:  # pragma: no branch
+            self._called_install_requirements = True
+            extra_requirements = self.gather_requirements()
+            if extra_requirements:  # pragma: no branch
+                install_requirements(extra_requirements)
+
+    async def a_install_requirements(self) -> None:
+        """Install the requirements for the flow asynchronously."""
+        if not self._called_install_requirements:  # pragma: no branch
+            self._called_install_requirements = True
+            extra_requirements = self.gather_requirements()
+            if extra_requirements:  # pragma: no branch
+                await a_install_requirements(extra_requirements)
 
 
 # noinspection PyUnresolvedReferences
@@ -32,39 +83,40 @@ def install_requirements(
     """
     requirements_string = ", ".join(extra_requirements)
     printer(f"Installing requirements: {requirements_string}")
-    pip_install = [sys.executable, "-m", "pip", "install"]
-    break_system_packages = ""
-    if not in_virtualenv():  # it should
-        # if not, let's try to install as user
-        # not sure if --break-system-packages is safe,
-        # but it might fail if we don't
-        break_system_packages = os.environ.get("PIP_BREAK_SYSTEM_PACKAGES", "")
-        os.environ["PIP_BREAK_SYSTEM_PACKAGES"] = "1"
-        if not is_root():
-            pip_install.append("--user")
-    if upgrade:
-        pip_install.append("--upgrade")
-    pip_install.extend(extra_requirements)
-    # pylint: disable=too-many-try-statements
+    ensure_pip()
+    pip_install, break_system_packages, install_location = _before_pip(
+        extra_requirements, upgrade
+    )
+
+    # pylint: disable=too-many-try-statements,broad-exception-caught
     try:
         with subprocess.Popen(
             pip_install,
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,
         ) as proc:
-            if proc.stdout:
+            if proc.stdout:  # pragma: no branch
                 for line in io.TextIOWrapper(proc.stdout, encoding="utf-8"):
-                    printer(strip_ansi(line.strip()))
-            if proc.stderr:
+                    stripped_line = strip_ansi(line.strip())
+                    if stripped_line:  # Only print non-empty lines
+                        printer(stripped_line)
+            if proc.stderr:  # pragma: no branch
                 for line in io.TextIOWrapper(proc.stderr, encoding="utf-8"):
-                    printer(strip_ansi(line.strip()))
+                    stripped_line = strip_ansi(line.strip())
+                    if stripped_line:  # Only print non-empty lines
+                        printer(stripped_line)
+
+            # Wait for process to complete and check return code
+            return_code = proc.wait()
+            if return_code != 0:
+                printer(
+                    f"Package installation failed with exit code {return_code}"
+                )
+
+    except Exception as e:
+        printer(f"Failed to install requirements: {e}")
     finally:
-        if not in_virtualenv():
-            # restore the old env var
-            if break_system_packages:
-                os.environ["PIP_BREAK_SYSTEM_PACKAGES"] = break_system_packages
-            else:
-                del os.environ["PIP_BREAK_SYSTEM_PACKAGES"]
+        _after_pip(break_system_packages, install_location)
 
 
 async def a_install_requirements(
@@ -83,66 +135,147 @@ async def a_install_requirements(
     printer : Callable[..., None]
         The printer function to use, defaults to print.
     """
+    pip_install, break_system_packages, install_location = _before_pip(
+        extra_requirements, upgrade=upgrade
+    )
     requirements_string = ", ".join(extra_requirements)
     printer(f"Installing requirements: {requirements_string}")
-    pip_install = [sys.executable, "-m", "pip", "install"]
-    break_system_packages = ""
-    if not in_virtualenv():
-        break_system_packages = os.environ.get("PIP_BREAK_SYSTEM_PACKAGES", "")
-        os.environ["PIP_BREAK_SYSTEM_PACKAGES"] = "1"
-        if not is_root():
-            pip_install.extend(["--user"])
-    if upgrade:
-        pip_install.append("--upgrade")
-    pip_install.extend(extra_requirements)
-    # pylint: disable=too-many-try-statements
+    # pylint: disable=too-many-try-statements,broad-exception-caught
     try:
         proc = await asyncio.create_subprocess_exec(
             *pip_install,
             stdout=asyncio.subprocess.PIPE,
             stderr=asyncio.subprocess.PIPE,
         )
-        if proc.stdout:
-            async for line in proc.stdout:
-                printer(strip_ansi(line.decode().strip()))
-        if proc.stderr:
-            async for line in proc.stderr:
-                printer(strip_ansi(line.decode().strip()))
+
+        async def _pump_stream(stream: asyncio.StreamReader | None) -> None:
+            if not stream:
+                return
+            async for raw in stream:
+                text = strip_ansi(raw.decode().rstrip())
+                if text:
+                    printer(text)
+
+        # Create tasks for concurrent execution
+        tasks: list[asyncio.Task[int | None]] = [
+            asyncio.create_task(_pump_stream(proc.stdout)),
+            asyncio.create_task(_pump_stream(proc.stderr)),
+            asyncio.create_task(proc.wait()),
+        ]
+        await asyncio.gather(*tasks, return_exceptions=True)
+        if proc.returncode != 0:
+            printer(
+                f"Package installation failed with exit code {proc.returncode}"
+            )
+
+    except Exception as e:
+        printer(f"Failed to install requirements: {e}")
     finally:
-        if not in_virtualenv():
-            if break_system_packages:
-                os.environ["PIP_BREAK_SYSTEM_PACKAGES"] = break_system_packages
-            else:
-                del os.environ["PIP_BREAK_SYSTEM_PACKAGES"]
+        _after_pip(break_system_packages, install_location)
 
 
-def install_waldiez(
-    upgrade: bool = True,
-    printer: Callable[..., None] = print,
-) -> None:
-    """Install Waldiez.
+def _before_pip(
+    packages: set[str],
+    upgrade: bool,
+) -> tuple[list[str], str, str | None]:
+    """Gather the pip command for installing requirements.
 
     Parameters
     ----------
-    upgrade : bool, optional
-        Whether to upgrade Waldiez, by default True.
-    printer : Callable[..., None]
-        The printer function to use, defaults to print.
+    packages : set[str]
+        The packages to install.
+    upgrade : bool
+        Whether to upgrade the packages.
+
+    Returns
+    -------
+    tuple[list[str], str, str | None]
+        The pip command, break_system_packages flag, and install location.
     """
-    install_requirements({"waldiez"}, upgrade, printer)
+    ensure_pip()
+    pip_install = [
+        get_python_executable(),
+        "-m",
+        "pip",
+        "install",
+        "--disable-pip-version-check",
+        "--no-input",
+    ]
+    install_location = get_pip_install_location()
+    break_system_packages = ""
 
+    if install_location:
+        pip_install += ["--target", install_location]
+    elif not in_virtualenv():  # it should  # pragma: no cover
+        # if not, let's try to install as user
+        # not sure if --break-system-packages is safe,
+        # but it might fail if we don't
+        break_system_packages = os.environ.get("PIP_BREAK_SYSTEM_PACKAGES", "")
+        os.environ["PIP_BREAK_SYSTEM_PACKAGES"] = "1"
+        if not is_root():
+            pip_install.append("--user")
 
-async def a_install_waldiez(
-    upgrade: bool = True,
-    printer: Callable[..., None] = print,
+    if upgrade:  # pragma: no cover
+        pip_install.append("--upgrade")
+
+    pip_install.extend(packages)
+    return pip_install, break_system_packages, install_location
+
+
+def _after_pip(
+    break_system_packages: str, install_location: str | None
 ) -> None:
-    """Install Waldiez asynchronously.
+    """Restore environment variables after pip installation.
 
     Parameters
     ----------
-    upgrade : bool, optional
-        Whether to upgrade Waldiez, by default True.
-    printer : Callable[..., None]
-        The printer function to use, defaults to print.
+    break_system_packages : str
+        The original value of PIP_BREAK_SYSTEM_PACKAGES.
+    install_location : str | None
+        The install location used (None if default).
+    """
+    if install_location is None and not in_virtualenv():  # pragma: no cover
+        # restore the old env var
+        if break_system_packages:
+            os.environ["PIP_BREAK_SYSTEM_PACKAGES"] = break_system_packages
+        else:
+            # Use pop to avoid KeyError if the key doesn't exist
+            os.environ.pop("PIP_BREAK_SYSTEM_PACKAGES", None)
+
+
+def dump_waldiez(waldiez: Waldiez, output_path: str | Path | None) -> Path:
+    """Dump waldiez flow to a file.
+
+    Parameters
+    ----------
+    waldiez : Waldiez
+        The waldiez instance to dump.
+    output_path : str | Path | None
+        Optional output path to determine the directory to save the flow to.
+
+    Returns
+    -------
+    Path
+        The path to the generated file.
     """
-    await a_install_requirements({"waldiez"}, upgrade, printer)
+    file_path = Path(output_path) if output_path else None
+    if file_path:
+        file_name = file_path.name
+        if not file_name.endswith(".waldiez"):
+            file_path.with_suffix(".waldiez")
+
+    else:
+        full_name = waldiez.name
+        file_name = safe_filename(full_name, "waldiez")
+    file_dir: Path
+    if file_path:
+        file_dir = file_path if file_path.is_dir() else file_path.parent
+    else:
+        file_dir = Path(tempfile.mkdtemp())
+    file_dir.mkdir(parents=True, exist_ok=True)
+    output_path = file_dir / file_name
+    with output_path.open(
+        "w", encoding="utf-8", errors="replace", newline="\n"
+    ) as f_open:
+        f_open.write(waldiez.model_dump_json())
+    return output_path

waldiez/running/protocol.py

@@ -51,98 +51,24 @@ class WaldiezRunnerProtocol(Protocol):
             The path to the temporary directory created for the run.
         """
 
-    def start(
-        self,
-        output_path: str | Path | None,
-        uploads_root: str | Path | None,
-        structured_io: bool | None = None,
-        skip_mmd: bool = False,
-        skip_timeline: bool = False,
-        dot_env: str | Path | None = None,
-        **kwargs: Any,
-    ) -> None:
-        """Start running the Waldiez flow in a non-blocking way.
-
-        To allow "stoping" it later.
-
-        Parameters
-        ----------
-        output_path : str | Path | None
-            The output path.
-        uploads_root : str | Path | None
-            The runtime uploads root.
-        structured_io : bool
-            Whether to use structured IO instead of the default 'input/print'.
-        skip_mmd : bool
-            Whether to skip generating the mermaid diagram.
-        skip_timeline : bool
-            Whether to skip generating the timeline JSON.
-        dot_env : str | Path | None
-            The path to the .env file, if any.
-        **kwargs : Any
-            Additional keyword arguments for the start method.
-
-        Raises
-        ------
-        RuntimeError
-            If the runner is already running.
-        """
-
-    async def a_start(
-        self,
-        output_path: str | Path | None,
-        uploads_root: str | Path | None,
-        structured_io: bool | None = None,
-        skip_mmd: bool = False,
-        skip_timeline: bool = False,
-        dot_env: str | Path | None = None,
-        **kwargs: Any,
-    ) -> None:
-        """Asynchronously start running the Waldiez flow in a non-blocking way.
-
-        To allow "stoping" it later.
-
-        Parameters
-        ----------
-        output_path :  str | Path | None
-            The output path.
-        uploads_root : str | Path | None
-            The runtime uploads root.
-        structured_io : bool
-            Whether to use structured IO instead of the default 'input/print'.
-        skip_mmd : bool
-            Whether to skip generating the mermaid diagram.
-        skip_timeline : bool
-            Whether to skip generating the timeline JSON.
-        dot_env : str | Path | None
-            The path to the .env file, if any.
-        **kwargs : Any
-            Additional keyword arguments for the start method.
-
-        Raises
-        ------
-        RuntimeError
-            If the runner is already running.
-        """
-
     def run(
         self,
-        output_path: str | Path | None,
-        uploads_root: str | Path | None,
+        output_path: str | Path | None = None,
+        uploads_root: str | Path | None = None,
         structured_io: bool | None = None,
         skip_mmd: bool = False,
         skip_timeline: bool = False,
         dot_env: str | Path | None = None,
         **kwargs: Any,
     ) -> list[dict[str, Any]]:
-        """Run the Waldiez flow in a blocking way.
+        """Run the Waldiez flow.
 
         Parameters
         ----------
         output_path : str | Path | None
             The output path, by default None.
         uploads_root : str | Path | None
-            The runtime uploads root.
+            The runtime uploads root, by default None.
         structured_io : bool
             Whether to use structured IO instead of the default 'input/print'.
         skip_mmd : bool
@@ -157,20 +83,26 @@ class WaldiezRunnerProtocol(Protocol):
         Returns
         -------
         list[dict[str, Any]]
-            The result of the run.
+            The results of the run.
+
+        Raises
+        ------
+        RuntimeError
+            If the runner is already running
+            or an error occurs during the run.
         """
 
     async def a_run(
         self,
-        output_path: str | Path | None,
-        uploads_root: str | Path | None,
+        output_path: str | Path | None = None,
+        uploads_root: str | Path | None = None,
         structured_io: bool | None = None,
         skip_mmd: bool = False,
         skip_timeline: bool = False,
         dot_env: str | Path | None = None,
         **kwargs: Any,
     ) -> list[dict[str, Any]]:
-        """Run the Waldiez flow.
+        """Run the Waldiez flow asynchronously.
 
         Parameters
         ----------
@@ -192,7 +124,13 @@ class WaldiezRunnerProtocol(Protocol):
         Returns
         -------
         list[dict[str, Any]]
-            The result of the run.
+            The results of the run.
+
+        Raises
+        ------
+        RuntimeError
+            If the runner is already running, the workflow is not async
+            or an error occurs during the run.
         """
 
     def after_run(
@@ -257,21 +195,3 @@ class WaldiezRunnerProtocol(Protocol):
         bool
             True if the runner is running, False otherwise.
         """
-
-    def stop(self) -> None:
-        """Stop the runner if it is running.
-
-        Raises
-        ------
-        RuntimeError
-            If the runner is not running.
-        """
-
-    async def a_stop(self) -> None:
-        """Asynchronously stop the runner if it is running.
-
-        Raises
-        ------
-        RuntimeError
-            If the runner is not running.
-        """

waldiez/running/run_results.py

@@ -10,5 +10,5 @@ class WaldiezRunResults(TypedDict):
     """Results of the Waldiez run."""
 
     results: list[dict[str, Any]]
-    exception: Exception | None
+    exception: BaseException | None
     completed: bool