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

waldiez/io/structured.py

@@ -0,0 +1,439 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+
+"""Structured I/O stream for JSON-based communication over stdin/stdout."""
+
+import json
+import queue
+import sys
+import threading
+from getpass import getpass
+from pathlib import Path
+from typing import Any
+from uuid import uuid4
+
+from autogen.events import BaseEvent  # type: ignore
+from autogen.io import IOStream  # type: ignore
+
+from .models import (
+    PrintMessage,
+    UserInputData,
+    UserInputRequest,
+    UserResponse,
+)
+from .utils import (
+    gen_id,
+    get_image,
+    is_json_dumped,
+    now,
+    try_parse_maybe_serialized,
+)
+
+
+class StructuredIOStream(IOStream):
+    """Structured I/O stream using stdin and stdout."""
+
+    uploads_root: Path | None = None
+
+    def __init__(
+        self,
+        timeout: float = 120,
+        uploads_root: Path | str | None = None,
+        is_async: bool = False,
+    ) -> None:
+        self.timeout = timeout
+        self.is_async = is_async
+        if uploads_root is not None:
+            self.uploads_root = Path(uploads_root).resolve()
+            if not self.uploads_root.exists():
+                self.uploads_root.mkdir(parents=True, exist_ok=True)
+
+    # noinspection PyMethodMayBeStatic
+    # pylint: disable=no-self-use
+    def print(self, *args: Any, **kwargs: Any) -> None:
+        """Structured print to stdout.
+
+        Parameters
+        ----------
+        args : Any
+            The data to print.
+        kwargs : Any
+        """
+        sep = kwargs.get("sep", " ")
+        end = kwargs.get("end", "\n")
+        message = sep.join(map(str, args))
+        is_dumped = is_json_dumped(message)
+        if is_dumped and end.endswith("\n"):
+            # If the message is already JSON-dumped,
+            # let's try not to double dump it
+            message = json.loads(message)
+            payload: dict[str, Any] = {
+                "type": "print",
+                "id": uuid4().hex,
+                "timestamp": now(),
+                "data": message,
+            }
+        else:
+            message += end
+            payload = PrintMessage(
+                id=uuid4().hex,
+                timestamp=now(),
+                data=message,
+            ).model_dump(mode="json")
+        flush = kwargs.get("flush", True)
+        payload_type = kwargs.get("type", "print")
+        payload["type"] = payload_type
+        print(json.dumps(payload), flush=flush)
+
+    def input(self, prompt: str = "", *, password: bool = False) -> str:
+        """Structured input from stdin.
+
+        Parameters
+        ----------
+        prompt : str, optional
+            The prompt to display. Defaults to "".
+        password : bool, optional
+            Whether to read a password. Defaults to False.
+
+        Returns
+        -------
+        str
+            The line read from the input stream.
+        """
+        request_id = uuid4().hex
+        prompt = prompt or "> "
+
+        self._send_input_request(prompt, request_id, password)
+        user_input_raw = self._read_user_input(prompt, password, request_id)
+        response = self._handle_user_input(user_input_raw, request_id)
+        user_response = response.to_string(
+            uploads_root=self.uploads_root,
+            base_name=request_id,
+        )
+        # 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???)
+        if self.is_async:
+            # let's make a coroutine to just return the user response
+            async def async_input() -> str:
+                """Asynchronous input method.
+
+                Returns
+                -------
+                str
+                    The line read from the input stream.
+                """
+                return user_response
+
+            return async_input()  # type: ignore
+        return user_response
+
+    # noinspection PyMethodMayBeStatic
+    # pylint: disable=no-self-use
+    def send(self, message: BaseEvent) -> None:
+        """Structured sending of a BaseEvent.
+
+        Parameters
+        ----------
+        message : BaseEvent
+            The message to send.
+        """
+        message_dump = message.model_dump(mode="json")
+        if message_dump.get("type") == "text":
+            content_block = message_dump.get("content")
+            if (
+                isinstance(content_block, dict)
+                and "content" in content_block
+                and isinstance(content_block["content"], str)
+            ):
+                inner_content = content_block["content"]
+                content_block["content"] = try_parse_maybe_serialized(
+                    inner_content
+                )
+        print(json.dumps(message_dump), flush=True, file=sys.stdout)
+
+    # noinspection PyMethodMayBeStatic
+    # pylint: disable=no-self-use
+    def _send_input_request(
+        self,
+        prompt: str,
+        request_id: str,
+        password: bool,
+    ) -> None:
+        payload = UserInputRequest(
+            request_id=request_id,
+            prompt=prompt,
+            password=password,
+        ).model_dump(mode="json")
+        print(json.dumps(payload), flush=True)
+
+    def _read_user_input(
+        self,
+        prompt: str,
+        password: bool,
+        request_id: str,
+    ) -> str:
+        input_queue: queue.Queue[str] = queue.Queue()
+
+        def read_input() -> None:
+            """Read user input from stdin."""
+            try:
+                user_input = (
+                    getpass(prompt).strip()
+                    if password
+                    else input(prompt).strip()
+                )
+                input_queue.put(user_input)
+            except EOFError:
+                input_queue.put("")
+
+        input_thread = threading.Thread(target=read_input, daemon=True)
+        input_thread.start()
+
+        try:
+            return input_queue.get(timeout=self.timeout)
+        except queue.Empty:
+            self._send_timeout_message(request_id)
+            return ""
+
+    def _send_timeout_message(self, request_id: str) -> None:
+        timeout_payload = {
+            "id": gen_id(),
+            "type": "timeout",
+            "request_id": request_id,
+            "timestamp": now(),
+            "data": f"No input received after {self.timeout} seconds.",
+        }
+        print(json.dumps(timeout_payload), flush=True, file=sys.stderr)
+
+    def _handle_user_input(
+        self, user_input_raw: str, request_id: str
+    ) -> UserResponse:
+        """Handle user input and return the appropriate response.
+
+        Parameters
+        ----------
+        user_input_raw : str
+            The raw user input string.
+        request_id : str
+            The request ID to match against.
+
+        Returns
+        -------
+        UserResponse
+            The structured user response.
+        """
+        user_input = self._load_user_input(user_input_raw)
+        if isinstance(user_input, str):
+            return UserResponse(data=user_input, request_id=request_id)
+        return self._parse_user_input(user_input, request_id)
+
+    @staticmethod
+    def _load_user_input(user_input_raw: str) -> str | dict[str, Any]:
+        """Load user input from a raw string.
+
+        Parameters
+        ----------
+        user_input_raw : str
+            The raw user input string.
+
+        Returns
+        -------
+        str | dict[str, Any]
+            The loaded user input, either as a string or a dictionary.
+        """
+        response: str | dict[str, Any] = user_input_raw
+        try:
+            # Attempt to parse the input as JSON
+            response = json.loads(user_input_raw)
+        except json.JSONDecodeError:
+            # If it's not valid JSON, return as is
+            # This allows for backwards compatibility with raw text input
+            return user_input_raw
+        if isinstance(response, str):
+            # double inner dumped?
+            try:
+                response = json.loads(response)
+            except json.JSONDecodeError:
+                # If it's not valid JSON, return as is
+                return response
+        if not isinstance(response, dict):
+            return str(response)
+        if "data" in response and isinstance(response["data"], str):
+            # double inner dumped?
+            try:
+                response["data"] = json.loads(response["data"])
+            except json.JSONDecodeError:
+                pass
+        return response
+
+    def _parse_user_input(
+        self, user_input: dict[str, Any], request_id: str
+    ) -> UserResponse:
+        """Parse user input and return the appropriate response.
+
+        Parameters
+        ----------
+        user_input : str
+            The raw user input string.
+        request_id : str
+            The request ID to match against.
+
+        Returns
+        -------
+        UserResponse
+            The structured user response.
+        """
+        # Load the user input
+        if user_input.get("request_id") == request_id:
+            # We have a valid response to our request
+            data = user_input.get("data")
+            if not data:
+                # let's check if text|image keys are sent (outside data)
+                if "image" in user_input or "text" in user_input:
+                    return UserResponse(
+                        request_id=request_id,
+                        data=self._format_multimedia_response(
+                            request_id=request_id, data=user_input
+                        ),
+                    )
+            if isinstance(data, list):
+                return self._handle_list_response(
+                    data,  # pyright: ignore
+                    request_id=request_id,
+                )
+            if not data or not isinstance(data, (str, dict)):
+                # No / invalid data provided in the response
+                return UserResponse(
+                    request_id=request_id,
+                    data="",
+                )
+            # Process different data types
+            if isinstance(data, str):
+                # double inner dumped?
+                data = self._load_user_input(data)
+            if isinstance(data, dict):
+                return UserResponse(
+                    data=self._format_multimedia_response(
+                        request_id=request_id,
+                        data=data,  # pyright: ignore
+                    ),
+                    request_id=request_id,
+                )
+            # For other types (numbers, bools ,...),
+            #  let's just convert to string
+            return UserResponse(
+                data=str(data), request_id=request_id
+            )  # pragma: no cover
+        # This response doesn't match our request_id, log and return empty
+        self._log_mismatched_response(request_id, user_input)
+        return UserResponse(
+            request_id=request_id,
+            data="",
+        )
+
+    def _handle_list_response(
+        self,
+        data: list[dict[str, Any]],
+        request_id: str,
+    ) -> UserResponse:
+        if len(data) == 0:  # pyright: ignore
+            # Empty list, return empty response
+            return UserResponse(
+                request_id=request_id,
+                data="",
+            )
+
+        input_data: list[UserInputData] = []
+        for entry in data:  # pyright: ignore
+            # pylint: disable=broad-exception-caught
+            try:
+                content = UserInputData.model_validate(entry)
+                input_data.append(content)
+            except Exception as error:  # pragma: no cover
+                print({"type": "error", "message": str(error)}, file=sys.stderr)
+                continue
+        if not input_data:  # pragma: no cover
+            # No valid data in the list, return empty response
+            return UserResponse(
+                request_id=request_id,
+                data="",
+            )
+        return UserResponse(
+            request_id=request_id,
+            data=input_data,
+        )
+
+    @staticmethod
+    def _log_mismatched_response(expected_id: str, response: Any) -> None:
+        """Log information about mismatched response IDs.
+
+        Parameters
+        ----------
+        expected_id : str
+            The request ID we were expecting
+        response : Any
+            The response received
+        """
+        # Create a log message
+        log_payload: dict[str, Any] = {
+            "type": "warning",
+            "id": uuid4().hex,
+            "timestamp": now(),
+            "data": {
+                "message": (
+                    "Received response with mismatched request_id. "
+                    f"Expected: {expected_id}"
+                ),
+                "details": {
+                    "expected_id": expected_id,
+                    "received": str(response)[:100]
+                    + ("..." if len(str(response)) > 100 else ""),
+                },
+            },
+        }
+        # Print to stderr to avoid interfering with stdout communication
+        print(json.dumps(log_payload), file=sys.stderr)
+
+    def _format_multimedia_response(
+        self,
+        data: dict[str, Any],
+        request_id: str | None = None,
+    ) -> str:
+        """Format a multimedia response dict into a string with image tags.
+
+        Parameters
+        ----------
+        data : dict[str, Any]
+            The data dictionary containing "image" and "text" keys.
+        request_id : str | None
+            The input request ID, if available.
+
+        Returns
+        -------
+        str
+            The formatted string with image tags and text.
+        """
+        result: list[str] = []
+        if "content" in data and isinstance(data["content"], dict):
+            return self._format_multimedia_response(
+                data=data["content"],  # pyright: ignore
+                request_id=request_id,
+            )
+
+        # Handle image if present
+        if "image" in data and data["image"]:
+            image_data = data["image"]
+            image = get_image(self.uploads_root, image_data, request_id)
+            img_tag = f"<img {image}>"
+            result.append(img_tag)
+
+        # Handle text if present
+        if "text" in data and data["text"]:
+            result.append(str(data["text"]))
+
+        # If neither image nor text, return empty string
+        if not result:
+            return ""
+        # Join with a space
+        return " ".join(result)

