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

waldiez/models/chat/chat_data.py

@@ -1,13 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Chat data model."""
 
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
-from pydantic import ConfigDict, Field, field_validator, model_validator
-from pydantic.alias_generators import to_camel
+from pydantic import Field, field_validator, model_validator
 from typing_extensions import Annotated, Self
 
-from ..common import WaldiezBase
-from .chat_message import WaldiezChatMessage, validate_message_dict
+from ..agents.swarm_agent import (
+    WaldiezSwarmAfterWork,
+    WaldiezSwarmOnConditionAvailable,
+)
+from ..common import WaldiezBase, check_function
+from .chat_message import (
+    CALLABLE_MESSAGE,
+    CALLABLE_MESSAGE_ARGS,
+    WaldiezChatMessage,
+)
 from .chat_nested import WaldiezChatNested
 from .chat_summary import WaldiezChatSummary
 
@@ -47,6 +56,10 @@ class WaldiezChatData(WaldiezBase):
         The real source of the chat (overrides the source).
     real_target : Optional[str]
         The real target of the chat (overrides the target).
+    max_rounds : int
+        Maximum number of conversation rounds (swarm).
+    after_work : Optional[WaldiezSwarmAfterWork]
+        The work to do after the chat (swarm).
 
     Functions
     ---------
@@ -60,13 +73,6 @@ class WaldiezChatData(WaldiezBase):
         Get the chat arguments to use in autogen.
     """
 
-    model_config = ConfigDict(
-        extra="forbid",
-        alias_generator=to_camel,
-        populate_by_name=True,
-        frozen=False,
-    )
-
     name: Annotated[
         str, Field(..., title="Name", description="The name of the chat.")
     ]
@@ -153,6 +159,14 @@ class WaldiezChatData(WaldiezBase):
             description="The maximum number of turns for the chat.",
         ),
     ]
+    prerequisites: Annotated[
+        List[str],
+        Field(
+            title="Prerequisites",
+            description="The prerequisites (chat ids) for the chat (if async).",
+            default_factory=list,
+        ),
+    ]
     silent: Annotated[
         Optional[bool],
         Field(
@@ -179,14 +193,90 @@ class WaldiezChatData(WaldiezBase):
             description="The real target of the chat (overrides the target).",
         ),
     ]
+    max_rounds: Annotated[
+        int,
+        Field(
+            20,
+            title="Max Rounds",
+            description="Maximum number of conversation rounds.(swarm)",
+        ),
+    ] = 20
+    after_work: Annotated[
+        Optional[WaldiezSwarmAfterWork],
+        Field(
+            None,
+            alias="afterWork",
+            title="After Work",
+            description="The work to do after the chat (swarm).",
+        ),
+    ] = None
+    context_variables: Annotated[
+        Optional[Dict[str, Any]],
+        Field(
+            None,
+            alias="contextVariables",
+            title="Context Variables",
+            description="The context variables to use in the chat.",
+        ),
+    ] = None
+    available: Annotated[
+        WaldiezSwarmOnConditionAvailable,
+        Field(
+            default_factory=WaldiezSwarmOnConditionAvailable,
+            title="Available",
+            description="The available condition for the chat.",
+        ),
+    ]
 
     _message_content: Optional[str] = None
+    _chat_id: int = 0
+    _prerequisites: List[int] = []
 
     @property
     def message_content(self) -> Optional[str]:
         """Get the message content."""
         return self._message_content
 
