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

waldiez/exporting/skills/utils.py

@@ -0,0 +1,281 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Skills/tools related string generation functions.
+
+Functions
+---------
+get_agent_skill_registration
+    Get an agent's skill registration string.
+export_skills
+    Get the skills content and secrets.
+"""
+
+from pathlib import Path
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
+from waldiez.models import WaldiezAgent, WaldiezSkill
+
+
+def get_agent_skill_registration(
+    caller_name: str,
+    executor_name: str,
+    skill_name: str,
+    skill_description: str,
+    string_escape: Callable[[str], str],
+) -> str:
+    """Get the agent skill string and secrets.
+
+    Parameters
+    ----------
+    caller_name : str
+        The name of the caller (agent).
+    executor_name : str
+        The name of the executor (agent).
+    skill_name : str
+        The name of the skill.
+    skill_description : str
+        The skill description.
+    string_escape : Callable[[str], str]
+        The string escape function.
+    Returns
+    -------
+    str
+        The agent skill string.
+
+    Example
+    -------
+    ```python
+    >>> get_agent_skill_registration(
+    ...     caller_name="agent1",
+    ...     executor_name="agent2",
+    ...     skill_name="skill1",
+    ...     skill_description="A skill that does something.",
+    ...     string_escape=lambda x: x.replace('"', '\\"').replace("\\n", "\\n"),
+    ... )
+    register_function(
+        skill1,
+        caller=agent1,
+        executor=agent2,
+        name="skill1",
+        description="A skill that does something.",
+    )
+    ```
+    """
+    skill_description = string_escape(skill_description)
+    content = f"""register_function(
+    {skill_name},
+    caller={caller_name},
+    executor={executor_name},
+    name="{skill_name}",
+    description="{skill_description}",
+)
+"""
+    return content
+
+
+def _write_skill_secrets(
+    flow_name: str,
+    skill: WaldiezSkill,
+    skill_name: str,
+    output_dir: Optional[Union[str, Path]],
+) -> None:
+    """Write the skill secrets to a file.
+
+    Parameters
+    ----------
+    skill : WaldiezSkill
+        The skill.
+    skill_name : str
+        The name of the skill.
+    output_dir : Optional[Union[str, Path]]
+        The output directory to save the secrets to.
+    """
+    if not skill.secrets or not output_dir:
+        return
+    if not isinstance(output_dir, Path):
+        output_dir = Path(output_dir)
+    secrets_file = output_dir / f"{flow_name}_{skill_name}_secrets.py"
+    first_line = f'"""Secrets for the skill: {skill_name}."""' + "\n"
+    with secrets_file.open("w", encoding="utf-8", newline="\n") as f:
+        f.write(first_line)
+        f.write("import os\n\n")
+        for key, value in skill.secrets.items():
+            f.write(f'os.environ["{key}"] = "{value}"' + "\n")
+
+
+def export_skills(
+    flow_name: str,
+    skills: List[WaldiezSkill],
+    skill_names: Dict[str, str],
+    output_dir: Optional[Union[str, Path]] = None,
+) -> Tuple[List[str], List[Tuple[str, str]], str]:
+    """Get the skills' contents and secrets.
+
+    If `output_dir` is provided, the contents are saved to that directory.
+
+    Parameters
+    ----------
+    flow_name : str
+        The name of the flow.
+    skills : List[WaldiezSkill]
+        The skills.
+    skill_names : Dict[str, str]
+        The skill names.
+    output_dir : Optional[Union[str, Path]]
+        The output directory to save the skills to.
+
+    Returns
+    -------
+    Tuple[Set[str], Set[Tuple[str, str]], str]
+        - The skill imports to use in the main file.
+        - The skill secrets to set as environment variables.
+        - The skills contents.
+    """
+    skill_imports: List[str] = []
+    skill_secrets: List[Tuple[str, str]] = []
+    skill_contents: str = ""
+    # if the skill.is_shared,
+    # its contents must be first (before the other skills)
+    shared_skill_contents = ""
+    for skill in skills:
+        skill_imports.append(get_skill_imports(flow_name, skill))
+        for key, value in skill.secrets.items():
+            skill_secrets.append((key, value))
+        _write_skill_secrets(
+            flow_name=flow_name,
+            skill=skill,
+            skill_name=skill_names[skill.id],
+            output_dir=output_dir,
+        )
+        skill_content = skill.get_content()
+        if not skill_content:
+            continue
+        if skill.is_shared:
+            shared_skill_contents += skill_content + "\n\n"
+        else:
+            skill_contents += skill_content + "\n\n"
+    skill_contents = shared_skill_contents + skill_contents
+    return (
+        skill_imports,
+        skill_secrets,
+        skill_contents.replace("\n\n\n", "\n\n"),
+    )
+
+
+def get_skill_imports(flow_name: str, skill: WaldiezSkill) -> str:
+    """Get the skill imports string.
+
+    Parameters
+    ----------
+    flow_name : str
+        The name of the flow.
+    skill : WaldiezSkill
+        The skill.
+    Returns
+    -------
+    str
+        The skill imports string.
+    """
+    if not skill.secrets:
+        return ""
+        # fmt: on
+    module_name = f"{flow_name}_{skill.name}"
+    ignore_noqa = "  # type: ignore # noqa"
+    return f"import {module_name}_secrets{ignore_noqa}" + "\n"
+
+
+def get_agent_skill_registrations(
+    agent: WaldiezAgent,
+    agent_names: Dict[str, str],
+    all_skills: List[WaldiezSkill],
+    skill_names: Dict[str, str],
+    string_escape: Callable[[str], str],
+) -> str:
+    """Get the agent skill registrations.
+
+    example output:
+
+    ```python
+    >>> string_escape = lambda x: x.replace('"', '\\"').replace("\\n", "\\n")
+    >>> agent = WaldiezAgent(
+    ...     id="wa-1",
+    ...     name="agent1",
+    ...     description="An agent that does something.",
+    ...     ...,
+    ...     skills=[
+    ...         WaldiezSkillLink(id="ws-1", executor_id="wa-2", ...),
+    ...         WaldiezSkillLink(id="ws-2", executor_id="wa-3", ...),
+    ...     ],
+    ... )
+    >>> agent_names = {"wa-1": "agent1", "wa-2": "agent2", "wa-3": "agent3"}
+    >>> all_skills = [
+    ...     WaldiezSkill(id="ws-1", ...),
+    ...     WaldiezSkill(id="ws-2", ...),
+    ...     WaldiezSkill(id="ws-3", ...),
+    ... ]
+    >>> skill_names = {"ws-1": "skill1", "ws-2": "skill2", "ws-3": "skill3"}
+    >>> get_agent_skill_registrations(
+    ...     agent=agent,
+    ...     agent_names=agent_names,
+    ...     all_skills=all_skills,
+    ...     skill_names=skill_names,
+    ...     string_escape=string_escape,
+    ... )
+
+    register_function(
+        skill1,
+        caller=agent1,
+        executor=agent2,
+        name="skill1",
+        description="A skill that does something.",
+    )
+    register_function(
+        skill2,
+        caller=agent1,
+        executor=agent3,
+        name="skill2",
+        description="A skill that does something.",
+    )
+    ```
+
+    Parameters
+    ----------
+    agent : WaldiezAgent
+        The agent.
+    agent_names : Dict[str, str]
+        A mapping of agent id to agent name.
+    all_skills : List[WaldiezSkill]
+        All the skills in the flow.
+    skill_names : Dict[str, str]
+        A mapping of skill id to skill name.
+    string_escape : Callable[[str], str]
+        The string escape function.
+    Returns
+    -------
+    str
+        The agent skill registrations.
+    """
+    if not agent.data.skills or not all_skills:
+        return ""
+    content = "\n"
+    for linked_skill in agent.data.skills:
+        waldiez_skill = next(
+            skill for skill in all_skills if skill.id == linked_skill.id
+        )
+        skill_name = skill_names[linked_skill.id]
+        skill_description = (
+            waldiez_skill.description or f"Description of {skill_name}"
+        )
+        caller_name = agent_names[agent.id]
+        executor_name = agent_names[linked_skill.executor_id]
+        content += (
+            get_agent_skill_registration(
+                caller_name=caller_name,
+                executor_name=executor_name,
+                skill_name=skill_name,
+                skill_description=skill_description,
+                string_escape=string_escape,
+            )
+            + "\n"
+        )
+    return content.replace("\n\n\n", "\n\n")

waldiez/models/__init__.py

@@ -1,10 +1,13 @@
-"""Waldiez models package.
-
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Waldiez models.
 - Agents (Users, Assistants, Group Managers, etc.).
 - Chat (Messages, Summaries, etc.).
 - Model (LLM config, API type, etc.).
 - Skill (Skills/Tools to be registered).
 - Flow (Flow of the conversation).
