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

waldiez/exporting/agent/agent_exporter.py

@@ -0,0 +1,279 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+# pylint: disable=too-many-return-statements,too-many-instance-attributes
+"""Export agents."""
+
+from pathlib import Path
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
+from waldiez.models import WaldiezAgent, WaldiezChat
+
+from ..base import (
+    AgentPosition,
+    AgentPositions,
+    BaseExporter,
+    ExporterMixin,
+    ExporterReturnType,
+    ExportPosition,
+    ImportPosition,
+)
+from .utils import (
+    get_agent_class_name,
+    get_agent_code_execution_config,
+    get_agent_imports,
+    get_group_manager_extras,
+    get_is_termination_message,
+    get_rag_user_extras,
+    get_swarm_extras,
+)
+
+
+class AgentExporter(BaseExporter, ExporterMixin):
+    """Agents exporter."""
+
+    def __init__(
+        self,
+        agent: WaldiezAgent,
+        agent_names: Dict[str, str],
+        model_names: Dict[str, str],
+        skill_names: Dict[str, str],
+        chats: Tuple[List[WaldiezChat], Dict[str, str]],
+        is_async: bool,
+        group_chat_members: List[WaldiezAgent],
+        for_notebook: bool,
+        arguments_resolver: Callable[[WaldiezAgent], List[str]],
+        output_dir: Optional[Union[str, Path]] = None,
+    ) -> None:
+        """Initialize the agents exporter.
+
+        Parameters
+        ----------
+        agent : WaldiezAgent
+            The agent to export.
+        agent_names : Dict[str, str]
+            The agent ids to names mapping.
+        model_names : Dict[str, str]
+            The model ids to names mapping.
+        skill_names : Dict[str, str]
+            The skill ids to names mapping.
+        all_chats : List[WaldiezChat]
+            All the chats in the flow.
+        chat_names : Dict[str, str]
+            The chat ids to names mapping.
+        is_async : bool
+            Whether the whole flow is async.
+        for_notebook : bool
+            Whether the exporter is for a notebook.
+        output_dir : Optional[Union[str, Path]], optional
+            The output directory, by default None
+        """
+        self.for_notebook = for_notebook
+        self.agent = agent
+        self.agent_names = agent_names
+        if output_dir is not None and not isinstance(output_dir, Path):
+            output_dir = Path(output_dir)
+        self.output_dir = output_dir
+        self.model_names = model_names
+        self.skill_names = skill_names
+        self.arguments_resolver = arguments_resolver
+        self.group_chat_members = group_chat_members
+        self.chats = chats
+        self.is_async = is_async
+        self._agent_name = agent_names[agent.id]
+        self._agent_class = get_agent_class_name(self.agent)
+        # content, argument, import
+        self._code_execution = get_agent_code_execution_config(
+            agent=self.agent,
+            agent_name=self._agent_name,
+            skill_names=self.skill_names,
+        )
+        # before_rag, retrieve_arg, rag_imports
+        self._rag = get_rag_user_extras(
+            agent=self.agent,
+            agent_name=self._agent_name,
+            model_names=self.model_names,
+            path_resolver=self.path_resolver,
+            serializer=self.serializer,
+        )
+        # before_manager, group_chat_arg
+        self._group_chat = get_group_manager_extras(
+            agent=self.agent,
+            agent_names=self.agent_names,
+            group_chat_members=self.group_chat_members,
+            serializer=self.serializer,
+        )
+        # before_agent, extra args, handoff_registrations
+        self._swarm = get_swarm_extras(
+            agent=self.agent,
+            agent_names=self.agent_names,
+            skill_names=self.skill_names,
+            chats=self.chats,
+            is_async=self.is_async,
+            serializer=self.serializer,
+            string_escape=self.string_escape,
+        )
+        # before_agent, termination_arg
+        self._termination = get_is_termination_message(
+            agent=self.agent, agent_name=self._agent_name
+        )
+
+    def get_imports(self) -> Optional[List[Tuple[str, ImportPosition]]]:
+        """Get the imports.
+
+        Returns
+        -------
+        Optional[Tuple[str, ImportPosition]]
+            The imports.
+        """
+        position = ImportPosition.THIRD_PARTY
+        # default imports based on the agent class.
+        agent_imports = get_agent_imports(self._agent_class)
+        # if code execution is enabled, update the imports.
+        if self._code_execution[2]:
+            agent_imports.add(self._code_execution[2])
+        # if RAG is enabled, update the imports.
+        if self._rag[2]:
+            agent_imports.update(self._rag[2])
+        # if the agent has skills, add the register_function import.
+        if self.agent.data.skills:
+            agent_imports.add("from autogen import register_function")
+        return [(import_string, position) for import_string in agent_imports]
+
+    def get_system_message_arg(self) -> str:
+        """Get the system message argument.
+
+        Returns
+        -------
+        str
+            The system message argument.
+        """
+        if not self.agent.data.system_message:
+            return ""
+        system_message = self.string_escape(self.agent.data.system_message)
+        return ",\n    system_message=" + f'"{system_message}"'
+
+    def get_before_export(
+        self,
+    ) -> Optional[List[Tuple[str, Union[ExportPosition, AgentPosition]]]]:
+        """Generate the content before the main export.
+
+        Returns
+        -------
+        Optional[List[Tuple[str, Union[ExportPosition, AgentPosition]]]]
+            The exported content before the main export and its position.
+        """
+        before_agent_string = ""
+        if self._code_execution[0] and self._code_execution[2]:
+            before_agent_string += self._code_execution[0]
+        if self._termination[1]:
+            before_agent_string += self._termination[1]
+        if self._group_chat[0]:
+            before_agent_string += self._group_chat[0]
+        if self._rag[0]:
+            before_agent_string += self._rag[0]
+        if self._swarm[0]:
+            before_agent_string += self._swarm[0]
+        if before_agent_string:
+            return [
+                (
+                    before_agent_string,
+                    AgentPosition(self.agent, AgentPositions.BEFORE),
+                )
+            ]
+        return None
+
+    def get_after_export(
+        self,
+    ) -> Optional[List[Tuple[str, Union[ExportPosition, AgentPosition]]]]:
+        """Generate the content after the main export.
+
+        Returns
+        -------
+        Optional[List[Tuple[str, Union[ExportPosition, AgentPosition]]]]
+            The exported content after the main export and its position.
+        """
+        after_agent_string = ""
+        if self._swarm[2]:
+            after_agent_string += self._swarm[2]
+        if after_agent_string:
+            return [
+                (
+                    after_agent_string,
+                    AgentPosition(self.agent, AgentPositions.AFTER_ALL),
+                )
+            ]
+        return None
+
+    def generate(self) -> Optional[str]:
+        """Generate the exported agent.
+
+        Returns
+        -------
+        Optional[str]
+            The exported agent.
+        """
+        agent = self.agent
+        agent_name = self._agent_name
+        agent_class = self._agent_class
+        retrieve_arg = self._rag[1]
+        group_chat_arg = self._group_chat[1]
+        is_termination = self._termination[0]
+        code_execution_arg = self._code_execution[1]
+        system_message_arg = self.get_system_message_arg()
+        default_auto_reply: str = "None"
+        if agent.data.agent_default_auto_reply:
+            default_auto_reply = (
+                f'"{self.string_escape(agent.data.agent_default_auto_reply)}"'
+            )
+        agent_str = f"""{agent_name} = {agent_class}(
+    name="{agent_name}",
+    description="{agent.description}"{system_message_arg},
+    human_input_mode="{agent.data.human_input_mode}",
+    max_consecutive_auto_reply={agent.data.max_consecutive_auto_reply},
+    default_auto_reply={default_auto_reply},
+    code_execution_config={code_execution_arg},
+    is_termination_msg={is_termination},{group_chat_arg}{retrieve_arg}
+"""
+        if self._swarm[1]:
+            agent_str += self._swarm[1]
+        # e.g. llm_config=...
+        other_args = self.arguments_resolver(agent)
+        if other_args:
+            agent_str += ",\n".join(other_args)
+        if not agent_str.endswith("\n"):
+            agent_str += "\n"
+        agent_str += ")"
+        return agent_str
+
+    def export(self) -> ExporterReturnType:
+        """Export the agent.
+
+        Returns
+        -------
+        ExporterReturnType
+            The exported agent.
+        """
+        agent_string = self.generate() or ""
+        is_group_manager = self.agent.agent_type == "group_manager"
+        after_export = self.get_after_export() or []
+        content: Optional[str] = agent_string
+        if is_group_manager and agent_string:
+            content = None
+            # make sure the group manager is defined
+            # after the rest of the agents.
+            # to avoid issues with (for example):
+            #  'group_manager_group_chat = GroupChat(
+            #    # assistant and rag_user should be defined first
+            # '    agents=[assistant, rag_user],
+            # '    enable_clear_history=True,
+            # ...
+            after_export.append(
+                (agent_string, AgentPosition(None, AgentPositions.AFTER_ALL, 0))
+            )
+        return {
+            "content": content,
+            "imports": self.get_imports(),
+            "environment_variables": [],
+            "before_export": self.get_before_export(),
+            "after_export": after_export,
+        }