+    def get_chat_id(self) -> int:
+        """Get the chat id.
+
+        Returns
+        -------
+        int
+            The chat id.
+        """
+        return self._chat_id
+
+    def set_chat_id(self, value: int) -> None:
+        """Set the chat id.
+
+        Parameters
+        ----------
+        value : int
+            The chat id.
+        """
+        self._chat_id = value
+
+    def get_prerequisites(self) -> List[int]:
+        """Get the chat prerequisites.
+
+        Returns
+        -------
+        List[int]
+            The chat prerequisites (if async).
+        """
+        return self._prerequisites
+
+    def set_prerequisites(self, value: List[int]) -> None:
+        """Set the chat prerequisites.
+
+        Parameters
+        ----------
+        value : List[int]
+            The chat prerequisites to set.
+        """
+        self._prerequisites = value
+
     @model_validator(mode="after")
     def validate_chat_data(self) -> Self:
         """Validate the chat data.
@@ -206,16 +296,17 @@ class WaldiezChatData(WaldiezBase):
                 self._message_content = None
             elif self.message.type == "string":
                 self._message_content = self.message.content
+            elif self.message.type == "method":
+                valid, error_or_body = check_function(
+                    self.message.content or "",
+                    CALLABLE_MESSAGE,
+                    CALLABLE_MESSAGE_ARGS,
+                )
+                if not valid:
+                    raise ValueError(error_or_body)
+                self._message_content = error_or_body
             else:
-                self._message_content = validate_message_dict(
-                    value={
-                        "type": self.message.type,
-                        "content": self.message.content,
-                        "use_carryover": self.message.use_carryover,
-                    },
-                    function_name="callable_message",
-                    skip_definition=True,
-                ).content
+                self._message_content = self.message.content
         return self
 
     @field_validator("message", mode="before")
@@ -242,28 +333,47 @@ class WaldiezChatData(WaldiezBase):
             return WaldiezChatMessage(
                 type="none", use_carryover=False, content=None, context={}
             )
-        if isinstance(value, str):
+        if isinstance(value, (str, int, float, bool)):
             return WaldiezChatMessage(
-                type="string", use_carryover=False, content=value, context={}
+                type="string",
+                use_carryover=False,
+                content=str(value),
+                context={},
             )
         if isinstance(value, dict):
-            return validate_message_dict(
-                value, function_name="callable_message"
-            )
+            return WaldiezChatMessage.model_validate(value)
         if isinstance(value, WaldiezChatMessage):
-            return validate_message_dict(
-                value={
-                    "type": value.type,
-                    "use_carryover": value.use_carryover,
-                    "content": value.content,
-                    "context": value.context,
-                },
-                function_name="callable_message",
-            )
+            return value
         return WaldiezChatMessage(
             type="none", use_carryover=False, content=None, context={}
         )
 
+    @field_validator("context_variables", mode="after")
+    @classmethod
+    def validate_context_variables(cls, value: Any) -> Optional[Dict[str, Any]]:
+        """Validate the context variables.
+
+        Parameters
+        ----------
+        value : Any
+            The context variables value.
+
+        Returns
+        -------
+        Optional[Dict[str, Any]]
+            The validated context variables value.
+
+        Raises
+        ------
+        ValueError
+            If the validation fails.
+        """
+        if value is None:
+            return None
+        if not isinstance(value, dict):
+            raise ValueError("Context variables must be a dictionary.")
+        return get_context_dict(value)
+
     @property
     def summary_args(self) -> Optional[Dict[str, Any]]:
         """Get the summary args."""
@@ -289,23 +399,19 @@ class WaldiezChatData(WaldiezBase):
         """
         extra_args: Dict[str, Any] = {}
         if isinstance(self.message, WaldiezChatMessage):
-            for key, value in self.message.context.items():
-                if str(value).lower() in ("none", "null"):
-                    extra_args[key] = None
-                elif str(value).isdigit():
-                    extra_args[key] = int(value)
-                elif str(value).replace(".", "").isdigit():
-                    try:
-                        extra_args[key] = float(value)
-                    except ValueError:  # pragma: no cover
-                        extra_args[key] = value
-                else:
-                    extra_args[key] = value
+            extra_args.update(get_context_dict(self.message.context))
         return extra_args
 
