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

waldiez/io/structured.py

@@ -1,16 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0.
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 
+# pyright: reportMissingTypeStubs=false,reportUnknownVariableType=false
+# pyright: reportUnknownArgumentType=false,reportArgumentType=false
+# pyright: reportUnknownMemberType=false,reportAttributeAccessIssue=false
+# pyright: reportPrivateImportUsage=false
+
 """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
@@ -23,6 +25,9 @@ from .models import (
     UserResponse,
 )
 from .utils import (
+    DEBUG_INPUT_PROMPT,
+    START_CHAT_PROMPT,
+    MessageType,
     gen_id,
     get_image,
     get_message_dump,
@@ -61,33 +66,57 @@ class StructuredIOStream(IOStream):
             The data to print.
         kwargs : Any
         """
-        sep = kwargs.get("sep", " ")
-        end = kwargs.get("end", "\n")
+        sep = str(kwargs.get("sep", " "))
+        end = str(kwargs.get("end", "\n"))
+        flush = bool(kwargs.get("flush", True))
+        payload_type = str(kwargs.get("type", "print"))
         message = sep.join(map(str, args))
-        is_dumped = is_json_dumped(message)
-        if is_dumped and end.endswith("\n"):
+        if len(args) == 1 and isinstance(args[0], dict):
+            message = args[0]  # type: ignore
+            payload_type_ = message.get("type", payload_type)  # type: ignore
+            if isinstance(payload_type_, str):
+                payload_type = payload_type_
+            is_dumped = True
+        else:
+            is_dumped, message = is_json_dumped(message)
+        if is_dumped:
             # 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,
+                "id": gen_id(),
                 "timestamp": now(),
-                "data": message,
+                # "data": message,
             }
+            if isinstance(message, dict):
+                payload.update(message)
+                end = ""
+            else:
+                payload["data"] = message
+            if "type" not in payload:
+                payload["type"] = payload_type
+        else:
+            print_message = PrintMessage(data=message)
+            payload = print_message.model_dump(mode="json", fallback=str)
+            payload["type"] = payload_type
+        dumped = json.dumps(payload, default=str, ensure_ascii=False) + end
+        file = kwargs.get("file", None)
+        if file and file in [
+            sys.stderr,
+            sys.__stderr__,
+            sys.stdout,
+            sys.__stdout__,
+        ]:
+            print(dumped, file=file, flush=flush)
         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, default=str), flush=flush)
-
-    def input(self, prompt: str = "", *, password: bool = False) -> str:
+            print(dumped, flush=flush)
+
+    def input(
+        self,
+        prompt: str = "",
+        *,
+        password: bool = False,
+        request_id: str | None = None,
+    ) -> str:
         """Structured input from stdin.
 
         Parameters
@@ -96,25 +125,36 @@ class StructuredIOStream(IOStream):
             The prompt to display. Defaults to "".
         password : bool, optional
             Whether to read a password. Defaults to False.
+        request_id : str, optional
+            The request id. If not provided, a new will be generated.
 
         Returns
         -------
         str
             The line read from the input stream.
         """
-        request_id = uuid4().hex
+        input_request_id = request_id or gen_id()
         prompt = prompt or ">"
         if not prompt or prompt in [">", "> "]:  # pragma: no cover
             # if the prompt is just ">" or "> ",
             # let's use a more descriptive one
-            prompt = "Enter your message to start the conversation: "
-
-        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)
+            prompt = START_CHAT_PROMPT
+        input_type = "chat"
+        if prompt.strip() == DEBUG_INPUT_PROMPT.strip():
+            input_type = "debug"
+        self._send_input_request(
+            prompt,
+            input_request_id,
+            password,
+            input_type=input_type,
+        )
+        user_input_raw = self._read_user_input(
+            prompt, password, input_request_id
+        )
+        response = self._handle_user_input(user_input_raw, input_request_id)
         user_response = response.to_string(
             uploads_root=self.uploads_root,
-            base_name=request_id,
+            base_name=input_request_id,
         )
         return user_response
 
@@ -140,11 +180,8 @@ class StructuredIOStream(IOStream):
                 content_block["content"] = try_parse_maybe_serialized(
                     inner_content
                 )
