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

waldiez/models/agents/swarm_agent/update_system_message.py

@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Update the agent's system message before they reply."""
+
+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
+
+WaldiezSwarmUpdateFunctionType = Literal["string", "callable"]
+
+CUSTOM_UPDATE_SYSTEM_MESSAGE = "custom_update_system_message"
+CUSTOM_UPDATE_SYSTEM_MESSAGE_ARGS = ["agent", "messages"]
+CUSTOM_UPDATE_SYSTEM_MESSAGE_TYPES = (
+    ["ConversableAgent", "List[Dict[str, Any]]"],
+    "str",
+)
+
+
+class WaldiezSwarmUpdateSystemMessage(WaldiezBase):
+    """Update the agent's system message before they reply.
+
+    Attributes
+    ----------
+    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.
+        If the `function_type` is 'callable', it should have the signature:
+        ```
+        def custom_update_system_message(
+            agent: ConversableAgent,
+            messages: List[Dict[str, Any]]
+        ) -> str:
+
+        ```
+    """
+
+    update_function_type: Annotated[
+        WaldiezSwarmUpdateFunctionType,
+        Field(
+            "string",
+            title="Function Type",
+            alias="updateFunctionType",
+            description=(
+                "The type of the update function. "
+                "Can be either 'string' or 'callable'."
+            ),
+        ),
+    ]
+
+    update_function: Annotated[
+        str,
+        Field(
+            ...,
+            title="Update Function",
+            alias="updateFunction",
+            description=(
+                "The string template or function definition to update "
+                "the agent's system message. Can be a string or a Callable.  "
+                "If the `update_function_type` is 'string', "
+                " it will be used as a template and substitute "
+                "the context variables. If `update_function_type` "
+                "is 'callable', it should have signature: "
+                "def custom_update_system_message("
+                "   agent: ConversableAgent, "
+                "   messages: List[Dict[str, Any]] "
+                ") -> str"
+            ),
+        ),
+    ]
+
+    _update_function: str = ""
+
+    def get_update_function(
+        self,
+        name_prefix: Optional[str] = None,
+        name_suffix: Optional[str] = None,
+    ) -> Tuple[str, str]:
+        """Get the update function.
+
+        Parameters
+        ----------
+        name_prefix : str, optional
+            The prefix of the name, by default None
+        name_suffix : str, optional
+            The suffix of the name, by default None
+
+        Returns
+        -------
+        Tuple[str, str]
+            The update function and the function name.
+
+        """
+        function_name = CUSTOM_UPDATE_SYSTEM_MESSAGE
+        if name_prefix:
+            function_name = f"{name_prefix}_{function_name}"
+        if name_suffix:
+            function_name = f"{function_name}_{name_suffix}"
+        return (
+            generate_function(
+                function_name=function_name,
+                function_args=CUSTOM_UPDATE_SYSTEM_MESSAGE_ARGS,
+                function_types=CUSTOM_UPDATE_SYSTEM_MESSAGE_TYPES,
+                function_body=self._update_function,
+            ),
+            function_name,
+        )
+
+    @model_validator(mode="after")
+    def validate_update_system_message(self) -> Self:
+        """Validate the update system message function.
+
+        Returns
+        -------
+        UpdateSystemMessage
+            The validated update system message.
+
+        Raises
+        ------
+        ValueError
+            If the type is callable and the function is invalid.
+            or if the function type is not 'string' or 'callable'.
+
+        """
+        self._update_function = self.update_function
+        if self.update_function_type == "callable":
+            valid, error_or_body = check_function(
+                code_string=self.update_function,
+                function_name=CUSTOM_UPDATE_SYSTEM_MESSAGE,
+                function_args=CUSTOM_UPDATE_SYSTEM_MESSAGE_ARGS,
+            )
+            if not valid or not error_or_body:
+                # pylint: disable=inconsistent-quotes
+                raise ValueError(
+                    f"Invalid custom method: {error_or_body or 'no content'}"
+                )
+            self._update_function = error_or_body
+        return self

waldiez/models/agents/user_proxy/__init__.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """User proxy agent model."""
 
 from .user_proxy import WaldiezUserProxy

waldiez/models/agents/user_proxy/user_proxy.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """User proxy agent model."""
 
 from pydantic import Field

waldiez/models/agents/user_proxy/user_proxy_data.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """User proxy agent data module."""
 
 from pydantic import Field

waldiez/models/chat/__init__.py

