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

waldiez/io/redis.py

@@ -4,6 +4,8 @@
 # flake8: noqa: E501
 # pylint: disable=too-many-try-statements,broad-exception-caught
 # pylint: disable=line-too-long,duplicate-code
+# pyright: reportMissingTypeStubs=false,reportUnknownArgumentType=false
+# pyright: reportUnknownMemberType=false
 
 """A Redis I/O stream for handling print and input messages."""
 
@@ -12,16 +14,10 @@ import logging
 import time
 import traceback as tb
 import uuid
+from collections.abc import Awaitable
 from pathlib import Path
 from types import TracebackType
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Awaitable,
-    Callable,
-    Optional,
-    Type,
-)
+from typing import TYPE_CHECKING, Any, Callable
 
 try:
     import redis
@@ -60,10 +56,10 @@ class RedisIOStream(IOStream):
     redis: Redis
     task_id: str
     input_timeout: int
-    on_input_request: Optional[Callable[[str, str, str], None]]
-    on_input_received: Optional[Callable[[str, str], None]]
+    on_input_request: Callable[[str, str, str], None] | None
+    on_input_response: Callable[[str, str], None] | None
     max_stream_size: int
-    output_stream: str
+    task_output_stream: str
     input_request_channel: str
     input_response_channel: str
 
@@ -73,8 +69,8 @@ class RedisIOStream(IOStream):
         task_id: str | None = None,
         input_timeout: int = 120,
         max_stream_size: int = 1000,
-        on_input_request: Optional[Callable[[str, str, str], None]] = None,
-        on_input_response: Optional[Callable[[str, str], None]] = None,
+        on_input_request: Callable[[str, str, str], None] | None = None,
+        on_input_response: Callable[[str, str], None] | None = None,
         redis_connection_kwargs: dict[str, Any] | None = None,
         uploads_root: Path | str | None = None,
     ) -> None:
@@ -127,7 +123,7 @@ class RedisIOStream(IOStream):
 
     def __exit__(
         self,
-        exc_type: Type[Exception] | None,
+        exc_type: type[Exception] | None,
         exc_value: Exception | None,
         traceback: TracebackType | None,
     ) -> None:
@@ -165,7 +161,7 @@ class RedisIOStream(IOStream):
         """
         LOG.debug("Sending print message: %s", payload)
         RedisIOStream.try_do(
-            self.redis.xadd,  # pyright: ignore
+            self.redis.xadd,
             self.task_output_stream,
             payload,
             maxlen=self.max_stream_size,
@@ -182,7 +178,7 @@ class RedisIOStream(IOStream):
         """
         LOG.debug("Sending print message: %s", payload)
         RedisIOStream.try_do(
-            self.redis.xadd,  # pyright: ignore
+            self.redis.xadd,
             self.common_output_stream,
             payload,
             maxlen=self.max_stream_size,
@@ -414,7 +410,7 @@ class RedisIOStream(IOStream):
         )
 
     @staticmethod
-    def _extract_message_data(data: Any) -> Optional[dict[str, Any]]:
+    def _extract_message_data(data: Any) -> dict[str, Any] | None:
         """Extract and parse the message data field."""
         message_data = data
 
@@ -431,7 +427,7 @@ class RedisIOStream(IOStream):
             LOG.error("Invalid message data format: %s", message_data)
             return None
 
-        return message_data  # pyright: ignore
+        return message_data  # pyright: ignore[reportUnknownVariableType]
 
     @staticmethod
     def _message_has_required_fields(message_data: dict[str, Any]) -> bool:
@@ -445,7 +441,7 @@ class RedisIOStream(IOStream):
     @staticmethod
     def _process_nested_data(
         message_data: dict[str, Any],
-    ) -> Optional[dict[str, Any]]:
+    ) -> dict[str, Any] | None:
         """Process nested JSON data if present."""
         # Create a copy to avoid modifying the original
         processed_data = message_data.copy()
@@ -467,7 +463,7 @@ class RedisIOStream(IOStream):
     @staticmethod
     def _create_user_response(
         message_data: dict[str, Any],
-    ) -> Optional["UserResponse"]:
+    ) -> UserResponse | None:
         """Create UserResponse object from validated data."""
         try:
             return UserResponse.model_validate(message_data)
