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

waldiez/exporting/chats/chats_exporter.py

@@ -1,240 +0,0 @@
-# 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=1 if self.for_notebook else 2,
-                    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=1 if self.for_notebook else 2,
-                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=1 if self.for_notebook else 2,
-            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/swarm.py

@@ -1,210 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0.
-# Copyright (c) 2024 - 2025 Waldiez and contributors.
-# flake8: noqa: E501
-"""Swarm chat exporting functions."""
-
-from typing import Callable, Dict, List, Optional, Tuple
-
-from waldiez.models import WaldiezAgent, WaldiezChat
-
-
-# pylint: disable=too-many-locals,line-too-long
-def export_swarm_chat(
-    chat: WaldiezChat,
-    agent_names: Dict[str, str],
-    chat_names: Dict[str, str],
-    sender: WaldiezAgent,
-    recipient: WaldiezAgent,
-    get_swarm_members: Callable[
-        [WaldiezAgent], Tuple[List[WaldiezAgent], WaldiezAgent | None]
-    ],
-    serializer: Callable[..., str],
-    string_escape: Callable[[str], str],
-    tabs: int,
-    is_async: bool,
-) -> Tuple[str, str]:
-    """Get the swarm chat message string.
-
-    Parameters
-    ----------
-    chat : WaldiezChat
-        The chat.
-    agent_names : Dict[str, str]
-        A mapping of agent id to agent name.
-    chat_names : Dict[str, str]
-        A mapping of chat id to chat name.
-    sender : WaldiezAgent
-        The sender.
-    recipient : WaldiezAgent
-        The recipient.
-    get_swarm_members: Callable[[WaldiezAgent], Tuple[List[WaldiezAgent], WaldiezAgent | None]],
-        The function to get the swarm members.
-    serializer : Callable[..., str]
-        The function to generate the string representation of an object.
-    string_escape : Callable[[str], str]
-        The function to escape the string.
-    tabs : int
-        The number of tabs to use for indentation.
-    is_async : bool
-        Whether the chat is asynchronous.
-
-    Returns
-    -------
-    Tuple[str, str]
-        The `initiate_swarm_chat` message string and additional methods string.
-
-    Raises
-    ------
-    ValueError
-        If neither the sender nor the recipient is a swarm agent.
-    """
-    # either the sender or the recipient can be the swarm agent
-    # one of them MUST be the swarm agent.
-    if sender.agent_type != "swarm" and recipient.agent_type != "swarm":
-        raise ValueError("One of the agents must be a swarm agent")
-    initial_agent = sender if sender.agent_type == "swarm" else recipient
-    tab = "    " * tabs if tabs > 0 else ""
-    swarm_members = get_swarm_members(initial_agent)
-    swarm_agents_string, user_agent_string = get_swarm_agents_strings(
-        swarm_members=swarm_members,
-        sender=sender,
-        recipient=recipient,
-        agent_names=agent_names,
-        user_agent=None,
-    )
-    messages_string = get_swarm_messages_string(
-        chat=chat, string_escape=string_escape
-    )
-    context_vars_string = (
-        serializer(chat.context_variables, tabs=tabs + 1)
-        if chat.context_variables
-        else "{}"
-    )
-    chat_name = chat_names[chat.id]
-    after_work_string, additional_methods = get_swarm_after_work_string(
-        chat=chat,
-        agent_names=agent_names,
-        name_suffix=chat_name,
-    )
-    agent_name = agent_names[initial_agent.id]
-    results_is = f"{tab}results, _, __ = "
-    initiate = "initiate_swarm_chat"
-    if is_async:
-        results_is += "await "
-        initiate = "a_initiate_swarm_chat"
-    initiate_chat = "\n" + f"{results_is}{initiate}(" + "\n"
-    initiate_chat += f"{tab}    initial_agent={agent_name}," + "\n"
-    initiate_chat += f"{tab}    agents=[{swarm_agents_string}]," + "\n"
-    initiate_chat += f"{tab}    messages={messages_string}," + "\n"
-    initiate_chat += f"{tab}    context_variables={context_vars_string}," + "\n"
-    initiate_chat += f"{tab}    user_agent={user_agent_string}," + "\n"
-    initiate_chat += f"{tab}    after_work={after_work_string}," + "\n"
-    initiate_chat += f"{tab}    max_rounds={chat.max_rounds}," + "\n"
-    initiate_chat += f"{tab})" + "\n"
-    return initiate_chat, additional_methods
-
-
-def get_swarm_agents_strings(
-    swarm_members: Tuple[List[WaldiezAgent], WaldiezAgent | None],
-    sender: WaldiezAgent,
-    recipient: WaldiezAgent,
-    agent_names: Dict[str, str],
-    user_agent: Optional[WaldiezAgent],
-) -> Tuple[str, str]:
-    """Get the swarm agent strings to use in `initiate_swarm_chat`.
-
-    Parameters
-    ----------
-    swarm_members : Tuple[List[WaldiezAgent], WaldiezAgent | None]
-        The swarm agents and the user agent.
-    sender : WaldiezAgent
-        The chat initiator.
-    recipient : WaldiezAgent
-        The chat recipient.
-    agent_names : Dict[str, str]
-        A mapping of agent id to agent name.
-    user_agent : Optional[WaldiezAgent]
-        The user agent.
-
-    Returns
-    -------
-    Tuple[str, str]
-        The swarm agents string and the user agent string.
-    """
-    swarm_agents, user_member = swarm_members
-    if user_agent is None and user_member is not None:
-        user_agent = user_member
-    if user_agent is None:  # pragma: no cover
-        if sender.agent_type in ("user", "rag_user"):
-            user_agent = sender
-        elif recipient.agent_type in ("user", "rag_user"):
-            user_agent = recipient
-    agents_string = ", ".join(
-        [f"{agent_names[agent.id]}" for agent in swarm_agents]
-    )
-    user_agent_string = "None"
-    if user_agent:
-        user_agent_string = agent_names[user_agent.id]
-    return agents_string, user_agent_string
-
-
-def get_swarm_messages_string(
-    chat: WaldiezChat,
-    string_escape: Callable[[str], str],
-) -> str:
-    """Get the swarm chat messages string to use in `initiate_swarm_chat`.
-
-    Parameters
-    ----------
-    chat : WaldiezChat
-        The chat.
-    string_escape : Callable[[str], str]
-        The string escape function.
-
-    Returns
-    -------
-    str
-        The swarm chat messages string.
-    """
-    chat_message = chat.message
-    if chat.message.type == "string" and chat_message.content:
-        messages_string = '[{"role": "user", "content": '
-        escaped_message = string_escape(chat_message.content)
-        messages_string += f'"{escaped_message}"'
-        messages_string += "}]"
-    else:
-        messages_string = '"Start the chat."'
-    return messages_string
-
-
-def get_swarm_after_work_string(
-    chat: WaldiezChat,
-    agent_names: Dict[str, str],
-    name_suffix: str,
-) -> Tuple[str, str]:
-    """Get the swarm after work string.
-
-    Parameters
-    ----------
-    chat : WaldiezChat
-        The chat.
-    agent_names : Dict[str, str]
-        A mapping of agent id to agent name.
-    name_suffix : str
-        The suffix to use for the function name.
-
-    Returns
-    -------
-    Tuple[str, str]
-        The after work string and the additional methods string.
-    """
-    if not chat.after_work:
-        return "AfterWork(AfterWorkOption.TERMINATE)", ""
-    additional_methods = ""
-    after_work_string, function_content = chat.after_work.get_recipient(
-        agent_names=agent_names,
-        name_suffix=name_suffix,
-    )
-    if chat.after_work.recipient_type == "callable" and function_content:
-        additional_methods = "\n" + f"{function_content}" + "\n"
-    return after_work_string, additional_methods