waldiez 0.5.3__py3-none-any.whl → 0.5.5__py3-none-any.whl

waldiez/exporting/flow/utils/common.py

@@ -4,8 +4,6 @@
 # flake8: noqa: E501
 """Common utils for the final generatio."""
 
-from pathlib import Path
-
 from waldiez.models import Waldiez
 
 from ...core import FILE_HEADER
@@ -16,7 +14,6 @@ from .linting import (
     get_pyright_ignore_comment,
 )
 
-RETURN_TYPE_HINT = "Union[ChatResult, list[ChatResult], dict[int, ChatResult]]"
 GENERATED_WITH = "🧩 generated with ❤️ by Waldiez."
 
 
@@ -105,21 +102,34 @@ Tags: {", ".join(tags)}
 """'''
 
 
-def main_doc_string() -> str:
+def main_doc_string(is_async: bool) -> str:
     """Generate the docstring for the main function.
 
+    Parameters
+    ----------
+    is_async : bool
+        Whether the main function is asynchronous.
+
     Returns
     -------
     str
         The docstring for the main function.
     """
+    return_type_hint = (
+        "AsyncRunResponseProtocol" if is_async else "RunResponseProtocol"
+    )
     return f'''"""Start chatting.
 
     Returns
     -------
-    {RETURN_TYPE_HINT}
+    {return_type_hint}
         The result of the chat session, which can be a single ChatResult,
         a list of ChatResults, or a dictionary mapping integers to ChatResults.
+
+    Raises
+    ------
+    RuntimeError
+        If the chat session fails.
     """'''
 
 
@@ -207,94 +217,3 @@ 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/importing.py

@@ -21,6 +21,7 @@ TYPING_IMPORT_NAMES = [
     "Annotated",
     "Any",
     "Callable",
+    "Coroutine",
     "Dict",
     "List",
     "Optional",
@@ -37,6 +38,9 @@ COMMON_AUTOGEN_IMPORTS = [
     "from autogen import ChatResult",
     "from autogen import GroupChat",
     "from autogen import runtime_logging",
+    "from autogen.events import BaseEvent",
+    "from autogen.io.run_response import RunResponseProtocol",
+    "from autogen.io.run_response import AsyncRunResponseProtocol",
 ]
 
 

waldiez/io/mqtt.py

@@ -2,9 +2,10 @@
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 
 # 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,too-many-locals
+# pylint: disable=too-many-try-statements,broad-exception-caught
+# pylint: disable=line-too-long,duplicate-code,unused-argument
+# pylint: disable=too-many-arguments,too-many-positional-arguments
+# pylint: disable=too-many-locals,too-many-instance-attributes
 
 """An MQTT I/O stream for handling print and input messages."""
 
@@ -351,9 +352,13 @@ class MqttIOStream(IOStream):
         }
         LOG.debug("MQTT log: %s", payload)
         print_message = PrintMessage(data=buf)
-        self._print_to_common_output(
-            payload=print_message.model_dump(mode="json")
-        )
+        try:
+            payload = print_message.model_dump(mode="json")
+        except Exception:
+            payload = print_message.model_dump(
+                serialize_as_any=True, mode="json", fallback=str
+            )
+        self._print_to_common_output(payload=payload)
 
     def _handle_input_response(self, payload: str) -> None:
         """Handle input response message."""
@@ -463,7 +468,12 @@ class MqttIOStream(IOStream):
             Additional keyword arguments.
         """
         print_message = PrintMessage.create(*args, **kwargs)
-        payload = print_message.model_dump(mode="json")
+        try:
+            payload = print_message.model_dump(mode="json")
+        except Exception:
+            payload = print_message.model_dump(
+                serialize_as_any=True, mode="json", fallback=str
+            )
         self._print(payload)
 
     def send(self, message: BaseMessage) -> None:
@@ -476,11 +486,16 @@ class MqttIOStream(IOStream):
         """
         try:
             message_dump = message.model_dump(mode="json")
-        except Exception as e:  # pragma: no cover
-            message_dump = {
-                "error": str(e),
-                "type": message.__class__.__name__,
-            }
+        except Exception:
+            try:
+                message_dump = message.model_dump(
+                    serialize_as_any=True, mode="json", fallback=str
+                )
+            except Exception as e:
+                message_dump = {
+                    "error": str(e),
+                    "type": message.__class__.__name__,
+                }
 
         message_type = message_dump.get("type", None)
         if not message_type:  # pragma: no cover
@@ -523,8 +538,12 @@ class MqttIOStream(IOStream):
             prompt=prompt,
             password=password,
         )
-
-        payload = input_request.model_dump(mode="json")
+        try:
+            payload = input_request.model_dump(mode="json")
+        except Exception:  # pragma: no cover
+            payload = input_request.model_dump(
+                serialize_as_any=True, mode="json", fallback=str
+            )
         payload["task_id"] = self.task_id
         payload["password"] = str(password).lower()
 

waldiez/io/redis.py

@@ -2,7 +2,8 @@
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 
 # flake8: noqa: E501
-# pylint: disable=too-many-try-statements,broad-exception-caught,line-too-long
+# pylint: disable=too-many-try-statements,broad-exception-caught
+# pylint: disable=line-too-long,duplicate-code
 
 """A Redis I/O stream for handling print and input messages."""
 
@@ -29,8 +30,8 @@ except ImportError as error:  # pragma: no cover
     raise ImportError(
         "Redis client not installed. Please install redis-py with `pip install redis`."
     ) from error
+from autogen.events import BaseEvent  # type: ignore
 from autogen.io import IOStream  # type: ignore
-from autogen.messages import BaseMessage  # type: ignore
 
 from .models import (
     PrintMessage,
@@ -39,7 +40,7 @@ from .models import (
     UserInputRequest,
     UserResponse,
 )
-from .utils import gen_id, now
+from .utils import gen_id, get_message_dump, now
 
 if TYPE_CHECKING:
     Redis = redis.Redis[bytes]
@@ -217,7 +218,12 @@ class RedisIOStream(IOStream):
             Additional keyword arguments.
         """
         print_message = PrintMessage.create(*args, **kwargs)
-        payload = print_message.model_dump(mode="json")
+        try:
+            payload = print_message.model_dump(mode="json")
+        except Exception:  # pragma: no cover
+            payload = print_message.model_dump(
+                serialize_as_any=True, mode="json", fallback=str
+            )
         self._print(payload)
 
     def input(
@@ -249,7 +255,12 @@ class RedisIOStream(IOStream):
             prompt=prompt,
             password=password,
         )
-        payload = input_request.model_dump(mode="json")
+        try:
+            payload = input_request.model_dump(mode="json")
+        except Exception:  # pragma: no cover
+            payload = input_request.model_dump(
+                serialize_as_any=True, mode="json", fallback=str
+            )
         payload["password"] = str(password).lower()
         payload["task_id"] = self.task_id
         LOG.debug("Requesting input via Pub/Sub: %s", payload)
@@ -278,7 +289,7 @@ class RedisIOStream(IOStream):
         self._print(payload)
         return user_input
 
-    def send(self, message: BaseMessage) -> None:
+    def send(self, message: BaseEvent) -> None:
         """Send a structured message to Redis.
 
         Parameters
@@ -286,13 +297,7 @@ class RedisIOStream(IOStream):
         message : dict[str, Any]
             The message to send.
         """
-        try:
-            message_dump = message.model_dump(mode="json")
-        except Exception as e:  # pragma: no cover
-            message_dump = {
-                "type": "error",
-                "error": str(e),
-            }
+        message_dump = get_message_dump(message)
         message_type = message_dump.get("type", None)
         if not message_type:
             message_type = message.__class__.__name__

waldiez/io/structured.py