-    def get_chat_args(self) -> Dict[str, Any]:
+    def get_chat_args(self, for_queue: bool) -> Dict[str, Any]:
         """Get the chat arguments to use in autogen.
 
+        Without the 'message' key.
+
+        Parameters
+        ----------
+        for_queue : bool
+            Whether to get the arguments for a chat queue.
+
         Returns
         -------
         Dict[str, Any]
@@ -323,4 +429,42 @@ class WaldiezChatData(WaldiezBase):
         if isinstance(self.silent, bool):
             args["silent"] = self.silent
         args.update(self._get_context_args())
+        if for_queue:
+            args["chat_id"] = self._chat_id
+        if self._prerequisites:
+            args["prerequisites"] = self._prerequisites
         return args
+
+
+def get_context_dict(context: Dict[str, Any]) -> Dict[str, Any]:
+    """Get the context dictionary.
+
+    Try to determine the type of the context variables.
+
+    Parameters
+    ----------
+    context : Dict[str, Any]
+        The context variables.
+
+    Returns
+    -------
+    Dict[str, Any]
+        The context variables with the detected types.
+    """
+    new_dict: Dict[str, Any] = {}
+    for key, value in context.items():
+        value_lower = str(value).lower()
+        if value_lower in ("none", "null"):
+            new_dict[key] = None
+        elif value_lower in ("true", "false"):
+            new_dict[key] = value.lower() == "true"
+        elif str(value).isdigit():
+            new_dict[key] = int(value)
+        elif str(value).replace(".", "").isdigit():
+            try:
+                new_dict[key] = float(value)
+            except ValueError:  # pragma: no cover
+                new_dict[key] = value
+        else:
+            new_dict[key] = value
+    return new_dict

waldiez/models/chat/chat_message.py

@@ -1,16 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Waldiez Message Model."""
 
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional
 
-from pydantic import Field
-from typing_extensions import Annotated, Literal
+from pydantic import Field, model_validator
+from typing_extensions import Annotated, Literal, Self
 
-from ..common import WaldiezBase, WaldiezMethodName, check_function
+from ..common import WaldiezBase, check_function
 
 WaldiezChatMessageType = Literal[
     "string", "method", "rag_message_generator", "none"
 ]
 
+CALLABLE_MESSAGE = "callable_message"
+CALLABLE_MESSAGE_ARGS = ["sender", "recipient", "context"]
+CALLABLE_MESSAGE_TYPES = (
+    ["ConversableAgent", "ConversableAgent", "Dict[str, Any]"],
+    "Union[Dict[str, Any], str]",
+)
+CALLABLE_MESSAGE_RAG_WITH_CARRYOVER_TYPES = (
+    ["RetrieveUserProxyAgent", "ConversableAgent", "Dict[str, Any]"],
+    "Union[Dict[str, Any], str]",
+)
+
 
 class WaldiezChatMessage(WaldiezBase):
     """
@@ -84,148 +97,130 @@ class WaldiezChatMessage(WaldiezBase):
         ),
     ]
 
+    _content_body: Optional[str] = None
 
-def validate_message_dict(
-    value: Dict[
-        Literal["type", "use_carryover", "content", "context"],
-        Union[Optional[str], Optional[bool], Optional[Dict[str, Any]]],
-    ],
-    function_name: WaldiezMethodName,
-    skip_definition: bool = False,
-) -> WaldiezChatMessage:
-    """Validate a message dict.
+    @property
+    def content_body(self) -> Optional[str]:
+        """Get the content body."""
+        return self._content_body
 
-    Check the provided message dict.
-    Depending on the type, the content is validated.
-    If the type is "method", the content is checked against the function name.
+    @model_validator(mode="after")
+    def validate_context_vars(self) -> Self:
+        """Try to detect bools nulls and numbers from the context values.
 
-    Parameters
-    ----------
-    value : dict
-        The message dict.
-    function_name : str (WaldiezMethodName)
-        The function name.
-    skip_definition : bool, optional
-        Skip the function definition in the content, by default False
+        Returns
+        -------
+        WaldiezChatMessage
+            The validated instance.
+        """
+        for key, value in self.context.items():
+            if isinstance(value, str):
+                if value.lower() == "true":
+                    self.context[key] = True
+                elif value.lower() == "false":
+                    self.context[key] = False
+                elif value.lower() in ["null", "none"]:
+                    self.context[key] = None
+                else:
+                    self.context[key] = self._number_or_string(value)
+        return self
 
