waldiez 0.4.8__py3-none-any.whl → 0.4.11__py3-none-any.whl

waldiez/running/patch_io_stream.py

@@ -0,0 +1,208 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+
+# pylint: disable=import-outside-toplevel, broad-exception-caught
+# pylint: disable=too-many-try-statements, too-complex
+# noqa: C901
+"""Patch ag2's IOStream if a flow is async.
+
+# let's keep an eye here:
+# https://github.com/ag2ai/ag2/blob/main/autogen/agentchat/conversable_agent.py#L2973
+# reply = await iostream.input(prompt) ???? (await???)
+"""
+
+import inspect
+import sys
+from typing import Any, Callable
+
+from asyncer import syncify
+
+
+def patch_io_stream(is_async: bool) -> None:
+    """Patch the IOStream to handle async flows.
+
+    Parameters
+    ----------
+    is_async : bool
+        Whether the flow is async or not.
+    """
+    if is_async:
+        patch_async_io_stream()
+    else:
+        patch_sync_io_stream()
+
+
+def patch_sync_io_stream() -> None:
+    """Patch the IOStream to handle async flows."""
+    from autogen.io import IOStream  # type: ignore
+
+    iostream = IOStream.get_default()
+    original_input = iostream.input
+
+    def _safe_input(prompt: str = "", *, password: bool = False) -> str:
+        """Async input method."""
+        try:
+            input_or_coro = original_input(prompt, password=password)
+            if inspect.iscoroutine(input_or_coro):
+
+                async def _async_input() -> str:
+                    reply = await input_or_coro
+                    return reply
+
+                return syncify(_async_input)()
+            return input_or_coro
+
+        except EOFError:
+            # Handle EOFError gracefully
+            return ""
+
+    iostream.input = _safe_input  # pyright: ignore
+    iostream.print = get_printer()  # pyright: ignore
+    IOStream.set_default(iostream)
+
+
+def patch_async_io_stream() -> None:
+    """Patch the IOStream to handle async flows."""
+    from autogen.io import IOStream  # pyright: ignore
+
+    iostream = IOStream.get_default()
+    original_input = iostream.input
+
+    async def _async_input(prompt: str = "", *, password: bool = False) -> str:
+        """Async input method."""
+        try:
+            input_or_coro = original_input(prompt, password=password)
+            if inspect.iscoroutine(input_or_coro):
+                reply = await input_or_coro
+            else:
+                reply = input_or_coro
+            return reply
+        except EOFError:
+            # Handle EOFError gracefully
+            return ""
+
+    iostream.input = _async_input  # pyright: ignore
+    iostream.print = get_printer()  # pyright: ignore
+    IOStream.set_default(iostream)
+
+
+def get_printer() -> Callable[..., None]:  # noqa: C901
+    """Get the printer function.
+
+    Returns
+    -------
+    Callable[..., None]
+        The printer function that handles Unicode encoding errors gracefully.
+    """
+    try:
+        from autogen.io import IOStream  # pyright: ignore
+
+        printer = IOStream.get_default().print
+    except ImportError:
+        # Fallback to standard print if autogen is not available
+        printer = print
+
+    def safe_printer(*args: Any, **kwargs: Any) -> None:  # noqa: C901
+        """Safe printer that handles Unicode encoding errors.
+
+        Parameters
+        ----------
+        *args : Any
+            Arguments to pass to the printer
+        **kwargs : Any
+            Keyword arguments to pass to the printer
+        """
+        try:
+            printer(*args, **kwargs)
+        except UnicodeEncodeError:
+            # First fallback: try to get a safe string representation
+            try:
+                msg, flush = get_what_to_print(*args, **kwargs)
+                # Convert problematic characters to safe representations
+                safe_msg = msg.encode("utf-8", errors="replace").decode("utf-8")
+                printer(safe_msg, end="", flush=flush)
+            except (UnicodeEncodeError, UnicodeDecodeError):
+                # Second fallback: use built-in print with safe encoding
+                try:
+                    # Convert args to safe string representations
+                    safe_args: list[str] = []
+                    for arg in args:
+                        try:
+                            safe_args.append(
+                                str(arg)
+                                .encode("utf-8", errors="replace")
+                                .decode("utf-8")
+                            )
+                        except (UnicodeEncodeError, UnicodeDecodeError):
+                            safe_args.append(repr(arg))
+
+                    # Use built-in print instead of the custom printer
+                    print(*safe_args, **kwargs)
+
+                except Exception:
+                    # Final fallback: write directly to stderr buffer
+                    try:
+                        error_msg = (
+                            "Could not print the message "
+                            "due to encoding issues.\n"
+                        )
+                        if hasattr(sys.stderr, "buffer"):
+                            sys.stderr.buffer.write(
+                                error_msg.encode("utf-8", errors="replace")
+                            )
+                            sys.stderr.buffer.flush()
+                        else:
+                            sys.stderr.write(error_msg)
+                            sys.stderr.flush()
+                    except Exception:
+                        # If even this fails, we're in a very bad state
+                        pass
+        except Exception as e:
+            # Handle any other unexpected errors
+            try:
+                error_msg = f"Unexpected error in printer: {str(e)}\n"
+                if hasattr(sys.stderr, "buffer"):
+                    sys.stderr.buffer.write(
+                        error_msg.encode("utf-8", errors="replace")
+                    )
+                    sys.stderr.buffer.flush()
+                else:
+                    sys.stderr.write(error_msg)
+                    sys.stderr.flush()
+            except Exception:
+                pass
+
+    return safe_printer
+
+
+def get_what_to_print(*args: Any, **kwargs: Any) -> tuple[str, bool]:
+    """Extract message and flush flag from print arguments.
+
+    Parameters
+    ----------
+    *args : Any
+        Arguments to print
+    **kwargs : Any
+        Keyword arguments for print function
+
+    Returns
+    -------
+    tuple[str, bool]
+        Message to print and flush flag
+    """
+    # Convert all args to strings and join with spaces (like print does)
+    msg = " ".join(str(arg) for arg in args)
+
+    # Handle sep parameter
+    sep = kwargs.get("sep", " ")
+    if len(args) > 1:
+        msg = sep.join(str(arg) for arg in args)
+
+    # Handle end parameter
+    end = kwargs.get("end", "\n")
+    msg += end
+
+    # Handle flush parameter
+    flush = kwargs.get("flush", False)
+
+    return msg, flush

