waldiez 0.1.0__py3-none-any.whl

waldiez/models/chat/chat_nested.py

@@ -0,0 +1,160 @@
+"""Nested chat model."""
+
+from typing import Any, Optional
+
+from pydantic import (
+    ConfigDict,
+    Field,
+    ValidationInfo,
+    field_validator,
+    model_validator,
+)
+from pydantic.alias_generators import to_camel
+from typing_extensions import Annotated, Self
+
+from ..common import WaldiezBase, WaldiezMethodName
+from .chat_message import WaldiezChatMessage, validate_message_dict
+
+
+class WaldiezChatNested(WaldiezBase):
+    """Nested chat class.
+
+    Attributes
+    ----------
+    message : WaldiezChatMessage
+        The message in a nested chat (sender -> recipient).
+    reply : WaldiezChatMessage
+        The reply in a nested chat (recipient -> sender).
+    """
+
+    model_config = ConfigDict(
+        extra="forbid",
+        alias_generator=to_camel,
+        populate_by_name=True,
+        frozen=False,
+    )
+
+    message: Annotated[
+        Optional[WaldiezChatMessage],
+        Field(
+            None,
+            title="Message",
+            description="The message in a nested chat (sender -> recipient).",
+        ),
+    ]
+    reply: Annotated[
+        Optional[WaldiezChatMessage],
+        Field(
+            None,
+            title="Reply",
+            description="The reply in a nested chat (recipient -> sender).",
+        ),
+    ]
+
+    _message_content: Optional[str] = None
+    _reply_content: Optional[str] = None
+
+    @property
+    def message_content(self) -> Optional[str]:
+        """Get the message content."""
+        return self._message_content
+
+    @property
+    def reply_content(self) -> Optional[str]:
+        """Get the reply content."""
+        return self._reply_content
+
+    @field_validator("message", "reply", mode="before")
+    @classmethod
+    def validate_message(
+        cls, value: Any, info: ValidationInfo
+    ) -> WaldiezChatMessage:
+        """Validate the message.
+
+        Parameters
+        ----------
+        value : Any
+            The value.
+        info : ValidationInfo
+            The validation info.
+
+        Returns
+        -------
+        WaldiezChatMessage
+            The validated message.
+
+        Raises
+        ------
+        ValueError
+            If the validation fails.
+        """
+        function_name: WaldiezMethodName = (
+            "nested_chat_message"
+            if info.field_name == "message"
+            else "nested_chat_reply"
+        )
+        if not value:
+            return WaldiezChatMessage(
+                type="none", use_carryover=False, content=None, context={}
+            )
+        if isinstance(value, str):
+            return WaldiezChatMessage(
+                type="string", use_carryover=False, content=value, context={}
+            )
+        if isinstance(value, dict):
+            return validate_message_dict(value, function_name=function_name)
+        if isinstance(value, WaldiezChatMessage):
+            return validate_message_dict(
+                {
+                    "type": value.type,
+                    "use_carryover": False,
+                    "content": value.content,
+                    "context": value.context,
+                },
+                function_name=function_name,
+            )
+        raise ValueError(f"Invalid message type: {type(value)}")
+
+    @model_validator(mode="after")
+    def validate_nested_chat(self) -> Self:
+        """Validate the nested chat.
+
+        Returns
+        -------
+        WaldiezChatNested
+            The validated nested chat.
+
+        Raises
+        ------
+        ValueError
+            If the validation fails.
+        """
+        if self.message is not None:
+            if self.message.type == "none":
+                self._message_content = ""
+            elif self.message.type == "string":
+                self._message_content = self.message.content
+            else:
+                self._message_content = validate_message_dict(
+                    value={
+                        "type": "method",
+                        "content": self.message.content,
+                    },
+                    function_name="nested_chat_message",
+                    skip_definition=True,
+                ).content
+        if self.reply is not None:
+            if self.reply.type == "none":
+                self._reply_content = ""
+            elif self.reply.type == "string":
+                self._reply_content = self.reply.content
+            else:
+                self._reply_content = validate_message_dict(
+                    value={
+                        "type": "method",
+                        "content": self.reply.content,
+                    },
+                    function_name="nested_chat_reply",
+                    skip_definition=True,
+                ).content
+        return self