waldiez/exporting/agent/utils/__init__.py

@@ -0,0 +1,23 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Utility functions for generating agent related strings."""
+
+from .agent_class_name import get_agent_class_name
+from .agent_imports import get_agent_imports
+from .code_execution import get_agent_code_execution_config
+from .group_manager import get_group_manager_extras
+from .rag_user import get_rag_user_extras
+from .swarm_agent import get_swarm_extras
+from .teachability import get_agent_teachability_string
+from .termination_message import get_is_termination_message
+
+__all__ = [
+    "get_agent_class_name",
+    "get_agent_imports",
+    "get_agent_code_execution_config",
+    "get_agent_teachability_string",
+    "get_group_manager_extras",
+    "get_is_termination_message",
+    "get_rag_user_extras",
+    "get_swarm_extras",
+]

waldiez/exporting/agent/utils/agent_class_name.py

@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Get the agent class name."""
+
+from waldiez.models import WaldiezAgent
+
+
+# pylint: disable=too-many-return-statements
+def get_agent_class_name(agent: WaldiezAgent) -> str:
+    """Get the agent class name.
+
+    Parameters
+    ----------
+    agent : WaldiezAgent
+        The agent.
+
+    Returns
+    -------
+    str
+        The agent class name.
+    """
+    if agent.data.is_multimodal:
+        return "MultimodalConversableAgent"
+    if agent.agent_type == "assistant":
+        return "AssistantAgent"
+    if agent.agent_type == "user":
+        return "UserProxyAgent"
+    if agent.agent_type == "manager":
+        return "GroupChatManager"
+    if agent.agent_type == "rag_user":
+        return "RetrieveUserProxyAgent"
+    if agent.agent_type == "swarm":
+        return "SwarmAgent"
+    return "ConversableAgent"  # pragma: no cover