-    Returns
-    -------
-    WaldiezChatMessage
-        The validated message.
+    @staticmethod
+    def _number_or_string(value: Any) -> Any:
+        try:
+            int_value = int(value)
+            if str(int_value) == value:
+                return int_value
+        except ValueError:
+            try:
+                float_value = float(value)
+                if str(float_value) == value:
+                    return float_value
+            except ValueError:
+                pass
+        return value
 
-    Raises
-    ------
-    ValueError
-        If the validation fails.
-    """
-    message_type, use_carryover, content, context = _get_message_args_from_dict(
-        value
-    )
-    if message_type == "string":
-        if not content or not isinstance(content, str):
-            content = ""
-        if use_carryover:
-            method_content = _get_last_carryover_method_content(content)
-            return WaldiezChatMessage(
-                type="method",
-                use_carryover=True,
-                content=method_content,
-                context=context,
-            )
-        return WaldiezChatMessage(
-            type="string",
-            use_carryover=False,
-            content=content,
-            context=context,
-        )
-    if message_type == "none":
-        return WaldiezChatMessage(
-            type="none",
-            use_carryover=False,
-            content=None,
-            context=context,
-        )
-    if message_type == "method":
-        if not content:
+    @model_validator(mode="after")
+    def validate_content(self) -> Self:
+        """Validate the content (if not a method).
+
+        Returns
+        -------
+        WaldiezChatMessage
+            The validated instance.
+
+        Raises
+        ------
+        ValueError
+            If the content is invalid.
+        """
+        content: Optional[str] = None
+        if self.type == "none":
+            content = "None"
+        if self.type == "method":
+            if not self.content:
+                raise ValueError(
+                    "The message content is required for the method type"
+                )
+            content = self.content
+        if self.type == "string":
+            if not self.content:
+                self.content = ""
+            if self.use_carryover:
+                content = get_last_carryover_method_content(
+                    text_content=self.content,
+                )
+            content = self.content
+        if self.type == "rag_message_generator":
+            if self.use_carryover:
+                content = get_last_carryover_method_content(
+                    text_content=self.content or "",
+                )
+            else:
+                content = RAG_METHOD_WITH_CARRYOVER_BODY
+                self.content = RAG_METHOD_WITH_CARRYOVER
+        self._content_body = content
+        return self
+
+    def validate_method(
+        self,
+        function_name: str,
+        function_args: List[str],
+    ) -> str:
+        """Validate a method.
+
+        Parameters
+        ----------
+        function_name : str
+            The method name.
+        function_args : List[str]
+            The expected method arguments.
+
+        Returns
+        -------
+        str
+            The validated method body.
+
+        Raises
+        ------
+        ValueError
+            If the validation fails.
+        """
+        if not self.content:
             raise ValueError(
                 "The message content is required for the method type"
             )
-        valid, error_or_content = check_function(
-            content, function_name, skip_type_hints=use_carryover
+        is_valid, error_or_body = check_function(
+            code_string=self.content,
+            function_name=function_name,
+            function_args=function_args,
         )
-        if not valid:
-            raise ValueError(error_or_content)
-        message_content = error_or_content if skip_definition else content
-        return WaldiezChatMessage(
-            type="method",
-            use_carryover=use_carryover,
-            content=message_content,
-            context=context,
-        )
-    if message_type == "rag_message_generator":
-        if use_carryover:
-            return WaldiezChatMessage(
-                type="method",
-                use_carryover=True,
-                content=RAG_METHOD_WITH_CARRYOVER,
-                context=context,
-            )
-        return WaldiezChatMessage(
-            type="rag_message_generator",
-            use_carryover=use_carryover,
-            content=None,
-            context=context,
-        )
-    raise ValueError("Invalid message type")  # pragma: no cover
-
+        if not is_valid:
+            raise ValueError(error_or_body)
+        return error_or_body
 
-def _get_message_args_from_dict(
-    value: Dict[
-        Literal["type", "use_carryover", "content", "context"],
-        Union[Optional[str], Optional[bool], Optional[Dict[str, Any]]],
-    ],
-) -> Tuple[str, bool, Optional[str], Dict[str, Any]]:
-    """Get the message args from a dict.
 
