waldiez 0.4.7__py3-none-any.whl → 0.4.9__py3-none-any.whl

waldiez/runner.py

@@ -2,7 +2,7 @@
 # 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.
+The flow is first converted to an autogen flow with agents, chats, and tools.
 We then chown to temporary directory, call the flow's `main()` and
 return the results. Before running the flow, any additional environment
 variables specified in the waldiez file are set.
@@ -10,14 +10,22 @@ variables specified in the waldiez file are set.
 
 # pylint: disable=import-outside-toplevel,reimported
 
+import gc
 import importlib.util
 import sys
 import tempfile
 from pathlib import Path
-from types import TracebackType
-from typing import TYPE_CHECKING, Dict, List, Optional, Set, Type, Union
+from types import ModuleType, TracebackType
+from typing import (
+    TYPE_CHECKING,
+    Optional,
+    Set,
+    Type,
+    Union,
+)
 
 from .exporter import WaldiezExporter
+from .io import StructuredIOStream
 from .models.waldiez import Waldiez
 from .running import (
     a_chdir,
@@ -25,13 +33,11 @@ from .running import (
     after_run,
     before_run,
     chdir,
-    get_printer,
     install_requirements,
     refresh_environment,
     reset_env_vars,
     set_env_vars,
 )
-from .utils import check_pysqlite3
 
 if TYPE_CHECKING:
     from autogen import ChatResult  # type: ignore
@@ -56,8 +62,8 @@ class WaldiezRunner:
         waldiez_file: Union[str, Path],
         name: Optional[str] = None,
         description: Optional[str] = None,
-        tags: Optional[List[str]] = None,
-        requirements: Optional[List[str]] = None,
+        tags: Optional[list[str]] = None,
+        requirements: Optional[list[str]] = None,
     ) -> "WaldiezRunner":
         """Create a WaldiezRunner instance from a file.
 
@@ -69,9 +75,9 @@ class WaldiezRunner:
             The name of the Waldiez, by default None.
         description : Optional[str], optional
             The description of the Waldiez, by default None.
-        tags : Optional[List[str]], optional
+        tags : Optional[list[str]], optional
             The tags of the Waldiez, by default None.
-        requirements : Optional[List[str]], optional
+        requirements : Optional[list[str]], optional
             The requirements of the Waldiez, by default None.
 
         Returns
@@ -132,6 +138,11 @@ class WaldiezRunner:
         """Get the Waldiez instance."""
         return self._waldiez
 
+    @property
+    def is_async(self) -> bool:
+        """Check if the workflow is async."""
+        return self.waldiez.is_async
+
     @property
     def running(self) -> bool:
         """Get the running status."""
@@ -148,34 +159,50 @@ class WaldiezRunner:
         extra_requirements = {
             req for req in self.waldiez.requirements if req not in sys.modules
         }
-        if self.waldiez.has_captain_agents:
-            check_pysqlite3()
         return extra_requirements
 
     def install_requirements(self) -> None:
         """Install the requirements for the flow."""
         self._called_install_requirements = True
-        printer = get_printer()
         extra_requirements = self.gather_requirements()
         if extra_requirements:
-            install_requirements(extra_requirements, printer)
-            refresh_environment()
+            install_requirements(extra_requirements)
 
     async def a_install_requirements(self) -> None:
         """Install the requirements for the flow asynchronously."""
         self._called_install_requirements = True
-        printer = get_printer()
         extra_requirements = self.gather_requirements()
         if extra_requirements:
-            await a_install_requirements(extra_requirements, printer)
-            refresh_environment()
+            await a_install_requirements(extra_requirements)
+
+    def _before_run(
+        self,
+        temp_dir: Path,
+        file_name: str,
+        module_name: str,
+    ) -> tuple[ModuleType, dict[str, str]]:
+        self._exporter.export(Path(file_name))
+        # unique_names = self._exporter.context.get_unique_names()
+        spec = importlib.util.spec_from_file_location(
+            module_name, temp_dir / file_name
+        )
+        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)
+        print("<Waldiez> - Starting workflow...")
+        print(self.waldiez.info.model_dump_json())
+        return module, old_vars
 
     def _run(
         self,
         output_path: Optional[Union[str, Path]],
         uploads_root: Optional[Union[str, Path]],
+        use_structured_io: bool = False,
         skip_mmd: bool = False,
-    ) -> Union["ChatResult", List["ChatResult"], Dict[int, "ChatResult"]]:
+    ) -> Union["ChatResult", list["ChatResult"], dict[int, "ChatResult"]]:
         """Run the Waldiez workflow.
 
         Parameters