waldiez/io/utils.py

@@ -0,0 +1,184 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+
+"""Utility functions for the waldiez.io package."""
+
+import ast
+import json
+import uuid
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Literal, Union
+
+from autogen.agentchat.contrib.img_utils import get_pil_image  # type: ignore
+
+MessageType = Literal[
+    "input_request",
+    "input_response",
+    "print",
+    "input",
+]
+"""Possible message types for the structured I/O stream."""
+
+MediaType = Union[
+    Literal["text"],
+    Literal["image"],
+    Literal["image_url"],
+    Literal["video"],
+    Literal["audio"],
+    Literal["file"],
+]
+"""Possible media types for the structured I/O stream."""
+
+
+def gen_id() -> str:
+    """Generate a unique identifier.
+
+    Returns
+    -------
+    str
+        A unique id
+    """
+    return str(uuid.uuid4())
+
+
+def now() -> str:
+    """Get the current time as an ISO string.
+
+    Returns
+    -------
+    str
+        The current time as an ISO string.
+    """
+    return datetime.now().isoformat()
+
+
+def detect_media_type(value: dict[str, Any]) -> MediaType:
+    """Detect mediatype from dict.
+
+    Either using the 'type' field or by checking the keys.
+
+    Parameters
+    ----------
+    value : dict[str, Any]
+        The input dictionary
+
+    Returns
+    -------
+    MediaType
+        The detected media type
+
+    Raises
+    ------
+    ValueError
+        If the media type is not valid or not found.
+    """
+    valid_types = (
+        "text",
+        "image",
+        "image_url",
+        "video",
+        "audio",
+        "file",
+    )
+    if "type" in value:
+        content_type = value["type"]
+        if content_type not in valid_types:
+            raise ValueError(f"Invalid media type: {content_type}")
+        return content_type
+    for valid_type in valid_types:
+        if valid_type in value:
+            return valid_type  # type: ignore
+    raise ValueError(f"No type in value: {value}.")
+
+
+def get_image(
+    uploads_root: Path | None,
+    image_data: str,
+    base_name: str | None = None,
+) -> str:
+    """Store the image data in a file and return the file path.
+
+    Parameters
+    ----------
+    uploads_root : Path | None
+        The root directory for storing images, optional.
+    image_data : str
+        The base64-encoded image data.
+    base_name : str | None
+        The base name for the image file, optional.
+
+    Returns
+    -------
+    str
+        The file path of the stored image.
+    """
+    if uploads_root:
+        # noinspection PyBroadException
+        # pylint: disable=broad-exception-caught
+        try:
+            pil_image = get_pil_image(image_data)
+        except BaseException:
+            return image_data
+        if not base_name:
+            base_name = uuid.uuid4().hex
+        file_name = f"{base_name}.png"
+        file_path = uploads_root / file_name
+        pil_image.save(file_path, format="PNG")
+        return str(file_path)
+    return image_data
+
+
+def is_json_dumped(value: str) -> bool:
+    """Check if a string is JSON-dumped.
+
+    Parameters
+    ----------
+    value : str
+        The string to check.
+
+    Returns
+    -------
+    bool
+        True if the string is JSON-dumped, False otherwise.
+    """
+    try:
+        parsed = json.loads(value)
+        # If we can parse it as JSON and it's not a string,
+        # we consider it JSON-dumped
+        return not isinstance(parsed, str)
+    except json.JSONDecodeError:
+        return False
+
+
+def try_parse_maybe_serialized(value: str) -> Any:
+    """Parse a string that may be JSON or Python serialized.
+
+    Returns the parsed object if successful, or the original string otherwise.
+
+    Parameters
+    ----------
+    value : str
+        The string to parse.
+
+    Returns
+    -------
+    Any
+        The parsed object or the original string if parsing fails.
+    """
+    for parser in (json.loads, ast.literal_eval):
+        # pylint: disable=broad-exception-caught, too-many-try-statements
+        try:
+            parsed: dict[str, Any] | list[Any] | str = parser(value)
+            # Normalize: if it's a single-item list of a string
+            # return the string
+            if (
+                isinstance(parsed, list)
+                and len(parsed) == 1
+                and isinstance(parsed[0], str)
+            ):
+                return parsed[0]
+            return parsed
+        except Exception:
+            pass  # Try next parser
+    return value  # Return original if all parsing fails