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

waldiez/models/flow/flow.py

@@ -1,29 +1,18 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Waldiez flow model."""
 
 import uuid
-from datetime import datetime, timezone
 from typing import List, Optional, Tuple
 
 from pydantic import Field, model_validator
 from typing_extensions import Annotated, Literal, Self
 
-from ..agents import WaldiezAgent
+from ..agents import WaldiezAgent, WaldiezSwarmAgent
 from ..chat import WaldiezChat
 from ..common import WaldiezBase, now
 from .flow_data import WaldiezFlowData
-
-
-def id_factory() -> str:
-    """Generate a unique ID.
-
-    Returns
-    -------
-    str
-        The unique ID.
-    """
-    now_td = datetime.now(timezone.utc)
-    now_str = now_td.strftime("%Y%m%d%H%M%S%f")
-    return f"{now_str}-{uuid.uuid4().hex}"
+from .utils import check_handoff_to_nested_chat, id_factory
 
 
 class WaldiezFlow(WaldiezBase):
@@ -138,6 +127,41 @@ class WaldiezFlow(WaldiezBase):
         List[Tuple[WaldiezChat, WaldiezAgent, WaldiezAgent]]
     ] = None
 
+    @property
+    def is_async(self) -> bool:
+        """Check if the flow is asynchronous.
+
+        Returns
+        -------
+        bool
+            True if the flow is asynchronous, False otherwise.
+        """
+        return self.data.is_async
+
+    @property
+    def has_shared_skill(self) -> bool:
+        """Check if the flow has a skill with shared variables (global).
+
+        Returns
+        -------
+        bool
+            True if the flow has shared skills, False otherwise.
+        """
+        return any(skill.is_shared for skill in self.data.skills)
+
+    @property
+    def is_swarm_flow(self) -> bool:
+        """Check if the flow is a swarm flow.
+
+        Returns
+        -------
+        bool
+            True if the flow is a swarm flow, False otherwise.
+        """
+        return any(
+            agent.agent_type == "swarm" for agent in self.data.agents.members
+        )
+
     @property
     def ordered_flow(
         self,
@@ -165,18 +189,9 @@ class WaldiezFlow(WaldiezBase):
         ValueError
             If the agent with the given ID is not found.
         """