waldiez/exporting/agent/utils/agent_imports.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Get the imports needed for the agent."""
+
+from typing import Set
+
+
+def get_agent_imports(agent_class: str) -> Set[str]:
+    """Get the imports needed for the agent.
+
+    Parameters
+    ----------
+    agent_class : str
+        The agent class name.
+
+    Returns
+    -------
+    Set[str]
+        The imports needed for the agent.
+    """
+    imports = set(["import autogen"])
+    if agent_class == "AssistantAgent":
+        imports.add("from autogen import AssistantAgent")
+    elif agent_class == "UserProxyAgent":
+        imports.add("from autogen import UserProxyAgent")
+    elif agent_class == "GroupChatManager":
+        imports.add("from autogen import GroupChatManager")
+    elif agent_class == "RetrieveUserProxyAgent":
+        imports.add(
+            "from autogen.agentchat.contrib.retrieve_user_proxy_agent "
+            "import RetrieveUserProxyAgent"
+        )
+    elif agent_class == "MultimodalConversableAgent":
+        imports.add(
+            "from autogen.agentchat.contrib.multimodal_conversable_agent "
+            "import MultimodalConversableAgent"
+        )
+    elif agent_class == "SwarmAgent":
+        imports.add(
+            "from autogen import "
+            "AFTER_WORK, "
+            "ON_CONDITION, "
+            "UPDATE_SYSTEM_MESSAGE, "
+            "AfterWorkOption, "
+            "SwarmAgent, "
+            "SwarmResult"
+        )
+    else:
+        imports.add("from autogen import ConversableAgent")
+    return imports

waldiez/exporting/{agents → agent/utils}/code_execution.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Code execution related functions for exporting agents."""
 
 from typing import Dict, Tuple
@@ -36,15 +38,15 @@ def get_agent_code_execution_config(
         if use_docker
         else "LocalCommandLineCodeExecutor"
     )
-    executor_content = f"{agent_name}_executor = {executor_class_name}(\n"
+    executor_content = f"{agent_name}_executor = {executor_class_name}(" + "\n"
     if agent.data.code_execution_config.work_dir:
         wok_dir = agent.data.code_execution_config.work_dir.replace(
             '"', '\\"'
         ).replace("\n", "\\n")
