Chats

def create_chats_exporter(
    all_agents: list[WaldiezAgent],
    agent_names: dict[str, str],
    all_chats: list[WaldiezChat],
    chat_names: dict[str, str],
    main_chats: list[WaldiezAgentConnection],
    root_group_manager: Optional[WaldiezGroupManager] = None,
    context: Optional[ExporterContext] = None,
    **kwargs: Any,
) -> ChatsExporter:
    """Create a chats exporter.

    Parameters
    ----------
    all_agents : list[WaldiezAgent]
        All agents involved in the chats.
    agent_names : dict[str, str]
        Mapping of agent IDs to their names.
    all_chats : list[WaldiezChat]
        All chats to be exported.
    chat_names : dict[str, str]
        Mapping of chat IDs to their names.
    main_chats : list[WaldiezAgentConnection]
        Main chats that are connections between agents.
    root_group_manager : Optional[WaldiezGroupManager], optional
        The root group manager for managing chat groups, if any.
    context : Optional[ExporterContext], optional
        Exporter context with dependencies, by default None
    **kwargs : Any
        Additional keyword arguments for the exporter.

    Returns
    -------
    ChatsExporter
        The created chats exporter.
    """
    if context is None:
        context = get_default_exporter_context()
    return ChatsExporter(
        all_agents=all_agents,
        agent_names=agent_names,
        all_chats=all_chats,
        chat_names=chat_names,
        main_chats=main_chats,
        root_group_manager=root_group_manager,
        context=context,
        **kwargs,
    )

def export_group_chats(
    agent_names: dict[str, str],
    manager: WaldiezGroupManager,
    initial_chat: str | None,
    tabs: int,
    is_async: bool,
) -> str:
    """Get the group chat string.

    Parameters
    ----------
    agent_names : dict[str, str]
        The agent names.
    manager : WaldiezGroupManager
        The group manager agent.
    initial_chat : str | None
        The initial chat to use if any.
    tabs : int
        The number of tabs for indentation.
    is_async : bool
        Whether the chat is asynchronous.

    Returns
    -------
    tuple[str, str]
        The group chat string and the import string.
    """
    tab = "    " * tabs
    initiate_group_chat = "initiate_group_chat"
    if is_async:
        initiate_group_chat = "a_initiate_group_chat"
    manager_name = agent_names[manager.id]
    pattern_name = f"{manager_name}_pattern"
    content = f"{tab}results, _, __ = {initiate_group_chat}(" + "\n"
    content += f"{tab}    pattern={pattern_name}," + "\n"
    if initial_chat:
        content += f"{tab}    messages={json.dumps(initial_chat)}," + "\n"
    else:
        content += f'{tab}    messages="",\n'
    content += f"{tab}    max_rounds={manager.data.max_round},\n"
    content += f"{tab})\n"
    return content