@@ -684,7 +680,7 @@ class RedisIOStream(IOStream):
         """
         for key in redis_client.scan_iter("task:*:output", count=100):
             RedisIOStream.try_do(
-                redis_client.xtrim,  # pyright: ignore
+                redis_client.xtrim,
                 key,
                 maxlen=maxlen,
                 approximate=approximate,
@@ -760,7 +756,7 @@ class RedisIOStream(IOStream):
         ):  # pragma: no branch
             before = await redis_client.xlen(key)
             await RedisIOStream.a_try_do(
-                redis_client.xtrim,  # pyright: ignore
+                redis_client.xtrim,
                 key,
                 maxlen=maxlen,
                 approximate=approximate,

waldiez/io/structured.py

@@ -1,12 +1,15 @@
 # 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
@@ -63,14 +66,16 @@ class StructuredIOStream(IOStream):
             The data to print.
         kwargs : Any
         """
-        sep = kwargs.get("sep", " ")
-        end = kwargs.get("end", "\n")
-        flush = kwargs.get("flush", True)
-        payload_type = kwargs.get("type", "print")
+        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))
         if len(args) == 1 and isinstance(args[0], dict):
-            message = args[0]  # pyright: ignore
-            payload_type = message.get("type", payload_type)  # pyright: ignore
+            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)
@@ -83,27 +88,35 @@ class StructuredIOStream(IOStream):
                 # "data": message,
             }
             if isinstance(message, dict):
-                payload.update(message)  # pyright: ignore
+                payload.update(message)
+                end = ""
             else:
                 payload["data"] = message
             if "type" not in payload:
                 payload["type"] = payload_type
         else:
