waldiez 0.5.8__py3-none-any.whl → 0.5.10__py3-none-any.whl

waldiez/running/utils.py

@@ -3,17 +3,44 @@
 """Common utilities for the waldiez runner."""
 
 import asyncio
+import functools
+import inspect
+import logging
 import os
 import re
 import subprocess
 import sys
+import threading
+import traceback
 
 # noinspection PyProtectedMember
 from asyncio.subprocess import Process
 from contextlib import asynccontextmanager, contextmanager
 from dataclasses import dataclass
+from getpass import getpass
 from pathlib import Path
-from typing import AsyncIterator, Iterator, Union
+from typing import (
+    Any,
+    AsyncIterator,
+    Callable,
+    Coroutine,
+    Generic,
+    Iterator,
+    TypeVar,
+    Union,
+    cast,
+)
+
+logger = logging.getLogger(__name__)
+T = TypeVar("T")
+
+
+@dataclass
+class _ResultContainer(Generic[T]):
+    """Container for thread execution results with proper typing."""
+
+    result: T | None = None
+    exception: BaseException | None = None
 
 
 @dataclass
@@ -134,3 +161,351 @@ async def create_async_subprocess(setup: ProcessSetup) -> Process:
         # stdin=asyncio.subprocess.PIPE,
         env={**os.environ},
     )