-        print(
-            json.dumps(message_dump, default=str),
-            flush=True,
-            file=sys.stdout,
-        )
+        message_dump["timestamp"] = now()
+        print(json.dumps(message_dump, default=str), flush=True)
 
     # noinspection PyMethodMayBeStatic
     # pylint: disable=no-self-use
@@ -152,14 +189,21 @@ class StructuredIOStream(IOStream):
         self,
         prompt: str,
         request_id: str,
-        password: bool,
+        password: bool = False,
+        input_type: str = "chat",
     ) -> None:
+        if input_type not in ("chat", "debug"):
+            input_type = "chat"
+        request_type = (
+            "debug_input_request" if input_type == "debug" else "input_request"
+        )
         payload = UserInputRequest(
+            type=request_type,  # type: ignore
             request_id=request_id,
             prompt=prompt,
             password=password,
         ).model_dump(mode="json")
-        print(json.dumps(payload), flush=True)
+        print(json.dumps(payload, default=str), flush=True)
 
     def _read_user_input(
         self,
@@ -167,33 +211,15 @@ class StructuredIOStream(IOStream):
         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 (
+                getpass(prompt).strip() if password else input(prompt).strip()
+            )
+        except EOFError:
             return ""
         except BaseException as e:
             self._send_error_message(request_id, str(e))
             return ""
-        finally:
-            input_thread.join(timeout=1)
 
     def _send_timeout_message(self, request_id: str) -> None:
         timeout_payload = {
@@ -294,14 +320,23 @@ class StructuredIOStream(IOStream):
         UserResponse
             The structured user response.
         """
-        # Load the user input
-        if user_input.get("request_id") == request_id:
+        response_type: MessageType
+        _response_type = user_input.get("type", "input_response")
+        if _response_type not in ("input_response", "debug_input_response"):
+            response_type = "input_response"
+        else:
+            response_type = _response_type
+        if (
+            user_input.get("request_id") == request_id
+            or response_type == "debug_input_response"
+        ):
             # 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(
+                        type=response_type,
                         request_id=request_id,
                         data=self._format_multimedia_response(
                             request_id=request_id, data=user_input
@@ -309,12 +344,14 @@ class StructuredIOStream(IOStream):
                     )
             if isinstance(data, list):
                 return self._handle_list_response(
-                    data,  # pyright: ignore
+                    data,
                     request_id=request_id,
+                    response_type=response_type,
                 )
             if not data or not isinstance(data, (str, dict)):
                 # No / invalid data provided in the response
                 return UserResponse(
+                    type=response_type,
                     request_id=request_id,
                     data="",
                 )
@@ -324,20 +361,24 @@ class StructuredIOStream(IOStream):
                 data = self._load_user_input(data)
             if isinstance(data, dict):
                 return UserResponse(
+                    type=response_type,
                     data=self._format_multimedia_response(
                         request_id=request_id,
-                        data=data,  # pyright: ignore
+                        data=data,
                     ),
                     request_id=request_id,
                 )
             # For other types (numbers, bools ,...),
             #  let's just convert to string
             return UserResponse(
-                data=str(data), request_id=request_id
+                type=response_type,
+                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(
+            type=response_type,
             request_id=request_id,
             data="",
         )
@@ -347,16 +388,18 @@ class StructuredIOStream(IOStream):
         self,
         data: list[dict[str, Any]],
         request_id: str,
+        response_type: MessageType,
     ) -> UserResponse:
-        if len(data) == 0:  # pyright: ignore
+        if len(data) == 0:
             # Empty list, return empty response
             return UserResponse(
+                type=response_type,
                 request_id=request_id,
                 data="",
             )
 
         input_data: list[UserInputData] = []
-        for entry in data:  # pyright: ignore
+        for entry in data:
             # pylint: disable=broad-exception-caught
             try:
                 content = UserInputData.model_validate(entry)
@@ -367,10 +410,12 @@ class StructuredIOStream(IOStream):
         if not input_data:  # pragma: no cover
             # No valid data in the list, return empty response
             return UserResponse(
+                type=response_type,
                 request_id=request_id,
                 data="",
             )
         return UserResponse(
+            type=response_type,
             request_id=request_id,
             data=input_data,
         )
@@ -387,9 +432,16 @@ class StructuredIOStream(IOStream):
             The response received
         """
         # Create a log message
+        got_id: str | None = None
+        if isinstance(response, dict):
+            got_id = response.get("request_id")
+        response_str = str(response)
+        message = response_str[:100] + (
+            "..." if len(response_str) > 100 else ""
+        )
         log_payload: dict[str, Any] = {
             "type": "warning",
-            "id": uuid4().hex,
+            "id": gen_id(),
             "timestamp": now(),
             "data": {
                 "message": (
@@ -398,8 +450,8 @@ class StructuredIOStream(IOStream):
                 ),
                 "details": {
                     "expected_id": expected_id,
-                    "received": str(response)[:100]
-                    + ("..." if len(str(response)) > 100 else ""),
+                    "received_id": got_id,
+                    "message": message,
                 },
             },
         }
@@ -428,7 +480,7 @@ class StructuredIOStream(IOStream):
         result: list[str] = []
         if "content" in data and isinstance(data["content"], dict):
             return self._format_multimedia_response(
-                data=data["content"],  # pyright: ignore
+                data=data["content"],
                 request_id=request_id,
             )
 

waldiez/io/utils.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0.
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 
+# pyright: reportMissingTypeStubs=false,reportDeprecated=false
+
 """Utility functions for the waldiez.io package."""
 
 import ast
@@ -14,9 +16,17 @@ from autogen.agentchat.contrib.img_utils import get_pil_image  # type: ignore
 from autogen.events import BaseEvent  # type: ignore
 from autogen.messages import BaseMessage  # type: ignore
 
+DEBUG_INPUT_PROMPT = (
+    # cspell: disable-next-line
+    "[Step] (c)ontinue, (r)un, (q)uit, (i)nfo, (h)elp, (st)ats: "
+)
+START_CHAT_PROMPT = "Enter your message to start the conversation: "
+
 MessageType = Literal[
     "input_request",
     "input_response",
+    "debug_input_request",
+    "debug_input_response",
     "print",
     "input",
 ]
@@ -131,26 +141,25 @@ def get_image(
     return image_data
 
 
-def is_json_dumped(value: str) -> bool:
-    """Check if a string is JSON-dumped.
+def is_json_dumped(value: Any) -> tuple[bool, Any]:
+    """Check if a value is JSON-dumped.
 
     Parameters
     ----------
-    value : str
-        The string to check.
+    value : Any
+        The value to check.
 
     Returns
     -------
     bool
-        True if the string is JSON-dumped, False otherwise.
+        True if the value is JSON-dumped, False otherwise.
     """
+    to_check = value.strip() if isinstance(value, str) else value
     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)
+        parsed = json.loads(to_check)
+        return True, parsed
     except json.JSONDecodeError:
-        return False
+        return False, value
 
 
 def try_parse_maybe_serialized(value: str) -> Any:

waldiez/io/ws.py

@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0.
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
+
 # pylint: disable=too-many-try-statements
+# pyright: reportMissingTypeStubs=false,reportReturnType=false
+
 """WebSocket IOStream implementation for AsyncIO."""
 
 import asyncio
@@ -30,12 +33,13 @@ from .utils import (
 LOG = logging.getLogger(__name__)
 
 if not is_websocket_available():  # pragma: no cover
-    raise ImportError(
+    _msg = (  # pylint: disable=invalid-name
         "WebSocket support requires "
         "either the 'websockets' or 'starlette' package. "
         "Please install one of them with "
         "`pip install websockets` or `pip install starlette`."
     )
+    raise ImportError(_msg)
 
 
 class AsyncWebsocketsIOStream(IOStream):
@@ -118,8 +122,8 @@ class AsyncWebsocketsIOStream(IOStream):
         end = kwargs.get("end", "\n")
         msg = sep.join(str(arg) for arg in args)
 
-        is_dumped = is_json_dumped(msg)
-        if is_dumped and end.endswith("\n"):  # pragma: no cover
+        is_dumped, msg = is_json_dumped(msg)
+        if is_dumped:  # pragma: no cover
             msg = json.loads(msg)
         else:
             msg = f"{msg}{end}"

waldiez/logger.py

@@ -3,6 +3,8 @@
 # pylint: disable=broad-exception-caught,too-many-try-statements
 """Waldiez logger."""
 
+from __future__ import annotations
+
 import inspect
 import logging
 import os
@@ -10,13 +12,15 @@ import re
 import string
 import threading
 import traceback
+from collections.abc import Mapping
 from datetime import datetime
 from enum import IntEnum
 from pathlib import Path
 from types import TracebackType
-from typing import Any, Callable, Mapping, Optional
+from typing import Any, Callable
 
 import click
+from typing_extensions import override
 
 HERE = Path(__file__).parent
 
@@ -41,7 +45,7 @@ class WaldiezLogger(logging.Logger):
        - logger.info("User {user} has {count} items", user="john", count=5)
     """
 
-    _instance: Optional["WaldiezLogger"] = None
+    _instance: WaldiezLogger | None = None
     _lock: threading.Lock = threading.Lock()
 
     _INTERNAL_METHODS = {
@@ -179,6 +183,7 @@ class WaldiezLogger(logging.Logger):
         return self._level_map.get(level.upper(), LogLevel.INFO)
 
     # pylint: disable=unused-argument
+    @override
     def log(
         self,
         level: int,
@@ -249,6 +254,7 @@ class WaldiezLogger(logging.Logger):
         level_int = self._get_level_number(level)
         self.log(level_int, msg, *args, **kwargs)
 
+    @override
     def debug(self, msg: Any, *args: Any, **kwargs: Any) -> None:
         """Log a debug message.
 
@@ -263,6 +269,7 @@ class WaldiezLogger(logging.Logger):
         """
         self.do_log(msg, *args, level="debug", **kwargs)
 
+    @override
     def info(self, msg: Any, *args: Any, **kwargs: Any) -> None:
         """Log an informational message.
 
@@ -291,6 +298,7 @@ class WaldiezLogger(logging.Logger):
         """
         self.do_log(msg, *args, level="success", **kwargs)
 
+    @override
     def warning(self, msg: Any, *args: Any, **kwargs: Any) -> None:
         """Log a warning message.
 
@@ -305,6 +313,7 @@ class WaldiezLogger(logging.Logger):
         """
         self.do_log(msg, *args, level="warning", **kwargs)
 
+    @override
     def error(self, msg: Any, *args: Any, **kwargs: Any) -> None:
         """Log an error message.
 
@@ -319,6 +328,7 @@ class WaldiezLogger(logging.Logger):
         """
         self.do_log(msg, *args, level="error", **kwargs)
 
+    @override
     def critical(self, msg: Any, *args: Any, **kwargs: Any) -> None:
         """Log a critical error message.
 
@@ -333,6 +343,7 @@ class WaldiezLogger(logging.Logger):
         """
         self.do_log(msg, *args, level="critical", **kwargs)
 
+    @override
     def exception(self, msg: Any, *args: Any, **kwargs: Any) -> None:
         """Log an exception message.
 
@@ -354,6 +365,7 @@ class WaldiezLogger(logging.Logger):
         if tb and "NoneType: None" not in tb:  # pragma: no branch
             click.echo(click.style(tb, fg="red", dim=True))
 
+    @override
     def setLevel(self, level: int | str) -> None:
         """
         Set the logging level.
@@ -388,10 +400,11 @@ class WaldiezLogger(logging.Logger):
             self._level = level_upper
             super().setLevel(self._level_map[level_upper].value)
         else:
-            raise ValueError(
+            msg = (
                 f"Invalid log level: {level}. "
                 f"Valid levels are: {list(self._level_map.keys())}"
             )
+            raise ValueError(msg)
 
     def get_level(self) -> str:
         """Get the current logging level.

waldiez/models/agents/__init__.py

@@ -1,5 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0.
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
+
+# pyright: reportImportCycles=false
+
 """Agent models."""
 
 from .agent import (