waldiez 0.2.1__py3-none-any.whl → 0.3.0__py3-none-any.whl

waldiez/models/waldiez.py

@@ -1,4 +1,6 @@
-"""Waldiez data class.
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Waldiez
 
 A Waldiez class contains all the information that is needed to generate
 and run an autogen workflow. It has the model/LLM configurations, the agent
@@ -204,6 +206,11 @@ class Waldiez:
         """Get the flow tags."""
         return self.flow.tags
 
+    @property
+    def is_async(self) -> bool:
+        """Check if the flow is asynchronous."""
+        return self.flow.is_async
+
     @property
     def requirements(self) -> List[str]:
         """Get the flow requirements."""
@@ -236,9 +243,9 @@ class Waldiez:
         for model in self.models:
             if model.data.api_type in models_with_additional_requirements:
                 requirements.add(
-                    f"pyautogen[{model.data.api_type}]==" f"{autogen_version}"
+                    f"pyautogen[{model.data.api_type}]=={autogen_version}"
                 )
-        return list(requirements)
+        return sorted(list(requirements))
 
     def get_flow_env_vars(self) -> List[Tuple[str, str]]:
         """Get the flow environment variables.
@@ -255,7 +262,7 @@ class Waldiez:
         return env_vars
 
     def get_group_chat_members(self, agent: WaldiezAgent) -> List[WaldiezAgent]:
-        """Get the chat members that connect to a group chat manger agent.
+        """Get the chat members that connect to a group chat manager agent.
 
         Parameters
         ----------
@@ -271,6 +278,23 @@ class Waldiez:
             return []
         return self.flow.get_group_chat_members(agent.id)
 
