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

waldiez/exporting/agent/utils/swarm_agent.py

@@ -0,0 +1,463 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+# pylint: disable=unused-argument
+"""Get the extras for a swarm agent."""
+
+from typing import Callable, Dict, List, Tuple
+
+from waldiez.exporting.chats.utils.nested import get_nested_chat_queue
+from waldiez.models import (
+    WaldiezAgent,
+    WaldiezChat,
+    WaldiezSwarmAfterWork,
+    WaldiezSwarmAgent,
+    WaldiezSwarmOnCondition,
+    WaldiezSwarmUpdateSystemMessage,
+)
+
+# SwarmAgent is a subclass of ConversableAgent.
+
+# Additional args:
+# functions (List[Callable]):
+#   -A list of functions to register with the agent.
+# update_agent_state_before_reply (List[Callable]):
+# - A list of functions, including UPDATE_SYSTEM_MESSAGEs,
+#   called to update the agent before it replies.
+
+# Additional methods:
+# register_hand_off(hand_offs: List[AfterWork|OnCondition]):
+
+
+def get_swarm_extras(
+    agent: WaldiezAgent,
+    agent_names: Dict[str, str],
+    skill_names: Dict[str, str],
+    chats: Tuple[List[WaldiezChat], Dict[str, str]],
+    is_async: bool,
+    serializer: Callable[..., str],
+    string_escape: Callable[[str], str],
+) -> Tuple[str, str, str]:
+    """Get the extras of a swarm agent.
+
+    Parameters
+    ----------
+    agent : WaldiezAgent
+        The agent to get the extras for.
+    agent_names : Dict[str, str]
+        A mapping of agent IDs to agent names.
+    skill_names : Dict[str, str]
+        A mapping of skill IDs to skill names.
+    chats : Tuple[List[WaldiezChat], Dict[str, str]]
+        The list of all chats and the mapping of chat IDs to chat names.
+    is_async : bool
+        Whether the chat is asynchronous.
+    serializer : Callable[..., str]
+        The serializer to get the string representation of an object.
+    string_escape : Callable[[str], str]
+        The function to escape the string quotes and newlines.
+    Returns
+    -------
+    Tuple[str, str, str]
+        The extras of the swarm agent:
+        the content before the agent,
+        the extra argument(s) for the agent initialization,
+        and the content after the agent.
+    """
+    args_string = ""
+    before_agent = ""
+    after_agent = ""
+    if agent.agent_type != "swarm" or not isinstance(agent, WaldiezSwarmAgent):
+        return args_string, before_agent, after_agent
+    args_string = get_function_arg(agent, skill_names)
+    before_reply = get_update_agent_state_before_reply_arg(
+        agent=agent,
+        agent_names=agent_names,
+        skill_names=skill_names,
+        string_escape=string_escape,
+    )
+    args_string += before_reply[0]
+    before_agent += before_reply[1]
+    before_registration, after_agent = get_agent_handoff_registrations(
+        agent=agent,
+        agent_names=agent_names,
+        all_chats=chats[0],
+        chat_names=chats[1],
+        is_async=is_async,
+        serializer=serializer,
+        string_escape=string_escape,
+    )
+    before_agent += before_registration
+    return before_agent, args_string, after_agent
+
+
+def get_function_arg(
+    agent: WaldiezSwarmAgent,
+    skill_names: Dict[str, str],
+) -> str:
+    """Get the function argument of a swarm agent.
+
+    Parameters
+    ----------
+    agent : WaldiezSwarmAgent
+        The swarm agent to get the function argument for.
+    skill_names : Dict[str, str]
+        A mapping of skill IDs to skill names.
+
+    Returns
+    -------
+    str
+        The function argument of the swarm agent.
+    """
+    tab = "    "
+    arg_string = f"{tab}functions=["
+    added_skills = False
+    for function in agent.data.functions:
+        skill_name = skill_names.get(function, "")
+        if skill_name:
+            arg_string += "\n" + f"{tab}{tab}{skill_name},"
+            added_skills = True
+    if added_skills:
+        arg_string += "\n" + tab
+    arg_string += "],\n"
+    return arg_string
+
+
+def get_update_agent_state_before_reply_arg(
+    agent: WaldiezSwarmAgent,
+    agent_names: Dict[str, str],
+    skill_names: Dict[str, str],
+    string_escape: Callable[[str], str],
+) -> Tuple[str, str]:
+    """Get the update_agent_state_before_reply argument of a swarm agent.
+
+    Parameters
+    ----------
+    agent : WaldiezSwarmAgent
+        The swarm agent to get the argument for.
+    agent_names : Dict[str, str]
+        A mapping of agent IDs to agent names.
+    skill_names : Dict[str, str]
+        A mapping of skill IDs to skill names.
+    string_escape : Callable[[str], str]
+        The function to escape the string quotes and newlines.
+
+    Returns
+    -------
+    Tuple[str, str]
+        The update_agent_state_before_reply argument of the swarm agent
+        and the content before the agent if any.
+    """
+    #     update_function_type : Literal["string", "callable"]
+    #     The type of the update function. Can be either a string or a callable.
+    # update_function : str
+    #     "The string template or function definition to update "
+    #          "the agent's system message. Can be a string or a Callable. "
+    #           "If the function_type is 'string' it will be used as a "
+    #            "template and substitute the context variables. "
+    #         ag2 checks for: vars = re.findall(r"\{(\w+)\}", function)
+    #     "If function_type is 'callable', it should have signature:
+    #  "def custom_update_system_message("
+    #     " agent: ConversableAgent, "
+    #     " messages: List[Dict[str, Any]]
+    #   ) -> str"
+    tab = "    "
+    before_agent = ""
+    arg_string = f"{tab}update_agent_state_before_reply=["
+    added_functions = False
+    # pylint: disable=line-too-long
+    for function in agent.data.update_agent_state_before_reply:
+        if isinstance(function, WaldiezSwarmUpdateSystemMessage):
+            added_functions = True
+            if function.update_function_type == "callable":
+                function_content, function_name = function.get_update_function(
+                    name_suffix=agent_names[agent.id],
+                )
+                arg_string += (
+                    "\n" + f"{tab}{tab}UPDATE_SYSTEM_MESSAGE({function_name}),"
+                )
+                before_agent += "\n" + function_content + "\n"
+            else:
+                escaped_function = string_escape(function.update_function)
+                arg_string += (
+                    "\n"
+                    + f'{tab}{tab}UPDATE_SYSTEM_MESSAGE("{escaped_function}"),'
+                )
+        else:
+            skill_name = skill_names.get(function, "")
+            if skill_name:
+                added_functions = True
+                arg_string += "\n" + f"{tab}{tab}{skill_name},"
+    if added_functions:
+        arg_string = arg_string + "\n" + tab
+    arg_string += "],\n"
+    return arg_string, before_agent
+
+
+def get_agent_handoff_registrations(
+    agent: WaldiezSwarmAgent,
+    agent_names: Dict[str, str],
+    all_chats: List[WaldiezChat],
+    chat_names: Dict[str, str],
+    is_async: bool,
+    serializer: Callable[..., str],
+    string_escape: Callable[[str], str],
+) -> Tuple[str, str]:
+    """Get the agent handoff registrations of a swarm agent.
+
+    Parameters
+    ----------
+    agent : WaldiezSwarmAgent
+        The swarm agent to get the agent handoff registrations for.
+    agent_names : Dict[str, str]
+        A mapping of agent IDs to agent names.
+    all_chats : List[WaldiezChat]
+        The list of all chats.
+    chat_names : Dict[str, str]
+        A mapping of chat IDs to chat names.
+    is_async : bool
+        Whether the chat is asynchronous.
+    serializer : Callable[..., str]
+        The serializer to get the string representation of an object.
+    string_escape : Callable[[str], str]
+        The function to escape the string quotes and newlines.
+
+    Returns
+    -------
+    Tuple[str, str]
+        the contents before and after the agent.
+    """
+    agent_name = agent_names[agent.id]
+    registrations = []
+    before_agent = ""
+    after_agent = ""
+    if not agent.handoffs:
+        return before_agent, after_agent
+    tab = "    "
+    after_agent = f"{agent_name}.register_hand_off(" + "\n" + f"{tab}[" + "\n"
+    for hand_off in agent.handoffs:
+        if isinstance(hand_off, WaldiezSwarmOnCondition):
+            registration, before_handoff = get_agent_on_condition_handoff(
+                agent=agent,
+                hand_off=hand_off,
+                agent_names=agent_names,
+                all_chats=all_chats,
+                chat_names=chat_names,
+                is_async=is_async,
+                serializer=serializer,
+                string_escape=string_escape,
+            )
+            if registration:
+                registrations.append(registration)
+            before_agent += before_handoff
+        elif isinstance(hand_off, WaldiezSwarmAfterWork):
+            registration, before_handoff = get_agent_after_work_handoff(
+                hand_off=hand_off,
+                agent_names=agent_names,
+                agent_name=agent_name,
+            )
+            registrations.append(registration)
+            before_agent += before_handoff
+    after_agent += "\n".join(registrations) + "\n" + f"{tab}]" + "\n" + ")"
+    return before_agent, after_agent
+
+
+def get_agent_after_work_handoff(
+    hand_off: WaldiezSwarmAfterWork,
+    agent_names: Dict[str, str],
+    agent_name: str,
+) -> Tuple[str, str]:
+    """Get the agent's after work hand off registration.
+
+    Parameters
+    ----------
+    hand_off : WaldiezSwarmAfterWork
+        The hand off to get the registration for.
+    agent_names : Dict[str, str]
+        A mapping of agent IDs to agent names.
+    agent_name : str
+        The name of the agent to register the hand off.
+
+    Returns
+    -------
+    Tuple[str, str]
+        The registration and the content before the agent.
+    """
+    before_agent = ""
+    tab = "    "
+    recipient_type = hand_off.recipient_type
+    recipient, function_content = hand_off.get_recipient(
+        agent_names=agent_names,
+        name_suffix=agent_name,
+    )
+    registration = f"{tab}{tab}{recipient},"
+    if recipient_type == "callable" and function_content:
+        before_agent += "\n" + function_content + "\n"
+    return registration, before_agent
+
+
+def get_agent_on_condition_handoff(
+    agent: WaldiezSwarmAgent,
+    hand_off: WaldiezSwarmOnCondition,
+    agent_names: Dict[str, str],
+    all_chats: List[WaldiezChat],
+    chat_names: Dict[str, str],
+    is_async: bool,
+    serializer: Callable[..., str],
+    string_escape: Callable[[str], str],
+) -> Tuple[str, str]:
+    """Get the agent's on condition hand off registration.
+
+    Parameters
+    ----------
+    agent : WaldiezSwarmAgent
+        The agent to get the registration for.
+    hand_off : WaldiezSwarmAfterWork
+        The hand off to get the registration for.
+    agent_names : Dict[str, str]
+        A mapping of agent IDs to agent names.
+    all_chats : List[WaldiezChat]
+        The list of all chats.
+    chat_names : Dict[str, str]
+        A mapping of chat IDs to chat names.
+    is_async : bool
+        Whether the chat is asynchronous.
+    serializer : Callable[..., str]
+        The serializer to get the string representation of an object.
+    string_escape : Callable[[str], str]
+        The function to escape the string quotes and newlines.
+
+    Returns
+    -------
+    Tuple[str, str]
+        The registration and the content before the agent.
+    """
+    before_agent = ""
+    registration = ""
+    available, available_function = hand_off.get_available(
+        name_suffix=agent_names[agent.id],
+    )
+    if available and not available_function:
+        available = f'"{string_escape(available)}"'
+    if hand_off.target_type == "agent":
+        recipient = agent_names[hand_off.target.id]
+        condition = (
+            string_escape(hand_off.condition) or f"Transfer to {recipient}"
+        )
+        results = _get_agent_on_condition_handoff_to_agent(
+            recipient=recipient,
+            available=available,
+            condition=condition,
+            available_function=available_function,
+        )
+        before_agent += results[0]
+        registration = results[1]
+    # else: # target_type == "nested_chat"
+    if hand_off.target_type == "nested_chat":
+        condition = _get_condition_string(
+            condition=hand_off.condition,
+            chat_id=hand_off.target.id,
+            all_chats=all_chats,
+            agent_names=agent_names,
+        )
+        results = _get_agent_on_condition_handoff_to_nested_chat(
+            agent=agent,
+            agent_names=agent_names,
+            condition=hand_off.condition,
+            available=available,
+            available_function=available_function,
+            all_chats=all_chats,
+            chat_names=chat_names,
+            is_async=is_async,
+            serializer=serializer,
+            string_escape=string_escape,
+        )
+        before_agent += results[0]
+        registration = results[1]
+    return registration, before_agent
+
+
+def _get_agent_on_condition_handoff_to_agent(
+    recipient: str,
+    available: str,
+    condition: str,
+    available_function: str,
+) -> Tuple[str, str]:
+    before_agent = ""
+    tab = "    "
+    on_condition = (
+        f"{tab}{tab}ON_CONDITION(" + "\n"
+        f"{tab}{tab}{tab}target={recipient}," + "\n"
+        f'{tab}{tab}{tab}condition="{condition}",' + "\n"
+    )
+    if available:
+        on_condition += f"{tab}{tab}{tab}available={available}," + "\n"
+    if available_function:
+        before_agent += "\n" + available_function + "\n"
+    on_condition += f"{tab}{tab}),"
+    return before_agent, on_condition
+
+
+# pylint: disable=too-many-locals
+def _get_agent_on_condition_handoff_to_nested_chat(
+    agent: WaldiezAgent,
+    agent_names: Dict[str, str],
+    condition: str,
+    available: str,
+    available_function: str,
+    all_chats: List[WaldiezChat],
+    chat_names: Dict[str, str],
+    is_async: bool,
+    serializer: Callable[..., str],
+    string_escape: Callable[[str], str],
+) -> Tuple[str, str]:
+    if not agent.data.nested_chats or not agent.data.nested_chats[0].messages:
+        return "", ""
+    chat_queue, extra_methods = get_nested_chat_queue(
+        nested_chat=agent.data.nested_chats[0],
+        agent=agent,
+        agent_names=agent_names,
+        chat_names=chat_names,
+        all_chats=all_chats,
+        serializer=serializer,
+        string_escape=string_escape,
+    )
+    if not chat_queue:
+        return "", ""
+    before_agent = ""
+    tab = "    "
+    chat_queue_var_name = f"{agent_names[agent.id]}_handoff_nested_chat_queue"
+    if extra_methods:
+        before_agent += "\n".join(extra_methods) + "\n"
+    before_agent += f"{chat_queue_var_name} = {chat_queue} " + "\n"
+    condition_string = string_escape(condition)
+    on_condition = (
+        f"{tab}{tab}ON_CONDITION(" + "\n"
+        f"{tab}{tab}{tab}target=" + "{\n"
+        f'{tab}{tab}{tab}{tab}"chat_queue": {chat_queue_var_name},' + "\n"
+        f'{tab}{tab}{tab}{tab}"config": None,' + "\n"
+        f'{tab}{tab}{tab}{tab}"reply_func_from_nested_chats": None,' + "\n"
+        f'{tab}{tab}{tab}{tab}"use_async": {is_async},' + "\n"
+        f"{tab}{tab}{tab}" + "},\n"
+        f'{tab}{tab}{tab}condition="{condition_string}",' + "\n"
+    )
+    if available:
+        on_condition += f"{tab}{tab}{tab}available={available}," + "\n"
+    if available_function:
+        before_agent += "\n" + available_function + "\n"
+    on_condition += f"{tab}{tab}),"
+    return before_agent, on_condition
+
+
+def _get_condition_string(
+    condition: str,
+    chat_id: str,
+    all_chats: List[WaldiezChat],
+    agent_names: Dict[str, str],
+) -> str:
+    if not condition:
+        chat = next((c for c in all_chats if c.id == chat_id), None)
+        if chat:
+            target_name = agent_names[chat.target]
+            return f"Transfer to {target_name}"
+        return "Transfer to the next agent"
+    return condition