waldiez/running/post_run.py

@@ -0,0 +1,121 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+
+# pylint: disable=unused-argument
+"""Utilities for running code."""
+
+import datetime
+import shutil
+from pathlib import Path
+from typing import Optional, Union
+
+from .gen_seq_diagram import generate_sequence_diagram
+
+
+def after_run(
+    temp_dir: Path,
+    output_file: Optional[Union[str, Path]],
+    flow_name: str,
+    uploads_root: Optional[Path] = None,
+    skip_mmd: bool = False,
+) -> None:
+    """Actions to perform after running the flow.
+
+    Parameters
+    ----------
+    temp_dir : Path
+        The temporary directory.
+    output_file : Optional[Union[str, Path]]
+        The output file.
+    flow_name : str
+        The flow name.
+    uploads_root : Optional[Path], optional
+        The runtime uploads root, by default None
+    skip_mmd : bool, optional
+        Whether to skip the mermaid sequence diagram generation,
+        by default, False
+    """
+    if isinstance(output_file, str):
+        output_file = Path(output_file)
+    mmd_dir = output_file.parent if output_file else Path.cwd()
+    if skip_mmd is False:
+        events_csv_path = temp_dir / "logs" / "events.csv"
+        if events_csv_path.exists():
+            print("Generating mermaid sequence diagram...")
+            mmd_path = temp_dir / f"{flow_name}.mmd"
+            generate_sequence_diagram(events_csv_path, mmd_path)
+            if (
+                not output_file
+                and mmd_path.exists()
+                and mmd_path != mmd_dir / f"{flow_name}.mmd"
+            ):
+                try:
+                    shutil.copyfile(mmd_path, mmd_dir / f"{flow_name}.mmd")
+                except BaseException:  # pylint: disable=broad-exception-caught
+                    pass
+    if output_file:
+        destination_dir = output_file.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
+        print(f"Copying the results to {destination_dir}")
+        copy_results(
+            temp_dir=temp_dir,
+            output_file=output_file,
+            destination_dir=destination_dir,
+        )
+    shutil.rmtree(temp_dir)
+
+
+def copy_results(
+    temp_dir: Path,
+    output_file: Path,
+    destination_dir: Path,
+) -> None:
+    """Copy the results to the output directory.
+
+    Parameters
+    ----------
+    temp_dir : Path
+        The temporary directory.
+    output_file : Path
+        The output file.
+    destination_dir : Path
+        The destination directory.
+    """
+    temp_dir.mkdir(parents=True, exist_ok=True)
+    output_dir = output_file.parent
+    for item in temp_dir.iterdir():
+        # skip cache files
+        if (
+            item.name.startswith("__pycache__")
+            or item.name.endswith(".pyc")
+            or item.name.endswith(".pyo")
+            or item.name.endswith(".pyd")
+            or item.name == ".cache"
+        ):
+            continue
+        if item.is_file():
+            # let's also copy the "tree of thoughts" image
+            # to the output directory
+            if item.name.endswith("tree_of_thoughts.png") or item.name.endswith(
+                "reasoning_tree.json"
+            ):
+                shutil.copy(item, output_dir / item.name)
+            shutil.copy(item, destination_dir)
+        else:
+            shutil.copytree(item, destination_dir / item.name)
+    if output_file.is_file():
+        if output_file.suffix == ".waldiez":
+            output_file = output_file.with_suffix(".py")
+        if output_file.suffix == ".py":
+            src = temp_dir / output_file.name
+            if src.exists():
+                dst = destination_dir / output_file.name
+                if dst.exists():
+                    dst.unlink()
+                shutil.copyfile(src, output_dir / output_file.name)

waldiez/running/pre_run.py

@@ -0,0 +1,105 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Actions to perform before running the flow."""
+
+import asyncio
+import io
+import os
+import subprocess
+import sys
+from typing import Callable
+
+from .environment import in_virtualenv, is_root
+from .utils import strip_ansi
+
+
+def install_requirements(
+    extra_requirements: set[str],
+    printer: Callable[..., None] = print,
+) -> None:
+    """Install the requirements.
+
+    Parameters
+    ----------
+    extra_requirements : set[str]
+        The extra requirements.
+    printer : Callable[..., None]
+        The printer function to use, defaults to print.
+    """
+    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")
+    pip_install.extend(extra_requirements)
+    # pylint: disable=too-many-try-statements
+    try:
+        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(strip_ansi(line.strip()))
+            if proc.stderr:
+                for line in io.TextIOWrapper(proc.stderr, encoding="utf-8"):
+                    printer(strip_ansi(line.strip()))
+    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"]
+
+
+async def a_install_requirements(
+    extra_requirements: set[str],
+    printer: Callable[..., None] = print,
+) -> None:
+    """Install the requirements asynchronously.
+
+    Parameters
+    ----------
+    extra_requirements : set[str]
+        The extra requirements.
+    printer : Callable[..., None]
+        The printer function to use, defaults to print.
+    """
+    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"])
+    pip_install.extend(extra_requirements)
+    # pylint: disable=too-many-try-statements
+    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()))
+    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"]