+
+
+async def input_async(prompt: str, *, password: bool = False) -> str:
+    """Asynchronous input function.
+
+    Parameters
+    ----------
+    prompt : str
+        The prompt to display to the user.
+    password : bool, optional
+        Whether to hide input (password mode), by default False.
+
+    Returns
+    -------
+    str
+        The user input.
+    """
+    if password:
+        try:
+            return await asyncio.to_thread(getpass, prompt)
+        except EOFError:
+            return ""
+    try:
+        return await asyncio.to_thread(input, prompt)
+    except EOFError:
+        return ""
+
+
+def input_sync(prompt: str, *, password: bool = False) -> str:
+    """Input function (synchronous).
+
+    Parameters
+    ----------
+    prompt : str
+        The prompt to display to the user.
+    password : bool, optional
+        Whether to hide input (password mode), by default False.
+
+    Returns
+    -------
+    str
+        The user input.
+    """
+    if password:
+        try:
+            return getpass(prompt)
+        except EOFError:
+            return ""
+    try:
+        return input(prompt)
+    except EOFError:
+        return ""
+
+
+# pylint: disable=import-outside-toplevel,too-complex
+def get_printer() -> Callable[..., None]:  # noqa: C901
+    """Get the printer function.
+
+    Returns
+    -------
+    Callable[..., None]
+        The printer function that handles Unicode encoding errors gracefully.
+    """
+    try:
+        # noinspection PyUnresolvedReferences
+        from autogen.io import IOStream  # type: ignore
+
+        printer = IOStream.get_default().print
+    except ImportError:  # pragma: no cover
+        # Fallback to standard print if autogen is not available
+        printer = print
+
+    # noinspection PyBroadException,TryExceptPass
+    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
+        """
+        # pylint: disable=broad-exception-caught,too-many-try-statements
+        try:
+            printer(*args, **kwargs)
+        except (UnicodeEncodeError, UnicodeDecodeError):
+            # 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,
+                        ):  # pragma: no cover
+                            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
+                    error_msg = (
+                        "Could not print the message due to encoding issues.\n"
+                    )
+                    to_sys_stderr(error_msg)
+        except Exception as e:
+            # Handle any other unexpected errors
+            traceback.print_exc()
+            error_msg = f"Unexpected error in printer: {str(e)}\n"
+            to_sys_stderr(error_msg)
+
+    return safe_printer
+
+
+def to_sys_stderr(msg: str) -> None:
+    """Write a message to sys.stderr.
+
+    Parameters
+    ----------
+    msg : str
+        The message to write to stderr.
+    """
+    # pylint: disable=broad-exception-caught
+    # noinspection TryExceptPass,PyBroadException
+    try:
+        if hasattr(sys.stderr, "buffer"):
+            sys.stderr.buffer.write(msg.encode("utf-8", errors="replace"))
+            sys.stderr.buffer.flush()
+        else:  # pragma: no cover
+            sys.stderr.write(msg)
+            sys.stderr.flush()
+    except Exception:  # pragma: no cover
+        pass
+
+
+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 isinstance(sep, bytes):  # pragma: no cover
+        sep = sep.decode("utf-8", errors="replace")
+    if len(args) > 1:
+        msg = sep.join(str(arg) for arg in args)
+
+    # Handle end parameter
+    end = kwargs.get("end", "\n")
+    if isinstance(end, bytes):  # pragma: no cover
+        end = end.decode("utf-8", errors="replace")
+    msg += end
+
+    # Handle flush parameter
+    flush = kwargs.get("flush", False)
+    # noinspection PyUnreachableCode
+    if not isinstance(flush, bool):  # pragma: no cover
+        flush = False
+
+    return msg, flush
+
+
+def is_async_callable(fn: Any) -> bool:
+    """Check if a function is async callable, including partials/callables.
+
+    Parameters
+    ----------
+    fn : Any
+        The function to check.
+
+    Returns
+    -------
+    bool
+        True if the function is async callable, False otherwise.
+    """
+    if isinstance(fn, functools.partial):
+        fn = fn.func
+    unwrapped = inspect.unwrap(fn)
+    return inspect.iscoroutinefunction(
+        unwrapped
+    ) or inspect.iscoroutinefunction(
+        getattr(unwrapped, "__call__", None),  # noqa: B004
+    )
+
+
+def syncify(
+    async_func: Callable[..., Coroutine[Any, Any, T]],
+    timeout: float | None = None,
+) -> Callable[..., T]:
+    """Convert an async function to a sync function.
+
+    This function handles the conversion of async functions to sync functions,
+    properly managing event loops and thread execution contexts.
+
+    Parameters
+    ----------
+    async_func : Callable[..., Coroutine[Any, Any, T]]
+        The async function to convert.
+    timeout : float | None, optional
+        The timeout for the sync function. Defaults to None.
+
+    Returns
+    -------
+    Callable[..., T]
+        The converted sync function.
+
+    Raises
+    ------
+    TimeoutError
+        If the async function times out.
+    RuntimeError
+        If there are issues with event loop management.
+    """
+
+    def _sync_wrapper(*args: Any, **kwargs: Any) -> T:
+        """Get the result of the async function."""
+        # pylint: disable=too-many-try-statements
+        try:
+            # Check if we're already in an event loop
+            asyncio.get_running_loop()
+            return _run_in_thread(async_func, args, kwargs, timeout)
+        except RuntimeError:
+            # No event loop running, we can use asyncio.run directly
+            logger.debug("No event loop running, using asyncio.run")
+
+        # Create a new event loop and run the coroutine
+        try:
+            if timeout is not None:
+                # Need to run with asyncio.run and wait_for inside
+                async def _with_timeout() -> T:
+                    return await asyncio.wait_for(
+                        async_func(*args, **kwargs), timeout=timeout
+                    )
+
+                return asyncio.run(_with_timeout())
+            return asyncio.run(async_func(*args, **kwargs))
+        except (asyncio.TimeoutError, TimeoutError) as e:
+            raise TimeoutError(
+                f"Async function timed out after {timeout} seconds"
+            ) from e
+
+    return _sync_wrapper
+
+
+def _run_in_thread(
+    async_func: Callable[..., Coroutine[Any, Any, T]],
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+    timeout: float | None,
+) -> T:
+    """Run async function in a separate thread.
+
+    Parameters
+    ----------
+    async_func : Callable[..., Coroutine[Any, Any, T]]
+        The async function to run.
+    args : tuple[Any, ...]
+        Positional arguments for the function.
+    kwargs : dict[str, Any]
+        Keyword arguments for the function.
+    timeout : float | None
+        Timeout in seconds.
+
+    Returns
+    -------
+    T
+        The result of the async function.
+
+    Raises
+    ------
+    TimeoutError
+        If the function execution times out.
+    RuntimeError
+        If thread execution fails unexpectedly.
+    """
+    result_container: _ResultContainer[T] = _ResultContainer()
+    finished_event = threading.Event()
+
+    def _thread_target() -> None:
+        """Target function for the thread."""
+        # pylint: disable=too-many-try-statements, broad-exception-caught
+        try:
+            if timeout is not None:
+
+                async def _with_timeout() -> T:
+                    return await asyncio.wait_for(
+                        async_func(*args, **kwargs), timeout=timeout
+                    )
+
+                result_container.result = asyncio.run(_with_timeout())
+            else:
+                result_container.result = asyncio.run(
+                    async_func(*args, **kwargs)
+                )
+        except (
+            BaseException
+        ) as e:  # Catch BaseException to propagate cancellations
+            result_container.exception = e
+        finally:
+            finished_event.set()
+
+    thread = threading.Thread(target=_thread_target, daemon=True)
+    thread.start()
+
+    # Wait for completion with timeout
+    timeout_buffer = 1.0  # 1 second buffer for cleanup
+    wait_timeout = timeout + timeout_buffer if timeout is not None else None
+
+    if not finished_event.wait(timeout=wait_timeout):  # pragma: no cover
+        raise TimeoutError(
+            f"Function execution timed out after {timeout} seconds"
+        )
+
+    thread.join(timeout=timeout_buffer)  # Give thread time to clean up
+
+    if result_container.exception is not None:
+        raise result_container.exception
+
+    # Use cast since we know the result should be T if no exception occurred
+    return cast(T, result_container.result)

waldiez/utils/version.py

@@ -30,7 +30,7 @@ def _get_waldiez_version_from_package_json() -> str | None:
 def _get_waldiez_version_from_version_py() -> str | None:
     """Get the Waldiez version from _version.py."""
     version_py_path = Path(__file__).parent.parent / "_version.py"
-    if version_py_path.exists():
+    if version_py_path.exists():  # pragma: no branch
         with open(version_py_path, "r", encoding="utf-8") as f:
             for line in f:
                 if line.startswith("__version__"):
@@ -48,14 +48,10 @@ def get_waldiez_version() -> str:
         The Waldiez version, or "dev" if not found.
     """
     w_version = _get_waldiez_version_from_importlib()