+    def get_swarm_members(
+        self, initial_agent: WaldiezAgent
+    ) -> Tuple[List[WaldiezAgent], Optional[WaldiezAgent]]:
+        """Get the chat members that connect to a swarm agent.
+
+        Parameters
+        ----------
+        initial_agent : WaldiezAgent
+            The initial agent.
+
+        Returns
+        -------
+        Tuple[List[WaldiezAgent], Optional[WaldiezAgent]]
+            The swarm agents and the user agent.
+        """
+        return self.flow.get_swarm_chat_members(initial_agent)
+
 
 def _get_flow(
     data: Dict[str, Any],

waldiez/runner.py

@@ -1,6 +1,6 @@
-"""Waldiez Flow runner.
-
-Run a waldiez flow.
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Run a waldiez flow.
 The flow is first converted to an autogen flow with agents, chats and skills.
 We then chown to temporary directory, call the flow's `main()` and
 return the results. Before running the flow, any additional environment
@@ -9,60 +9,34 @@ variables specified in the waldiez file are set.
 
 # pylint: disable=import-outside-toplevel,reimported
 
-import datetime
 import importlib.util
-import io
-import os
-import shutil
-import site
-import subprocess  # nosemgrep # nosec
 import sys
 import tempfile
-import warnings
-from contextlib import contextmanager
 from pathlib import Path
 from types import TracebackType
-from typing import (
-    TYPE_CHECKING,
-    Callable,
-    Dict,
-    Iterator,
-    List,
-    Optional,
-    Tuple,
-    Type,
-    Union,
-)
+from typing import TYPE_CHECKING, List, Optional, Type, Union
+
+from asyncer import syncify
 
 from .exporter import WaldiezExporter
 from .models.waldiez import Waldiez
+from .running import (
+    a_chdir,
+    a_install_requirements,
+    after_run,
+    before_run,
+    chdir,
+    get_printer,
+    install_requirements,
+    refresh_environment,
+    reset_env_vars,
+    set_env_vars,
+)
 
 if TYPE_CHECKING:
     from autogen import ChatResult  # type: ignore
 
 
-@contextmanager
-def _chdir(to: Union[str, Path]) -> Iterator[None]:
-    """Change the current working directory in a context.
-
-    Parameters
-    ----------
-    to : Union[str, Path]
-        The directory to change to.
-
-    Yields
-    ------
-    Iterator[None]
-        The context manager.
-    """
-    old_cwd = str(os.getcwd())
-    os.chdir(to)
-    try:
-        yield
-    finally:
-        os.chdir(old_cwd)
-
-
 class WaldiezRunner:
     """Waldiez runner class."""
 
@@ -127,6 +101,12 @@ class WaldiezRunner:
         """Enter the context manager."""
         return self
 
+    async def __aenter__(
+        self,
+    ) -> "WaldiezRunner":
+        """Enter the context manager asynchronously."""
+        return self
+
     def __exit__(
         self,
         exc_type: Type[BaseException],
@@ -137,6 +117,16 @@ class WaldiezRunner:
         if self._running:
             self._running = False
 
+    async def __aexit__(
+        self,
+        exc_type: Type[BaseException],
+        exc_value: BaseException,
+        traceback: TracebackType,
+    ) -> None:
+        """Exit the context manager asynchronously."""
+        if self._running:
+            self._running = False
+
     @property
     def waldiez(self) -> Waldiez:
         """Get the Waldiez instance."""
@@ -155,78 +145,25 @@ class WaldiezRunner:
             req for req in self.waldiez.requirements if req not in sys.modules
         )
         if extra_requirements:
-            requirements_string = ", ".join(extra_requirements)
-            printer(f"Installing requirements: {requirements_string}")
-            pip_install = [sys.executable, "-m", "pip", "install"]
-            if not in_virtualenv():
-                pip_install.append("--user")
-            pip_install.extend(extra_requirements)
-            with subprocess.Popen(
-                pip_install,
-                stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE,
-            ) as proc:
-                if proc.stdout:
-                    for line in io.TextIOWrapper(proc.stdout, encoding="utf-8"):
-                        printer(line.strip())
-                if proc.stderr:
-                    for line in io.TextIOWrapper(proc.stderr, encoding="utf-8"):
-                        printer(line.strip())
+            install_requirements(extra_requirements, printer)
             refresh_environment()
 
-    @staticmethod
-    def _after_run(
-        temp_dir: Path,
-        output_path: Optional[Union[str, Path]],
-        printer: Callable[..., None],
-    ) -> None:
-        if output_path:
-            destination_dir = Path(output_path).parent
-            destination_dir = (
-                destination_dir
-                / "waldiez_out"
-                / datetime.datetime.now().strftime("%Y%m%d%H%M%S")
-            )
-            destination_dir.mkdir(parents=True, exist_ok=True)
-            # copy the contents of the temp dir to the destination dir
-            printer(f"Copying the results to {destination_dir}")
-            for item in temp_dir.iterdir():
-                # skip cache files
-                if (
-                    item.name.startswith("__pycache__")
-                    or item.name.endswith(".pyc")
-                    or item == ".cache"
-                ):
-                    continue
-                if item.is_file():
-                    shutil.copy(item, destination_dir)
-                else:
-                    shutil.copytree(item, destination_dir / item.name)
-        shutil.rmtree(temp_dir)
-
-    def _set_env_vars(self) -> Dict[str, str]:
-        """Set environment variables and return the old ones (if any)."""
-        old_vars: Dict[str, str] = {}
-        for var_key, var_value in self.waldiez.get_flow_env_vars():
-            if var_key:
-                current = os.environ.get(var_key, "")
-                old_vars[var_key] = current
-                os.environ[var_key] = var_value
-        return old_vars
-
-    @staticmethod
-    def _reset_env_vars(old_vars: Dict[str, str]) -> None:
-        """Reset the environment variables."""
-        for var_key, var_value in old_vars.items():
-            if not var_value:
-                os.environ.pop(var_key, "")
-            else:
-                os.environ[var_key] = var_value
+    async def a_install_requirements(self) -> None:
+        """Install the requirements for the flow asynchronously."""
+        self._called_install_requirements = True
+        printer = get_printer()
+        extra_requirements = set(
+            req for req in self.waldiez.requirements if req not in sys.modules
+        )
+        if extra_requirements:
+            await a_install_requirements(extra_requirements, printer)
+            refresh_environment()
 
-    def _do_run(
+    def _run(
         self,
         output_path: Optional[Union[str, Path]],
         uploads_root: Optional[Union[str, Path]],
+        skip_mmd: bool = False,
     ) -> Union["ChatResult", List["ChatResult"]]:
         """Run the Waldiez workflow.
 
@@ -236,12 +173,16 @@ class WaldiezRunner:
             The output path.
         uploads_root : Optional[Union[str, Path]]
             The runtime uploads root.
-
+        skip_mmd : bool
+            Whether to skip the Mermaid diagram generation.
         Returns
         -------
         Union[ChatResult, List[ChatResult]]
             The result(s) of the chat(s).
         """
+        temp_dir = Path(tempfile.mkdtemp())
+        file_name = before_run(output_path, uploads_root)
+        module_name = file_name.replace(".py", "")
         if not self._called_install_requirements:
             self.install_requirements()
         else:
@@ -253,22 +194,52 @@ class WaldiezRunner:
             "you might need to restart the kernel."
         )
         results: Union["ChatResult", List["ChatResult"]] = []
-        if not uploads_root:
-            uploads_root = Path(tempfile.mkdtemp())
-        else:
-            uploads_root = Path(uploads_root)
-        if not uploads_root.exists():
-            uploads_root.mkdir(parents=True)
-        temp_dir = Path(tempfile.mkdtemp())
-        file_name = "flow.py" if not output_path else Path(output_path).name
-        if file_name.endswith((".json", ".waldiez")):
-            file_name = file_name.replace(".json", ".py").replace(
-                ".waldiez", ".py"
+        with chdir(to=temp_dir):
+            self._exporter.export(Path(file_name))
+            spec = importlib.util.spec_from_file_location(
+                module_name, temp_dir / file_name
             )
-        if not file_name.endswith(".py"):
-            file_name += ".py"
+            if not spec or not spec.loader:
+                raise ImportError("Could not import the flow")
+            sys.path.insert(0, str(temp_dir))
+            old_vars = set_env_vars(self.waldiez.get_flow_env_vars())
+            module = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(module)
+            printer("<Waldiez> - Starting workflow...")
+            results = module.main()
+            sys.path.pop(0)
+            reset_env_vars(old_vars)
+        after_run(
+            temp_dir=temp_dir,
+            output_path=output_path,
+            printer=printer,
+            flow_name=self.waldiez.name,
+            skip_mmd=skip_mmd,
+        )
+        return results
+
+    async def _a_run(
+        self,
+        output_path: Optional[Union[str, Path]],
+        uploads_root: Optional[Union[str, Path]],
+        skip_mmd: bool = False,
+    ) -> Union["ChatResult", List["ChatResult"]]:
+        """Run the Waldiez workflow asynchronously."""
+        temp_dir = Path(tempfile.mkdtemp())
+        file_name = before_run(output_path, uploads_root)
         module_name = file_name.replace(".py", "")
-        with _chdir(to=temp_dir):
+        if not self._called_install_requirements:
+            await self.a_install_requirements()
+        else:
+            refresh_environment()
+        printer = get_printer()
+        printer(
+            "Requirements installed.\n"
+            "NOTE: If new packages were added and you are using Jupyter, "
+            "you might need to restart the kernel."
+        )
+        results: Union["ChatResult", List["ChatResult"]] = []
+        async with a_chdir(to=temp_dir):
             self._exporter.export(Path(file_name))
             spec = importlib.util.spec_from_file_location(
                 module_name, temp_dir / file_name
@@ -276,20 +247,27 @@ class WaldiezRunner:
             if not spec or not spec.loader:
                 raise ImportError("Could not import the flow")
             sys.path.insert(0, str(temp_dir))
-            old_vars = self._set_env_vars()
+            old_vars = set_env_vars(self.waldiez.get_flow_env_vars())
             module = importlib.util.module_from_spec(spec)
             spec.loader.exec_module(module)
             printer("<Waldiez> - Starting workflow...")
-            results = module.main()
+            results = await module.main()
             sys.path.pop(0)
-            self._reset_env_vars(old_vars)
-        self._after_run(temp_dir, output_path, printer)
+            reset_env_vars(old_vars)
+        after_run(
+            temp_dir=temp_dir,
+            output_path=output_path,
+            printer=printer,
+            flow_name=self.waldiez.name,
+            skip_mmd=skip_mmd,
+        )
         return results
 
     def run(
         self,
         output_path: Optional[Union[str, Path]] = None,
         uploads_root: Optional[Union[str, Path]] = None,
+        skip_mmd: bool = False,
     ) -> Union["ChatResult", List["ChatResult"]]:
         """Run the Waldiez workflow.
 
@@ -299,6 +277,8 @@ class WaldiezRunner:
             The output path, by default None.
         uploads_root : Optional[Union[str, Path]], optional
             The uploads root, to get user-uploaded files, by default None.
+        skip_mmd : bool, optional
+            Whether to skip the Mermaid diagram generation, by default False.
 
         Returns
         -------
@@ -310,114 +290,48 @@ class WaldiezRunner:
         RuntimeError
             If the workflow is already running.
         """
+        if self.waldiez.is_async:
+            return syncify(self.a_run, raise_sync_error=False)(
+                output_path, uploads_root
+            )
         if self._running is True:
             raise RuntimeError("Workflow already running")
         self._running = True
         file_path = output_path or self._file_path
         try:
-            return self._do_run(file_path, uploads_root)
+            return self._run(file_path, uploads_root, skip_mmd)
         finally:
             self._running = False
 
+    async def a_run(
+        self,
+        output_path: Optional[Union[str, Path]] = None,
+        uploads_root: Optional[Union[str, Path]] = None,
+    ) -> Union["ChatResult", List["ChatResult"]]:
+        """Run the Waldiez workflow asynchronously.
 
-def in_virtualenv() -> bool:
-    """Check if we are inside a virtualenv.
-
-    Returns
-    -------
-    bool
-        True if inside a virtualenv, False otherwise.
-    """
-    return hasattr(sys, "real_prefix") or (
-        hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix
-    )
-
-
-def refresh_environment() -> None:
-    """Refresh the environment."""
-    # backup the default IOStream
-    from autogen.io import IOStream  # type: ignore
-
-    default_io_stream = IOStream.get_default()
-    site.main()
-    # pylint: disable=import-outside-toplevel
-    modules_to_reload = [
-        mod for mod in sys.modules if mod.startswith("autogen")
-    ]
-    for mod in modules_to_reload:
-        del sys.modules[mod]
-    warnings.filterwarnings(
-        "ignore", module="flaml", message="^.*flaml.automl is not available.*$"
-    )
-    import autogen
-    from autogen.io import IOStream
-
-    importlib.reload(autogen)
-    # restore the default IOStream
-    IOStream.set_global_default(default_io_stream)
-
-
-def get_printer() -> Callable[..., None]:
-    """Get the printer function.
-
-    Returns
-    -------
-    Callable[..., None]
-        The printer function.
-    """
-    from autogen.io import IOStream
+        Parameters
+        ----------
+        output_path : Optional[Union[str, Path]], optional
+            The output path, by default None.
+        uploads_root : Optional[Union[str, Path]], optional
+            The uploads root, to get user-uploaded files, by default None.
 
-    printer = IOStream.get_default().print
+        Returns
+        -------
+        Union[ChatResult, List[ChatResult]]
+            The result(s) of the chat(s).
 
-    def safe_printer(*args: object, **kwargs: object) -> None:
+        Raises
+        ------
+        RuntimeError
+            If the workflow is already running.
+        """
+        if self._running is True:
+            raise RuntimeError("Workflow already running")
+        self._running = True
+        file_path = output_path or self._file_path
         try:
-            printer(*args, **kwargs)
-        except UnicodeEncodeError:
-            # pylint: disable=too-many-try-statements
-            try:
-                msg, flush = get_what_to_print(*args, **kwargs)
-                printer(msg, end="", flush=flush)
-            except UnicodeEncodeError:
-                sys.stdout = io.TextIOWrapper(
-                    sys.stdout.buffer, encoding="utf-8"
-                )
-                sys.stderr = io.TextIOWrapper(
-                    sys.stderr.buffer, encoding="utf-8"
-                )
-                try:
-                    printer(*args, **kwargs)
-                except UnicodeEncodeError:
-                    sys.stderr.write(
-                        "Could not print the message due to encoding issues.\n"
-                    )
-
-    return safe_printer
-
-
-def get_what_to_print(*args: object, **kwargs: object) -> Tuple[str, bool]:
-    """Get what to print.
-
-    Parameters
-    ----------
-    args : object
-        The arguments.
-    kwargs : object
-        The keyword arguments.
-
-    Returns
-    -------
-    Tuple[str, bool]
-        The message and whether to flush.
-    """
-    sep = kwargs.get("sep", " ")
-    if not isinstance(sep, str):
-        sep = " "
-    end = kwargs.get("end", "\n")
-    if not isinstance(end, str):
-        end = "\n"
-    flush = kwargs.get("flush", False)
-    if not isinstance(flush, bool):
-        flush = False
-    msg = sep.join(str(arg) for arg in args) + end
-    utf8_msg = msg.encode("utf-8", errors="replace").decode("utf-8")
-    return utf8_msg, flush
+            return await self._a_run(file_path, uploads_root)
+        finally:
+            self._running = False

waldiez/running/__init__.py

@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Common utility functions."""
+
+from .environment import (
+    in_virtualenv,
+    refresh_environment,
+    reset_env_vars,
+    set_env_vars,
+)
+from .running import (
+    a_chdir,
+    a_install_requirements,
+    after_run,
+    before_run,
+    chdir,
+    get_printer,
+    install_requirements,
+)
+
+__all__ = [
+    "a_chdir",
+    "a_install_requirements",
+    "after_run",
+    "before_run",
+    "chdir",
+    "get_printer",
+    "in_virtualenv",
+    "install_requirements",
+    "refresh_environment",
+    "reset_env_vars",
+    "set_env_vars",
+]

waldiez/running/environment.py

@@ -0,0 +1,83 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+# pylint: disable=import-outside-toplevel,reimported
+"""Environment related utilities."""
+
+import importlib.util
+import os
+import site
+import sys
+import warnings
+from typing import Dict, List, Tuple
+
+
+def in_virtualenv() -> bool:
+    """Check if we are inside a virtualenv.
+
+    Returns
+    -------
+    bool
+        True if inside a virtualenv, False otherwise.
+    """
+    return hasattr(sys, "real_prefix") or (
+        hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix
+    )
+
+
+def refresh_environment() -> None:
+    """Refresh the environment."""
+    # backup the default IOStream
+    from autogen.io import IOStream  # type: ignore
+
+    default_io_stream = IOStream.get_default()
+    site.main()
+    # pylint: disable=import-outside-toplevel
+    modules_to_reload = [mod for mod in sys.modules if "autogen" in mod]
+    for mod in modules_to_reload:
+        del sys.modules[mod]
+    warnings.filterwarnings(
+        "ignore", module="flaml", message="^.*flaml.automl is not available.*$"
+    )
+    import autogen  # type: ignore
+    from autogen.io import IOStream
+
+    importlib.reload(autogen)
+    # restore the default IOStream
+    IOStream.set_global_default(default_io_stream)
+
+
+def set_env_vars(flow_env_vars: List[Tuple[str, str]]) -> Dict[str, str]:
+    """Set environment variables and return the old ones (if any).
+
+    Parameters
+    ----------
+    flow_env_vars : List[Tuple[str, str]]
+        The environment variables to set.
+
+    Returns
+    -------
+    Dict[str, str]
+        The old environment variables.
+    """
+    old_vars: Dict[str, str] = {}
+    for var_key, var_value in flow_env_vars:
+        if var_key:
+            current = os.environ.get(var_key, "")
+            old_vars[var_key] = current
+            os.environ[var_key] = var_value
+    return old_vars
+
+
+def reset_env_vars(old_vars: Dict[str, str]) -> None:
+    """Reset the environment variables.
+
+    Parameters
+    ----------
+    old_vars : Dict[str, str]
+        The old environment variables.
+    """
+    for var_key, var_value in old_vars.items():
+        if not var_value:
+            os.environ.pop(var_key, "")
+        else:
+            os.environ[var_key] = var_value