waldiez 0.2.2__py3-none-any.whl → 0.3.1__py3-none-any.whl

waldiez/exporting/chats/chats_exporter.py

@@ -0,0 +1,240 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+# flake8: noqa E501
+# pylint: disable=line-too-long
+"""Chats exporter."""
+
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
+from waldiez.models import WaldiezAgent, WaldiezChat
+
+from ..base import (
+    AgentPosition,
+    AgentPositions,
+    BaseExporter,
+    ExporterMixin,
+    ExporterReturnType,
+    ExportPosition,
+    ExportPositions,
+    ImportPosition,
+)
+from .utils import (
+    export_nested_chat_registration,
+    export_sequential_chat,
+    export_single_chat,
+    export_swarm_chat,
+)
+
+
+class ChatsExporter(BaseExporter, ExporterMixin):
+    """Chats exporter."""
+
+    _chat_string: Optional[str]
+    _before_chat: Optional[str]
+    _generated: bool
+
+    def __init__(
+        self,
+        get_swarm_members: Callable[
+            [WaldiezAgent], Tuple[List[WaldiezAgent], Optional[WaldiezAgent]]
+        ],
+        all_agents: List[WaldiezAgent],
+        agent_names: Dict[str, str],
+        all_chats: List[WaldiezChat],
+        chat_names: Dict[str, str],
+        main_chats: List[Tuple[WaldiezChat, WaldiezAgent, WaldiezAgent]],
+        for_notebook: bool,
+        is_async: bool,
+    ):
+        """Initialize the chats exporter.
+
+        Parameters
+        ----------
+        get_swarm_members : Callable[
+                [WaldiezAgent],
+                Tuple[List[WaldiezAgent], Optional[WaldiezAgent]]
+            ]
+            The function to use to resolve the swarm members.
+        all_agents : List[WaldiezAgent]
+            All the agents in the flow.
+        agent_names : Dict[str, str]
+            A mapping of agent id to agent name.
+        all_chats : List[WaldiezChat]
+            All the chats in the flow.
+        chat_names : Dict[str, str]
+            A mapping of chat id to chat name.
+        main_chats : List[Tuple[WaldiezChat, WaldiezAgent, WaldiezAgent]]
+            The main chats in the flow.
+        for_notebook : bool
+            Whether the export is for a notebook.
+        is_async : bool
+            Whether the chat is asynchronous.
+        """
+        self.all_agents = all_agents
+        self.agent_names = agent_names
+        self.main_chats = main_chats
+        self.all_chats = all_chats
+        self.chat_names = chat_names
+        self.get_swarm_members = get_swarm_members
+        self.for_notebook = for_notebook
+        self.is_async = is_async
+        self._chat_string = None
+        self._before_chat = None
+        self._generated = False
+
+    def _export_chats(self) -> None:
+        """Export the chats content."""
+        if len(self.main_chats) == 1:
+            main_chat = self.main_chats[0]
+            chat, sender, recipient = main_chat
+            if sender.agent_type == "swarm" or recipient.agent_type == "swarm":
+                self._chat_string, self._before_chat = export_swarm_chat(
+                    get_swarm_members=self.get_swarm_members,
+                    chat=chat,
+                    agent_names=self.agent_names,
+                    chat_names=self.chat_names,
+                    sender=sender,
+                    recipient=recipient,
+                    serializer=self.serializer,
+                    string_escape=self.string_escape,
+                    tabs=0 if self.for_notebook else 1,
+                    is_async=self.is_async,
+                )
+                return
+            self._chat_string, self._before_chat = export_single_chat(
+                sender=sender,
+                recipient=recipient,
+                chat=chat,
+                agent_names=self.agent_names,
+                chat_names=self.chat_names,
+                serializer=self.serializer,
+                string_escape=self.string_escape,
+                tabs=0 if self.for_notebook else 1,
+                is_async=self.is_async,
+            )
+            return
+        self._chat_string, self._before_chat = export_sequential_chat(
+            main_chats=self.main_chats,
+            agent_names=self.agent_names,
+            chat_names=self.chat_names,
+            serializer=self.serializer,
+            string_escape=self.string_escape,
+            tabs=0 if self.for_notebook else 1,
+            is_async=self.is_async,
+        )
+
+    def get_imports(self) -> Optional[List[Tuple[str, ImportPosition]]]:
+        """Get the imports string.
+
+        Returns
+        -------
+        str
+            The imports string.
+        """
+        if len(self.main_chats) == 1:
+            _, sender, recipient = self.main_chats[0]
+            if sender.agent_type == "swarm" or recipient.agent_type == "swarm":
+                import_string = "from autogen import initiate_swarm_chat"
+                if self.is_async:
+                    import_string = "from autogen import a_initiate_swarm_chat"
+                return [(import_string, ImportPosition.THIRD_PARTY)]
+            # no additional imports, it is `sender.initiate_chat(....)`
+            return None
+        if self.is_async:
+            import_string = (
+                "from autogen.agentchat.chat import a_initiate_chats"
+            )
+        else:
+            import_string = "from autogen.agentchat.chat import initiate_chats"
+        return [(import_string, ImportPosition.THIRD_PARTY)]
+
+    def generate(self) -> str:
+        """Generate the chats content.
+
+        Returns
+        -------
+        str
+            The chats content.
+        """
+        if self._generated is False:
+            self._export_chats()
+            self._generated = True
+        return self._chat_string or ""
+
+    def get_before_export(
+        self,
+    ) -> Optional[List[Tuple[str, Union[ExportPosition, AgentPosition]]]]:
+        """Generate the content before the main export.
+
+        Returns
+        -------
+        Optional[List[Tuple[str, Union[ExportPosition, AgentPosition]]]]
+            The exported content before the main export and its position.
+        """
+        before: List[Tuple[str, Union[ExportPosition, AgentPosition]]] = []
+        if self._generated is False:
+            self._export_chats()
+            self._generated = True
+        if self._before_chat:
+            before.append(
+                (self._before_chat, ExportPosition(ExportPositions.CHATS))
+            )
+        return before
+
+    def get_after_export(
+        self,
+    ) -> Optional[List[Tuple[str, Union[ExportPosition, AgentPosition]]]]:
+        """Generate the content after the main export.
+
+        Returns
+        -------
+        Optional[List[Tuple[str, Union[ExportPosition, AgentPosition]]]]
+            The exported content after the main export and its position.
+        """
+        after: List[Tuple[str, Union[ExportPosition, AgentPosition]]] = []
+        # not per agent, we might have references to agents not yet defined.
+        # let's use one string for all nested chat registrations
+        nested_chat_registrations = ""
+        for agent in self.all_agents:
+            if agent.agent_type != "swarm":
+                registration_string = export_nested_chat_registration(
+                    agent=agent,
+                    all_chats=self.all_chats,
+                    chat_names=self.chat_names,
+                    agent_names=self.agent_names,
+                    string_escape=self.string_escape,
+                    serializer=self.serializer,
+                    is_async=self.is_async,
+                )
+                if registration_string:
+                    nested_chat_registrations += "\n" + registration_string
+        if nested_chat_registrations:
+            # let's place it before the chats (after all agents are defined)
+            after.append(
+                (
+                    nested_chat_registrations,
+                    AgentPosition(None, AgentPositions.AFTER_ALL, 2),
+                )
+            )
+        return after
+
+    def export(self) -> ExporterReturnType:
+        """Export the chats.
+
+        Returns
+        -------
+        ExporterReturnType
+            The exported chats, the imports, the before export strings,
+            the after export strings, and the environment variables.
+        """
+        exported_string = self.generate()
+        imports = self.get_imports()
+        before_export = self.get_before_export()
+        after_export = self.get_after_export()
+        return {
+            "content": exported_string,
+            "imports": imports,
+            "before_export": before_export,
+            "after_export": after_export,
+            "environment_variables": None,
+        }