-    if not w_version:
+    if not w_version:  # pragma: no branch
         w_version = _get_waldiez_version_from_version_py()
     if not w_version:
         w_version = _get_waldiez_version_from_package_json()
     if not w_version:
         w_version = "dev"
     return w_version
-
-
-if __name__ == "__main__":
-    print(get_waldiez_version())

waldiez/ws/__init__.py

@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+# pylint: disable=duplicate-code
+"""Waldiez WebSocket server module."""
+
+import typer
+from typer.models import CommandInfo
+
+from .cli import serve
+from .client_manager import ClientManager
+from .errors import (
+    ErrorCode,
+    ErrorHandler,
+    MessageHandlingError,
+    MessageParsingError,
+    OperationTimeoutError,
+    ServerOverloadError,
+    UnsupportedActionError,
+    WaldiezServerError,
+)
+from .server import WaldiezWsServer, run_server
+from .session_manager import SessionManager
+from .utils import (
+    ConnectionManager,
+    HealthChecker,
+    ServerHealth,
+    get_available_port,
+    is_port_available,
+    test_server_connection,
+)
+
+
+def add_ws_app(app: typer.Typer) -> None:
+    """Add WebSocket server commands to the CLI.
+
+    Parameters
+    ----------
+    app : typer.Typer
+        The Typer application instance.
+    """
+    app.registered_commands.append(
+        CommandInfo(
+            name="ws",
+            help="Start the Waldiez WebSocket server.",
+            callback=serve,
+        )
+    )
+
+
+__all__ = [
+    "WaldiezWsServer",
+    "run_server",
+    "ClientManager",
+    "ConnectionManager",
+    "HealthChecker",
+    "ServerHealth",
+    "test_server_connection",
+    "ErrorHandler",
+    "ErrorCode",
+    "MessageParsingError",
+    "MessageHandlingError",
+    "UnsupportedActionError",
+    "ServerOverloadError",
+    "SessionManager",
+    "OperationTimeoutError",
+    "WaldiezServerError",
+    "get_available_port",
+    "is_port_available",
+    "add_ws_app",
+]

