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

waldiez/exporting/chats/utils/nested.py

@@ -3,7 +3,7 @@
 # pylint: disable=too-many-locals
 """Nested chats exporting."""
 
-from typing import Callable, Dict, List, Optional, Tuple
+from typing import Any, Callable, Optional
 
 from waldiez.models import (
     WaldiezAgent,
@@ -12,16 +12,15 @@ from waldiez.models import (
     WaldiezChat,
 )
 
-from .common import update_summary_chat_args
+# 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],
+    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.
@@ -30,16 +29,14 @@ def export_nested_chat_registration(
     ----------
     agent : WaldiezAgent
         The agent.
-    all_chats : List[WaldiezChat]
+    all_chats : list[WaldiezChat]
         All the chats in the flow.
-    chat_names : Dict[str, str]
+    chat_names : dict[str, str]
         The chat names.
-    agent_names : Dict[str, str]
+    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.
 
@@ -51,7 +48,7 @@ def export_nested_chat_registration(
     if not agent.data.nested_chats:
         return ""
     content = ""
-    extra_contents = []
+    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):
@@ -65,7 +62,6 @@ def export_nested_chat_registration(
             chat_names=chat_names,
             all_chats=all_chats,
             serializer=serializer,
-            string_escape=string_escape,
         )
         if not chat_queue:  # pragma: no cover
             continue
@@ -75,9 +71,9 @@ def export_nested_chat_registration(
             if use_suffix
             else f"{agent_name}_chat_queue"
         )
-        content += f"{var_name} = {chat_queue}" + "\n"
+        content += f"{var_name}: list[dict[str, Any]] = {chat_queue}" + "\n"
         content += f"""
-{agent_name}.register_nested_chats(
+{agent_name}.register_nested_chats(  # pyright: ignore
     trigger={trigger_names},
     chat_queue={var_name},
     use_async={is_async},
@@ -100,7 +96,7 @@ def export_nested_chat_registration(
 
 def get_nested_chat_trigger_agent_names(
     nested_chat: WaldiezAgentNestedChat,
-    agent_names: Dict[str, str],
+    agent_names: dict[str, str],
 ) -> str:
     """Get the trigger agent names for the nested chat.
 
@@ -108,7 +104,7 @@ def get_nested_chat_trigger_agent_names(
     ----------
     nested_chat : WaldiezAgentNestedChat
         The nested chat.
-    agent_names : Dict[str, str]
+    agent_names : dict[str, str]
         A mapping of agent id to agent name.
 
     Returns
@@ -126,11 +122,10 @@ def get_nested_chat_message_string(
     waldiez_chat: WaldiezChat,
     message: WaldiezAgentNestedChatMessage,
     agent: WaldiezAgent,
-    agent_names: Dict[str, str],
-    chat_names: Dict[str, str],
+    agent_names: dict[str, str],
+    chat_names: dict[str, str],
     serializer: Callable[..., str],
-    string_escape: Callable[[str], str],
-) -> Tuple[str, Optional[str]]:
+) -> tuple[str, Optional[str]]:
     """Get the nested chat message string.
 
     Parameters
@@ -141,18 +136,16 @@ def get_nested_chat_message_string(
         The message.
     agent : WaldiezAgent
         The agent.
-    agent_names : Dict[str, str]
+    agent_names : dict[str, str]
         A mapping of agent id to agent name.
-    chat_names : Dict[str, str]
+    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
     -------
-    Tuple[str, Optional[str]]
+    tuple[str, Optional[str]]
         The message string and the method name if the message is a method.
     """
     sender_name: Optional[str] = None
@@ -163,8 +156,8 @@ 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(for_queue=True)
-    chat_dict = update_summary_chat_args(chat_dict, string_escape)
+    chat_dict: dict[str, Any] = waldiez_chat.get_chat_args(for_queue=True)
+    # chat_dict = update_summary_chat_args(chat_dict)
     chat_dict["recipient"] = recipient_name
     if sender_name:
         chat_dict["sender"] = sender_name
@@ -172,7 +165,6 @@ def get_nested_chat_message_string(
         chat=waldiez_chat,
         is_reply=message.is_reply,
         chat_names=chat_names,
-        string_escape=string_escape,
     )
     chat_dict["message"] = message_value
     message_dict_str = serializer(chat_dict, tabs=1)
@@ -180,7 +172,7 @@ def get_nested_chat_message_string(
         # it's not a string, its the name of the function
         message_dict_str = message_dict_str.replace(
             f': "{message_value}"', f": {message_value}"
-        )
+        ).replace(f'"{message_value}"', f"{message_value}")
     if sender_name:
         message_dict_str = message_dict_str.replace(
             f': "{sender_name}"', f": {sender_name}"
@@ -195,12 +187,11 @@ def get_nested_chat_message_string(
 def get_nested_chat_queue(
     nested_chat: WaldiezAgentNestedChat,
     agent: WaldiezAgent,
-    agent_names: Dict[str, str],
-    chat_names: Dict[str, str],
-    all_chats: List[WaldiezChat],
+    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]]:
+) -> tuple[str, list[str]]:
     """Get the nested chat queue.
 
     Parameters
@@ -209,24 +200,22 @@ def get_nested_chat_queue(
         The nested chat.
     agent : WaldiezAgent
         The agent.
-    agent_names : Dict[str, str]
+    agent_names : dict[str, str]
         A mapping of agent id to agent name.
-    chat_names : Dict[str, str]
+    chat_names : dict[str, str]
         A mapping of chat id to chat name.
-    all_chats : List[WaldiezChat]
+    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]]
+    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 = []
+    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)
@@ -237,7 +226,6 @@ def get_nested_chat_queue(
             agent_names=agent_names,
             chat_names=chat_names,
             serializer=serializer,
-            string_escape=string_escape,
         )
         if message_source:
             message_methods_to_include.append(message_source)
@@ -251,9 +239,8 @@ def get_nested_chat_queue(
 def get_chat_nested_string(
     chat: WaldiezChat,
     is_reply: bool,
-    chat_names: Dict[str, str],
-    string_escape: Callable[[str], str],
-) -> Tuple[str, Optional[str]]:
+    chat_names: dict[str, str],
+) -> tuple[str, Optional[str]]:
     """Get the nested chat message.
 
     Parameters
@@ -262,14 +249,12 @@ def get_chat_nested_string(
         The chat.
     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.
-    string_escape : Callable[[str], str]
-        The function to escape the string.
+    chat_names : dict[str, str]
+        A mapping of chat id to chat name..
 
     Returns
     -------
-    Tuple[str, Optional[str]]
+    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.
@@ -282,7 +267,7 @@ 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 string_escape(message.content), None
+        return message.content, None
     chat_name = chat_names[chat.id]
     if is_reply:
         function_content, function_name = chat.get_nested_chat_message_function(

waldiez/exporting/chats/utils/sequential.py

@@ -2,49 +2,48 @@
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Utility functions for exporting sequential chats."""
 
-from typing import Callable, Dict, List, Tuple
+from typing import Callable
 
 from waldiez.models import (
-    WaldiezAgent,
-    WaldiezChat,
+    WaldiezAgentConnection,
     WaldiezChatMessage,
-    WaldiezRagUser,
+    WaldiezRagUserProxy,
 )
 
-from .common import get_chat_message_string, update_summary_chat_args
+from .common import get_chat_message_string
 
 
 def export_sequential_chat(
-    main_chats: List[Tuple[WaldiezChat, WaldiezAgent, WaldiezAgent]],
-    chat_names: Dict[str, str],
-    agent_names: Dict[str, str],
+    main_chats: list[WaldiezAgentConnection],
+    chat_names: dict[str, str],
+    agent_names: dict[str, str],
     serializer: Callable[..., str],
-    string_escape: Callable[[str], str],
     tabs: int,
     is_async: bool,
-) -> Tuple[str, str]:
+    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[Tuple[WaldiezChat, WaldiezAgent, WaldiezAgent]]
+    main_chats : list[WaldiezAgentConnection]
         The main chats.
-    chat_names : Dict[str, str]
+    chat_names : dict[str, str]
         A mapping of chat id to chat name.
-    agent_names : Dict[str, str]
+    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.
-    string_escape : Callable[[str], str]
-        The string escape function.
     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]
+    tuple[str, str]
         The main chats content and additional methods string if any.
 
     Example
@@ -117,16 +116,14 @@ def export_sequential_chat(
     content = "\n"
     additional_methods_string = ""
     content += _get_initiate_chats_line(tab, is_async)
-    for chat, sender, recipient in main_chats:
+    for connection in main_chats:
         chat_string, additional_methods = _get_chat_dict_string(
-            chat=chat,
             chat_names=chat_names,
-            sender=sender,
-            recipient=recipient,
+            connection=connection,
             agent_names=agent_names,
             serializer=serializer,
-            string_escape=string_escape,
             tabs=tabs + 1,
+            skip_cache=skip_cache,
         )
         additional_methods_string += additional_methods
         content += "\n" + f"{tab}    {chat_string}"
@@ -134,17 +131,14 @@ def export_sequential_chat(
     return content, additional_methods_string
 
 
-# pylint: disable=too-many-locals
 def _get_chat_dict_string(
-    chat: WaldiezChat,
-    sender: WaldiezAgent,
-    recipient: WaldiezAgent,
-    chat_names: Dict[str, str],
-    agent_names: Dict[str, str],
+    connection: WaldiezAgentConnection,
+    chat_names: dict[str, str],
+    agent_names: dict[str, str],
     serializer: Callable[..., str],
-    string_escape: Callable[[str], str],
     tabs: int,
-) -> Tuple[str, str]:
+    skip_cache: bool,
+) -> tuple[str, str]:
     """Get a chat dictionary string.
 
     If the chat message is a separate method and not a string or a lambda,
@@ -153,62 +147,50 @@ def _get_chat_dict_string(
 
     Parameters
     ----------
-    chat : WaldiezChat
-        The chat.
-    sender : WaldiezAgent
-        The sender.
-    recipient : WaldiezAgent
-        The recipient.
-    chat_names : Dict[str, str]
+    connection : WaldiezAgentConnection
+        The connection object containing the chat and agents.
+    chat_names : dict[str, str]
         A mapping of chat id to chat name.
-    agent_names : Dict[str, str]
+    agent_names : dict[str, str]
         A mapping of agent id to agent name.
     serializer : Callable[[str], str]
         The function to serialize the dictionaries or lists.
-    string_escape : Callable[[str], str]
-        The function to escape the string.
     tabs : int
         The number of tabs to use for indentation.
+    skip_cache : bool
+        Whether to skip the cache argument.
 
     Returns
     -------
-    Tuple[str, str]
+    tuple[str, str]
         The chat dictionary string and additional methods string if any.
     """
     tab = "    " * tabs
-    chat_args = chat.get_chat_args(for_queue=True, sender=sender)
-    chat_args = update_summary_chat_args(chat_args, string_escape)
-    chat_string = "{"
-    chat_string += "\n" + f'{tab}    "sender": {agent_names[sender.id]},'
-    chat_string += "\n" + f'{tab}    "recipient": {agent_names[recipient.id]},'
-    chat_string += "\n" + f'{tab}    "cache": cache,'
-    additional_methods_string = ""
-    for key, value in chat_args.items():
-        if isinstance(value, str):
-            chat_string += "\n" + f'{tab}    "{key}": "{value}",'
-        elif isinstance(value, dict):
-            chat_string += (
-                "\n" + f'{tab}    "{key}": {serializer(value, tabs=tabs + 1)},'
-            )
-        else:
-            chat_string += "\n" + f'{tab}    "{key}": {value},'
+    chat = connection["chat"]
+    sender = connection["source"]
+    chat_string = _get_chat_strig_start(
+        connection=connection,
+        agent_names=agent_names,
+        serializer=serializer,
+        tabs=tabs,
+        skip_cache=skip_cache,
+    )
     if (
-        sender.agent_type == "rag_user"
-        and isinstance(sender, WaldiezRagUser)
+        sender.is_rag_user
+        and isinstance(sender, WaldiezRagUserProxy)
         and chat.message.type == "rag_message_generator"
     ):
         message = f"{agent_names[sender.id]}.message_generator"
         chat_string += "\n" + f'{tab}    "message": {message},'
         chat_string += "\n" + tab + "},"
-        return chat_string, additional_methods_string
+        return chat_string, ""
+    additional_methods_string = ""
     message, method_content = get_chat_message_string(
         sender=sender,
         chat=chat,
         chat_names=chat_names,
-        string_escape=string_escape,
     )
     if message and isinstance(chat.data.message, WaldiezChatMessage):
-        message = string_escape(message)
         if chat.data.message.type == "method":
             if method_content:
                 additional_methods_string += "\n" + method_content
@@ -219,6 +201,35 @@ def _get_chat_dict_string(
     return chat_string, additional_methods_string
 
 
+def _get_chat_strig_start(
+    connection: WaldiezAgentConnection,
+    agent_names: dict[str, str],
+    serializer: Callable[..., str],
+    tabs: int,
+    skip_cache: bool,
+) -> str:
+    tab = "    " * tabs
+    chat = connection["chat"]
+    sender = connection["source"]
+    recipient = connection["target"]
+    chat_args = chat.get_chat_args(for_queue=True, sender=sender)
+    # chat_args = update_summary_chat_args(chat_args)
+    chat_string = "{"
+    chat_string += "\n" + f'{tab}    "sender": {agent_names[sender.id]},'
+    chat_string += "\n" + f'{tab}    "recipient": {agent_names[recipient.id]},'
+    if skip_cache is False:
+        chat_string += "\n" + f'{tab}    "cache": cache,'
+    # additional_methods_string = ""
+    for key, value in chat_args.items():
+        if isinstance(value, (dict, str)):
+            chat_string += (
+                "\n" + f'{tab}    "{key}": {serializer(value, tabs=tabs + 1)},'
+            )
+        else:
+            chat_string += "\n" + f'{tab}    "{key}": {value},'
+    return chat_string
+
+
 def _get_initiate_chats_line(
     tab: str,
     is_async: bool,