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

waldiez/models/chat/chat_nested.py

@@ -1,19 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """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 pydantic import Field, field_validator, model_validator
 from typing_extensions import Annotated, Self
 
-from ..common import WaldiezBase, WaldiezMethodName
-from .chat_message import WaldiezChatMessage, validate_message_dict
+from ..common import WaldiezBase
+from .chat_message import WaldiezChatMessage
+
+NESTED_CHAT_MESSAGE = "nested_chat_message"
+NESTED_CHAT_REPLY = "nested_chat_reply"
+NESTED_CHAT_ARGS = ["recipient", "messages", "sender", "config"]
+NESTED_CHAT_TYPES = (
+    [
+        "ConversableAgent",
+        "List[Dict[str, Any]]",
+        "ConversableAgent",
+        "Dict[str, Any]",
+    ],
+    "Union[Dict[str, Any], str]",
+)
 
 
 class WaldiezChatNested(WaldiezBase):
@@ -27,13 +35,6 @@ class WaldiezChatNested(WaldiezBase):
         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(
@@ -66,17 +67,13 @@ class WaldiezChatNested(WaldiezBase):
 
     @field_validator("message", "reply", mode="before")
     @classmethod
-    def validate_message(
-        cls, value: Any, info: ValidationInfo
-    ) -> WaldiezChatMessage:
+    def validate_message(cls, value: Any) -> WaldiezChatMessage:
         """Validate the message.
 
         Parameters
         ----------
         value : Any
             The value.
-        info : ValidationInfo
-            The validation info.
 
         Returns
         -------
@@ -88,11 +85,6 @@ class WaldiezChatNested(WaldiezBase):
         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={}
@@ -102,17 +94,9 @@ class WaldiezChatNested(WaldiezBase):
                 type="string", use_carryover=False, content=value, context={}
             )
         if isinstance(value, dict):
-            return validate_message_dict(value, function_name=function_name)
+            return WaldiezChatMessage.model_validate(value)
         if isinstance(value, WaldiezChatMessage):
-            return validate_message_dict(
-                {
-                    "type": value.type,
-                    "use_carryover": False,
-                    "content": value.content,
-                    "context": value.context,
-                },
-                function_name=function_name,
-            )
+            return value
         raise ValueError(f"Invalid message type: {type(value)}")
 
     @model_validator(mode="after")
@@ -134,27 +118,23 @@ class WaldiezChatNested(WaldiezBase):
                 self._message_content = ""
             elif self.message.type == "string":
                 self._message_content = self.message.content
+            elif self.message.type == "method":
+                self._message_content = self.message.validate_method(
+                    function_name=NESTED_CHAT_MESSAGE,
+                    function_args=NESTED_CHAT_ARGS,
+                )
             else:
-                self._message_content = validate_message_dict(
-                    value={
-                        "type": "method",
-                        "content": self.message.content,
-                    },
-                    function_name="nested_chat_message",
-                    skip_definition=True,
-                ).content
+                self._message_content = self.message.content_body
         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
+            elif self.reply.type == "method":
+                self._reply_content = self.reply.validate_method(
+                    function_name=NESTED_CHAT_REPLY,
+                    function_args=NESTED_CHAT_ARGS,
+                )
             else:
-                self._reply_content = validate_message_dict(
-                    value={
-                        "type": "method",
-                        "content": self.reply.content,
-                    },
-                    function_name="nested_chat_reply",
-                    skip_definition=True,
-                ).content
+                self._reply_content = self.reply.content_body
         return self

waldiez/models/chat/chat_summary.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Waldiez chat summary options."""
 
 from typing import Any, Dict, Optional

waldiez/models/common/__init__.py

@@ -1,13 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """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,
+    generate_function,
+    get_function,
     parse_code_string,
 )
 
@@ -29,10 +30,9 @@ def now() -> str:
 
 __all__ = [
     "WaldiezBase",
-    "METHOD_ARGS",
-    "METHOD_TYPE_HINTS",
-    "WaldiezMethodName",
     "now",
     "check_function",
+    "get_function",
+    "generate_function",
     "parse_code_string",
 ]

waldiez/models/common/base.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Base class to inherit from."""
 
 from typing import Any, Dict
@@ -19,7 +21,8 @@ class WaldiezBase(BaseModel):
         alias_generator=to_camel,
         # allow passing either `skill_id` or `skillId`
         populate_by_name=True,
-        frozen=True,
+        # allow setting any attribute after initialization
+        frozen=False,
     )
 
     def model_dump(self, **kwargs: Any) -> Dict[str, Any]:

waldiez/models/common/method_utils.py

@@ -1,48 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Function related utilities."""
 
-# flake8: noqa E501
 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]",