waldiez/ws/__main__.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Main entry point for Waldiez WebSocket server."""
+
+import sys  # pragma: no cover
+from pathlib import Path  # pragma: no cover
+
+try:  # pragma: no cover
+    from .cli import app
+except ImportError:  # pragma: no cover
+    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
+    from waldiez.ws.cli import app
+
+if __name__ == "__main__":
+    app()

waldiez/ws/_file_handler.py

@@ -0,0 +1,201 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+
+"""Files related request handler."""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from typing import Any
+
+from waldiez.exporter import WaldiezExporter
+from waldiez.models import Waldiez
+
+from .models import (
+    ConvertWorkflowRequest,
+    ConvertWorkflowResponse,
+    SaveFlowRequest,
+    SaveFlowResponse,
+)
+
+
+class FileRequestHandler:
+    """Handles file-related requests."""
+
+    @staticmethod
+    def handle_save_flow_request(
+        msg: SaveFlowRequest,
+        workspace_dir: Path,
+        client_id: str,
+        logger: logging.Logger,
+    ) -> dict[str, Any]:
+        """Handle save flow request.
+
+        Parameters
+        ----------
+        msg : SaveFlowRequest
+            The save flow request message.
+        workspace_dir : Path
+            The workspace directory.
+        client_id : str
+            The client ID.
+        logger : logging.Logger
+            The logger instance.
+
+        Returns
+        -------
+        dict[str, Any]
+            The response dictionary.
+        """
+        filename = msg.filename or f"waldiez_{client_id}.waldiez"
+        try:
+            output_path = resolve_output_path(
+                filename,
+                workspace_dir=workspace_dir,
+                expected_ext="waldiez",
+            )
+        except ValueError as exc:
+            logger.error("Error resolving output path: %s", exc)
+            return SaveFlowResponse.fail(
+                error=f"Invalid output path: {exc}",
+                file_path=filename,
+            ).model_dump(mode="json")
+        # pylint: disable=too-many-try-statements
+        try:
+            if output_path.exists() and not msg.force_overwrite:
+                return SaveFlowResponse.fail(
+                    error=f"File exists: {output_path}",
+                    file_path=str(output_path.relative_to(workspace_dir)),
+                ).model_dump(mode="json")
+
+            # Parent dir already created by resolve_output_path
+            output_path.write_text(msg.flow_data, encoding="utf-8")
+
+            return SaveFlowResponse.ok(
+                file_path=str(output_path.relative_to(workspace_dir))
+            ).model_dump(mode="json")
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            logger.error("Error saving flow: %s", e)
+            return SaveFlowResponse.fail(error=str(e)).model_dump(mode="json")
+
+    @staticmethod
+    def handle_convert_workflow_request(
+        msg: ConvertWorkflowRequest,
+        client_id: str,
+        workspace_dir: Path,
+        logger: logging.Logger,
+    ) -> dict[str, Any]:
+        """Handle a convert workflow request.
+
+        Parameters
+        ----------
+        msg : ConvertWorkflowRequest
+            The convert workflow request message.
+        client_id : str
+            The client ID.
+        workspace_dir : Path
+            The workspace directory.
+        logger : logging.Logger
+            The logger instance.
+
+        Returns
+        -------
+        dict[str, Any]
+            The response dictionary.
+        """
+        target_format = (msg.target_format or "").strip().lower()
+        if target_format not in {"py", "ipynb"}:
+            return ConvertWorkflowResponse.fail(
+                error=f"Unsupported target format: {target_format}",
+                target_format=target_format,
+            ).model_dump(mode="json")
+
+        try:
+            waldiez_data = Waldiez.from_dict(json.loads(msg.flow_data))
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            return ConvertWorkflowResponse.fail(
+                error=f"Invalid flow_data: {e}",
+                target_format=target_format,
+            ).model_dump(mode="json")
+
+        try:
+            # Use normalized target_format for default name
+            filename = msg.output_path or f"waldiez_{client_id}.{target_format}"
+            output_path = resolve_output_path(
+                filename,
+                workspace_dir=workspace_dir,
+                expected_ext=target_format,
+            )
+        except ValueError as exc:
+            logger.error("Error resolving output path: %s", exc)
+            return ConvertWorkflowResponse.fail(
+                error=f"Invalid output path: {exc}",
+                target_format=target_format,
+            ).model_dump(mode="json")
+
+        try:
+            exporter = WaldiezExporter(waldiez_data)
+            exporter.export(path=output_path, force=True, structured_io=True)
+
+            return ConvertWorkflowResponse.ok(
+                target_format=target_format,
+                output_path=str(output_path.relative_to(workspace_dir)),
+            ).model_dump(mode="json")
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            logger.error("Error converting workflow: %s", e)
+            return ConvertWorkflowResponse.fail(
+                error=str(e), target_format=target_format
+            ).model_dump(mode="json")
+
+
+def resolve_output_path(
+    filename: str,
+    workspace_dir: Path,
+    expected_ext: str | None = None,
+) -> Path:
+    """
+    Resolve output path inside the workspace.
+
+    Parameters
+    ----------
+    filename : str
+        Provided filename (may be relative or absolute).
+    workspace_dir : Path
+        The workspace directory to resolve the output path against.
+    expected_ext : str | None
+        If provided, ensure the filename ends with this extension.
+
+    Returns
+    -------
+    Path
+        Resolved absolute path, with parent directories created.
+
+    Raises
+    ------
+    ValueError
+        If the output path is outside the workspace.
+    """
+    # Normalize workspace_dir to an absolute path
+    workspace_dir = workspace_dir.resolve()
+
+    output_path = Path(filename)
+    if not output_path.is_absolute():
+        output_path = workspace_dir / output_path
+
+    if expected_ext and output_path.suffix != f".{expected_ext}":
+        output_path = output_path.with_suffix(f".{expected_ext}")
+
+    output_path = output_path.resolve()
+
+    # Ensure output_path is a subpath of workspace_dir
+    try:
+        output_path.relative_to(workspace_dir)
+    except ValueError as exc:
+        raise ValueError(
+            f"Output path {output_path} is outside workspace {workspace_dir}"
+        ) from exc
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    return output_path