waldiez/exporting/{agents → agent/utils}/teachability.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Exporting teachability data for agents."""
 
 from typing import Dict
@@ -23,15 +25,17 @@ def get_agent_teachability_string(
     str
         The teachability string
     """
-    if not agent.data.teachability.enabled:
+    if not agent.data.teachability or not agent.data.teachability.enabled:
         return ""
     agent_name = agent_names[agent.id]
     teachability = agent.data.teachability
-    content = f"{agent_name}_teachability = teachability.Teachability(\n"
-    content += f"    verbosity={teachability.verbosity},\n"
-    content += f"    reset_db={teachability.reset_db},\n"
-    content += f"    recall_threshold={teachability.recall_threshold},\n"
-    content += f"    max_num_retrievals={teachability.max_num_retrievals},\n"
+    content = f"{agent_name}_teachability = teachability.Teachability(" + "\n"
+    content += f"    verbosity={teachability.verbosity}," + "\n"
+    content += f"    reset_db={teachability.reset_db}," + "\n"
+    content += f"    recall_threshold={teachability.recall_threshold}," + "\n"
+    content += (
+        f"    max_num_retrievals={teachability.max_num_retrievals}," + "\n"
+    )
     content += ")\n\n\n"
     content += f"{agent_name}_teachability.add_to_agent({agent_name})"
     return content