-            print_message = PrintMessage(data=message)  # pyright: ignore
+            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
-        if kwargs.get("file") and kwargs["file"] in [
+        file = kwargs.get("file", None)
+        if file and file in [
             sys.stderr,
             sys.__stderr__,
             sys.stdout,
             sys.__stdout__,
         ]:
-            print(dumped, file=kwargs["file"], flush=flush)
+            print(dumped, file=file, flush=flush)
         else:
             print(dumped, flush=flush)
 
-    def input(self, prompt: str = "", *, password: bool = False) -> str:
+    def input(
+        self,
+        prompt: str = "",
+        *,
+        password: bool = False,
+        request_id: str | None = None,
+    ) -> str:
         """Structured input from stdin.
 
         Parameters
@@ -112,13 +125,15 @@ 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 = gen_id()
+        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 "> ",
@@ -129,15 +144,17 @@ class StructuredIOStream(IOStream):
             input_type = "debug"
         self._send_input_request(
             prompt,
-            request_id,
+            input_request_id,
             password,
             input_type=input_type,
         )
-        user_input_raw = self._read_user_input(prompt, password, request_id)
-        response = self._handle_user_input(user_input_raw, request_id)
+        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
 
@@ -163,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
@@ -197,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 = {
@@ -330,7 +326,10 @@ class StructuredIOStream(IOStream):
             response_type = "input_response"
         else:
             response_type = _response_type
-        if user_input.get("request_id") == request_id:
+        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:
@@ -345,7 +344,7 @@ class StructuredIOStream(IOStream):
                     )
             if isinstance(data, list):
                 return self._handle_list_response(
-                    data,  # pyright: ignore
+                    data,
                     request_id=request_id,
                     response_type=response_type,
                 )
@@ -365,7 +364,7 @@ class StructuredIOStream(IOStream):
                     type=response_type,
                     data=self._format_multimedia_response(
                         request_id=request_id,
-                        data=data,  # pyright: ignore
+                        data=data,
                     ),
                     request_id=request_id,
                 )
@@ -391,7 +390,7 @@ class StructuredIOStream(IOStream):
         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,
@@ -400,7 +399,7 @@ class StructuredIOStream(IOStream):
             )
 
         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)
@@ -435,8 +434,8 @@ class StructuredIOStream(IOStream):
         # Create a log message
         got_id: str | None = None
         if isinstance(response, dict):
-            got_id = response.get("request_id")  # pyright: ignore
-        response_str = str(response)  # pyright: ignore
+            got_id = response.get("request_id")
+        response_str = str(response)
         message = response_str[:100] + (
             "..." if len(response_str) > 100 else ""
         )
@@ -481,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
@@ -15,6 +17,7 @@ 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: "

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):

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 (

waldiez/models/agents/agent/agent.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0.
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
+
 # pylint: disable=too-many-public-methods
+# pyright:  reportArgumentType=false
 """Base agent class to be inherited by all agents."""
 
 import warnings
@@ -71,7 +73,7 @@ class WaldiezAgent(WaldiezBase):
             title="Type",
             description="The type of the 'node' in a graph.",
         ),
-    ] = "agent"
+    ]
     agent_type: Annotated[
         WaldiezAgentType,
         Field(
@@ -93,7 +95,7 @@ class WaldiezAgent(WaldiezBase):
             title="Description",
             description="The description of the agent",
         ),
-    ] = "Agent's description"
+    ]
     tags: Annotated[
         list[str],
         Field(
@@ -101,7 +103,7 @@ class WaldiezAgent(WaldiezBase):
             description="Tags of the agent",
             default_factory=list,
         ),
-    ] = []
+    ]
     requirements: Annotated[
         list[str],
         Field(
@@ -109,7 +111,7 @@ class WaldiezAgent(WaldiezBase):
             description="Python requirements for the agent",
             default_factory=list,
         ),
-    ] = []
+    ]
     created_at: Annotated[
         str,
         Field(
@@ -131,7 +133,7 @@ class WaldiezAgent(WaldiezBase):
         Field(
             title="Data",
             description="The data (properties) of the agent",
-            default_factory=WaldiezAgentData,  # pyright: ignore
+            default_factory=WaldiezAgentData,
         ),
     ]
 
@@ -139,7 +141,7 @@ class WaldiezAgent(WaldiezBase):
         list[WaldiezHandoff],
         Field(
             init=False,  # this is not a field in the constructor
-            default_factory=list[WaldiezHandoff],
+            default_factory=list,
             title="Handoffs",
             description=(
                 "A list of handoffs (target ids) to register. "
@@ -148,10 +150,10 @@ class WaldiezAgent(WaldiezBase):
         ),
     ] = []
 
-    _checked_nested_chats: Annotated[bool, Field(init=False, default=False)] = (
+    _checked_nested_chats: Annotated[bool, Field(default=False, init=False)] = (
         False
     )
-    _checked_handoffs: Annotated[bool, Field(init=False, default=False)] = False
+    _checked_handoffs: Annotated[bool, Field(default=False, init=False)] = False
 
     @property
     def args_to_skip(self) -> list[str]:
@@ -190,10 +192,11 @@ class WaldiezAgent(WaldiezBase):
             If handoffs have not been gathered yet.
         """
         if not self._checked_handoffs:
-            raise RuntimeError(
+            msg = (
                 "Handoffs have not been gathered yet. "
                 "Call gather_handoffs() first."
             )
+            raise RuntimeError(msg)
         return self._handoffs
 
     @field_validator("agent_type")
@@ -392,7 +395,7 @@ class WaldiezAgent(WaldiezBase):
         return class_name  # pragma: no cover
 
     @property
-    def ag2_imports(self) -> set[str]:
+    def ag2_imports(self) -> set[str]:  # pylint: disable=too-complex
         """Return the AG2 imports of the agent.
 
         Returns
@@ -413,25 +416,28 @@ class WaldiezAgent(WaldiezBase):
             case "UserProxyAgent":
                 imports.add("from autogen import UserProxyAgent")
             case "RetrieveUserProxyAgent":
-                imports.add(
+                _imp = (
                     "from autogen.agentchat.contrib.retrieve_user_proxy_agent "
                     "import RetrieveUserProxyAgent"
                 )
+                imports.add(_imp)
             case "MultimodalConversableAgent":
-                imports.add(
+                _imp = (
                     "from "
                     "autogen.agentchat.contrib.multimodal_conversable_agent "
                     "import MultimodalConversableAgent"
                 )
+                imports.add(_imp)
             case "ReasoningAgent":
                 imports.add(
                     "from autogen.agents.experimental import ReasoningAgent"
                 )
             case "CaptainAgent":
-                imports.add(
+                _imp = (
                     "from autogen.agentchat.contrib.captainagent "
                     "import CaptainAgent"
                 )
+                imports.add(_imp)
             case "GroupChatManager":  # pragma: no branch
                 imports.add("from autogen import GroupChat")
                 imports.add("from autogen.agentchat import GroupChatManager")
@@ -443,10 +449,11 @@ class WaldiezAgent(WaldiezBase):
             case "ConversableAgent":
                 imports.add("from autogen import ConversableAgent")
             case _:  # pragma: no cover
-                raise ValueError(
+                msg = (
                     f"Unknown agent class: {agent_class}. "
                     "Please implement the imports for this class."
                 )
+                raise ValueError(msg)
         return imports
 
     def validate_linked_tools(
@@ -678,8 +685,8 @@ class WaldiezAgent(WaldiezBase):
             nested_chat = WaldiezAgentNestedChat(
                 triggered_by=triggered_by,
                 messages=messages,
-                condition=chat_with_condition.condition,  # pyright: ignore
-                available=chat_with_condition.available,  # pyright: ignore
+                condition=chat_with_condition.condition,
+                available=chat_with_condition.available,
             )
             self.data.nested_chats.append(nested_chat)