@@ -24,6 +24,7 @@ from .models import (
 from .utils import (
     gen_id,
     get_image,
+    get_message_dump,
     is_json_dumped,
     now,
     try_parse_maybe_serialized,
@@ -83,7 +84,7 @@ class StructuredIOStream(IOStream):
         flush = kwargs.get("flush", True)
         payload_type = kwargs.get("type", "print")
         payload["type"] = payload_type
-        print(json.dumps(payload), flush=flush)
+        print(json.dumps(payload, default=str), flush=flush)
 
     def input(self, prompt: str = "", *, password: bool = False) -> str:
         """Structured input from stdin.
@@ -133,7 +134,7 @@ class StructuredIOStream(IOStream):
         return user_response
 
     # noinspection PyMethodMayBeStatic
-    # pylint: disable=no-self-use
+    # pylint: disable=no-self-use,broad-exception-caught
     def send(self, message: BaseEvent) -> None:
         """Structured sending of a BaseEvent.
 
@@ -142,7 +143,7 @@ class StructuredIOStream(IOStream):
         message : BaseEvent
             The message to send.
         """
-        message_dump = message.model_dump(mode="json")
+        message_dump = get_message_dump(message)
         if message_dump.get("type") == "text":
             content_block = message_dump.get("content")
             if (
@@ -154,7 +155,11 @@ class StructuredIOStream(IOStream):
                 content_block["content"] = try_parse_maybe_serialized(
                     inner_content
                 )
-        print(json.dumps(message_dump), flush=True, file=sys.stdout)
+        print(
+            json.dumps(message_dump, default=str),
+            flush=True,
+            file=sys.stdout,
+        )
 
     # noinspection PyMethodMayBeStatic
     # pylint: disable=no-self-use

waldiez/io/utils.py

@@ -11,6 +11,7 @@ from pathlib import Path
 from typing import Any, Literal, Union
 
 from autogen.agentchat.contrib.img_utils import get_pil_image  # type: ignore
+from autogen.events import BaseEvent  # type: ignore
 
 MessageType = Literal[
     "input_request",
@@ -182,3 +183,34 @@ def try_parse_maybe_serialized(value: str) -> Any:
         except Exception:
             pass  # Try next parser
     return value  # Return original if all parsing fails
+
+
+def get_message_dump(
+    message: BaseEvent,
+) -> dict[str, Any]:
+    """Get the message dump.
+
+    Parameters
+    ----------
+    message : BaseEvent
+        The message to dump.
+
+    Returns
+    -------
+    dict[str, Any]
+        The dumped message.
+    """
+    # pylint: disable=broad-exception-caught
+    try:
+        message_dump = message.model_dump(mode="json")
+    except Exception:
+        try:
+            message_dump = message.model_dump(
+                serialize_as_any=True, mode="json", fallback=str
+            )
+        except Exception as e:
+            message_dump = {
+                "error": str(e),
+                "type": message.__class__.__name__,
+            }
+    return message_dump

waldiez/io/ws.py

@@ -19,7 +19,12 @@ from ._ws import (
     is_websocket_available,
 )
 from .models import UserResponse
-from .utils import is_json_dumped, now, try_parse_maybe_serialized
+from .utils import (
+    get_message_dump,
+    is_json_dumped,
+    now,
+    try_parse_maybe_serialized,
+)
 
 LOG = logging.getLogger(__name__)
 
@@ -137,7 +142,7 @@ class AsyncWebsocketsIOStream(IOStream):
         message : BaseEvent
             The message to send.
         """
-        message_dump = message.model_dump(mode="json")
+        message_dump = get_message_dump(message)
 
         if message_dump.get("type") == "text":
             content_block = message_dump.get("content")
@@ -150,6 +155,7 @@ class AsyncWebsocketsIOStream(IOStream):
                 content_block["content"] = try_parse_maybe_serialized(
                     inner_content
                 )
+                message_dump["content"] = content_block
 
         json_dump = json.dumps(message_dump, ensure_ascii=False)
 

waldiez/models/__init__.py

@@ -29,6 +29,9 @@ from .agents import (
     WaldiezCaptainAgent,
     WaldiezCaptainAgentData,
     WaldiezCaptainAgentLibEntry,
+    WaldiezDocAgent,
+    WaldiezDocAgentData,
+    WaldiezDocAgentQueryEngine,
     WaldiezGroupManager,
     WaldiezGroupManagerData,
     WaldiezGroupManagerSpeakers,
@@ -150,6 +153,8 @@ __all__ = [
     "WaldiezContextBasedTransition",
     "WaldiezContextStrLLMCondition",
     "WaldiezDefaultCondition",
+    "WaldiezDocAgent",
+    "WaldiezDocAgentData",
     "WaldiezExpressionContextCondition",
     "WaldiezFlow",
     "WaldiezFlowData",
@@ -171,6 +176,7 @@ __all__ = [
     "WaldiezModelAWS",
     "WaldiezModelData",
     "WaldiezModelPrice",
+    "WaldiezDocAgentQueryEngine",
     "WaldiezRagUserProxy",
     "WaldiezRagUserProxyChunkMode",
     "WaldiezRagUserProxyData",

waldiez/models/agents/__init__.py

@@ -25,6 +25,11 @@ from .captain import (
     WaldiezCaptainAgentData,
     WaldiezCaptainAgentLibEntry,
 )
+from .doc_agent import (
+    WaldiezDocAgent,
+    WaldiezDocAgentData,
+    WaldiezDocAgentQueryEngine,
+)
 from .extra_requirements import (
     get_captain_agent_extra_requirements,
     get_retrievechat_extra_requirements,
@@ -99,6 +104,8 @@ __all__ = [
     "WaldiezCaptainAgent",
     "WaldiezCaptainAgentData",
     "WaldiezCaptainAgentLibEntry",
+    "WaldiezDocAgent",
+    "WaldiezDocAgentData",
     "WaldiezRagUserProxy",
     "WaldiezRagUserProxyData",
     "WaldiezRagUserProxyModels",
@@ -107,6 +114,7 @@ __all__ = [
     "WaldiezReasoningAgentReasonConfig",
     "WaldiezUserProxy",
     "WaldiezUserProxyData",
+    "WaldiezDocAgentQueryEngine",
     "WaldiezRagUserProxyRetrieveConfig",
     "WaldiezRagUserProxyTask",
     "WaldiezRagUserProxyChunkMode",