-        executor_content += f'    work_dir="{wok_dir}",\n'
+        executor_content += f'    work_dir="{wok_dir}",' + "\n"
     if agent.data.code_execution_config.timeout:
         executor_content += (
-            f"    timeout={agent.data.code_execution_config.timeout},\n"
+            f"    timeout={agent.data.code_execution_config.timeout}," + "\n"
         )
     if use_docker is False and agent.data.code_execution_config.functions:
         function_names = []
@@ -53,15 +55,11 @@ def get_agent_code_execution_config(
             function_names.append(skill_name)
         if function_names:
             # pylint: disable=inconsistent-quotes
+            function_names_string = ", ".join(function_names)
             executor_content += (
-                f"    functions=[{', '.join(function_names)}],\n"
+                f"    functions=[{function_names_string}]," + "\n"
             )
     executor_content += ")\n\n"
-    # if (
-    #     executor_content
-    #     == f"{agent_name}_executor = {executor_class_name}(\n)\n\n"
-    # ):
-    #     # empty executor?
-    #     return "", "False", ""
     executor_arg = f'{{"executor": {agent_name}_executor}}'
-    return executor_content, executor_arg, executor_class_name
+    the_import = f"from autogen.coding import {executor_class_name}"
+    return executor_content, executor_arg, the_import

waldiez/exporting/{agents → agent/utils}/group_manager.py

@@ -1,16 +1,17 @@
-"""Export group manger and group chat to string."""
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Export group manager and group chat to string."""
 
-from typing import Dict, List, Optional, Tuple
+from typing import Callable, Dict, List, Optional, Tuple
 
 from waldiez.models import WaldiezAgent, WaldiezGroupManager
 
-from ..utils import get_method_string, get_object_string
-
 
 def get_group_manager_extras(
     agent: WaldiezAgent,
     group_chat_members: List[WaldiezAgent],
     agent_names: Dict[str, str],
+    serializer: Callable[..., str],
 ) -> Tuple[str, str]:
     """Get the group manager extra string and custom selection method if any.
 
@@ -22,6 +23,8 @@ def get_group_manager_extras(
         The group members.
     agent_names : Dict[str, str]
         The agent names.
+    serializer : Callable[..., str]
+        The serializer function.
 
     Returns
     -------
@@ -34,12 +37,17 @@ def get_group_manager_extras(
     custom_speaker_selection: Optional[str] = None
     if agent.agent_type == "manager" and isinstance(agent, WaldiezGroupManager):
         group_chat_string, group_chat_name, custom_speaker_selection = (
-            _get_group_manager_extras(agent, group_chat_members, agent_names)
+            _get_group_manager_extras(
+                agent=agent,
+                group_members=group_chat_members,
+                agent_names=agent_names,
+                serializer=serializer,
+            )
         )
         if group_chat_name:
-            group_chat_arg = f"\n    groupchat={group_chat_name},"
+            group_chat_arg = "\n" + f"    groupchat={group_chat_name},"
         if custom_speaker_selection:
-            before_agent_string += f"{custom_speaker_selection}\n\n"
+            before_agent_string += f"{custom_speaker_selection}" + "\n"
         if group_chat_string:
             before_agent_string += group_chat_string
     return before_agent_string, group_chat_arg
@@ -49,6 +57,7 @@ def _get_group_manager_extras(
     agent: WaldiezGroupManager,
     group_members: List[WaldiezAgent],
     agent_names: Dict[str, str],
+    serializer: Callable[..., str],
 ) -> Tuple[str, str, Optional[str]]:
     """Get the group manager extra string and custom selection method if any.
 
@@ -60,6 +69,8 @@ def _get_group_manager_extras(
         The group members.
     agent_names : Dict[str, str]
         The agent names.
+    serializer : Callable[..., str]
+        The serializer function.
 
     Returns
     -------
@@ -88,25 +99,19 @@ def _get_group_manager_extras(
         group_chat_string += f"    max_round={agent.data.max_round}," + "\n"
     if agent.data.admin_name:
         group_chat_string += f'    admin_name="{agent.data.admin_name}",' + "\n"
-    extra_group_chat_string, method_name_and_content = (
-        _get_group_chat_speakers_string(agent, agent_names)
+    extra_group_chat_string, custom_selection_method = (
+        _get_group_chat_speakers_string(agent, agent_names, serializer)
     )
-    custom_selection_method: Optional[str] = None
     group_chat_string += extra_group_chat_string
     group_chat_string += ")\n\n"
-    if method_name_and_content:
-        method_name, method_content = method_name_and_content
-        custom_selection_method = get_method_string(
-            "custom_speaker_selection",
-            method_name,
-            method_content,
-        )
     return group_chat_string, group_chat_name, custom_selection_method
 
 
 def _get_group_chat_speakers_string(
-    agent: WaldiezGroupManager, agent_names: Dict[str, str]
-) -> Tuple[str, Optional[Tuple[str, str]]]:
+    agent: WaldiezGroupManager,
+    agent_names: Dict[str, str],
+    serializer: Callable[..., str],
+) -> Tuple[str, Optional[str]]:
     """Get the group chat speakers string.
 
     Parameters
@@ -115,16 +120,18 @@ def _get_group_chat_speakers_string(
         The agent.
     agent_names : Dict[str, str]
         The agent names.
+    serializer : Callable[..., str]
+        The serializer function.
 
     Returns
     -------
     str
         The group chat speakers string.
-    Optional[Tuple[str, str]]
-        The custom selection method name and content if any.
+    Optional[str]
+        The custom custom for speaker selection if any.
     """
     speakers_string = ""
-    method_name_and_content: Optional[Tuple[str, str]] = None
+    function_content: Optional[str] = None
     if agent.data.speakers.max_retries_for_selecting is not None:
         speakers_string += (
             "    max_retries_for_selecting_speaker="
@@ -139,12 +146,14 @@ def _get_group_chat_speakers_string(
         )
     else:
         agent_name = agent_names[agent.id]
-        method_name = f"custom_speaker_selection_method_{agent_name}"
-        method_name_and_content = (
-            method_name,
-            agent.data.speakers.custom_method_string or "",
+        function_content, function_name = (
+            agent.data.speakers.get_custom_method_function(
+                name_suffix=agent_name
+            )
+        )
+        speakers_string += (
+            f"    speaker_selection_method={function_name}," + "\n"
         )
-        speakers_string += f"    speaker_selection_method={method_name}," "\n"
     # selection_mode == "repeat":
     if agent.data.speakers.selection_mode == "repeat":
         speakers_string += _get_speakers_selection_repeat_string(
@@ -156,10 +165,12 @@ def _get_group_chat_speakers_string(
         and agent.data.speakers.allowed_or_disallowed_transitions
     ):
         speakers_string += _get_speakers_selection_transition_string(
-            agent, agent_names
+            agent=agent,
+            agent_names=agent_names,
+            serializer=serializer,
         )
     speakers_string = speakers_string.replace('"None"', "None")
-    return speakers_string, method_name_and_content
+    return speakers_string, function_content
 
 
 def _get_speakers_selection_repeat_string(
@@ -168,9 +179,8 @@ def _get_speakers_selection_repeat_string(
     speakers_string = ""
     if isinstance(agent.data.speakers.allow_repeat, bool):
         speakers_string += (
-            "    allow_repeat_speaker="
-            f"{agent.data.speakers.allow_repeat},"
-            "\n"
+            f"    allow_repeat_speaker={agent.data.speakers.allow_repeat},"
+            + "\n"
         )
     elif isinstance(agent.data.speakers.allow_repeat, list):
         # get the names of the agents
@@ -178,12 +188,14 @@ def _get_speakers_selection_repeat_string(
             agent_names[agent_id]
             for agent_id in agent.data.speakers.allow_repeat
         )
-        speakers_string += f"    allow_repeat=[{allow_repeat}]," "\n"
+        speakers_string += f"    allow_repeat=[{allow_repeat}]," + "\n"
     return speakers_string
 
 
 def _get_speakers_selection_transition_string(
-    agent: WaldiezGroupManager, agent_names: Dict[str, str]
+    agent: WaldiezGroupManager,
+    agent_names: Dict[str, str],
+    serializer: Callable[..., str],
 ) -> str:
     speakers_string = ""
     allowed_or_disallowed_speaker_transitions = {}
@@ -194,7 +206,7 @@ def _get_speakers_selection_transition_string(
         allowed_or_disallowed_speaker_transitions[agent_names[agent_id]] = [
             agent_names[transition] for transition in transitions
         ]
-    transitions_string = get_object_string(
+    transitions_string = serializer(
         allowed_or_disallowed_speaker_transitions, 1
     )
     transitions_string = transitions_string.replace('"', "").replace("'", "")

waldiez/exporting/{agents → agent/utils}/rag_user/__init__.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """RAG User Agent related string generation."""
 
 from .rag_user import get_rag_user_extras, get_rag_user_retrieve_config_str