def export_nested_chat_registration(
    agent: WaldiezAgent,
    all_chats: list[WaldiezChat],
    chat_names: dict[str, str],
    agent_names: dict[str, str],
    serializer: Callable[..., 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.
    is_async : bool
        Whether the chat is asynchronous.

    Returns
    -------
    str
        The nested chat string.
    """
    if not agent.data.nested_chats:
        return ""
    content = ""
    extra_contents: list[str] = []
    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,
        )
        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}: list[dict[str, Any]] = {chat_queue}" + "\n"
        content += f"""
{agent_name}.register_nested_chats(  # pyright: ignore
    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 export_sequential_chat(
    main_chats: list[WaldiezAgentConnection],
    chat_names: dict[str, str],
    agent_names: dict[str, str],
    serializer: Callable[..., str],
    tabs: int,
    is_async: bool,
    skip_cache: bool,
) -> tuple[str, str]:
    r"""Get the chats content, when there are more than one chats in the flow.

    Parameters
    ----------
    main_chats : list[WaldiezAgentConnection]
        The main chats.
    chat_names : dict[str, str]
        A mapping of chat id to chat name.
    agent_names : dict[str, str]
        A mapping of agent id to agent name.
    serializer : Callable[..., str]
        The serializer function to escape quotes in a string.
    tabs : int
        The number of tabs to use for indentation.
    is_async : bool
        Whether the chat is asynchronous.
    skip_cache : bool
        Whether to skip the cache argument.

    Returns
    -------
    tuple[str, str]
        The main chats content and additional methods string if any.

    Example
    -------
    ```python
    >>> from waldiez.models import (
    ...     WaldiezAgent,
    ...     WaldiezChat,
    ...     WaldiezChatData,
    ...     WaldiezChatMessage,
    ... )
    >>> chat1 = WaldiezChat(
    ...     id="wc-1",
    ...     name="chat1",
    ...     description="A chat between two agents.",
    ...     tags=["chat", "chat1"],
    ...     requirements=[],
    ...     data=WaldiezChatData(
    ...         sender="wa-1",
    ...         recipient="wa-2",
    ...         position=0,
    ...         message=WaldiezChatMessage(
    ...             type="string",
    ...             content="Hello, how are you?",
    ...         ),
    ...     ),
    ... )
    >>> chat2 = WaldiezChat(
    ...     id="wc-2",
    ...     name="chat2",
    ...     description="A chat between two agents.",
    ...     tags=["chat", "chat2"],
    ...     requirements=[],
    ...     data=WaldiezChatData(
    ...         sender="wa-2",
    ...         recipient="wa-1",
    ...         position=1,
    ...         message=WaldiezChatMessage(
    ...             type="string",
    ...             content="I am good, thank you. How about you?",
    ...         ),
    ...     ),
    ... )
    >>> agent_names = {"wa-1": "agent1", "wa-2": "agent2"}
    >>> chat_names = {"wc-1": "chat1", "wc-2": "chat2"}
    >>> serializer = lambda x: x.replace('"', "\"").replace("\n", "\\n")
    >>>  export_sequential_chat(
    ...     main_chats=[(chat1, agent1, agent2), (chat2, agent2, agent1)],
    ...     chat_names=chat_names,
    ...     agent_names=agent_names,
    ...     serializer=serializer,
    ...     tabs=0,
    ...     is_async=False,
    ... )
    results = initiate_chats([
        {
            "sender": agent1,
            "recipient": agent2,
            "message": "Hello, how are you?",
        },
        {
            "sender": agent2,
            "recipient": agent1,
            "message": "I am good, thank you. How about you?",
        },
    ])
    ```
    """
    tab = "    " * tabs if tabs > 0 else ""
    content = "\n"
    additional_methods_string = ""
    content += _get_initiate_chats_line(tab, is_async)
    for connection in main_chats:
        chat_string, additional_methods = _get_chat_dict_string(
            chat_names=chat_names,
            connection=connection,
            agent_names=agent_names,
            serializer=serializer,
            tabs=tabs + 1,
            skip_cache=skip_cache,
        )
        additional_methods_string += additional_methods
        content += "\n" + f"{tab}    {chat_string}"
    content += "\n" + "    " * tabs + "])\n"
    return content, additional_methods_string

def export_single_chat(
    sender: WaldiezAgent,
    recipient: WaldiezAgent,
    chat: WaldiezChat,
    agent_names: dict[str, str],
    chat_names: dict[str, str],
    serializer: Callable[[str], str],
    tabs: int,
    is_async: bool,
    skip_cache: bool,
    tab_length: int = 4,
) -> tuple[str, str]:
    """Get the chat string when there is only one chat in the flow.

    Parameters
    ----------
    sender : WaldiezAgent
        The sender.
    recipient : WaldiezAgent
        The recipient.
    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.
    serializer : Callable[[str], str]
        The serializer function to escape quotes in a string.
    tabs : int
        The number of tabs to use for indentation.
    is_async : bool
        Whether the chat is asynchronous.
    skip_cache : bool
        Whether to skip the cache argument.
    tab_length : int, optional
        The length of the tab string, by default 4.

    Returns
    -------
    tuple[str, str]
        The chat string and additional methods string if any

    Example
    -------
    ```python
    >>> from waldiez.models import WaldiezAgent, WaldiezChat, WaldiezChatData, WaldiezChatMessage
    >>> chat = WaldiezChat(
    ...     id="wc-1",
    ...     name="chat1",
    ...     description="A chat between two agents.",
    ...     tags=["chat", "chat1"],
    ...     requirements=[],
    ...     data=WaldiezChatData(
    ...         sender="wa-1",
    ...         recipient="wa-2",
    ...         message=WaldiezChatMessage(
    ...             type="string",
    ...             content="Hello, how are you?",
    ...         ),
    ...     ),
    ... )
    >>> agent_names = {"wa-1": "agent1", "wa-2": "agent2"}
    >>> chat_names = {"wc-1": "chat1"}
    >>> export_single_chat_string(
    ...     sender=agent1,
    ...     recipient=agent2,
    ...     chat=chat,
    ...     agent_names=agent_names,
    ...     chat_names=chat_names,
    ...     tabs=0,
    ... )
    agent1.initiate_chat(
        agent2,
        message="Hello, how are you?",
        cache=cache,
    )
    ```
    """
    tab = " " * tab_length * tabs if tabs > 0 else ""
    chat_args = chat.get_chat_args(for_queue=False, sender=sender)
    # chat_args = update_summary_chat_args(chat_args, skip_summary=True)
    if not chat_args:
        return get_empty_simple_chat_string(
            chat=chat,
            sender=sender,
            recipient=recipient,
            agent_names=agent_names,
            tab=tab,
            is_async=is_async,
            skip_cache=skip_cache,
        )
    return get_simple_chat_string(
        chat=chat,
        chat_args=chat_args,
        sender=sender,
        recipient=recipient,
        agent_names=agent_names,
        chat_names=chat_names,
        serializer=serializer,
        tabs=tabs,
        is_async=is_async,
        skip_cache=skip_cache,
    )

def get_nested_chat_queue(
    nested_chat: WaldiezAgentNestedChat,
    agent: WaldiezAgent,
    agent_names: dict[str, str],
    chat_names: dict[str, str],
    all_chats: list[WaldiezChat],
    serializer: Callable[..., str],
) -> tuple[str, list[str]]:
    """Get the nested chat queue.

    Parameters
    ----------
    nested_chat : WaldiezAgentNestedChat
        The nested chat.
    agent : WaldiezAgent
        The agent.
    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.
    all_chats : list[WaldiezChat]
        All the chats in the flow.
    serializer : Callable[..., str]
        The serializer to use to escape quotes in a string.

    Returns
    -------
    tuple[str, list[str]]
        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: list[str] = []
    chat_messages_str = "[\n"
    for message in nested_chat.messages:
        waldiez_chat = next(chat for chat in all_chats if chat.id == message.id)
        message_str, message_source = get_nested_chat_message_string(
            waldiez_chat=waldiez_chat,
            message=message,
            agent=agent,
            agent_names=agent_names,
            chat_names=chat_names,
            serializer=serializer,
        )
        if message_source:
            message_methods_to_include.append(message_source)
        chat_messages_str += f"    {message_str}," + "\n"
    chat_messages_str += "]"
    if chat_messages_str == "[\n]":
        return "", message_methods_to_include
    return chat_messages_str, message_methods_to_include