@@ -1,16 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Waldiez chat related models."""
 
 from .chat import WaldiezChat
 from .chat_data import WaldiezChatData
 from .chat_message import (
+    CALLABLE_MESSAGE,
+    CALLABLE_MESSAGE_ARGS,
+    CALLABLE_MESSAGE_RAG_WITH_CARRYOVER_TYPES,
+    CALLABLE_MESSAGE_TYPES,
     WaldiezChatMessage,
     WaldiezChatMessageType,
-    validate_message_dict,
 )
-from .chat_nested import WaldiezChatNested
+from .chat_nested import (
+    NESTED_CHAT_ARGS,
+    NESTED_CHAT_MESSAGE,
+    NESTED_CHAT_REPLY,
+    NESTED_CHAT_TYPES,
+    WaldiezChatNested,
+)
 from .chat_summary import WaldiezChatSummary, WaldiezChatSummaryMethod
 
 __all__ = [
+    "CALLABLE_MESSAGE",
+    "CALLABLE_MESSAGE_ARGS",
+    "CALLABLE_MESSAGE_TYPES",
+    "CALLABLE_MESSAGE_RAG_WITH_CARRYOVER_TYPES",
+    "NESTED_CHAT_MESSAGE",
+    "NESTED_CHAT_REPLY",
+    "NESTED_CHAT_ARGS",
+    "NESTED_CHAT_TYPES",
     "WaldiezChat",
     "WaldiezChatData",
     "WaldiezChatMessage",
@@ -18,5 +37,4 @@ __all__ = [
     "WaldiezChatNested",
     "WaldiezChatSummary",
     "WaldiezChatSummaryMethod",
-    "validate_message_dict",
 ]

waldiez/models/chat/chat.py

@@ -1,15 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Waldiez chat model."""
 
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional, Tuple
 
 from pydantic import Field
 from typing_extensions import Annotated
 
-from ..agents import WaldiezAgent, WaldiezRagUser
-from ..common import WaldiezBase
+from ..agents import WaldiezAgent, WaldiezRagUser, WaldiezSwarmAfterWork
+from ..common import WaldiezBase, generate_function
 from .chat_data import WaldiezChatData
-from .chat_message import WaldiezChatMessage
-from .chat_nested import WaldiezChatNested
+from .chat_message import (
+    CALLABLE_MESSAGE,
+    CALLABLE_MESSAGE_ARGS,
+    CALLABLE_MESSAGE_RAG_WITH_CARRYOVER_TYPES,
+    CALLABLE_MESSAGE_TYPES,
+    RAG_METHOD_WITH_CARRYOVER_BODY,
+    WaldiezChatMessage,
+)
+from .chat_nested import (
+    NESTED_CHAT_ARGS,
+    NESTED_CHAT_MESSAGE,
+    NESTED_CHAT_REPLY,
+    NESTED_CHAT_TYPES,
+    WaldiezChatNested,
+)
 
 
 class WaldiezChat(WaldiezBase):
@@ -63,6 +78,11 @@ class WaldiezChat(WaldiezBase):
         """Get the name."""
         return self.data.name
 
+    @property
+    def description(self) -> str:
+        """Get the description."""
+        return self.data.description
+
     @property
     def source(self) -> str:
         """Get the source."""
@@ -77,6 +97,11 @@ class WaldiezChat(WaldiezBase):
             return self.data.real_target
         return self.data.target
 
+    @property
+    def order(self) -> int:
+        """Get the order."""
+        return self.data.order
+
     @property
     def nested_chat(self) -> WaldiezChatNested:
         """Get the nested chat."""
@@ -101,22 +126,204 @@ class WaldiezChat(WaldiezBase):
         """Get the message content."""
         return self.data.message_content
 