@@ -184,12 +211,14 @@ class WaldiezRunner:
             The output path.
         uploads_root : Optional[Union[str, Path]]
             The runtime uploads root.
+        use_structured_io : bool
+            Whether to use structured IO instead of the default 'input/print'.
         skip_mmd : bool
             Whether to skip the Mermaid diagram generation.
 
         Returns
         -------
-        Union[ChatResult, List[ChatResult]]
+        Union[ChatResult, list[ChatResult]]
             The result(s) of the chat(s).
         """
         temp_dir = Path(tempfile.mkdtemp())
@@ -197,95 +226,99 @@ class WaldiezRunner:
         module_name = file_name.replace(".py", "")
         if not self._called_install_requirements:
             self.install_requirements()
-        else:
-            refresh_environment()
-        printer = get_printer()
-        printer(
+        refresh_environment()
+        print(
             "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"], Dict[int, "ChatResult"]
+            "ChatResult", list["ChatResult"], dict[int, "ChatResult"]
         ] = []
         with chdir(to=temp_dir):
-            self._exporter.export(Path(file_name))
-            spec = importlib.util.spec_from_file_location(
-                module_name, temp_dir / file_name
+            module, old_vars = self._before_run(
+                temp_dir=temp_dir,
+                file_name=file_name,
+                module_name=module_name,
             )
-            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()
-            printer("<Waldiez> - Workflow finished")
+            if use_structured_io:
+                stream = StructuredIOStream(
+                    timeout=120,  # 2 minutes
+                    uploads_root=uploads_root,
+                )
+                with StructuredIOStream.set_default(stream):
+                    results = module.main()
+            else:
+                results = module.main()
+            print("<Waldiez> - Workflow finished")
             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,
         )
+        gc.collect()
+        refresh_environment()
         return results
 
     async def _a_run(
         self,
         output_path: Optional[Union[str, Path]],
         uploads_root: Optional[Union[str, Path]],
+        use_structured_io: bool = False,
         skip_mmd: bool = False,
-    ) -> Union["ChatResult", List["ChatResult"], Dict[int, "ChatResult"]]:
+    ) -> Union["ChatResult", list["ChatResult"], dict[int, "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", "")
         if not self._called_install_requirements:
             await self.a_install_requirements()
-        else:
-            refresh_environment()
-        printer = get_printer()
-        printer(
+        refresh_environment()
+        print(
             "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"], Dict[int, "ChatResult"]
+            "ChatResult", list["ChatResult"], dict[int, "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
+            module, old_vars = self._before_run(
+                temp_dir=temp_dir,
+                file_name=file_name,
+                module_name=module_name,
             )
-            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 = await module.main()
+            if use_structured_io:
+                stream = StructuredIOStream(
+                    timeout=120,  # 2 minutes
+                    uploads_root=uploads_root,
+                )
+                with StructuredIOStream.set_default(stream):
+                    results = await module.main()
+            else:
+                results = await 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,
         )
+        gc.collect()
+        refresh_environment()
         return results
 
     def run(
         self,
         output_path: Optional[Union[str, Path]] = None,
         uploads_root: Optional[Union[str, Path]] = None,
+        use_structured_io: bool = False,
         skip_mmd: bool = False,
-    ) -> Union["ChatResult", List["ChatResult"], Dict[int, "ChatResult"]]:
+    ) -> Union["ChatResult", list["ChatResult"], dict[int, "ChatResult"]]:
         """Run the Waldiez workflow.
 
         Parameters
@@ -294,12 +327,15 @@ 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.
+        use_structured_io : bool, optional
+            Whether to use structured IO instead of the default 'input/print',
+            by default False.
         skip_mmd : bool, optional
             Whether to skip the Mermaid diagram generation, by default False.
 
         Returns
         -------
-        Union["ChatResult", List["ChatResult"], Dict[int, "ChatResult"]]
+        Union["ChatResult", list["ChatResult"], dict[int, "ChatResult"]]
             The result(s) of the chat(s).
 
         Raises
