waldiez 0.4.9__py3-none-any.whl → 0.5.0__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

@@ -1,5 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0.
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
+
+# pylint: disable=unused-argument
 """Utilities for running code."""
 
 import datetime
@@ -12,8 +14,9 @@ from .gen_seq_diagram import generate_sequence_diagram
 
 def after_run(
     temp_dir: Path,
-    output_path: Optional[Union[str, 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.
@@ -22,17 +25,19 @@ def after_run(
     ----------
     temp_dir : Path
         The temporary directory.
-    output_path : Optional[Union[str, Path]]
-        The output path.
+    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_path, str):
-        output_path = Path(output_path)
-    mmd_dir = output_path.parent if output_path else Path.cwd()
+    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():
@@ -40,7 +45,7 @@ def after_run(
             mmd_path = temp_dir / f"{flow_name}.mmd"
             generate_sequence_diagram(events_csv_path, mmd_path)
             if (
-                not output_path
+                not output_file
                 and mmd_path.exists()
                 and mmd_path != mmd_dir / f"{flow_name}.mmd"
             ):
@@ -48,8 +53,8 @@ def after_run(
                     shutil.copyfile(mmd_path, mmd_dir / f"{flow_name}.mmd")
                 except BaseException:  # pylint: disable=broad-exception-caught
                     pass
-    if output_path:
-        destination_dir = output_path.parent
+    if output_file:
+        destination_dir = output_file.parent
         destination_dir = (
             destination_dir
             / "waldiez_out"
@@ -60,8 +65,7 @@ def after_run(
         print(f"Copying the results to {destination_dir}")
         copy_results(
             temp_dir=temp_dir,
-            output_path=output_path,
-            output_dir=output_path.parent,
+            output_file=output_file,
             destination_dir=destination_dir,
         )
     shutil.rmtree(temp_dir)
@@ -69,8 +73,7 @@ def after_run(
 
 def copy_results(
     temp_dir: Path,
-    output_path: Path,
-    output_dir: Path,
+    output_file: Path,
     destination_dir: Path,
 ) -> None:
     """Copy the results to the output directory.
@@ -79,14 +82,13 @@ def copy_results(
     ----------
     temp_dir : Path
         The temporary directory.
-    output_path : Path
-        The output path.
-    output_dir : Path
-        The output 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 (
@@ -107,13 +109,13 @@ def copy_results(
             shutil.copy(item, destination_dir)
         else:
             shutil.copytree(item, destination_dir / item.name)
-    if output_path.is_file():
-        if output_path.suffix == ".waldiez":
-            output_path = output_path.with_suffix(".py")
-        if output_path.suffix == ".py":
-            src = temp_dir / output_path.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_path.name
+                dst = destination_dir / output_file.name
                 if dst.exists():
                     dst.unlink()
-                shutil.copyfile(src, output_dir / output_path.name)
+                shutil.copyfile(src, output_dir / output_file.name)

waldiez/running/pre_run.py

@@ -7,54 +7,10 @@ import io
 import os
 import subprocess
 import sys
-import tempfile
-from pathlib import Path
-from typing import Callable, Optional, Union
+from typing import Callable
 
 from .environment import in_virtualenv, is_root
-from .util import strip_ansi
-
-
-def before_run(
-    output_path: Optional[Union[str, Path]],
-    uploads_root: Optional[Union[str, Path]],
-) -> str:
-    """Actions to perform before running the flow.
-
-    Parameters
-    ----------
-    output_path : Optional[Union[str, Path]]
-        The output path.
-    uploads_root : Optional[Union[str, Path]]
-        The runtime uploads root.
-
-    Returns
-    -------
-    str
-        The file name.
-    """
-    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)
-    output_dir = Path.cwd()
-    if output_path and isinstance(output_path, str):
-        output_path = Path(output_path)
-    if output_path:
-        if output_path.is_dir():
-            output_dir = output_path
-        else:
-            output_dir = output_path.parent if output_path else Path.cwd()
-    if not output_dir.exists():
-        output_dir.mkdir(parents=True, exist_ok=True)
-    file_name = Path(output_path).name if output_path else "waldiez_flow.py"
-    if file_name.endswith((".json", ".waldiez")):
-        file_name = file_name.replace(".json", ".py").replace(".waldiez", ".py")
-    if not file_name.endswith(".py"):
-        file_name += ".py"
-    return file_name
+from .utils import strip_ansi
 
 
 def install_requirements(

waldiez/running/protocol.py

@@ -0,0 +1,281 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+
+"""Waldiez Runner protocol."""
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Protocol, Union, runtime_checkable
+
+if TYPE_CHECKING:
+    from autogen import ChatResult  # type: ignore[import-untyped]
+
+
+@runtime_checkable
+class WaldiezRunnerProtocol(Protocol):
+    """Waldiez Runner protocol."""
+
+    def before_run(
+        self,
+        output_file: Path,
+        uploads_root: Path | None,
+    ) -> Path:  # pyright: ignore
+        """Actions to perform before running the flow.
+
+        Parameters
+        ----------
+        output_file : Path
+            The output file.
+        uploads_root : Path | None
+            The runtime uploads root.
+
+        Returns
+        -------
+        Path
+            The path to the temporary directory created for the run.
+        """
+
+    async def a_before_run(
+        self,
+        output_file: Path,
+        uploads_root: Path | None,
+    ) -> Path:  # pyright: ignore
+        """Asynchronously perform actions before running the flow.
+
+        Parameters
+        ----------
+        output_file : Path
+            The output file.
+        uploads_root : Path | None
+            The runtime uploads root.
+
+        Returns
+        -------
+        Path
+            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_patch_io: bool | None = None,
+        skip_mmd: bool = False,
+    ) -> 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_patch_io : bool | None
+            Whether to skip patching I/O, by default None.
+            If None, it will use the value from the context.
+        skip_mmd : bool
+            Whether to skip generating the mermaid diagram.
+
+        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_patch_io: bool | None = None,
+        skip_mmd: bool = False,
+    ) -> 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_patch_io : bool | None
+            Whether to skip patching I/O, by default None.
+            If None, it will use the value from the context.
+        skip_mmd : bool
+            Whether to skip generating the mermaid diagram.
+
+        Raises
+        ------
+        RuntimeError
+            If the runner is already running.
+        """
+
+    def run(
+        self,
+        output_path: str | Path | None,
+        uploads_root: str | Path | 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
+        """Run the Waldiez flow in a blocking way.
+
+        Parameters
+        ----------
+        output_path : str | Path | None
+            The output path, by default None.
+        uploads_root : str | Path | None
+            The runtime uploads root.
+        structured_io : bool
+            Whether to use structured IO instead of the default 'input/print'.
+        threaded : bool | None
+            Whether to run the flow in a separate thread.
+        skip_patch_io : bool
+            Whether to skip patching I/O, by default None.
+            If None, it will use the value from the context.
+        skip_mmd : bool
+            Whether to skip generating the mermaid diagram.
+
+        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.
+        """
+
+    async def a_run(
+        self,
+        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,
+    ) -> Union[
+        "ChatResult",
+        list["ChatResult"],
+        dict[int, "ChatResult"],
+    ]:  # pyright: ignore
+        """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.
+        structured_io : bool
+            Whether to use structured IO instead of the default 'input/print'.
+        skip_patch_io : bool
+            Whether to skip patching I/O, by default None.
+            If None, it will use the value from the context.
+        skip_mmd : bool
+            Whether to skip generating the mermaid diagram.
+
+        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.
+        """
+
+    def after_run(
+        self,
+        results: Union[
+            "ChatResult",
+            list["ChatResult"],
+            dict[int, "ChatResult"],
+        ],
+        output_file: Path,
+        uploads_root: Path | None,
+        temp_dir: Path,
+        skip_mmd: bool,
+    ) -> None:
+        """Actions to perform after running the flow.
+
+        Parameters
+        ----------
+        results : Union[ChatResult, list[ChatResult], dict[int, ChatResult]]
+            The results of the run, which can be a single ChatResult,
+            a list of ChatResults,
+            or a dictionary mapping indices to ChatResults.
+        output_file : Path
+            The path to the output file.
+        uploads_root : Path | None
+            The runtime uploads root.
+        temp_dir : Path
+            The path to the temporary directory.
+        skip_mmd : bool
+            Whether to skip generating the mermaid diagram.
+        """
+
+    async def a_after_run(
+        self,
+        results: Union[
+            "ChatResult",
+            list["ChatResult"],
+            dict[int, "ChatResult"],
+        ],
+        output_file: Path,
+        uploads_root: Path | None,
+        temp_dir: Path,
+        skip_mmd: bool,
+    ) -> None:
+        """Asynchronously perform actions after running the flow.
+
+        Parameters
+        ----------
+        results : Union[ChatResult, list[ChatResult], dict[int, ChatResult]]
+            The results of the run, which can be a single ChatResult,
+            a list of ChatResults,
+            or a dictionary mapping indices to ChatResults.
+        output_file : Path
+            The path to the output file.
+        uploads_root : Path | None
+            The runtime uploads root.
+        temp_dir : Path
+            The path to the temporary directory.
+        skip_mmd : bool
+            Whether to skip generating the mermaid diagram.
+        """
+
+    def is_running(self) -> bool:  # pyright: ignore
+        """Check if the runner is currently running.
+
+        Returns
+        -------
+        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

@@ -0,0 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+
+"""Waldiez run results module."""
+
+from typing import TYPE_CHECKING, TypedDict, Union
+
+if TYPE_CHECKING:
+    from autogen import ChatResult  # type: ignore[import-untyped]
+
+
+class WaldiezRunResults(TypedDict):
+    """Results of the Waldiez run."""
+
+    results: Union[
+        "ChatResult",
+        list["ChatResult"],
+        dict[int, "ChatResult"],
+        None,
+    ]
+    exception: Exception | None
+    completed: bool