-    Parameters
-    ----------
-    value : dict
-        The message dict.
-
-    Returns
-    -------
-    tuple
-        The message type, content, and context.
-
-    Raises
-    ------
-    ValueError
-        If the message type is invalid.
-    """
-    message_type = value.get("type")
-    if not isinstance(message_type, str) or message_type not in (
-        "string",
-        "method",
-        "rag_message_generator",
-        "none",
-    ):
-        raise ValueError("Invalid message type")
-    use_carryover = value.get("use_carryover", False)
-    if not isinstance(use_carryover, bool):
-        use_carryover = False
-    content = value.get("content", "")
-    if not isinstance(content, str):
-        content = ""
-    context: Dict[str, Any] = {}
-    context_value = value.get("context")
-    if isinstance(context_value, dict):
-        context = context_value
-    if not isinstance(context, dict):  # pragma: no cover
-        context = {}
-    return message_type, use_carryover, content, context
-
-
-def _get_last_carryover_method_content(text_content: str) -> str:
+def get_last_carryover_method_content(text_content: str) -> str:
     """Get the last carryover method content.
 
     Parameters
@@ -238,8 +233,6 @@ def _get_last_carryover_method_content(text_content: str) -> str:
         The last carryover method content.
     """
     method_content = '''
-def callable_message(sender, recipient, context):
-    # type: (ConversableAgent, ConversableAgent, dict) -> Union[dict, str]
     """Get the message to send using the last carryover.
 
     Parameters
@@ -248,17 +241,22 @@ def callable_message(sender, recipient, context):
         The source agent.
     recipient : ConversableAgent
         The target agent.
-    context : dict
+    context : Dict[str, Any]
         The context.
 
     Returns
     -------
-    Union[dict, str]
+    Union[Dict[str, Any], str]
         The message to send using the last carryover.
     """
     carryover = context.get("carryover", "")
     if isinstance(carryover, list):
         carryover = carryover[-1]
+    if not isinstance(carryover, str):
+        if isinstance(carryover, list):
+            carryover = carryover[-1]
+        elif isinstance(carryover, dict):
+            carryover = carryover.get("content", "")
     if not isinstance(carryover, str):
         carryover = ""'''
     if text_content:
@@ -273,9 +271,7 @@ def callable_message(sender, recipient, context):
     return method_content
 
 
-RAG_METHOD_WITH_CARRYOVER = '''
-def callable_message(sender, recipient, context):
-    # type: (RetrieveUserProxyAgent, ConversableAgent, dict) -> Union[dict, str]
+RAG_METHOD_WITH_CARRYOVER_BODY = '''
     """Get the message using the RAG message generator method.
 
     Parameters
@@ -284,17 +280,22 @@ def callable_message(sender, recipient, context):
         The source agent.
     recipient : ConversableAgent
         The target agent.
-    context : dict
+    context : Dict[str, Any]
         The context.
 
     Returns
     -------
-    Union[dict, str]
+    Union[Dict[str, Any], str]
         The message to send using the last carryover.
     """
     carryover = context.get("carryover", "")
     if isinstance(carryover, list):
         carryover = carryover[-1]
+    if not isinstance(carryover, str):
+        if isinstance(carryover, list):
+            carryover = carryover[-1]
+        elif isinstance(carryover, dict):
+            carryover = carryover.get("content", "")
     if not isinstance(carryover, str):
         carryover = ""
     message = sender.message_generator(sender, recipient, context)
@@ -302,3 +303,11 @@ def callable_message(sender, recipient, context):
         message += carryover
     return message
 '''
+RAG_METHOD_WITH_CARRYOVER = (
+    "def callable_message(\n"
+    "    sender: RetrieveUserProxyAgent,\n"
+    "    recipient: ConversableAgent,\n"
+    "    context: Dict[str, Any],\n"
+    ") -> Union[Dict[str, Any], str]:"
+    f"{RAG_METHOD_WITH_CARRYOVER_BODY}"
+)