@@ -316,6 +352,7 @@ class WaldiezRunner:
                     self._a_run,
                     output_path,
                     uploads_root,
+                    use_structured_io,
                     skip_mmd,
                 )
         if self._running is True:
@@ -323,7 +360,12 @@ class WaldiezRunner:
         self._running = True
         file_path = output_path or self._file_path
         try:
-            return self._run(file_path, uploads_root, skip_mmd)
+            return self._run(
+                file_path,
+                uploads_root=uploads_root,
+                use_structured_io=use_structured_io,
+                skip_mmd=skip_mmd,
+            )
         finally:
             self._running = False
 
@@ -331,7 +373,9 @@ class WaldiezRunner:
         self,
         output_path: Optional[Union[str, Path]] = None,
         uploads_root: Optional[Union[str, Path]] = None,
-    ) -> Union["ChatResult", List["ChatResult"]]:
+        use_structured_io: bool = False,
+        skip_mmd: bool = False,
+    ) -> Union["ChatResult", list["ChatResult"], dict[int, "ChatResult"]]:
         """Run the Waldiez workflow asynchronously.
 
         Parameters
@@ -340,10 +384,15 @@ 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.
+        use_structured_io : bool, optional
+            Whether to use structured IO instead of the default 'input/print',
+            by default False.
+        skip_mmd : bool, optional
+            Whether to skip the Mermaid diagram generation, by default False.
 
         Returns
         -------
-        Union[ChatResult, List[ChatResult]]
+        Union[ChatResult, list[ChatResult]], dict[int, ChatResult]
             The result(s) of the chat(s).
 
         Raises
@@ -356,6 +405,11 @@ class WaldiezRunner:
         self._running = True
         file_path = output_path or self._file_path
         try:
-            return await self._a_run(file_path, uploads_root)
+            return await self._a_run(
+                file_path,
+                uploads_root=uploads_root,
+                use_structured_io=use_structured_io,
+                skip_mmd=skip_mmd,
+            )
         finally:
             self._running = False

waldiez/running/__init__.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0.
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
-"""Common utility functions."""
+"""Running related functions."""
 
 from .environment import (
     in_virtualenv,
@@ -9,23 +9,29 @@ from .environment import (
     reset_env_vars,
     set_env_vars,
 )
-from .running import (
-    a_chdir,
+from .post_run import after_run
+from .pre_run import (
     a_install_requirements,
-    after_run,
     before_run,
-    chdir,
-    get_printer,
     install_requirements,
 )
+from .util import (
+    a_chdir,
+    chdir,
+    create_async_subprocess,
+    create_sync_subprocess,
+    strip_ansi,
+)
 
 __all__ = [
     "a_chdir",
     "a_install_requirements",
     "after_run",
     "before_run",
+    "create_async_subprocess",
+    "create_sync_subprocess",
+    "strip_ansi",
     "chdir",
-    "get_printer",
     "in_virtualenv",
     "is_root",
     "install_requirements",

waldiez/running/environment.py

@@ -3,12 +3,9 @@
 # pylint: disable=import-outside-toplevel,reimported
 """Environment related utilities."""
 
-import importlib
 import os
-import site
 import sys
-import warnings
-from typing import Dict, Generator, List, Tuple
+from typing import Generator
 
 
 def in_virtualenv() -> bool:
@@ -25,7 +22,7 @@ def in_virtualenv() -> bool:
     )
 
 
-def is_root() -> bool:
+def is_root() -> bool:  # pragma: no cover  # os specific
     """Check if the script is running as root/administrator.
 
     Returns