+- Methods (Method names, arguments, hints, etc.).
+- Waldiez (Main class to hold the flow).
 """
 
 from .agents import (
@@ -34,6 +37,15 @@ from .agents import (
     WaldiezRagUserTask,
     WaldiezRagUserVectorDb,
     WaldiezRagUserVectorDbConfig,
+    WaldiezSwarmAfterWork,
+    WaldiezSwarmAfterWorkOption,
+    WaldiezSwarmAfterWorkRecipientType,
+    WaldiezSwarmAgent,
+    WaldiezSwarmAgentData,
+    WaldiezSwarmOnCondition,
+    WaldiezSwarmOnConditionAvailable,
+    WaldiezSwarmOnConditionTarget,
+    WaldiezSwarmUpdateSystemMessage,
     WaldiezUserProxy,
     WaldiezUserProxyData,
 )
@@ -45,7 +57,6 @@ from .chat import (
     WaldiezChatSummary,
     WaldiezChatSummaryMethod,
 )
-from .common import METHOD_ARGS, METHOD_TYPE_HINTS, WaldiezMethodName
 from .flow import WaldiezFlow, WaldiezFlowData
 from .model import (
     WaldiezModel,
@@ -53,14 +64,12 @@ from .model import (
     WaldiezModelData,
     WaldiezModelPrice,
 )
-from .skill import WaldiezSkill, WaldiezSkillData
+from .skill import SHARED_SKILL_NAME, WaldiezSkill, WaldiezSkillData
 from .waldiez import Waldiez
 
 # pylint: disable=duplicate-code
 __all__ = [
-    "METHOD_ARGS",
-    "METHOD_TYPE_HINTS",
-    "WaldiezMethodName",
+    "SHARED_SKILL_NAME",
     "Waldiez",
     "WaldiezAgent",
     "WaldiezAgentCodeExecutionConfig",
@@ -104,4 +113,13 @@ __all__ = [
     "WaldiezRagUserVectorDb",
     "WaldiezRagUserVectorDbConfig",
     "WaldiezRagUserModels",
+    "WaldiezSwarmAgent",
+    "WaldiezSwarmAfterWork",
+    "WaldiezSwarmAfterWorkOption",
+    "WaldiezSwarmAfterWorkRecipientType",
+    "WaldiezSwarmAgentData",
+    "WaldiezSwarmOnCondition",
+    "WaldiezSwarmOnConditionTarget",
+    "WaldiezSwarmOnConditionAvailable",
+    "WaldiezSwarmUpdateSystemMessage",
 ]

waldiez/models/agents/__init__.py

@@ -1,6 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Agent models."""
 
 from .agent import (
+    IS_TERMINATION_MESSAGE,
+    IS_TERMINATION_MESSAGE_ARGS,
+    IS_TERMINATION_MESSAGE_TYPES,
     WaldiezAgent,
     WaldiezAgentCodeExecutionConfig,
     WaldiezAgentData,
@@ -14,6 +19,9 @@ from .agent import (
 from .agents import WaldiezAgents
 from .assistant import WaldiezAssistant, WaldiezAssistantData
 from .group_manager import (
+    CUSTOM_SPEAKER_SELECTION,
+    CUSTOM_SPEAKER_SELECTION_ARGS,
+    CUSTOM_SPEAKER_SELECTION_TYPES,
     WaldiezGroupManager,
     WaldiezGroupManagerData,
     WaldiezGroupManagerSpeakers,
@@ -22,6 +30,15 @@ from .group_manager import (
     WaldiezGroupManagerSpeakersTransitionsType,
 )
 from .rag_user import (
+    CUSTOM_EMBEDDING_FUNCTION,
+    CUSTOM_EMBEDDING_FUNCTION_ARGS,
+    CUSTOM_EMBEDDING_FUNCTION_TYPES,
+    CUSTOM_TEXT_SPLIT_FUNCTION,
+    CUSTOM_TEXT_SPLIT_FUNCTION_ARGS,
+    CUSTOM_TEXT_SPLIT_FUNCTION_TYPES,
+    CUSTOM_TOKEN_COUNT_FUNCTION,
+    CUSTOM_TOKEN_COUNT_FUNCTION_ARGS,
+    CUSTOM_TOKEN_COUNT_FUNCTION_TYPES,
     WaldiezRagUser,
     WaldiezRagUserChunkMode,
     WaldiezRagUserData,
@@ -31,9 +48,53 @@ from .rag_user import (
     WaldiezRagUserVectorDb,
     WaldiezRagUserVectorDbConfig,
 )
+from .swarm_agent import (
+    CUSTOM_AFTER_WORK,
+    CUSTOM_AFTER_WORK_ARGS,
+    CUSTOM_AFTER_WORK_TYPES,
+    CUSTOM_ON_CONDITION_AVAILABLE,
+    CUSTOM_ON_CONDITION_AVAILABLE_ARGS,
+    CUSTOM_ON_CONDITION_AVAILABLE_TYPES,
+    CUSTOM_UPDATE_SYSTEM_MESSAGE,
+    CUSTOM_UPDATE_SYSTEM_MESSAGE_ARGS,
+    CUSTOM_UPDATE_SYSTEM_MESSAGE_TYPES,
+    WaldiezSwarmAfterWork,
+    WaldiezSwarmAfterWorkOption,
+    WaldiezSwarmAfterWorkRecipientType,
+    WaldiezSwarmAgent,
+    WaldiezSwarmAgentData,
+    WaldiezSwarmOnCondition,
+    WaldiezSwarmOnConditionAvailable,
+    WaldiezSwarmOnConditionTarget,
+    WaldiezSwarmUpdateSystemMessage,
+)
 from .user_proxy import WaldiezUserProxy, WaldiezUserProxyData
 
 __all__ = [
+    "IS_TERMINATION_MESSAGE",
+    "IS_TERMINATION_MESSAGE_ARGS",
+    "IS_TERMINATION_MESSAGE_TYPES",
+    "CUSTOM_AFTER_WORK",
+    "CUSTOM_AFTER_WORK_ARGS",
+    "CUSTOM_AFTER_WORK_TYPES",
+    "CUSTOM_EMBEDDING_FUNCTION",
+    "CUSTOM_EMBEDDING_FUNCTION_ARGS",
+    "CUSTOM_EMBEDDING_FUNCTION_TYPES",
+    "CUSTOM_ON_CONDITION_AVAILABLE",
+    "CUSTOM_ON_CONDITION_AVAILABLE_ARGS",
+    "CUSTOM_ON_CONDITION_AVAILABLE_TYPES",
+    "CUSTOM_SPEAKER_SELECTION",
+    "CUSTOM_SPEAKER_SELECTION_ARGS",
+    "CUSTOM_SPEAKER_SELECTION_TYPES",
+    "CUSTOM_TEXT_SPLIT_FUNCTION",
+    "CUSTOM_TEXT_SPLIT_FUNCTION_ARGS",
+    "CUSTOM_TEXT_SPLIT_FUNCTION_TYPES",
+    "CUSTOM_TOKEN_COUNT_FUNCTION",
+    "CUSTOM_TOKEN_COUNT_FUNCTION_ARGS",
+    "CUSTOM_TOKEN_COUNT_FUNCTION_TYPES",
+    "CUSTOM_UPDATE_SYSTEM_MESSAGE",
+    "CUSTOM_UPDATE_SYSTEM_MESSAGE_ARGS",
+    "CUSTOM_UPDATE_SYSTEM_MESSAGE_TYPES",
     "WaldiezAgent",
     "WaldiezAgentType",
     "WaldiezAgents",
@@ -62,4 +123,13 @@ __all__ = [
     "WaldiezRagUserChunkMode",
     "WaldiezRagUserVectorDb",
     "WaldiezRagUserVectorDbConfig",
+    "WaldiezSwarmAgent",
+    "WaldiezSwarmAgentData",
+    "WaldiezSwarmAfterWork",
+    "WaldiezSwarmAfterWorkOption",
+    "WaldiezSwarmAfterWorkRecipientType",
+    "WaldiezSwarmOnCondition",
+    "WaldiezSwarmOnConditionTarget",
+    "WaldiezSwarmOnConditionAvailable",
+    "WaldiezSwarmUpdateSystemMessage",
 ]

waldiez/models/agents/agent/__init__.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Base agent class to be inherited by all other agents."""
 
 from .agent import WaldiezAgent, WaldiezAgentType
@@ -6,9 +8,17 @@ from .code_execution import WaldiezAgentCodeExecutionConfig
 from .linked_skill import WaldiezAgentLinkedSkill
 from .nested_chat import WaldiezAgentNestedChat, WaldiezAgentNestedChatMessage
 from .teachability import WaldiezAgentTeachability
-from .termination_message import WaldiezAgentTerminationMessage
+from .termination_message import (
+    IS_TERMINATION_MESSAGE,
+    IS_TERMINATION_MESSAGE_ARGS,
+    IS_TERMINATION_MESSAGE_TYPES,
+    WaldiezAgentTerminationMessage,
+)
 
 __all__ = [
+    "IS_TERMINATION_MESSAGE",
+    "IS_TERMINATION_MESSAGE_ARGS",
+    "IS_TERMINATION_MESSAGE_TYPES",
     "WaldiezAgent",
     "WaldiezAgentCodeExecutionConfig",
     "WaldiezAgentData",

waldiez/models/agents/agent/agent.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Base agent class to be inherited by all agents."""
 
 from typing import List
@@ -9,7 +11,7 @@ from ...common import WaldiezBase, now
 from .agent_data import WaldiezAgentData
 from .code_execution import WaldiezAgentCodeExecutionConfig
 
-WaldiezAgentType = Literal["user", "assistant", "manager", "rag_user"]
+WaldiezAgentType = Literal["user", "assistant", "manager", "rag_user", "swarm"]
 
 
 class WaldiezAgent(WaldiezBase):
@@ -21,7 +23,7 @@ class WaldiezAgent(WaldiezBase):
         The ID of the agent.
     type : Literal["agent"]
         The type of the "node" in a graph: "agent"
-    agent_type : Literal["user", "assistant", "manager", "rag_user"]
+    agent_type : Literal["user", "assistant", "manager", "rag_user", "swarm"]
         The type of the agent
     name: str
         The name of the agent.
@@ -59,11 +61,14 @@ class WaldiezAgent(WaldiezBase):
         ),
     ]
     agent_type: Annotated[
-        Literal["user", "assistant", "manager", "rag_user"],
+        Literal["user", "assistant", "manager", "rag_user", "swarm"],
         Field(
             ...,
             title="Agent type",
-            description="The type of the agent: user, assistant, group manager",
+            description=(
+                "The type of the agent: user, assistant, group manager, "
+                "rag_user, or swarm"
+            ),
         ),
     ]
     name: Annotated[

waldiez/models/agents/agent/agent_data.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Common data structures for agents."""
 
 from typing import List, Optional, Union
@@ -117,7 +119,7 @@ class WaldiezAgentData(WaldiezBase):
         ),
     ]
     teachability: Annotated[
-        WaldiezAgentTeachability,
+        Optional[WaldiezAgentTeachability],
         Field(
             title="Teachability",
             description="The agent teachability configuration.",

waldiez/models/agents/agent/code_execution.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Waldiez Agent Code Execution Configuration."""
 
 from typing import List, Optional

waldiez/models/agents/agent/linked_skill.py

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

waldiez/models/agents/agent/nested_chat.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Waldiez Agent Nested Chat."""
 
 from typing import List

waldiez/models/agents/agent/teachability.py

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

waldiez/models/agents/agent/termination_message.py

@@ -1,12 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Waldiez Agent Termination Message Check."""
 
-from typing import List, Optional
+from typing import List, Optional, Tuple
 
-from pydantic import ConfigDict, Field, model_validator
-from pydantic.alias_generators import to_camel
+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, generate_function
+
+IS_TERMINATION_MESSAGE = "is_termination_message"
+IS_TERMINATION_MESSAGE_ARGS = ["message"]
+IS_TERMINATION_MESSAGE_TYPES = (
+    ["Dict[str, Any]"],
+    "bool",
+)
 
 
 class WaldiezAgentTerminationMessage(WaldiezBase):
@@ -34,13 +42,6 @@ class WaldiezAgentTerminationMessage(WaldiezBase):
         Validate the termination message configuration.
     """
 
-    model_config = ConfigDict(
-        extra="forbid",
-        alias_generator=to_camel,
-        populate_by_name=True,
-        frozen=False,
-    )
-
     type: Annotated[
         Literal["none", "keyword", "method"],
         Field(
@@ -90,6 +91,40 @@ class WaldiezAgentTerminationMessage(WaldiezBase):
 
     _string: str = "None"
 
+    def get_termination_function(
+        self,
+        name_prefix: Optional[str] = None,
+        name_suffix: Optional[str] = None,
+    ) -> Tuple[str, str]:
+        """Get the termination function.
+
+        Parameters
+        ----------
+        name_prefix : str
+            The function name prefix.
+        name_suffix : str
+            The function name suffix.
+
+        Returns
+        -------
+        Tuple[str, str]
+            The termination function and the function name.
+        """
+        function_name = "is_termination_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=IS_TERMINATION_MESSAGE_ARGS,
+                function_types=IS_TERMINATION_MESSAGE_TYPES,
+                function_body=self.string,
+            ),
+            function_name,
+        )
+
     @property
     def string(self) -> str:
         """Get the value of the termination message.
@@ -112,9 +147,10 @@ class WaldiezAgentTerminationMessage(WaldiezBase):
             raise ValueError(
                 "Method content is required for method termination."
             )
-        expected_function_name: WaldiezMethodName = "is_termination_message"
         valid, error_or_content = check_function(
-            self.method_content, expected_function_name
+            self.method_content,
+            function_name=IS_TERMINATION_MESSAGE,
+            function_args=IS_TERMINATION_MESSAGE_ARGS,
         )
         if not valid:
             raise ValueError(error_or_content)