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

waldiez/models/agents/swarm_agent/after_work.py

@@ -0,0 +1,178 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Swarm after work model.
+
+Handles the next step in the conversation when an
+agent doesn't suggest a tool call or a handoff.
+
+"""
+
+# pylint: disable=line-too-long
+
+from typing import Dict, Optional, Tuple
+
+from pydantic import Field, model_validator
+from typing_extensions import Annotated, Literal, Self
+
+from ...common import WaldiezBase, check_function, generate_function
+
+WaldiezSwarmAfterWorkRecipientType = Literal["agent", "option", "callable"]
+WaldiezSwarmAfterWorkOption = Literal[
+    "TERMINATE", "REVERT_TO_USER", "STAY", "SWARM_MANAGER"
+]
+
+
+CUSTOM_AFTER_WORK = "custom_after_work"
+CUSTOM_AFTER_WORK_ARGS = ["last_speaker", "messages", "groupchat"]
+CUSTOM_AFTER_WORK_TYPES = (
+    ["SwarmAgent", "List[Dict[str, Any]]", "GroupChat"],
+    "Union[AfterWorkOption, SwarmAgent, str]",
+)
+
+
+class WaldiezSwarmAfterWork(WaldiezBase):
+    """Swarm after work.
+
+
+    Attributes
+    ----------
+    recipient : str
+        The agent_id to hand off to, an AfterWork option,
+        or the custom after work method.
+        If it is an AfterWork option, it can be one of
+        ('TERMINATE', 'REVERT_TO_USER', 'STAY', 'SWARM_MANAGER').
+
+    recipient_type : WaldiezSwarmAfterWorkRecipientType
+        The type of recipient.
+        Can be 'agent', 'option', or 'callable'.
+        If 'agent', the recipient is a SwarmAgent.
+        If 'option', the recipient is an AfterWorkOption :
+            ('TERMINATE', 'REVERT_TO_USER', 'STAY', 'SWARM_MANAGER').
+        If 'callable', it should have the signature:
+        def custom_after_work(
+            last_speaker: SwarmAgent,
+            messages: List[dict],
+            groupchat: GroupChat,
+        ) -> Union[AfterWorkOption, SwarmAgent, str]:
+
+    """
+
+    recipient: Annotated[
+        str,
+        Field(
+            "TERMINATE",
+            title="Recipient",
+            description=(
+                "The agent_id to hand off to, an AfterWork option, "
+                "or the custom after work method. "
+                "If it is an AfterWork option, it can be one of "
+                "('TERMINATE', 'REVERT_TO_USER', 'STAY')"
+            ),
+        ),
+    ]
+    recipient_type: Annotated[
+        WaldiezSwarmAfterWorkRecipientType,
+        Field(
+            "option",
+            alias="recipientType",
+            title="Recipient Type",
+            description=(
+                "The type of recipient. "
+                "Can be 'agent', 'option', or 'callable'. "
+                "If 'agent', the recipient is a SwarmAgent.  "
+                "If 'option', the recipient is an AfterWorkOption :"
+                "    ('TERMINATE', 'REVERT_TO_USER', 'STAY', 'SWARM_MANAGER'). "
+                "If 'callable', it should have the signature: "
+                "def custom_after_work("
+                "    last_speaker: SwarmAgent,"
+                "    messages: List[Dict[str, Any]],"
+                "    groupchat: GroupChat,"
+                ") -> Union[AfterWorkOption, SwarmAgent, str]:"
+            ),
+        ),
+    ]
+
+    _recipient_string: str = ""
+
+    def get_recipient(
+        self,
+        agent_names: Dict[str, str],
+        name_prefix: Optional[str] = None,
+        name_suffix: Optional[str] = None,
+    ) -> Tuple[str, str]:
+        """Get the recipient string.
+
+        Parameters
+        ----------
+        agent_names : Dict[str, str]
+            A mapping of agent id to agent name.
+        name_prefix : Optional[str], optional
+            The prefix for the function name, by default None.
+        name_suffix : Optional[str], optional
+            The suffix for the function name, by default None.
+
+        Returns
+        -------
+        Tuple[str, str]
+            The recipient string and the function content if applicable.
+        """
+        if self.recipient_type == "option":
+            return f"AFTER_WORK(AfterWorkOption.{self.recipient})", ""
+        if self.recipient_type == "agent":
+            # the the recipient is passed as the agent name
+            # (and not its id), care should be taken to ensure
+            # the all the agents in the flow have unique names
+            agent_instance = agent_names.get(self.recipient, self.recipient)
+            return f"AFTER_WORK({agent_instance})", ""
+
+        function_name = CUSTOM_AFTER_WORK
+        if name_prefix:
+            function_name = f"{name_prefix}_{function_name}"
+        if name_suffix:
+            function_name = f"{function_name}_{name_suffix}"
+        return (
+            f"AFTER_WORK({function_name})",
+            generate_function(
+                function_name=function_name,
+                function_args=CUSTOM_AFTER_WORK_ARGS,
+                function_body=self._recipient_string,
+                function_types=CUSTOM_AFTER_WORK_TYPES,
+            ),
+        )
+
+    @model_validator(mode="after")
+    def validate_recipient(self) -> Self:
+        """Validate the recipient.
+
+        Returns
+        -------
+        WaldiezSwarmAfterWork
+            The validated after work model.
+
+        Raises
+        ------
+        ValueError
+            If the validation fails.
+        """
+        self._recipient_string = self.recipient
+        if self.recipient_type == "callable":
+            is_valid, error_or_body = check_function(
+                code_string=self.recipient,
+                function_name=CUSTOM_AFTER_WORK,
+                function_args=CUSTOM_AFTER_WORK_ARGS,
+            )
+            if not is_valid or not error_or_body:
+                # pylint: disable=inconsistent-quotes
+                raise ValueError(
+                    f"Invalid custom method: {error_or_body or 'no content'}"
+                )
+            self._recipient_string = error_or_body
+        elif self.recipient_type == "option":
+            if self.recipient not in [
+                "TERMINATE",
+                "REVERT_TO_USER",
+                "STAY",
+                "SWARM_MANAGER",
+            ]:
+                raise ValueError("Invalid option.")
+        return self

waldiez/models/agents/swarm_agent/on_condition.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Swarm condition model for handoff."""
+
+from typing import Optional, Tuple
+
+from pydantic import Field
+from typing_extensions import Annotated, Literal
+
+from ...common import WaldiezBase
+from .on_condition_available import WaldiezSwarmOnConditionAvailable
+from .on_condition_target import WaldiezSwarmOnConditionTarget
+
+WaldiezSwarmOnConditionTargetType = Literal["agent", "nested_chat"]
+
+
+class WaldiezSwarmOnCondition(WaldiezBase):
+    """Swarm condition to handle handoff.
+
+    Attributes
+    ----------
+    target : WaldiezSwarmOnConditionTarget
+        The agent or nested chat configuration to hand off to.
+
+    target_type: Literal["agent", "nested_chat"]
+        The type of the target. Can be either 'agent' or 'nested_chat'.
+        Default is 'agent'.
+
+    condition : str
+        The condition for transitioning to the target agent
+
+    available: str, optional
+        Optional condition to determine if this ON_CONDITION is available.
+        Can be a Callable or a string.  If a string, it will look up the
+        value of the context variable with that name, which should be a bool.
+
+    available_check_type : Literal["string", "callable", "none"]
+        The type of the `available` property to check. Default is "none".
+    """
+
+    target: Annotated[
+        WaldiezSwarmOnConditionTarget,
+        Field(
+            title="Target",
+            description=(
+                "The agent or nested chat configuration to hand off to."
+            ),
+        ),
+    ]
+    target_type: Annotated[
+        WaldiezSwarmOnConditionTargetType,
+        Field(
+            "agent",
+            alias="targetType",
+            title="Target Type",
+            description=(
+                "The type of the target. "
+                "Can be either 'agent' or 'nested_chat'.Default is 'agent'."
+            ),
+        ),
+    ] = "agent"
+    condition: Annotated[
+        str,
+        Field(
+            ...,
+            title="Condition",
+            description="The condition for transitioning to the target agent",
+        ),
+    ]
+    available: Annotated[
+        WaldiezSwarmOnConditionAvailable,
+        Field(
+            default_factory=WaldiezSwarmOnConditionAvailable,
+            title="Available",
+            description=(
+                "Optional condition to determine if this ON_CONDITION "
+                "is available."
+            ),
+        ),
+    ]
+
+    def get_available(
+        self,
+        name_prefix: Optional[str] = None,
+        name_suffix: Optional[str] = None,
+    ) -> Tuple[str, str]:
+        """Get the available string.
+
+        Parameters
+        ----------
+        name_prefix : str, optional
+            The prefix to add to the function name. Default is None.
+        name_suffix : str, optional
+            The suffix to add to the function name. Default is None.
+        Returns
+        -------
+        Tuple[str, str]
+            The available string or function name and code if available.
+        """
+        return self.available.get_available(
+            name_prefix,
+            name_suffix,
+        )