waldiez/models/chat/chat_summary.py

@@ -0,0 +1,110 @@
+"""Waldiez chat summary options."""
+
+from typing import Any, Dict, Optional
+
+from pydantic import (
+    Field,
+    FieldSerializationInfo,
+    field_serializer,
+    field_validator,
+)
+from typing_extensions import Annotated, Literal
+
+from ..common import WaldiezBase
+
+WaldiezChatSummaryMethod = Literal[
+    "reflectionWithLlm",
+    "lastMsg",
+    "reflection_with_llm",
+    "last_msg",
+]
+
+
+class WaldiezChatSummary(WaldiezBase):
+    """Llm summary method options.
+
+    Attributes
+    ----------
+    method : Optional[WaldiezChatSummaryMethod]
+        The method to use for the LLM summary. Defaults to "last_msg".
+    prompt : str
+        The prompt for the LLM summary method.
+    args : Optional[Dict[str, Any]]
+        The additional arguments for the LLM summary method, by default None.
+    """
+
+    method: Annotated[
+        Optional[WaldiezChatSummaryMethod],
+        Field(
+            "last_msg",
+            title="Method",
+            description="The method to use for the LLM summary.",
+        ),
+    ]
+    prompt: Annotated[
+        str,
+        Field(
+            "",
+            title="Prompt",
+            description="The prompt for the LLM summary method.",
+        ),
+    ]
+    args: Annotated[
+        Dict[str, str],
+        Field(
+            title="Arguments",
+            description="The additional arguments for the LLM summary method.",
+            default_factory=dict,
+        ),
+    ]
+
+    @field_validator("method", mode="before")
+    @classmethod
+    def validate_summary_method(
+        cls, value: Optional[WaldiezChatSummaryMethod]
+    ) -> Optional[WaldiezChatSummaryMethod]:
+        """Validate the summary method.
+
+        Parameters
+        ----------
+        value : Optional[WaldiezChatSummaryMethod]
+            The passed WaldiezChatSummaryMethod
+
+        Returns
+        -------
+        Optional[WaldiezChatSummaryMethod]
+            The validated message summary method
+        """
+        if str(value).lower() == "none":
+            return None
+        if value == "lastMsg":
+            return "last_msg"
+        if value == "reflectionWithLlm":
+            return "reflection_with_llm"
+        return value
+
+    @field_serializer("method")
+    @classmethod
+    def serialize_summary_method(
+        cls, value: Any, info: FieldSerializationInfo
+    ) -> Any:
+        """Serialize summary method.
+
+        Parameters
+        ----------
+        value : Any
+            The value to serialize.
+        info : FieldSerializationInfo
+            The serialization info.
+
+        Returns
+        -------
+        Any
+            The serialized value.
+        """
+        if info.by_alias is True:
+            if value == "reflection_with_llm":
+                return "reflectionWithLlm"
+            if value == "last_msg":
+                return "lastMsg"
+        return value

waldiez/models/common/__init__.py

@@ -0,0 +1,38 @@
+"""Common utils for all models."""
+
+from datetime import datetime, timezone
+
+from .base import WaldiezBase
+from .method_utils import (
+    METHOD_ARGS,
+    METHOD_TYPE_HINTS,
+    WaldiezMethodName,
+    check_function,
+    parse_code_string,
+)
+
+
+def now() -> str:
+    """Get the current date and time in UTC.
+
+    Returns
+    -------
+    str
+        The current date and time in UTC.
+    """
+    return (
+        datetime.now(tz=timezone.utc)
+        .isoformat(timespec="milliseconds")
+        .replace("+00:00", "Z")
+    )
+
+
+__all__ = [
+    "WaldiezBase",
+    "METHOD_ARGS",
+    "METHOD_TYPE_HINTS",
+    "WaldiezMethodName",
+    "now",
+    "check_function",
+    "parse_code_string",
+]