-}
+from typing import List, Optional, Tuple
+
+import parso
+import parso.python
+import parso.tree
 
 
 def parse_code_string(
@@ -65,28 +30,27 @@ def parse_code_string(
     try:
         tree = ast.parse(code_string)
     except SyntaxError as e:
-        return f"SyntaxError: {e}, in \n{code_string}", None
+        return f"SyntaxError: {e}, in " + "\n" + f"{code_string}", None
     except BaseException as e:  # pragma: no cover
-        return f"Invalid code: {e}, in \n{code_string}", None
+        return f"Invalid code: {e}, in " + "\n" + f"{code_string}", None
     return None, tree
 
 
 def check_function(
     code_string: str,
-    function_name: WaldiezMethodName,
-    skip_type_hints: bool = False,
+    function_name: str,
+    function_args: List[str],
 ) -> 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.
-
+        The code string to check.
+    function_name : str
+        The expected method name.
+    function_args : List[str]
+        The expected method arguments.
     Returns
     -------
     Tuple[bool, str]
@@ -96,24 +60,19 @@ def check_function(
     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(
+    return _validate_function_body(
         tree,
         code_string,
         function_name,
-        expected_method_args,
-        skip_type_hints=skip_type_hints,
+        function_args,
     )
 
 
-def _get_function_body(
+def _validate_function_body(
     tree: ast.Module,
     code_string: str,
-    function_name: WaldiezMethodName,
-    method_args: List[str],
-    skip_type_hints: bool = False,
+    function_name: str,
+    function_args: List[str],
 ) -> Tuple[bool, str]:
     """Get the function body.
 
@@ -121,13 +80,12 @@ def _get_function_body(
     ----------
     tree : ast.Module
         The ast module.
-    code_string : str
-        The code string.
-    function_name : WaldiezMethodName
-        The expected function name.
-    method_args : List[str]
+    function_body : str
+        The function body.
+    function_name : str
+        The expected method name.
+    function_args : List[str]
         The expected method arguments.
-
     Returns
     -------
     Tuple[bool, str]
@@ -138,29 +96,151 @@ def _get_function_body(
         if isinstance(node, ast.FunctionDef):
             if node.name != function_name:
                 continue
-            if len(node.args.args) != len(method_args):
+            if len(node.args.args) != len(function_args):
                 return (
                     False,
-                    f"Invalid number of arguments in function {node.name}",
+                    (
+                        f"Invalid number of arguments, in function {node.name},"
+                        f" expected: {len(function_args)},"
+                        f" got: {len(node.args.args)} :("
+                    ),
                 )
-            for arg, expected_arg in zip(node.args.args, method_args):
+            for arg, expected_arg in zip(node.args.args, function_args):
                 if arg.arg != expected_arg:
                     return (
                         False,
-                        f"Invalid argument name in function {node.name}",
+                        (
+                            f"Invalid argument name: {arg.arg}"
+                            f" 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}"
-                )
+            if not node.body:
+                return False, "No body found in the function"
+            function_body = _get_function_body(code_string, node)
             return True, function_body
     error_msg = (
-        f"No function with name `{function_name}`"
-        f" and arguments `{method_args}` found"
+        f"No method with name `{function_name}`"
+        f" and arguments `{function_args}` found"
     )
     return False, error_msg
+
+
+def get_function(
+    code_string: str,
+    function_name: str,
+) -> str:
+    """Get the function signature and body.
+
+    Parameters
+    ----------
+    code_string : str
+        The code string.
+    function_name : str
+        The function name.
+
+    Returns
+    -------
+    str
+        The function signature and body.
+    """
+    try:
+        tree = parso.parse(code_string)  # type: ignore
+    except BaseException:  # pylint: disable=broad-except
+        return ""
+    for node in tree.iter_funcdefs():
+        if node.name.value == function_name:
+            return node.get_code()
+    return ""
+
+
+def _get_function_body(
+    code_string: str,
+    node: ast.FunctionDef,
+) -> str:
+    """Get the function body, including docstring and comments inside.
+
+    Parameters
+    ----------
+    code_string : str
+        The code string.
+    node : ast.FunctionDef
+        The function node.
+
+    Returns
+    -------
+    str
+        The function body.
+
+    Raises
+    ------
+    ValueError
+        If no body found in the function.
+    """
+    lines = code_string.splitlines()
+    signature_start_line = node.lineno - 1
+    body_start_line = node.body[0].lineno - 1
+    signature_end_line = signature_start_line
+    for i in range(signature_start_line, body_start_line):
+        if ")" in lines[i]:
+            signature_end_line = i
+            break
+    function_body_lines = lines[signature_end_line + 1 :]
+    last_line = function_body_lines[-1]
+    if not last_line.strip() and len(function_body_lines) > 1:
+        function_body_lines = function_body_lines[:-1]
+    function_body = "\n".join(function_body_lines)
+    while function_body.startswith("\n"):
+        function_body = function_body[1:]
+    while function_body.endswith("\n"):
+        function_body = function_body[:-1]
+    return function_body
+
+
+def generate_function(
+    function_name: str,
+    function_args: List[str],
+    function_types: Tuple[List[str], str],
+    function_body: str,
+    types_as_comments: bool = False,
+) -> str:
+    """Generate a function.
+
+    Parameters
+    ----------
+    function_name : str
+        The function name.
+    function_args : List[str]
+        The function arguments.
+    function_types : Tuple[List[str], str]
+        The function types.
+    function_body : str
+        The function body.
+    types_as_comments : bool, optional
+        Include the type hints as comments (or in the function signature)
+        (default is False).
+    Returns
+    -------
+    str
+        The generated function.
+    """
+    function_string = f"def {function_name}("
+    if not function_args:
+        function_string += ")"
+    else:
+        function_string += "\n"
+        for arg, arg_type in zip(function_args, function_types[0]):
+            if types_as_comments:
+                function_string += f"    {arg},  # type: {arg_type}" + "\n"
+            else:
+                function_string += f"    {arg}: {arg_type}," + "\n"
+        function_string += ")"
+    if types_as_comments:
+        function_string += ":\n"
+        function_string += "    # type: (...) -> " + function_types[1]
+    else:
+        function_string += " -> " + function_types[1] + ":"
+    function_string += "\n" if not function_body.startswith("\n") else ""
+    function_string += f"{function_body}"
+    if not function_string.endswith("\n"):
+        function_string += "\n"
+    return function_string

waldiez/models/flow/__init__.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Waldiez flow related models."""
 
 from .flow import WaldiezFlow