waldiez 0.4.8__py3-none-any.whl → 0.4.11__py3-none-any.whl

waldiez/exporting/flow/orchestrator.py

@@ -23,6 +23,7 @@ from .utils import (
     generate_header,
     get_after_run_content,
     get_np_no_nep50_handle,
+    get_set_io_stream,
     get_sqlite_out,
     get_start_logging,
     get_stop_logging,
@@ -182,6 +183,17 @@ class ExportOrchestrator:
             order=ContentOrder.CLEANUP.value + 1,  # after imports
             skip_strip=True,  # keep newlines
         )
+        if not self.config.skip_patch_io:
+            merged_result.add_content(
+                get_set_io_stream(
+                    use_structured_io=self.config.structured_io,
+                    is_async=self.waldiez.is_async,
+                    uploads_root=self.config.uploads_root,
+                ),
+                position=ExportPosition.IMPORTS,  # after imports, before models
+                order=ContentOrder.CLEANUP.value + 2,  # after start logging
+            )
+        # merged_result.add_content
         merged_result.add_content(
             get_sqlite_out(is_async=self.waldiez.is_async),
             position=ExportPosition.AGENTS,

waldiez/exporting/flow/utils/__init__.py

@@ -6,6 +6,7 @@ from .common import (
     generate_header,
     get_after_run_content,
     get_np_no_nep50_handle,
+    get_set_io_stream,
 )
 from .importing import get_sorted_imports, get_the_imports_string
 from .logging import get_sqlite_out, get_start_logging, get_stop_logging
@@ -19,4 +20,5 @@ __all__ = [
     "get_start_logging",
     "get_stop_logging",
     "get_sqlite_out",
+    "get_set_io_stream",
 ]

waldiez/exporting/flow/utils/common.py

@@ -4,6 +4,8 @@
 # flake8: noqa: E501
 """Common utils for the final generatio."""
 
+from pathlib import Path
+
 from waldiez.models import Waldiez
 
 from ...core import FILE_HEADER
@@ -151,7 +153,7 @@ def get_after_run_content(
         if agent.is_reasoning:
             agent_name = agent_names[agent.id]
             content += f"""
-{space}# pylint: disable=broad-except,too-many-try-statements
+{space}# pylint: disable=broad-exception-caught,too-many-try-statements
 {space}try:
 {space}{tab}{agent_name}.visualize_tree()
 {space}{tab}if os.path.exists("tree_of_thoughts.png"):
@@ -160,8 +162,9 @@ def get_after_run_content(
 {space}except BaseException:
 {space}{tab}pass
 {space}# save the tree to json
+{space}# pylint: disable=protected-access
 {space}try:
-{space}{tab}data = {agent_name}._root.to_dict()  # pylint: disable=protected-access  # pyright: ignore
+{space}{tab}data = {agent_name}._root.to_dict()  # pyright: ignore
 {space}{tab}with open("{agent_name}_reasoning_tree.json", "w", encoding="utf-8") as f:
 {space}{tab}{tab}json.dump(data, f)
 {space}except BaseException:
@@ -204,3 +207,94 @@ if not hasattr(np, "_no_pep50_warning"):
     setattr(np, "_no_pep50_warning", _np_no_nep50_warning)  # noqa
 '''
     return content
+
+
+def get_set_io_stream(
+    use_structured_io: bool,
+    is_async: bool,
+    uploads_root: Path | None,
+) -> str:
+    """Get the content to set structured IO.
+
+    Parameters
+    ----------
+    use_structured_io : bool
+        Whether to use structured IO or not.
+    is_async : bool
+        Whether the flow is async or not.
+    uploads_root : Path | None
+        The uploads root, to get user-uploaded files, by default None.
+
+    Returns
+    -------
+    str
+        The content to set structured IO.
+    """
+    if use_structured_io:
+        return get_set_structured_io_stream(
+            is_async=is_async,
+            uploads_root=uploads_root,
+        )
+    return get_patch_default_io_stream(is_async)
+
+
+def get_set_structured_io_stream(
+    is_async: bool,
+    uploads_root: Path | None,
+) -> str:
+    """Get the content to set structured IO.
+
+    Parameters
+    ----------
+    is_async : bool
+        Whether the flow is async or not.
+    uploads_root : Path | None
+        The uploads root, to get user-uploaded files, by default None.
+
+    Returns
+    -------
+    str
+        The content to set structured IO.
+    """
+    upload_root_arg = f'r"{uploads_root}"' if uploads_root else "None"
+    return f"""
+# set structured IO
+try:
+    # pylint: disable=import-outside-toplevel
+    from autogen.io import IOStream
+    from waldiez.io import StructuredIOStream
+    stream = StructuredIOStream(
+        is_async={is_async},
+        uploads_root={upload_root_arg}
+    )
+    IOStream.set_default(stream)
+except BaseException:  # pylint: disable=broad-exception-caught
+    # allow running the flow without structured IO
+    pass
+"""
+
+
+def get_patch_default_io_stream(is_async: bool) -> str:
+    """Get the content to patch the default IO stream.
+
+    Parameters
+    ----------
+    is_async : bool
+        Whether the flow is async or not.
+
+    Returns
+    -------
+    str
+        The content to patch the default IO stream.
+    """
+    # copy from waldiez/running/patch_io_stream.py
+    return f"""
+# patch the default IOStream
+try:
+    # pylint: disable=import-outside-toplevel
+    from waldiez.running.patch_io_stream import patch_io_stream
+    patch_io_stream(is_async={is_async})
+except BaseException:  # pylint: disable=broad-exception-caught
+    # allow running the flow without patching the IOStream
+    pass
+"""

waldiez/exporting/flow/utils/logging.py

@@ -193,7 +193,7 @@ def get_async_sqlite_out() -> str:
         query = f"SELECT * FROM {table}"  # nosec
         try:
             cursor = await conn.execute(query)
-        except BaseException:  # pylint: disable=broad-except
+        except BaseException:  # pylint: disable=broad-exception-caught
             await conn.close()
             return
         rows = await cursor.fetchall()
@@ -210,6 +210,7 @@ def get_async_sqlite_out() -> str:
             await file.write(json.dumps(data, indent=4, ensure_ascii=False)
     ```
     """
+    # fmt: off
     content = "\n\n"
     content += "async def get_sqlite_out(dbname: str, table: str, csv_file: str) -> None:\n"
     content += '    """Convert a sqlite table to csv and json files.\n\n'
@@ -226,7 +227,7 @@ def get_async_sqlite_out() -> str:
     content += '    query = f"SELECT * FROM {table}"  # nosec\n'
     content += "    try:\n"
     content += "        cursor = await conn.execute(query)\n"
-    content += "    except BaseException:  # pylint: disable=broad-except\n"
+    content += "    except BaseException:  # pylint: disable=broad-exception-caught\n"
     content += "        await conn.close()\n"
     content += "        return\n"
     content += "    rows = await cursor.fetchall()\n"
@@ -235,10 +236,7 @@ def get_async_sqlite_out() -> str:
     content += "    data = [dict(zip(column_names, row)) for row in rows]\n"
     content += "    await cursor.close()\n"
     content += "    await conn.close()\n"
-    content += (
-        '    async with aiofiles.open(csv_file, "w", newline="", '
-        'encoding="utf-8") as file:\n'
-    )
+    content += '    async with aiofiles.open(csv_file, "w", newline="", encoding="utf-8") as file:\n'
     content += '        csv_writer = AsyncDictWriter(file, fieldnames=column_names, dialect="unix")\n'
     content += "        await csv_writer.writeheader()\n"
     content += "        await csv_writer.writerows(data)\n"
@@ -246,6 +244,7 @@ def get_async_sqlite_out() -> str:
     content += '    async with aiofiles.open(json_file, "w", encoding="utf-8") as file:\n'
     content += "        await file.write(json.dumps(data, indent=4, ensure_ascii=False))\n"
     content += "\n"
+    # fmt: on
     return content
 
 

waldiez/io/mqtt.py

@@ -4,7 +4,7 @@
 # flake8: noqa: E501
 # pylint: disable=too-many-try-statements,broad-exception-caught,
 # pylint: disable=line-too-long,unused-argument,too-many-instance-attributes
-# pylint: disable=too-many-arguments,too-many-positional-arguments
+# pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals
 
 """An MQTT I/O stream for handling print and input messages."""
 
@@ -82,6 +82,7 @@ class MqttIOStream(IOStream):
         broker_port: int = 1883,
         task_id: str | None = None,
         input_timeout: int = 120,
+        connect_timeout: int = 10,
         max_retain_messages: int = 1000,
         on_input_request: Optional[Callable[[str, str, str], None]] = None,
         on_input_response: Optional[Callable[[str, str], None]] = None,
@@ -104,6 +105,8 @@ class MqttIOStream(IOStream):
             An ID to use for the topics. If not provided, a random UUID will be generated.
         input_timeout : int, optional
             The time to wait for user input in seconds, by default 120.
+        connect_timeout : int, optional
+            The time to wait for MQTT connection in seconds, by default 10.
         on_input_request : Optional[Callable[[str, str, str], None]], optional
             Callback for input request, by default None
             parameters: prompt, request_id, task_id
@@ -129,6 +132,7 @@ class MqttIOStream(IOStream):
         self.broker_port = broker_port
         self.task_id = task_id or uuid.uuid4().hex
         self.input_timeout = input_timeout
+        self.connect_timeout = connect_timeout
         self.on_input_request = on_input_request
         self.on_input_response = on_input_response
         self.max_retain_messages = max_retain_messages
@@ -191,7 +195,7 @@ class MqttIOStream(IOStream):
             self.client.loop_start()
 
             # Wait for connection
-            timeout = 10  # seconds
+            timeout = self.connect_timeout  # seconds
             start_time = time.time()
             while (
                 not self.client.is_connected()
@@ -280,7 +284,7 @@ class MqttIOStream(IOStream):
             while reconnect_count < MQTT_MAX_RECONNECT_COUNT:
                 LOG.info("Reconnecting in %d seconds...", reconnect_delay)
                 time.sleep(reconnect_delay)
-                # pylint: disable=broad-except
+                # pylint: disable=broad-exception-caught
                 try:
                     client.reconnect()
                 except Exception as err:

waldiez/io/structured.py

@@ -36,9 +36,13 @@ class StructuredIOStream(IOStream):
     uploads_root: Path | None = None
 
     def __init__(
-        self, timeout: float = 120, uploads_root: Path | str | None = None
+        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():
@@ -97,7 +101,11 @@ class StructuredIOStream(IOStream):
             The line read from the input stream.
         """
         request_id = uuid4().hex
-        prompt = prompt or "> "
+        prompt = prompt or ">"
+        if not prompt or prompt in [">", "> "]:
+            # 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)
@@ -106,6 +114,22 @@ class StructuredIOStream(IOStream):
             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

waldiez/models/common/method_utils.py

@@ -73,7 +73,7 @@ def parse_code_string(
         If valid, None and the ast module.
         If invalid, the error message and None.
     """
-    # pylint: disable=broad-except
+    # pylint: disable=broad-exception-caught
     try:
         tree = ast.parse(code_string)
     except SyntaxError as e:

waldiez/models/tool/tool.py

@@ -137,9 +137,10 @@ class WaldiezTool(WaldiezBase):
             raise FileNotFoundError(f"File not found: {resolved}")
         with resolved.open("r", encoding="utf-8") as file:
             data_string = file.read()
+            # pylint: disable=broad-exception-caught
             try:
                 data_dict = json.loads(data_string)
-            except BaseException as exc:  # pylint: disable=broad-except
+            except BaseException as exc:
                 raise ValueError(f"Invalid WaldiezTool/JSON: {exc}") from exc
             return WaldiezTool.model_validate(data_dict)