waldiez/models/agents/swarm_agent/on_condition_available.py

@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Waldiez Swarm Agent ON_CONDITION Available Model."""
+
+from typing import Optional, Tuple
+
+from pydantic import Field, model_validator
+from typing_extensions import Annotated, Literal, Self
+
+from ...common import WaldiezBase, check_function, generate_function
+
+WaldiezSwarmOnConditionAvailableCheckType = Literal[
+    "string", "callable", "none"
+]
+
+CUSTOM_ON_CONDITION_AVAILABLE = "custom_on_condition_available"
+CUSTOM_ON_CONDITION_AVAILABLE_ARGS = ["agent", "message"]
+CUSTOM_ON_CONDITION_AVAILABLE_TYPES = (
+    ["ConversableAgent", "Dict[str, Any]"],
+    "bool",
+)
+
+# In ag2 it is used as:
+#
+# if on_condition.available is not None:
+#     if isinstance(on_condition.available, Callable):
+#         is_available = on_condition.available(
+#           agent, next(iter(agent.chat_messages.values()))
+#         )
+#     elif isinstance(on_condition.available, str):
+#         is_available = agent.get_context(on_condition.available) or False
+
+
+class WaldiezSwarmOnConditionAvailable(WaldiezBase):
+    """Swarm condition availability check."""
+
+    type: Annotated[
+        WaldiezSwarmOnConditionAvailableCheckType,
+        Field(
+            "none",
+            alias="availableCheckType",
+            title="Available Check Type",
+            description=("The type of the `available` property to check. "),
+        ),
+    ] = "none"
+    value: Annotated[
+        Optional[str],
+        Field(
+            None,
+            title="Available",
+            description=(
+                "Optional condition to determine if this ON_CONDITION "
+                "is available. Can be a Callable or a string.  If a string, "
+                " it will look up the value of the context variable with that "
+                "name, which should be a bool."
+            ),
+        ),
+    ]
+
+    _available_string: str = ""
+
+    @property
+    def available_string(self) -> str:
+        """Get the available string.
+
+        Returns
+        -------
+        str
+            The available string.
+        """
+        return self._available_string
+
+    def get_available(
+        self,
+        name_prefix: Optional[str] = None,
+        name_suffix: Optional[str] = None,
+    ) -> Tuple[str, str]:
+        """Get the available string.
+
+        Parameters
+        ----------
+        name_prefix : str, optional
+            The prefix to add to the function name. Default is None.
+        name_suffix : str, optional
+            The suffix to add to the function name. Default is None.
+        Returns
+        -------
+        Tuple[str, str]
+            The available string or function name and code if available.
+        """
+        if self.type == "none" or not self.value:
+            return "", ""
+        if self.type == "string":
+            return self._available_string, ""
+        function_name = CUSTOM_ON_CONDITION_AVAILABLE
+        if name_prefix:
+            function_name = f"{name_prefix}_{function_name}"
+        if name_suffix:
+            function_name = f"{function_name}_{name_suffix}"
+        return function_name, generate_function(
+            function_name=function_name,
+            function_args=CUSTOM_ON_CONDITION_AVAILABLE_ARGS,
+            function_types=CUSTOM_ON_CONDITION_AVAILABLE_TYPES,
+            function_body=self._available_string,
+        )
+
+    @model_validator(mode="after")
+    def validate_available(self) -> Self:
+        """Validate the available check.
+
+        Returns
+        -------
+        Self
+            The swarm agent's on condition available model.
+
+        Raises
+        ------
+        ValueError
+            If the available check fails.
+        """
+        if self.type == "callable":
+            if not self.value:
+                raise ValueError(
+                    "A callable is expected, but no value was provided."
+                )
+            is_valid, error_or_body = check_function(
+                code_string=self.value,
+                function_name=CUSTOM_ON_CONDITION_AVAILABLE,
+                function_args=CUSTOM_ON_CONDITION_AVAILABLE_ARGS,
+            )
+            if not is_valid:
+                raise ValueError(error_or_body)
+            self._available_string = error_or_body
+        if self.type == "string":
+            if not self.value:
+                raise ValueError(
+                    "A string is expected, but no value was provided."
+                )
+            self._available_string = self.value
+        return self

waldiez/models/agents/swarm_agent/on_condition_target.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Swarm condition model for handoff."""
+
+from pydantic import Field
+from typing_extensions import Annotated
+
+from ...common import WaldiezBase
+
+
+class WaldiezSwarmOnConditionTarget(WaldiezBase):
+    """Swarm condition target.
+
+    If the condition's target is "agent", the id refers to the agent's ID.
+    If the condition's target is "nested_chat", the id refers to the edge's ID.
+
+    Attributes
+    ----------
+    id : str
+        The ID of the target agent or edge.
+    order : int
+        The order of the target agent or edge.
+    """
+
+    id: Annotated[
+        str,
+        Field(
+            ...,
+            title="ID",
+            description="The ID of the target agent or edge.",
+        ),
+    ]
+    order: Annotated[
+        int,
+        Field(
+            ...,
+            title="Order",
+            description="The order of the target agent or edge.",
+        ),
+    ]