waldiez/models/common/base.py

@@ -0,0 +1,63 @@
+"""Base class to inherit from."""
+
+from typing import Any, Dict
+
+from pydantic import BaseModel, ConfigDict
+from pydantic.alias_generators import to_camel
+
+
+class WaldiezBase(BaseModel):
+    """Base model class to inherit from.
+
+    It contains the default configuration for all models.
+    It also `model_dumps` by alias by default.
+    """
+
+    model_config = ConfigDict(
+        extra="forbid",
+        # treat `skillId` as `skill_id`
+        alias_generator=to_camel,
+        # allow passing either `skill_id` or `skillId`
+        populate_by_name=True,
+        frozen=True,
+    )
+
+    def model_dump(self, **kwargs: Any) -> Dict[str, Any]:
+        """Dump the model to a dictionary.
+
+        Parameters
+        ----------
+        **kwargs : Any
+            Additional keyword arguments.
+
+        Returns
+        -------
+        Dict[str, Any]
+            The dictionary representation of the model.
+        """
+        by_alias = kwargs.pop("by_alias", None)
+        if by_alias is None:
+            by_alias = True
+        if not isinstance(by_alias, bool):
+            by_alias = True
+        return super().model_dump(by_alias=by_alias, **kwargs)
+
+    def model_dump_json(self, **kwargs: Any) -> str:
+        """Dump the model to a JSON string.
+
+        Parameters
+        ----------
+        **kwargs : Any
+            Additional keyword arguments.
+
+        Returns
+        -------
+        str
+            The JSON string.
+        """
+        by_alias = kwargs.pop("by_alias", None)
+        if by_alias is None:
+            by_alias = True
+        if not isinstance(by_alias, bool):
+            by_alias = True
+        return super().model_dump_json(by_alias=by_alias, **kwargs)

waldiez/models/common/method_utils.py