+    @property
+    def context_variables(self) -> Dict[str, Any]:
+        """Get the context variables."""
+        return self.data.context_variables or {}
+
+    @property
+    def max_rounds(self) -> int:
+        """Get the max rounds for swarm chat."""
+        return self.data.max_rounds
+
+    @property
+    def after_work(self) -> Optional[WaldiezSwarmAfterWork]:
+        """Get the after work."""
+        return self.data.after_work
+
+    @property
+    def chat_id(self) -> int:
+        """Get the chat ID."""
+        return self.data.get_chat_id()
+
+    @property
+    def prerequisites(self) -> List[int]:
+        """Get the chat prerequisites."""
+        return self.data.get_prerequisites()
+
+    def set_chat_id(self, value: int) -> None:
+        """Set the chat ID.
+
+        Parameters
+        ----------
+        value : int
+            The chat ID.
+        """
+        self.data.set_chat_id(value)
+
+    def set_prerequisites(self, value: List[int]) -> None:
+        """Set the chat prerequisites.
+
+        Parameters
+        ----------
+        value : List[int]
+            The chat prerequisites.
+        """
+        self.data.set_prerequisites(value)
+
+    def get_message_function(
+        self,
+        name_prefix: Optional[str] = None,
+        name_suffix: Optional[str] = None,
+        is_rag: bool = False,
+    ) -> Tuple[str, str]:
+        """Get the message function.
+
+        Parameters
+        ----------
+        name_prefix : str
+            The function name prefix.
+        name_suffix : str
+            The function name suffix.
+        is_rag : bool
+            If the message is from a RAG user.
+
+        Returns
+        -------
+        Tuple[str, str]
+            The message function and the function name.
+        """
+        if self.message.type in ("string", "none") or (
+            not self.message_content and is_rag is False
+        ):
+            return "", ""
+        function_types = CALLABLE_MESSAGE_TYPES
+        function_name = CALLABLE_MESSAGE
+        if name_prefix:
+            function_name = f"{name_prefix}_{function_name}"
+        if name_suffix:
+            function_name = f"{function_name}_{name_suffix}"
+        if is_rag and self.message.type == "rag_message_generator":
+            function_types = CALLABLE_MESSAGE_RAG_WITH_CARRYOVER_TYPES
+            return (
+                generate_function(
+                    function_name=function_name,
+                    function_args=CALLABLE_MESSAGE_ARGS,
+                    function_types=function_types,
+                    function_body=self.message.content_body
+                    or RAG_METHOD_WITH_CARRYOVER_BODY,
+                ),
+                function_name,
+            )
+        return (
+            generate_function(
+                function_name=function_name,
+                function_args=CALLABLE_MESSAGE_ARGS,
+                function_types=function_types,
+                function_body=self.message_content or "",
+            ),
+            function_name,
+        )
+
+    def get_nested_chat_message_function(
+        self,
+        name_prefix: Optional[str] = None,
+        name_suffix: Optional[str] = None,
+    ) -> Tuple[str, str]:
+        """Get the nested chat message function.
+
+        Parameters
+        ----------
+        name_prefix : str
+            The function name prefix.
+        name_suffix : str
+            The function name suffix.
+
+        Returns
+        -------
+        Tuple[str, str]
+            The nested chat message function and the function name.
+        """
+        if (
+            not self.nested_chat.message
+            or self.nested_chat.message.type in ("string", "none")
+            or not self.nested_chat.message_content
+        ):
+            return "", ""
+        function_name = NESTED_CHAT_MESSAGE
+        if name_prefix:
+            function_name = f"{name_prefix}_{function_name}"
+        if name_suffix:
+            function_name = f"{function_name}_{name_suffix}"
+        return (
+            generate_function(
+                function_name=function_name,
+                function_args=NESTED_CHAT_ARGS,
+                function_types=NESTED_CHAT_TYPES,
+                function_body=self.nested_chat.message_content,
+            ),
+            function_name,
+        )
+
+    def get_nested_chat_reply_function(
+        self,
+        name_prefix: Optional[str] = None,
+        name_suffix: Optional[str] = None,
+    ) -> Tuple[str, str]:
+        """Get the nested chat reply function.
+
+        Parameters
+        ----------
+        name_prefix : str
+            The function name prefix.
+        name_suffix : str
+            The function name suffix.
+
+        Returns
+        -------
+        Tuple[str, str]
+            The nested chat reply function and the function name.
+        """
+        if (
+            not self.nested_chat.reply
+            or self.nested_chat.reply.type in ("string", "none")
+            or not self.nested_chat.reply_content
+        ):
+            return "", ""
+        function_name = NESTED_CHAT_REPLY
+        if name_prefix:
+            function_name = f"{name_prefix}_{function_name}"
+        if name_suffix:
+            function_name = f"{function_name}_{name_suffix}"
+        return (
+            generate_function(
+                function_name=function_name,
+                function_args=NESTED_CHAT_ARGS,
+                function_types=NESTED_CHAT_TYPES,
+                function_body=self.nested_chat.reply_content,
+            ),
+            function_name,
+        )
+
     def get_chat_args(
         self,
+        for_queue: bool,
         sender: Optional[WaldiezAgent] = None,
     ) -> Dict[str, Any]:
         """Get the chat arguments to use in autogen.
 
         Parameters
         ----------
-        sender : WaldiezAgent
+        for_queue : bool
+            Whether to get the chat arguments for a chat queue.
+        sender : WaldiezAgent, optional
             The sender agent, to check if it's a RAG user.
         Returns
         -------
         dict
             The chat arguments.
         """
-        args_dict = self.data.get_chat_args()
+        args_dict = self.data.get_chat_args(for_queue)
         if (
             isinstance(sender, WaldiezRagUser)
             and sender.agent_type == "rag_user"
@@ -127,3 +334,30 @@ class WaldiezChat(WaldiezBase):
             if isinstance(n_results, int) and n_results > 0:
                 args_dict["n_results"] = n_results
         return args_dict
+
+    def model_dump(self, **kwargs: Any) -> Dict[str, Any]:
+        """Dump the model to a dict including the chat attributes.
+
+        Parameters
+        ----------
+        kwargs : Any
+            The keyword arguments.
+        Returns
+        -------
+        Dict[str, Any]
+            The model dump with the chat attributes.
+        """
+        dump = super().model_dump(**kwargs)
+        dump["name"] = self.name
+        dump["description"] = self.description
+        dump["source"] = self.source
+        dump["target"] = self.target
+        dump["nested_chat"] = self.nested_chat.model_dump()
+        dump["message"] = self.message.model_dump()
+        dump["message_content"] = self.message_content
+        dump["context_variables"] = self.context_variables
+        dump["max_rounds"] = self.max_rounds
+        dump["after_work"] = (
+            self.after_work.model_dump() if self.after_work else None
+        )
+        return dump