waldiez/exporting/chats/utils/__init__.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Utils for exporting chats."""
+
+from .nested import export_nested_chat_registration
+from .sequential import export_sequential_chat
+from .single_chat import export_single_chat
+from .swarm import export_swarm_chat
+
+__all__ = [
+    "export_nested_chat_registration",
+    "export_sequential_chat",
+    "export_single_chat",
+    "export_swarm_chat",
+]

waldiez/exporting/chats/utils/common.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Common utilities for exporting chats."""
+
+from typing import Any, Callable, Dict, Optional, Tuple
+
+from waldiez.models import WaldiezAgent, WaldiezChat
+
+
+def update_summary_chat_args(
+    chat_args: Dict[str, Any],
+    string_escape: Callable[[str], str],
+) -> Dict[str, Any]:
+    """Escape quotes in the summary args if they are strings.
+
+    Parameters
+    ----------
+    chat_args : Dict[str, Any]
+        The chat arguments.
+    string_escape : Callable[[str], str]
+        The function to escape the string
+
+    Returns
+    -------
+    Dict[str, Any]
+        The chat arguments with the summary prompt escaped.
+    """
+    if "summary_args" in chat_args and isinstance(
+        chat_args["summary_args"], dict
+    ):
+        for key, value in chat_args["summary_args"].items():
+            if isinstance(value, str):
+                chat_args["summary_args"][key] = string_escape(value)
+    return chat_args
+
+
+def get_chat_message_string(
+    sender: WaldiezAgent,
+    chat: WaldiezChat,
+    chat_names: Dict[str, str],
+    string_escape: Callable[[str], str],
+) -> Tuple[str, Optional[str]]:
+    """Get the agent's message as a string.
+
+    Parameters
+    ----------
+    sender : WaldiezAgent
+        The sender.
+    chat : WaldiezChat
+        The chat.
+    chat_names : Dict[str, str]
+        A mapping of chat id to chat name with all the chats in the flow.
+    string_escape : Callable[[str], str]
+        The function to escape the string.
+
+    Returns
+    -------
+    Tuple[str, Optional[str]]
+        If the message is a string, the message content and None.
+        If the message is a method, the method name and the method content.
+        If the message is None, 'None' and None.
+    """
+    if (
+        not chat.message
+        or chat.message.type == "none"
+        or chat.message.content is None
+        or chat.message_content is None
+    ):
+        return "None", None
+    if chat.message.type == "string":
+        return string_escape(chat.message.content), None
+
+    is_rag_with_carryover = (
+        sender.agent_type == "rag_user" and chat.message.use_carryover
+    )
+    chat_name = chat_names[chat.id]
+    function_content, function_name = chat.get_message_function(
+        name_suffix=chat_name,
+        is_rag=is_rag_with_carryover,
+    )
+    return function_name, function_content