@@ -0,0 +1,165 @@
+"""Function related utilities."""
+
+import ast
+from typing import Dict, List, Literal, Optional, Tuple
+
+WaldiezMethodName = Literal[
+    "callable_message",  # Chat
+    "is_termination_message",  # Agent
+    "nested_chat_message",  # Agents' NestedChat
+    "nested_chat_reply",  # Agents' NestedChat
+    "custom_speaker_selection",  # GroupChat
+    "custom_embedding_function",  # RAG
+    "custom_token_count_function",  # RAG
+    "custom_text_split_function",  # RAG
+]
+
+METHOD_ARGS: Dict[WaldiezMethodName, List[str]] = {
+    "callable_message": ["sender", "recipient", "context"],
+    "is_termination_message": ["message"],
+    "nested_chat_message": ["recipient", "messages", "sender", "config"],
+    "nested_chat_reply": ["recipient", "messages", "sender", "config"],
+    "custom_speaker_selection": ["last_speaker", "groupchat"],
+    "custom_embedding_function": [],
+    "custom_token_count_function": ["text", "model"],
+    "custom_text_split_function": [
+        "text",
+        "max_tokens",
+        "chunk_mode",
+        "must_break_at_empty_line",
+        "overlap",
+    ],
+}
+
+# pylint: disable=line-too-long
+METHOD_TYPE_HINTS: Dict[WaldiezMethodName, str] = {
+    "callable_message": "# type: (ConversableAgent, ConversableAgent, dict) -> Union[dict, str]",
+    "is_termination_message": "# type: (dict) -> bool",
+    "nested_chat_message": "# type: (ConversableAgent, list[dict], ConversableAgent, dict) -> Union[dict, str]",
+    "nested_chat_reply": "# type: (ConversableAgent, list[dict], ConversableAgent, dict) -> Union[dict, str]",
+    "custom_speaker_selection": "# type: (ConversableAgent, GroupChat) -> Union[Agent, str, None]",
+    "custom_embedding_function": "# type: () -> Callable[..., Any]",
+    "custom_token_count_function": "# type: (str, str) -> int",
+    "custom_text_split_function": "# type: (str, int, str, bool, int) -> List[str]",
+}
+
+
+def parse_code_string(
+    code_string: str,
+) -> Tuple[Optional[str], Optional[ast.Module]]:
+    """Parse the code string.
+
+    Parameters
+    ----------
+    code_string : str
+        The code string.
+
+    Returns
+    -------
+    Tuple[Optional[str], Optional[ast.Module]]
+        If valid, None and the ast module.
+        If invalid, the error message and None.
+    """
+    # pylint: disable=broad-except
+    try:
+        tree = ast.parse(code_string)
+    except SyntaxError as e:
+        return f"SyntaxError: {e}, in \n{code_string}", None
+    except BaseException as e:  # pragma: no cover
+        return f"Invalid code: {e}, in \n{code_string}", None
+    return None, tree
+
+
+def check_function(
+    code_string: str,
+    function_name: WaldiezMethodName,
+    skip_type_hints: bool = False,
+) -> Tuple[bool, str]:
+    """Check the function.
+
+    Parameters
+    ----------
+    code_string : str
+        The code string.
+    function_name : WaldiezMethodName
+        The expected function name.
+    skip_type_hints : bool, optional
+        Whether to skip type hints in the function body, by default False.
+
+    Returns
+    -------
+    Tuple[bool, str]
+        If valid, True and the function body (only), no extra lines.
+        If invalid, False and the error message.
+    """
+    error, tree = parse_code_string(code_string)
+    if error is not None or tree is None:
+        return False, error or "Invalid code"
+    if function_name not in METHOD_ARGS:
+        return False, f"Invalid function name: {function_name}"
+    expected_method_args = METHOD_ARGS[function_name]
+    return _get_function_body(
+        tree,
+        code_string,
+        function_name,
+        expected_method_args,
+        skip_type_hints=skip_type_hints,
+    )
+
+
+def _get_function_body(
+    tree: ast.Module,
+    code_string: str,
+    function_name: WaldiezMethodName,
+    method_args: List[str],
+    skip_type_hints: bool = False,
+) -> Tuple[bool, str]:
+    """Get the function body.
+
+    Parameters
+    ----------
+    tree : ast.Module
+        The ast module.
+    code_string : str
+        The code string.
+    function_name : WaldiezMethodName
+        The expected function name.
+    method_args : List[str]
+        The expected method arguments.
+
+    Returns
+    -------
+    Tuple[bool, str]
+        If valid, True and the function body (only), no extra lines.
+        If invalid, False and the error message.
+    """
+    for node in ast.walk(tree):
+        if isinstance(node, ast.FunctionDef):
+            if node.name != function_name:
+                continue
+            if len(node.args.args) != len(method_args):
+                return (
+                    False,
+                    f"Invalid number of arguments in function {node.name}",
+                )
+            for arg, expected_arg in zip(node.args.args, method_args):
+                if arg.arg != expected_arg:
+                    return (
+                        False,
+                        f"Invalid argument name in function {node.name}",
+                    )
+            function_body_lines = code_string.splitlines()[
+                node.lineno - 1 : node.end_lineno
+            ]
+            function_body = "\n".join(function_body_lines[1:])
+            if not skip_type_hints:
+                # add type hints after the function definition
+                function_body = (
+                    f"    {METHOD_TYPE_HINTS[function_name]}\n{function_body}"
+                )
+            return True, function_body
+    error_msg = (
+        f"No function with name `{function_name}`"
+        f" and arguments `{method_args}` found"
+    )
+    return False, error_msg

waldiez/models/flow/__init__.py

@@ -0,0 +1,9 @@
+"""Waldiez flow related models."""
+
+from .flow import WaldiezFlow
+from .flow_data import WaldiezFlowData
+
+__all__ = [
+    "WaldiezFlow",
+    "WaldiezFlowData",
+]