waldiez/exporting/{agents → agent/utils}/termination_message.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Get the `is_termination_message` check for the agent."""
 
 from typing import Tuple
@@ -33,13 +35,10 @@ def get_is_termination_message(
     if agent.data.termination.type == "keyword":
         return agent.data.termination.string, ""
     if agent.data.termination.type == "method":
-        method_name = f"is_termination_message_{agent_name}"
-        content = (
-            "\n\n"
-            + f"def is_termination_message_{agent_name}(message):"
-            + "\n"
-            + f"{agent.data.termination.string}"
-            + "\n\n"
+        content, function_name = (
+            agent.data.termination.get_termination_function(
+                name_suffix=agent_name
+            )
         )
-        return method_name, content
+        return function_name, "\n\n" + content + "\n"
     raise ValueError(f"Invalid termination type: {agent.data.termination.type}")

waldiez/exporting/base/__init__.py

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Base classes, mixins, and utilities for exporting data.
+
+Each exporter should inherit from the `BaseExporter` class and implement the
+`export` method. The `export` method should return the exported content as an
+instance of the `ExporterReturnType` typed dictionary.
+"""
+
+from .agent_position import AgentPosition, AgentPositions
+from .base_exporter import BaseExporter, ExporterReturnType
+from .export_position import ExportPosition, ExportPositions
+from .import_position import ImportPosition
+from .mixin import ExporterMixin
+
+__all__ = [
+    "AgentPosition",
+    "AgentPositions",
+    "BaseExporter",
+    "ExporterMixin",
+    "ExportPosition",
+    "ExportPositions",
+    "ExporterReturnType",
+    "ImportPosition",
+]