waldiez/exporting/chats/{nested.py → utils/nested.py}

@@ -1,6 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+# pylint: disable=too-many-locals
 """Nested chats exporting."""
 
-from typing import Dict, List, Optional, Tuple
+from typing import Callable, Dict, List, Optional, Tuple
 
 from waldiez.models import (
     WaldiezAgent,
@@ -9,8 +12,90 @@ from waldiez.models import (
     WaldiezChat,
 )
 
-from ..utils import get_escaped_string, get_object_string
-from .helpers import escape_summary_args_quotes
+from .common import update_summary_chat_args
+
+
+def export_nested_chat_registration(
+    agent: WaldiezAgent,
+    all_chats: List[WaldiezChat],
+    chat_names: Dict[str, str],
+    agent_names: Dict[str, str],
+    serializer: Callable[..., str],
+    string_escape: Callable[[str], str],
+    is_async: bool,
+) -> str:
+    """Get the nested chat string.
+
+    Parameters
+    ----------
+    agent : WaldiezAgent
+        The agent.
+    all_chats : List[WaldiezChat]
+        All the chats in the flow.
+    chat_names : Dict[str, str]
+        The chat names.
+    agent_names : Dict[str, str]
+        The agent names.
+    serializer : Callable[..., str]
+        The serializer to use to escape quotes in a string.
+    string_escape : Callable[[str], str]
+        The string escape function.
+    is_async : bool
+        Whether the chat is asynchronous.
+
+    Returns
+    -------
+    str
+        The nested chat string.
+    """
+    if not agent.data.nested_chats:
+        return ""
+    content = ""
+    extra_contents = []
+    agent_name = agent_names[agent.id]
+    use_suffix = len(agent.data.nested_chats) > 1
+    for index, entry in enumerate(agent.data.nested_chats):
+        trigger_names = get_nested_chat_trigger_agent_names(
+            nested_chat=entry, agent_names=agent_names
+        )
+        chat_queue, extra_methods = get_nested_chat_queue(
+            nested_chat=entry,
+            agent=agent,
+            agent_names=agent_names,
+            chat_names=chat_names,
+            all_chats=all_chats,
+            serializer=serializer,
+            string_escape=string_escape,
+        )
+        if not chat_queue:  # pragma: no cover
+            continue
+        extra_contents.extend(extra_methods)
+        var_name = (
+            f"{agent_name}_chat_queue_{index}"
+            if use_suffix
+            else f"{agent_name}_chat_queue"
+        )
+        content += f"{var_name} = {chat_queue}" + "\n"
+        content += f"""
+{agent_name}.register_nested_chats(
+    trigger={trigger_names},
+    chat_queue={var_name},
+    use_async={is_async},
+    ignore_async_in_sync_chat=True,
+)
+"""
+    functions_string = "\n".join(sorted(extra_contents))
+    if functions_string:
+        functions_string = functions_string + "\n"
+    content = f"{functions_string}{content}"
+    return (
+        content.replace('"None"', "None")
+        .replace("'None'", "None")
+        .replace('"False"', "False")
+        .replace("'False'", "False")
+        .replace("'True'", "True")
+        .replace('"True"', "True")
+    )
 
 
 def get_nested_chat_trigger_agent_names(
@@ -32,7 +117,8 @@ def get_nested_chat_trigger_agent_names(
         The trigger agent names.
     """
     agents = [agent_names[agent_id] for agent_id in nested_chat.triggered_by]
-    trigger_string = f'{[", ".join(agents)]}'
+    agents_string = [", ".join(agents)]
+    trigger_string = f"{agents_string}"
     return trigger_string.replace("'", '"')
 
 
@@ -42,6 +128,8 @@ def get_nested_chat_message_string(
     agent: WaldiezAgent,
     agent_names: Dict[str, str],
     chat_names: Dict[str, str],
+    serializer: Callable[..., str],
+    string_escape: Callable[[str], str],
 ) -> Tuple[str, Optional[str]]:
     """Get the nested chat message string.
 
@@ -57,6 +145,10 @@ def get_nested_chat_message_string(
         A mapping of agent id to agent name.
     chat_names : Dict[str, str]
         A mapping of chat id to chat name.
+    serializer : Callable[..., str]
+        The function to serialize the chat arguments.
+    string_escape : Callable[[str], str]
+        The function to escape the string.
 
     Returns
     -------
@@ -71,16 +163,19 @@ def get_nested_chat_message_string(
     if sender_id != agent.id:
         sender_name = agent_names[sender_id]
     recipient_name = agent_names[recipient_id]
-    chat_dict = waldiez_chat.get_chat_args()
-    chat_dict = escape_summary_args_quotes(chat_dict)
+    chat_dict = waldiez_chat.get_chat_args(for_queue=True)
+    chat_dict = update_summary_chat_args(chat_dict, string_escape)
     chat_dict["recipient"] = recipient_name
     if sender_name:
         chat_dict["sender"] = sender_name
     message_value, message_source = get_chat_nested_string(
-        chat=waldiez_chat, is_reply=message.is_reply, chat_names=chat_names
+        chat=waldiez_chat,
+        is_reply=message.is_reply,
+        chat_names=chat_names,
+        string_escape=string_escape,
     )
     chat_dict["message"] = message_value
-    message_dict_str = get_object_string(chat_dict, tabs=1)
+    message_dict_str = serializer(chat_dict, tabs=1)
     if message_source:
         # it's not a string, its the name of the function
         message_dict_str = message_dict_str.replace(
@@ -103,6 +198,8 @@ def get_nested_chat_queue(
     agent_names: Dict[str, str],
     chat_names: Dict[str, str],
     all_chats: List[WaldiezChat],
+    serializer: Callable[..., str],
+    string_escape: Callable[[str], str],
 ) -> Tuple[str, List[str]]:
     """Get the nested chat queue.
 
@@ -118,11 +215,16 @@ def get_nested_chat_queue(
         A mapping of chat id to chat name.
     all_chats : List[WaldiezChat]
         All the chats in the flow.
+    serializer : Callable[..., str]
+        The serializer to use to escape quotes in a string.
+    string_escape : Callable[[str], str]
+        The function to escape the string.
 
     Returns
     -------
     Tuple[str, List[str]]
-        The nested chat queue and the methods to include.
+        The nested chat queue and the methods to include
+        (methods: message string and method name if the message is a method).
     """
     message_methods_to_include = []
     chat_messages_str = "[\n"
@@ -134,6 +236,8 @@ def get_nested_chat_queue(
             agent=agent,
             agent_names=agent_names,
             chat_names=chat_names,
+            serializer=serializer,
+            string_escape=string_escape,
         )
         if message_source:
             message_methods_to_include.append(message_source)
@@ -148,6 +252,7 @@ def get_chat_nested_string(
     chat: WaldiezChat,
     is_reply: bool,
     chat_names: Dict[str, str],
+    string_escape: Callable[[str], str],
 ) -> Tuple[str, Optional[str]]:
     """Get the nested chat message.
 
@@ -158,7 +263,9 @@ def get_chat_nested_string(
     is_reply : bool
         Whether to use the nested chat's reply message or not.
     chat_names : Dict[str, str]
-        A mapping of chat id to chat name
+        A mapping of chat id to chat name.
+    string_escape : Callable[[str], str]
+        The function to escape the string.
 
     Returns
     -------
@@ -175,82 +282,14 @@ def get_chat_nested_string(
     if not message or message.type == "none" or message.content is None:
         return "None", None
     if message.type == "string":
-        return get_escaped_string(message.content), None
+        return string_escape(message.content), None
     chat_name = chat_names[chat.id]
-    method_args = "recipient, messages, sender, config"
-    function_name = "nested_chat_reply" if is_reply else "nested_chat_message"
-    new_function_name = f"{function_name}_{chat_name}"
-    function_def = f"\ndef {new_function_name}({method_args}):"
-    attribute_name = "reply_content" if is_reply else "message_content"
-    function_content = getattr(chat.data.nested_chat, attribute_name)
-    return new_function_name, function_def + "\n" + function_content + "\n"
-
-
-def export_nested_chat(
-    agent: WaldiezAgent,
-    all_chats: List[WaldiezChat],
-    chat_names: Dict[str, str],
-    agent_names: Dict[str, str],
-) -> str:
-    """Get the nested chat string.
-
-    Parameters
-    ----------
-    agent : WaldiezAgent
-        The agent.
-    all_chats : List[WaldiezChat]
-        All the chats in the flow.
-    chat_names : Dict[str, str]
-        The chat names.
-    agent_names : Dict[str, str]
-        The agent names.
-
-    Returns
-    -------
-    str
-        The nested chat string.
-    """
-    if not agent.data.nested_chats:
-        return ""
-    content = ""
-    extra_contents = []
-    agent_name = agent_names[agent.id]
-    use_suffix = len(agent.data.nested_chats) > 1
-    for index, entry in enumerate(agent.data.nested_chats):
-        trigger_names = get_nested_chat_trigger_agent_names(
-            nested_chat=entry, agent_names=agent_names
-        )
-        chat_queue, extra_methods = get_nested_chat_queue(
-            nested_chat=entry,
-            agent=agent,
-            agent_names=agent_names,
-            chat_names=chat_names,
-            all_chats=all_chats,
+    if is_reply:
+        function_content, function_name = chat.get_nested_chat_message_function(
+            name_suffix=chat_name
         )
-        if not chat_queue:
-            continue
-        extra_contents.extend(extra_methods)
-        var_name = (
-            f"{agent_name}_chat_queue_{index}"
-            if use_suffix
-            else f"{agent_name}_chat_queue"
+    else:
+        function_content, function_name = chat.get_nested_chat_message_function(
+            name_suffix=chat_name
         )
-        content += f"\n{var_name} = {chat_queue}" + "\n"
-        content += f"""\n
-{agent_name}.register_nested_chats(
-    trigger={trigger_names},
-    chat_queue={var_name},
-)\n
-"""
-    functions_string = "\n".join(sorted(extra_contents))
-    if functions_string:
-        functions_string = functions_string + "\n"
-    content = f"{functions_string}{content}"
-    return (
-        content.replace('"None"', "None")
-        .replace("'None'", "None")
-        .replace('"False"', "False")
-        .replace("'False'", "False")
-        .replace("'True'", "True")
-        .replace('"True"', "True")
-    )
+    return function_name, function_content