waldiez/models/agents/swarm_agent/swarm_agent.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Swarm agent."""
+
+from typing import List, Union
+
+from pydantic import Field
+from typing_extensions import Annotated, Literal
+
+from ..agent import WaldiezAgent, WaldiezAgentNestedChat
+from .after_work import WaldiezSwarmAfterWork
+from .on_condition import WaldiezSwarmOnCondition
+from .swarm_agent_data import WaldiezSwarmAgentData
+from .update_system_message import WaldiezSwarmUpdateSystemMessage
+
+
+class WaldiezSwarmAgent(WaldiezAgent):
+    """Swarm agent.
+
+    It extends a user agent and has swarm related parameters.
+
+    Attributes
+    ----------
+    agent_type : Literal["swarm"]
+        The agent type: 'swarm' for a swarm agent.
+    data : WaldiezSwarmAgentData
+        The swarm agent's data.
+        See `WaldiezSwarmAgentData` for more info.
+    """
+
+    agent_type: Annotated[
+        Literal["swarm"],
+        Field(
+            "swarm",
+            title="Agent type",
+            description="The agent type: 'swarm' for a swarm agent.",
+            alias="agentType",
+        ),
+    ]
+
+    data: Annotated[
+        WaldiezSwarmAgentData,
+        Field(
+            title="Data",
+            description="The swarm agent's data",
+            default_factory=WaldiezSwarmAgentData,
+        ),
+    ]
+
+    @property
+    def functions(self) -> List[str]:
+        """Get the functions that the agent can use.
+
+        Returns
+        -------
+        List[str]
+            The functions that the agent can use.
+        """
+        return self.data.functions
+
+    @property
+    def update_agent_state_before_reply(
+        self,
+    ) -> List[Union[str, WaldiezSwarmUpdateSystemMessage]]:
+        """Get the functions that update the agent's state before it replies.
+
+        Returns
+        -------
+        List[str]
+            The functions that update the agent's state before it replies.
+        """
+        return self.data.update_agent_state_before_reply
+
+    @property
+    def handoffs(
+        self,
+    ) -> List[Union[WaldiezSwarmOnCondition, WaldiezSwarmAfterWork]]:
+        """Get the hand offs to register.
+
+        Returns
+        -------
+        List[str]
+            The hand offs to register.
+        """
+        return self.data.handoffs
+
+    @property
+    def nested_chats(self) -> List[WaldiezAgentNestedChat]:
+        """Get the nested chats.
+
+        Returns
+        -------
+        List[WaldiezChat]
+            The nested chats.
+        """
+        return self.data.nested_chats
+
+    @property
+    def is_initial(self) -> bool:
+        """Check if the agent is the initial agent.
+
+        Returns
+        -------
+        bool
+            Whether the agent is the initial agent.
+        """
+        return self.data.is_initial

waldiez/models/agents/swarm_agent/swarm_agent_data.py

@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+# https://docs.ag2.ai/docs/reference/agentchat/contrib/swarm_agent
+"""Swarm agent data."""
+
+from typing import List, Union
+
+from pydantic import Field, model_validator
+from typing_extensions import Annotated, Self
+
+from ..agent import WaldiezAgentData
+from .after_work import WaldiezSwarmAfterWork
+from .on_condition import WaldiezSwarmOnCondition
+from .update_system_message import WaldiezSwarmUpdateSystemMessage
+
+
+# flake8: noqa: E501
+# pylint: disable=line-too-long
+class WaldiezSwarmAgentData(WaldiezAgentData):
+    """Swarm agent data.
+
+    Attributes
+    ----------
+    is_initial: bool
+        Whether the agent is the initial agent.
+    functions : List[str]
+        A list of functions (skill ids) to register with the agent.
+
+    update_agent_state_before_reply : List[str]
+        A list of functions, including `UPDATE_SYSTEM_MESSAGE`,
+        called to update the agent's state before it replies. Each function
+        is called when the agent is selected and before it speaks.
+
+    handoffs : List[Union[WaldiezSwarmOnCondition, WaldiezSwarmAfterWork]]
+        A list of hand offs to register.
+
+    Notes
+    -----
+    Each agent should have at most one `AfterWork` and (if any) it should be
+    at the end the list of hand offs.
+    """
+
+    is_initial: Annotated[
+        bool,
+        Field(
+            False,
+            title="Is Initial",
+            alias="isInitial",
+            description=("Whether the agent is the initial agent."),
+        ),
+    ]
+    functions: Annotated[
+        List[str],
+        Field(
+            title="Functions",
+            description=(
+                "A list of functions (skill ids) to register with the agent."
+            ),
+            default_factory=list,
+        ),
+    ]
+    update_agent_state_before_reply: Annotated[
+        List[Union[str, WaldiezSwarmUpdateSystemMessage]],
+        Field(
+            title="Update Agent State Before Reply",
+            alias="updateAgentStateBeforeReply",
+            description=(
+                "A list of functions, including UPDATE_SYSTEM_MESSAGEs,"
+                "called to update the agent's state before it replies. "
+                " Each function is called when the agent is selected "
+                "and before it speaks. If not an UPDATE_SYSTEM_MESSAGE, "
+                "it should be a skill id."
+            ),
+            default_factory=list,
+        ),
+    ]
+    handoffs: Annotated[
+        List[Union[WaldiezSwarmOnCondition, WaldiezSwarmAfterWork]],
+        Field(
+            title="Handoffs",
+            description=(
+                "A list of hand offs to register. "
+                "There should only be at most one `AfterWork` per agent"
+                "And (if any) it should be at the end of the list."
+            ),
+            default_factory=list,
+        ),
+    ]
+
+    @model_validator(mode="after")
+    def validate_handoffs(self) -> Self:
+        """Validate the hand offs.
+
+        Returns
+        -------
+        Self
+            The swarm agent data.
+
+        Raises
+        ------
+        ValueError
+            If there are more than one `AfterWork`s.
+        """
+        after_works: List[WaldiezSwarmAfterWork] = [
+            hand_off
+            for hand_off in self.handoffs
+            if isinstance(hand_off, WaldiezSwarmAfterWork)
+        ]
+        if len(after_works) > 1:
+            raise ValueError(
+                "Each agent should have at most one `AfterWork` "
+                "and (if any) it should be at the end of the list."
+            )
+        on_conditions: List[WaldiezSwarmOnCondition] = [
+            hand_off
+            for hand_off in self.handoffs
+            if isinstance(hand_off, WaldiezSwarmOnCondition)
+        ]
+        on_conditions = sorted(on_conditions, key=lambda x: x.target.order)
+        handoffs = on_conditions + after_works
+        if after_works and after_works[0] != handoffs[-1]:
+            handoffs.remove(after_works[0])
+            handoffs.append(after_works[0])
+        self.handoffs = handoffs
+        return self