@@ -33,13 +30,15 @@ def is_root() -> bool:
     bool
         True if running as root/administrator, False otherwise.
     """
-    # pylint: disable=import-outside-toplevel,line-too-long
+    # pylint: disable=import-outside-toplevel,line-too-long,no-member
     if os.name == "nt":
+        # noinspection PyBroadException
+        # pylint: disable=broad-exception-caught
         try:
             import ctypes
 
             return ctypes.windll.shell32.IsUserAnAdmin() != 0  # type: ignore[unused-ignore,attr-defined]  # noqa: E501
-        except Exception:  # pylint: disable=broad-exception-caught
+        except Exception:
             return False
     else:
         return os.getuid() == 0
@@ -47,75 +46,57 @@ def is_root() -> bool:
 
 def refresh_environment() -> None:
     """Refresh the environment."""
-    with warnings.catch_warnings():
-        warnings.filterwarnings(
-            "ignore",
-            module="flaml",
-            message="^.*flaml.automl is not available.*$",
-        )
-        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]
-        import autogen  # type: ignore
-        from autogen.io import IOStream
-
-        importlib.reload(autogen)
-        # restore the default IOStream
-        IOStream.set_global_default(default_io_stream)
-        # reload any other modules that may have been affected
-        for mod in modules_to_reload:
-            if mod not in sys.modules:
-                importlib.import_module(mod)
-        # a swarm chat without a user agent
-        # creates a new user (this has a default code execution with docker)
-        # captain also generates new agents that also have
-        # default code execution with docker
-        # temp (until we handle/detect docker setup)
-        os.environ["AUTOGEN_USE_DOCKER"] = "0"
-        # we might get:
-        # module 'numpy' has no attribute '_no_nep50_warning'
-        # in autogen/agentchat/contrib/captainagent/tool_retriever.py
-        os.environ["NEP50_DEPRECATION_WARNING"] = "0"
-        os.environ["NEP50_DISABLE_WARNING"] = "1"
-        os.environ["NPY_PROMOTION_STATE"] = "weak"
-        import numpy as np
-
-        if not hasattr(np, "_no_pep50_warning"):
-            import contextlib
-
-            @contextlib.contextmanager
-            def _np_no_nep50_warning() -> Generator[None, None, None]:
-                """Avoid no_nep50 warning.
-
-                Yields
-                ------
-                None
-                    Dummy value.
-                """
-                yield
-
-            setattr(np, "_no_pep50_warning", _np_no_nep50_warning)  # noqa
-
-
-def set_env_vars(flow_env_vars: List[Tuple[str, str]]) -> Dict[str, str]:
+    # a group chat without a user agent
+    # creates a new user (this has a default code execution with docker)
+    # captain also generates new agents that also have
+    # default code execution with docker
+    # temp (until we handle/detect docker setup)
+    os.environ["AUTOGEN_USE_DOCKER"] = "0"
+    try_handle_the_np_thing()
+
+
+def try_handle_the_np_thing() -> None:
+    """Try to handle the numpy deprecation warning."""
+    # we might get:
+    # module 'numpy' has no attribute '_no_nep50_warning'
+    # (sentnence_transformers?)
+    # in autogen/agentchat/contrib/captainagent/tool_retriever.py
+    os.environ["NEP50_DEPRECATION_WARNING"] = "0"
+    os.environ["NEP50_DISABLE_WARNING"] = "1"
+    os.environ["NPY_PROMOTION_STATE"] = "weak"
+    import numpy as np
+
+    if not hasattr(np, "_no_pep50_warning"):
+        import contextlib
+
+        @contextlib.contextmanager
+        def _np_no_nep50_warning() -> Generator[None, None, None]:
+            """Avoid no_nep50 warning.
+
+            Yields
+            ------
+            None
+                Nothing.
+            """
+            yield  # pragma: no cover
+
+        setattr(np, "_no_pep50_warning", _np_no_nep50_warning)  # noqa
+
+
+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]]
+    flow_env_vars : list[tuple[str, str]]
         The environment variables to set.
 
     Returns
     -------
-    Dict[str, str]
+    dict[str, str]
         The old environment variables.
     """
-    old_vars: Dict[str, str] = {}
+    old_vars: dict[str, str] = {}
     for var_key, var_value in flow_env_vars:
         if var_key:
             current = os.environ.get(var_key, "")
@@ -124,12 +105,12 @@ def set_env_vars(flow_env_vars: List[Tuple[str, str]]) -> Dict[str, str]:
     return old_vars
 
 
-def reset_env_vars(old_vars: Dict[str, str]) -> None:
+def reset_env_vars(old_vars: dict[str, str]) -> None:
     """Reset the environment variables.
 
     Parameters
     ----------
-    old_vars : Dict[str, str]
+    old_vars : dict[str, str]
         The old environment variables.
     """
     for var_key, var_value in old_vars.items():