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

waldiez/exporting/chats/utils/swarm.py

@@ -0,0 +1,207 @@
+# 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 "AFTER_WORK(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

waldiez/exporting/flow/__init__.py

@@ -1,5 +1,7 @@
-"""Export the entire flow to string."""
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Flow exporter."""
 
-from .flow import export_flow
+from .flow_exporter import FlowExporter
 
-__all__ = ["export_flow"]
+__all__ = ["FlowExporter"]