waldiez/exporting/base/agent_position.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Agent position generation."""
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional
+
+from waldiez.models import WaldiezAgent
+
+
+class AgentPositions(Enum):
+    """Agent positions.
+
+    Attributes
+    ----------
+    BEFORE_ALL: int
+        Before all agents.
+    BEFORE: int
+        Before the agent.
+    AS_ARGUMENT: int
+        As an argument of the agent's initialization.
+    AFTER: int
+        After the agent.
+    AFTER_ALL: int
+        After all agents.
+    """
+
+    BEFORE_ALL = 0
+    BEFORE = 1
+    AS_ARGUMENT = 2
+    AFTER = 3
+    AFTER_ALL = 4
+
+
+POSITIONS_WITHOUT_AGENT = [
+    AgentPositions.BEFORE_ALL,
+    AgentPositions.AFTER_ALL,
+]
+
+
+@dataclass(order=True, frozen=True, slots=True)
+class AgentPosition:
+    """Agent position.
+
+    Attributes
+    ----------
+    agent: Optional[WaldiezAgent]
+        The agent.
+    position: AgentPositions
+        The position.
+    order: int
+        The order of the agent position.
+
+    Raises
+    ------
+    ValueError
+        If the position is not "BEFORE_ALL" or "AFTER_ALL"
+        and the agent is not provided.
+    """
+
+    agent: Optional[WaldiezAgent]
+    position: AgentPositions
+    order: int = 0
+
+    def __post_init__(self) -> None:
+        """Post initialization.
+
+        Raises
+        ------
+        ValueError
+            If the agent is not provided for the given position.
+        """
+        if self.agent is None and self.position not in POSITIONS_WITHOUT_AGENT:
+            raise ValueError("Agent must be provided for the given position.")