-        for user in self.data.agents.users:
-            if user.id == agent_id:
-                return user
-        for assistant in self.data.agents.assistants:
-            if assistant.id == agent_id:
-                return assistant
-        for manager in self.data.agents.managers:
-            if manager.id == agent_id:
-                return manager
-        for rag_user in self.data.agents.rag_users:
-            if rag_user.id == agent_id:
-                return rag_user
+        for agent in self.data.agents.members:
+            if agent.id == agent_id:
+                return agent
         raise ValueError(f"Agent with ID {agent_id} not found.")
 
     def _get_flow_order(
@@ -187,10 +202,12 @@ class WaldiezFlow(WaldiezBase):
         # we only keep the ones with order >=0
         # and sort them by this property
         ordered_flow: List[Tuple[WaldiezChat, WaldiezAgent, WaldiezAgent]] = []
-        sorted_chats_by_order = sorted(
-            self.data.chats, key=lambda chat: chat.data.order
-        )
-        for chat in sorted_chats_by_order:
+        # if swarm, we only keep the first chat
+        if self.is_swarm_flow:
+            ordered_flow = self._get_swarm_flow()
+        if ordered_flow:
+            return ordered_flow
+        for chat in self.data.chats:
             if chat.data.order < 0:
                 continue
             source = self.get_agent_by_id(chat.source)
@@ -204,6 +221,51 @@ class WaldiezFlow(WaldiezBase):
                 ordered_flow.append((chat, source, target))
         return ordered_flow
 
+    def _get_swarm_flow(
+        self,
+    ) -> List[Tuple[WaldiezChat, WaldiezAgent, WaldiezAgent]]:
+        # valid "first" chat:
+        # - source is a user|rag_user and target is a swarm
+        # - source is a swarm and target is a  (and source.is_initial)
+        valid_chats: List[WaldiezChat] = []
+        for chat in self.data.chats:
+            target = self.get_agent_by_id(chat.target)
+            source = self.get_agent_by_id(chat.source)
+            if (
+                source.agent_type in ["user", "rag_user"]
+                and target.agent_type == "swarm"
+            ):
+                return [(chat, source, target)]
+            if source.agent_type == "swarm" and target.agent_type == "swarm":
+                if isinstance(source, WaldiezSwarmAgent) and source.is_initial:
+                    return [(chat, source, target)]
+                valid_chats.append(chat)
+        if not valid_chats:
+            return []
+        first_chat: Optional[WaldiezChat] = None
+        # first check the order
+        by_order = sorted(
+            filter(lambda edge: edge.data.order >= 0, valid_chats),
+            key=lambda edge: edge.data.order,
+        )
+        if not by_order:
+            # let's order by position
+            by_position = sorted(
+                valid_chats,
+                key=lambda chat: chat.data.position,
+            )
+            if by_position:
+                first_chat = by_position[0]
+            else:
+                first_chat = valid_chats[0]
+        else:
+            first_chat = by_order[0]
+        if first_chat:
+            source = self.get_agent_by_id(first_chat.source)
+            target = self.get_agent_by_id(first_chat.target)
+            return [(first_chat, source, target)]
+        return []
+
     def get_agent_connections(
         self, agent_id: str, all_chats: bool = True
     ) -> List[str]:
@@ -222,7 +284,7 @@ class WaldiezFlow(WaldiezBase):
         List[str]
             The list of agent ids that the agent with the given ID connects to.
         """
-        connections = []
+        connections: List[str] = []
         if all_chats:
             for chat in self.data.chats:
                 if chat.source == agent_id:
@@ -255,9 +317,81 @@ class WaldiezFlow(WaldiezBase):
         agent = self.get_agent_by_id(group_manager_id)
         if agent.agent_type != "manager":
             return []
-        connections = self.get_agent_connections(group_manager_id)
+        connections = self.get_agent_connections(
+            group_manager_id,
+            all_chats=True,
+        )
         return [self.get_agent_by_id(member_id) for member_id in connections]
 
+    def get_initial_swarm_agent(self) -> Optional[WaldiezAgent]:
+        """Get the initial swarm agent.
+
+        Returns
+        -------
+        Optional[WaldiezAgent]
+            The initial swarm agent if found, None otherwise.
+        """
+        fallback_agent = None
+        for chat in self.data.chats:
+            source_agent = self.get_agent_by_id(chat.source)
+            target_agent = self.get_agent_by_id(chat.target)
+            if (
+                target_agent.agent_type == "swarm"
+                and source_agent.agent_type != "swarm"
+            ):
+                return target_agent
+            if (
+                source_agent.agent_type == "swarm"
+                and target_agent.agent_type == "swarm"
+            ):
+                fallback_agent = source_agent
+                break
+        for swarm_agent in self.data.agents.swarm_agents:
+            if swarm_agent.is_initial:
+                return swarm_agent
+        return fallback_agent
+
+    def get_swarm_chat_members(
+        self,
+        initial_agent: WaldiezAgent,
+    ) -> Tuple[List[WaldiezAgent], Optional[WaldiezAgent]]:
+        """Get the swarm chat members.
+
+        Parameters
+        ----------
+        initial_agent : WaldiezAgent
+            The initial agent.
+
+        Returns
+        -------
+        Tuple[List[WaldiezAgent], Optional[WaldiezAgent]]
+            The list of swarm chat members and the user agent if any.
+        """
+        if initial_agent.agent_type != "swarm":
+            return [], None
+        members: List[WaldiezAgent] = [initial_agent]
+        user_agent: Optional[WaldiezAgent] = None
+        visited_agents = set()
+        visited_agents.add(initial_agent.id)
+        connections = self.get_agent_connections(
+            initial_agent.id,
+            all_chats=True,
+        )
+        while connections:
+            agent_id = connections.pop()
+            if agent_id in visited_agents:
+                continue
+            agent = self.get_agent_by_id(agent_id)
+            visited_agents.add(agent_id)
+            if agent.agent_type == "swarm":
+                members.append(agent)
+                connections.extend(
+                    self.get_agent_connections(agent_id, all_chats=True)
+                )
+            if agent.agent_type in ["user", "rag_user"] and not user_agent:
+                user_agent = agent
+        return members, user_agent
+
     def _validate_agent_connections(self) -> None:
         for agent in self.data.agents.members:
             if not any(
@@ -268,13 +402,6 @@ class WaldiezFlow(WaldiezBase):
                     f"Agent {agent.id} ({agent.name}) "
                     "does not connect to any other node."
                 )
-            # already covered above
-            # if agent.agent_type == "manager":
-            #     chat_member_ids = self.get_agent_connections(agent.id)
-            #     if not chat_member_ids:
-            #         raise ValueError(
-            #             f"Manager's {agent.id} group chat has no members."
-            #         )
 
     @model_validator(mode="after")
     def validate_flow(self) -> Self:
@@ -288,6 +415,8 @@ class WaldiezFlow(WaldiezBase):
         - all the managers have at least one member in the chat group
         - the ordered flow (chats with position >=0) is not empty
         - all agents' code execution config functions exist in the flow skills
+        - if swarm flow, there is at least one swarm agent
+        - if swarm flow, there is an initial swarm agent
 
         Returns
         -------
@@ -313,4 +442,11 @@ class WaldiezFlow(WaldiezBase):
             raise ValueError("Skill IDs must be unique.")
         self.data.agents.validate_flow(model_ids, skills_ids)
         self._validate_agent_connections()
+        if self.is_swarm_flow:
+            for swarm_agent in self.data.agents.swarm_agents:
+                check_handoff_to_nested_chat(
+                    swarm_agent,
+                    all_agents=list(self.data.agents.members),
+                    all_chats=self.data.chats,
+                )
         return self

waldiez/models/flow/flow_data.py

@@ -1,9 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Waldiez flow data."""
 
 from typing import Any, Dict, List
 
-from pydantic import Field
-from typing_extensions import Annotated
+from pydantic import Field, model_validator
+from typing_extensions import Annotated, Self
 
 from ..agents import WaldiezAgents
 from ..chat import WaldiezChat
@@ -36,6 +38,8 @@ class WaldiezFlowData(WaldiezBase):
         The skills of the flow. See `WaldiezSkill`.
     chats : List[WaldiezChat]
         The chats of the flow. See `WaldiezChat`.
+    is_async : bool
+        Whether the flow is asynchronous or not.
     """
 
     # the ones below (nodes,edges, viewport) we ignore
@@ -97,3 +101,60 @@ class WaldiezFlowData(WaldiezBase):
             default_factory=list,
         ),
     ]
+    is_async: Annotated[
+        bool,
+        Field(
+            False,
+            description="Whether the flow is asynchronous or not",
+            title="Is Async",
+        ),
+    ]
+
+    @model_validator(mode="after")
+    def validate_flow_chats(self) -> Self:
+        """Validate the flow chats.
+
+        Returns
+        -------
+        WaldiezFlowData
+            The flow data.
+
+        Raises
+        ------
+        ValueError
+            If there is a chat with a prerequisite that does not exist.
+        """
+        self.chats = sorted(self.chats, key=lambda x: x.order)
+        # in async, ag2 uses the "chat_id" field (and it must be an int):
+        # ```
+        #    prerequisites = []
+        #    for chat_info in chat_queue:
+        #        if "chat_id" not in chat_info:
+        #            raise ValueError(
+        #               "Each chat must have a unique id for "
+        #               "async multi-chat execution."
+        #            )
+        #     chat_id = chat_info["chat_id"]
+        #     pre_chats = chat_info.get("prerequisites", [])
+        #     for pre_chat_id in pre_chats:
+        #         if not isinstance(pre_chat_id, int):
+        #             raise ValueError("Prerequisite chat id is not int.")
+        #         prerequisites.append((chat_id, pre_chat_id))
+        #    return prerequisites
+        # ```
+        id_to_chat_id: Dict[str, int] = {}
+        for index, chat in enumerate(self.chats):
+            id_to_chat_id[chat.id] = index
+            chat.set_chat_id(index)
+        # also update the chat prerequisites
+        #  we have ids(str), not chat_ids(int)
+        for chat in self.chats:
+            chat_prerequisites = []
+            for chat_id in chat.data.prerequisites:
+                if chat_id not in id_to_chat_id:
+                    raise ValueError(
+                        f"Chat with id {chat_id} not found in the flow."
+                    )
+                chat_prerequisites.append(id_to_chat_id[chat_id])
+            chat.set_prerequisites(chat_prerequisites)
+        return self

waldiez/models/flow/utils.py

@@ -0,0 +1,172 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Helpers for the flow model."""
+
+import uuid
+from datetime import datetime, timezone
+from typing import List
+
+from ..agents import (
+    WaldiezAgent,
+    WaldiezAgentNestedChat,
+    WaldiezAgentNestedChatMessage,
+    WaldiezSwarmAgent,
+    WaldiezSwarmOnCondition,
+)
+from ..chat import WaldiezChat
+
+
+def id_factory() -> str:
+    """Generate a unique ID.
+
+    Returns
+    -------
+    str
+        The unique ID.
+    """
+    now_td = datetime.now(timezone.utc)
+    now_str = now_td.strftime("%Y%m%d%H%M%S%f")
+    return f"{now_str}-{uuid.uuid4().hex}"
+
+
+def check_handoff_to_nested_chat(
+    agent: WaldiezSwarmAgent,
+    all_agents: List[WaldiezAgent],
+    all_chats: List[WaldiezChat],
+) -> None:
+    """Check the handoffs to a nested chat.
+
+    If we have one and the agent does not have nested_chats,
+    we should generate them with the `handoff.target.id`
+    as the first (chat's) message.
+
+    Parameters
+    ----------
+    agent : WaldiezSwarmAgent
+        The swarm agent.
+    all_agents : List[WaldiezAgent]
+        All agents.
+    all_chats : List[WaldiezChat]
+        All chats.
+
+    Raises
+    ------
+    ValueError
+        If the agent has a handoff to a nested chat,
+        but no chat found with it as a source.
+    """
+    # pylint: disable=too-complex
+    for handoff in agent.handoffs:
+        if not isinstance(handoff, WaldiezSwarmOnCondition):
+            continue
+        is_nested_chat = handoff.target_type == "nested_chat"
+        if is_nested_chat:
+            # check if the agent already has nested_chats
+            # but only to get the order (and the first chat/message)
+            # either way, we must include all the connections that
+            # are swarm => non-swarm
+            # if we have the orders, ok, else we get them from the
+            # edges positions
+            all_connections = sorted(
+                get_nested_chat_swarm_connections(agent, all_agents, all_chats),
+                key=lambda x: x.data.order,
+            )
+            if not all_connections:
+                raise ValueError(
+                    f"Agent {agent.name} has a handoff to a nested chat, "
+                    "but no chat found with it as a source."
+                )
+            agent_nested_chats = agent.nested_chats
+            if not agent_nested_chats or not agent_nested_chats[0].messages:
+                agent.data.nested_chats = [
+                    WaldiezAgentNestedChat(
+                        triggered_by=[],
+                        messages=[
+                            WaldiezAgentNestedChatMessage(
+                                id=chat.id,
+                                is_reply=False,
+                            )
+                            for chat in all_connections
+                        ],
+                    )
+                ]
+                break
+            agent.data.nested_chats = merge_nested_chat_messages(
+                agent_nested_chats[0].messages, all_connections
+            )
+
+
+def get_nested_chat_swarm_connections(
+    agent: WaldiezSwarmAgent,
+    all_agents: List[WaldiezAgent],
+    all_chats: List[WaldiezChat],
+) -> List[WaldiezChat]:
+    """Get the nested chat connections.
+
+    Parameters
+    ----------
+    agent : WaldiezSwarmAgent
+        The swarm agent.
+    all_agents : Iterator[WaldiezAgent]
+        All agents.
+    all_chats : List[WaldiezChat]
+        All chats.
+
+    Returns
+    -------
+    List[WaldiezAgentNestedChat]
+        The nested chat connections.
+    """
+    connections_with_non_swarm_targets = []
+    for chat in all_chats:
+        if chat.source != agent.id:
+            continue
+        target_agent = next(
+            (a for a in all_agents if a.id == chat.target), None
+        )
+        if not target_agent or target_agent.agent_type != "swarm":
+            connections_with_non_swarm_targets.append(chat)
+    return connections_with_non_swarm_targets
+
+
+def merge_nested_chat_messages(
+    agent_nested_chat_messages: List[WaldiezAgentNestedChatMessage],
+    all_connections: List[WaldiezChat],
+) -> List[WaldiezAgentNestedChat]:
+    """Merge the nested chat messages.
+
+    Parameters
+    ----------
+    all_connections : List[WaldiezChat]
+        The connections.
+    agent_nested_chat_messages : List[WaldiezAgentNestedChatMessage]
+        The agent's nested chat messages.
+
+    Returns
+    -------
+    List[WaldiezAgentNestedChat]
+        The merged nested chat with all the messages.
+    """
+    nested_chat = WaldiezAgentNestedChat(triggered_by=[], messages=[])
+    chat_ids_added: List[str] = []
+    for message in agent_nested_chat_messages:
+        chat = next((c for c in all_connections if c.id == message.id), None)
+        if chat and chat.id not in chat_ids_added:
+            nested_chat.messages.append(
+                WaldiezAgentNestedChatMessage(
+                    id=chat.id,
+                    is_reply=False,
+                )
+            )
+            chat_ids_added.append(chat.id)
+    for chat in all_connections:
+        if chat.id not in chat_ids_added:
+            nested_chat.messages.append(
+                WaldiezAgentNestedChatMessage(
+                    id=chat.id,
+                    is_reply=False,
+                )
+            )
+            chat_ids_added.append(chat.id)
+    nested_chat.messages.sort(key=lambda x: chat_ids_added.index(x.id))
+    return [nested_chat]

waldiez/models/model/__init__.py

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

waldiez/models/model/model.py

@@ -1,4 +1,6 @@
-"""Waldiez model(llm) model."""
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Waldiez model model."""
 
 import os
 from typing import Any, Dict, List, Optional
@@ -102,6 +104,27 @@ class WaldiezModel(WaldiezBase):
         Field(..., title="Data", description="The data of the model."),
     ]
 
+    @property
+    def api_key_env_key(self) -> str:
+        """Get the model's api key environment key to check.
+
+        - openai: 'OPENAI_API_KEY',
+        - azure: 'AZURE_API_KEY',
+        - google: 'GOOGLE_GEMINI_API_KEY',
+        - anthropic: 'ANTHROPIC_API_KEY',
+        - mistral: 'MISTRAL_API_KEY',
+        - groq: 'GROQ_API_KEY',
+        - together: 'TOGETHER_API_KEY',
+        - nim: 'NIM_API_KEY',
+        - other: 'OPENAI_API_KEY'
+        """
+        env_key = "OPENAI_API_KEY"
+        if self.data.api_type == "google":
+            env_key = "GOOGLE_GEMINI_API_KEY"
+        elif self.data.api_type not in ["openai", "other"]:
+            env_key = f"{self.data.api_type.upper()}_API_KEY"
+        return env_key
+
     @property
     def api_key(self) -> str:
         """Get the model's api key.
@@ -120,11 +143,7 @@ class WaldiezModel(WaldiezBase):
         """
         if self.data.api_key:
             return self.data.api_key
-        env_key = "OPENAI_API_KEY"
-        if self.data.api_type == "google":
-            env_key = "GOOGLE_GEMINI_API_KEY"
-        elif self.data.api_type not in ["openai", "other"]:
-            env_key = f"{self.data.api_type.upper()}_API_KEY"
+        env_key = self.api_key_env_key
         api_key = os.environ.get(env_key, "")
         return api_key
 

waldiez/models/model/model_data.py

@@ -1,6 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+# flake8: noqa E501
 """Waldiez Model Data."""
 
-# flake8: noqa E501
 from typing import Dict, Optional
 
 from pydantic import Field

waldiez/models/skill/__init__.py

@@ -1,9 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Waldiez Skill related models."""
 
-from .skill import WaldiezSkill
+from .skill import SHARED_SKILL_NAME, WaldiezSkill
 from .skill_data import WaldiezSkillData
 
 __all__ = [
+    "SHARED_SKILL_NAME",
     "WaldiezSkill",
     "WaldiezSkillData",
 ]

waldiez/models/skill/skill.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Waldiez Skill model."""
 
 from typing import Dict, List
@@ -5,9 +7,11 @@ from typing import Dict, List
 from pydantic import Field, model_validator
 from typing_extensions import Annotated, Literal, Self
 
-from ..common import WaldiezBase, now, parse_code_string
+from ..common import WaldiezBase, get_function, now, parse_code_string
 from .skill_data import WaldiezSkillData
 
+SHARED_SKILL_NAME = "waldiez_shared"
+
 
 class WaldiezSkill(WaldiezBase):
     """Waldiez Skill.
@@ -93,6 +97,30 @@ class WaldiezSkill(WaldiezBase):
         ),
     ]
 
+    @property
+    def is_shared(self) -> bool:
+        """Check if the skill is shared.
+
+        Returns
+        -------
+        bool
+            True if the skill is shared, False otherwise.
+        """
+        return self.name == SHARED_SKILL_NAME
+
+    def get_content(self) -> str:
+        """Get the content of the skill.
+
+        Returns
+        -------
+        str
+            The content of the skill.
+        """
+        if self.name == SHARED_SKILL_NAME:
+            # the whole content (globals)
+            return self.data.content
+        return get_function(self.data.content, self.name)
+
     @model_validator(mode="after")
     def validate_data(self) -> Self:
         """Validate the data.
@@ -109,7 +137,7 @@ class WaldiezSkill(WaldiezBase):
             If the skill content is invalid.
         """
         search = f"def {self.name}("
-        if search not in self.data.content:
+        if self.name != SHARED_SKILL_NAME and search not in self.data.content:
             raise ValueError(
                 f"The skill name '{self.name}' is not in the content."
             )

waldiez/models/skill/skill_data.py

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