waldiez 0.3.11__py3-none-any.whl → 0.4.0__py3-none-any.whl

waldiez/exporting/flow/utils/importing_utils.py

@@ -36,21 +36,6 @@ COMMON_AUTOGEN_IMPORTS = [
 ]
 
 
-def get_standard_imports() -> str:
-    """Get the standard imports.
-
-    Returns
-    -------
-    str
-        The standard imports.
-    """
-    builtin_imports = BUILTIN_IMPORTS.copy()
-    imports_string = "\n".join(builtin_imports) + "\n"
-    typing_imports = "from typing import " + ", ".join(TYPING_IMPORTS)
-    imports_string += typing_imports
-    return imports_string
-
-
 def sort_imports(
     all_imports: List[Tuple[str, ImportPosition]],
 ) -> Tuple[List[str], List[str], List[str], List[str], bool]:
@@ -66,10 +51,10 @@ def sort_imports(
     Tuple[List[str], List[str], List[str], List[str], bool]
         The sorted imports and a flag if we got `import autogen`.
     """
-    builtin_imports = []
-    third_party_imports = []
-    local_imports = []
-    autogen_imports = COMMON_AUTOGEN_IMPORTS.copy()
+    builtin_imports: List[str] = []
+    third_party_imports: List[str] = []
+    local_imports: List[str] = []
+    autogen_imports: List[str] = COMMON_AUTOGEN_IMPORTS.copy()
     got_import_autogen = False
     for import_string, position in all_imports:
         if "import autogen" in import_string:
@@ -85,11 +70,22 @@ def sort_imports(
         elif position == ImportPosition.LOCAL:
             local_imports.append(import_string)
     autogen_imports = list(set(autogen_imports))
+    third_party_imports = ensure_np_import(third_party_imports)
+    sorted_builtins = sorted(
+        [imp for imp in builtin_imports if imp.startswith("import ")]
+    ) + sorted([imp for imp in builtin_imports if imp.startswith("from ")])
+    sorted_third_party = sorted(
+        [imp for imp in third_party_imports if imp.startswith("import ")]
+    ) + sorted([imp for imp in third_party_imports if imp.startswith("from ")])
+    sorted_locals = sorted(
+        [imp for imp in local_imports if imp.startswith("import ")]
+    ) + sorted([imp for imp in local_imports if imp.startswith("from ")])
+
     return (
-        sorted(builtin_imports),
+        sorted_builtins,
         sorted(autogen_imports),
-        sorted(third_party_imports),
-        sorted(local_imports),
+        sorted_third_party,
+        sorted_locals,
         got_import_autogen,
     )
 
@@ -150,6 +146,27 @@ def get_the_imports_string(
     return final_string.replace("\n\n\n", "\n\n")  # avoid too many newlines
 
 
+def ensure_np_import(third_party_imports: List[str]) -> List[str]:
+    """Ensure numpy is imported.
+
+    Parameters
+    ----------
+    third_party_imports : List[str]
+        The third party imports.
+
+    Returns
+    -------
+    List[str]
+        The third party imports with numpy.
+    """
+    if (
+        not third_party_imports
+        or "import numpy as np" not in third_party_imports
+    ):
+        third_party_imports.append("import numpy as np")
+    return third_party_imports
+
+
 def gather_imports(
     model_imports: Optional[List[Tuple[str, ImportPosition]]],
     skill_imports: Optional[List[Tuple[str, ImportPosition]]],
@@ -174,13 +191,14 @@ def gather_imports(
     Tuple[str, ImportPosition]
         The gathered imports.
     """
-    imports_string = get_standard_imports()
-    all_imports: List[Tuple[str, ImportPosition]] = [
-        (
-            imports_string,
-            ImportPosition.BUILTINS,
+    all_imports: List[Tuple[str, ImportPosition]] = []
+    for import_statement in BUILTIN_IMPORTS:
+        all_imports.append(
+            (
+                import_statement,
+                ImportPosition.BUILTINS,
+            )
         )
-    ]
     if model_imports:
         all_imports.extend(model_imports)
     if skill_imports:
@@ -189,4 +207,21 @@ def gather_imports(
         all_imports.extend(chat_imports)
     if agent_imports:
         all_imports.extend(agent_imports)
-    return list(set(all_imports))
+    # let's try to avoid this:
+    # from typing import Annotated
+    # from typing import Annotated, Any, Callable, Dict, ...Union
+    all_typing_imports = TYPING_IMPORTS.copy()
+    final_imports: List[Tuple[str, ImportPosition]] = []
+    for import_statement, import_position in all_imports:
+        if import_statement.startswith("from typing"):
+            to_import = import_statement.split("import")[1].strip()
+            if to_import:
+                all_typing_imports.append(to_import)
+        else:
+            final_imports.append((import_statement, import_position))
+    unique_typing_imports = list(set(all_typing_imports))
+    one_typing_import = "from typing import " + ", ".join(
+        sorted(unique_typing_imports)
+    )
+    final_imports.insert(1, (one_typing_import, ImportPosition.BUILTINS))
+    return list(set(final_imports))

waldiez/exporting/skills/skills_exporter.py

@@ -89,12 +89,18 @@ class SkillsExporter(BaseExporter, ExporterMixin):
         Tuple[str, int]
             The exported imports and the position of the imports.
         """
-        if not self.skill_imports:
-            return []
         imports: List[Tuple[str, ImportPosition]] = []
-        for skill_import in self.skill_imports:
-            if (skill_import, ImportPosition.LOCAL) not in imports:
-                imports.append((skill_import, ImportPosition.LOCAL))
+        if not self.skill_imports:
+            return imports
+        # standard imports
+        for import_statement in self.skill_imports[0]:
+            imports.append((import_statement, ImportPosition.BUILTINS))
+        # third party imports
+        for import_statement in self.skill_imports[1]:
+            imports.append((import_statement, ImportPosition.THIRD_PARTY))
+        # secrets/local imports
+        for import_statement in self.skill_imports[2]:
+            imports.append((import_statement, ImportPosition.LOCAL))
         return imports
 
     def get_before_export(
@@ -156,11 +162,12 @@ class SkillsExporter(BaseExporter, ExporterMixin):
             the before export strings, the after export strings,
             and the environment variables.
         """
+        content = self.generate()
         imports = self.get_imports()
         after_export = self.get_after_export()
         environment_variables = self.get_environment_variables()
         result: ExporterReturnType = {
-            "content": self.generate(),
+            "content": content,
             "imports": imports,
             "before_export": None,
             "after_export": after_export,

waldiez/exporting/skills/utils.py

@@ -94,6 +94,7 @@ def _write_skill_secrets(
         return
     if not isinstance(output_dir, Path):
         output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
     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:
@@ -108,7 +109,7 @@ def export_skills(
     skills: List[WaldiezSkill],
     skill_names: Dict[str, str],
     output_dir: Optional[Union[str, Path]] = None,
-) -> Tuple[List[str], List[Tuple[str, str]], str]:
+) -> Tuple[Tuple[List[str], List[str], 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.
@@ -126,19 +127,26 @@ def export_skills(
 
     Returns
     -------
-    Tuple[Set[str], Set[Tuple[str, str]], str]
+    Tuple[Tuple[List[str], List[str], List[str]], List[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_imports: Tuple[List[str], List[str], 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))
+        standard_skill_imports, third_party_skill_imports = skill.get_imports()
+        if standard_skill_imports:
+            skill_imports[0].extend(standard_skill_imports)
+        if third_party_skill_imports:
+            skill_imports[1].extend(third_party_skill_imports)
+        secrets_import = get_skill_secrets_import(flow_name, skill)
+        if secrets_import:
+            skill_imports[2].append(secrets_import)
         for key, value in skill.secrets.items():
             skill_secrets.append((key, value))
         _write_skill_secrets(
@@ -153,8 +161,14 @@ def export_skills(
         if skill.is_shared:
             shared_skill_contents += skill_content + "\n\n"
         else:
+            if skill.is_interop:
+                skill_content += _add_interop_extras(
+                    skill=skill, skill_names=skill_names
+                )
             skill_contents += skill_content + "\n\n"
     skill_contents = shared_skill_contents + skill_contents
+    # remove dupes from imports if any and sort them
+    skill_imports = _sort_imports(skill_imports)
     return (
         skill_imports,
         skill_secrets,
@@ -162,8 +176,77 @@ def export_skills(
     )
 
 
-def get_skill_imports(flow_name: str, skill: WaldiezSkill) -> str:
-    """Get the skill imports string.
+def _add_interop_extras(
+    skill: WaldiezSkill,
+    skill_names: Dict[str, str],
+) -> str:
+    """Add the interop conversion.
+
+    Parameters
+    ----------
+    skill : WaldiezSkill
+        The skill
+    skill_names : Dict[str, str]
+        The skill names.
+
+    Returns
+    -------
+    str
+        The extra content to convert the tool.
+    """
+    skill_name = skill_names[skill.id]
+    interop_instance = f"ag2_{skill_name}_interop = Interoperability()" + "\n"
+    extra_content = (
+        f"ag2_{skill_name} = "
+        f"ag2_{skill_name}_interop.convert_tool("
+        f"tool={skill_name}, "
+        f'type="{skill.skill_type}")'
+    )
+    return "\n" + interop_instance + extra_content
+
+
+def _sort_imports(
+    skill_imports: Tuple[List[str], List[str], List[str]],
+) -> Tuple[List[str], List[str], List[str]]:
+    """Sort the imports.
+
+    Parameters
+    ----------
+    skill_imports : Tuple[List[str], List[str], List[str]]
+        The skill imports.
+
+    Returns
+    -------
+    Tuple[List[str], List[str], List[str]]
+        The sorted skill imports.
+    """
+
+    # "from x import y" and "import z"
+    # the "import a" should be first (and sorted)
+    # then the "from b import c" (and sorted)
+    standard_lib_imports = skill_imports[0]
+    third_party_imports = skill_imports[1]
+    secrets_imports = skill_imports[2]
+
+    sorted_standard_lib_imports = sorted(
+        [imp for imp in standard_lib_imports if imp.startswith("import ")]
+    ) + sorted([imp for imp in standard_lib_imports if imp.startswith("from ")])
+
+    sorted_third_party_imports = sorted(
+        [imp for imp in third_party_imports if imp.startswith("import ")]
+    ) + sorted([imp for imp in third_party_imports if imp.startswith("from ")])
+
+    sorted_secrets_imports = sorted(secrets_imports)
+
+    return (
+        sorted_standard_lib_imports,
+        sorted_third_party_imports,
+        sorted_secrets_imports,
+    )
+
+
+def get_skill_secrets_import(flow_name: str, skill: WaldiezSkill) -> str:
+    """Get the skill secrets import string.
 
     Parameters
     ----------
@@ -263,6 +346,9 @@ def get_agent_skill_registrations(
             skill for skill in all_skills if skill.id == linked_skill.id
         )
         skill_name = skill_names[linked_skill.id]
+        if waldiez_skill.is_interop:
+            # the name of the the converted to ag2 tool
+            skill_name = f"ag2_{skill_name}"
         skill_description = (
             waldiez_skill.description or f"Description of {skill_name}"
         )

waldiez/models/__init__.py

@@ -23,6 +23,9 @@ from .agents import (
     WaldiezAgentType,
     WaldiezAssistant,
     WaldiezAssistantData,
+    WaldiezCaptainAgent,
+    WaldiezCaptainAgentData,
+    WaldiezCaptainAgentLibEntry,
     WaldiezGroupManager,
     WaldiezGroupManagerData,
     WaldiezGroupManagerSpeakers,
@@ -87,6 +90,9 @@ __all__ = [
     "WaldiezAgentType",
     "WaldiezAssistant",
     "WaldiezAssistantData",
+    "WaldiezCaptainAgent",
+    "WaldiezCaptainAgentData",
+    "WaldiezCaptainAgentLibEntry",
     "WaldiezChat",
     "WaldiezChatData",
     "WaldiezChatSummary",

waldiez/models/agents/__init__.py

@@ -18,6 +18,15 @@ from .agent import (
 )
 from .agents import WaldiezAgents
 from .assistant import WaldiezAssistant, WaldiezAssistantData
+from .captain_agent import (
+    WaldiezCaptainAgent,
+    WaldiezCaptainAgentData,
+    WaldiezCaptainAgentLibEntry,
+)
+from .extra_requirements import (
+    get_captain_agent_extra_requirements,
+    get_retrievechat_extra_requirements,
+)
 from .group_manager import (
     CUSTOM_SPEAKER_SELECTION,
     CUSTOM_SPEAKER_SELECTION_ARGS,
@@ -77,6 +86,8 @@ from .swarm_agent import (
 from .user_proxy import WaldiezUserProxy, WaldiezUserProxyData
 
 __all__ = [
+    "get_retrievechat_extra_requirements",
+    "get_captain_agent_extra_requirements",
     "IS_TERMINATION_MESSAGE",
     "IS_TERMINATION_MESSAGE_ARGS",
     "IS_TERMINATION_MESSAGE_TYPES",
@@ -113,6 +124,9 @@ __all__ = [
     "WaldiezAgentNestedChatMessage",
     "WaldiezAgentTeachability",
     "WaldiezAgentTerminationMessage",
+    "WaldiezCaptainAgent",
+    "WaldiezCaptainAgentData",
+    "WaldiezCaptainAgentLibEntry",
     "WaldiezGroupManager",
     "WaldiezGroupManagerData",
     "WaldiezGroupManagerSpeakers",

waldiez/models/agents/agent/__init__.py

@@ -2,8 +2,9 @@
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Base agent class to be inherited by all other agents."""
 
-from .agent import WaldiezAgent, WaldiezAgentType
+from .agent import WaldiezAgent
 from .agent_data import WaldiezAgentData
+from .agent_type import WaldiezAgentType
 from .code_execution import WaldiezAgentCodeExecutionConfig
 from .linked_skill import WaldiezAgentLinkedSkill
 from .nested_chat import WaldiezAgentNestedChat, WaldiezAgentNestedChatMessage

waldiez/models/agents/agent/agent.py

@@ -2,19 +2,16 @@
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Base agent class to be inherited by all agents."""
 
-from typing import List
+from typing import List, Set
 
 from pydantic import Field
 from typing_extensions import Annotated, Literal
 
 from ...common import WaldiezBase, now
 from .agent_data import WaldiezAgentData
+from .agent_type import WaldiezAgentType
 from .code_execution import WaldiezAgentCodeExecutionConfig
 
-WaldiezAgentType = Literal[
-    "user", "assistant", "manager", "rag_user", "swarm", "reasoning"
-]
-
 
 class WaldiezAgent(WaldiezBase):
     """Waldiez Agent.
@@ -25,9 +22,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", "swarm", "reasoning"
-        ]
+    agent_type : WaldiezAgentType
         The type of the agent
     name: str
         The name of the agent.
@@ -65,9 +60,7 @@ class WaldiezAgent(WaldiezBase):
         ),
     ]
     agent_type: Annotated[
-        Literal[
-            "user", "assistant", "manager", "rag_user", "swarm", "reasoning"
-        ],
+        WaldiezAgentType,
         Field(
             ...,
             title="Agent type",
@@ -129,6 +122,73 @@ class WaldiezAgent(WaldiezBase):
         ),
     ]
 
+    @property
+    def ag2_class(self) -> str:
+        """Return the AG2 class of the agent."""
+        class_name = "ConversableAgent"
+        if self.data.is_multimodal:
+            return "MultimodalConversableAgent"
+        if self.agent_type == "assistant":
+            class_name = "AssistantAgent"
+        if self.agent_type == "user":
+            class_name = "UserProxyAgent"
+        if self.agent_type == "manager":
+            class_name = "GroupChatManager"
+        if self.agent_type == "rag_user":
+            class_name = "RetrieveUserProxyAgent"
+        if self.agent_type == "swarm":
+            class_name = "SwarmAgent"
+        if self.agent_type == "reasoning":
+            class_name = "ReasoningAgent"
+        if self.agent_type == "captain":
+            class_name = "CaptainAgent"
+        return class_name
+
+    @property
+    def ag2_imports(self) -> Set[str]:
+        """Return the AG2 imports of the agent."""
+        agent_class = self.ag2_class
+        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 "
+                "register_hand_off, "
+                "AfterWork, "
+                "OnCondition, "
+                "UpdateSystemMessage, "
+                "AfterWorkOption, "
+                "SwarmResult"
+            )
+        elif agent_class == "ReasoningAgent":
+            imports.add(
+                "from autogen.agentchat.contrib.reasoning_agent "
+                "import ReasoningAgent, visualize_tree"
+            )
+        elif agent_class == "CaptainAgent":
+            imports.add(
+                "from autogen.agentchat.contrib.captainagent "
+                "import CaptainAgent"
+            )
+        else:  # pragma: no cover
+            imports.add("import ConversableAgent")
+        return imports
+
     def validate_linked_skills(
         self, skill_ids: List[str], agent_ids: List[str]
     ) -> None:

waldiez/models/agents/agent/agent_type.py

@@ -0,0 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Waldiez Agent types."""
+
+from typing_extensions import Literal
+
+# pylint: disable=line-too-long
+# fmt: off
+WaldiezAgentType = Literal["user", "assistant", "manager", "rag_user", "swarm", "reasoning", "captain"]  # noqa: E501
+"""Possible types of a Waldiez Agent: user, assistant, manager, rag_user, swarm, reasoning, captain."""  # noqa: E501
+# fmt: on

waldiez/models/agents/agents.py

@@ -10,7 +10,8 @@ from typing_extensions import Annotated, Self
 from ..common import WaldiezBase
 from .agent import WaldiezAgent
 from .assistant import WaldiezAssistant
-from .group_manager.group_manager import WaldiezGroupManager
+from .captain_agent import WaldiezCaptainAgent
+from .group_manager import WaldiezGroupManager
 from .rag_user import WaldiezRagUser
 from .reasoning import WaldiezReasoningAgent
 from .swarm_agent import WaldiezSwarmAgent
@@ -80,6 +81,14 @@ class WaldiezAgents(WaldiezBase):
             default_factory=list,
         ),
     ]
+    captain_agents: Annotated[
+        List[WaldiezCaptainAgent],
+        Field(
+            title="Captain Agents.",
+            description="Captain agents",
+            default_factory=list,
+        ),
+    ]
 
     @property
     def members(self) -> Iterator[WaldiezAgent]:
@@ -96,6 +105,7 @@ class WaldiezAgents(WaldiezBase):
         yield from self.reasoning_agents
         yield from self.swarm_agents
         yield from self.managers
+        yield from self.captain_agents
 
     @model_validator(mode="after")
     def validate_agents(self) -> Self:

waldiez/models/agents/captain_agent/__init__.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Captain agent model."""
+
+from .captain_agent import WaldiezCaptainAgent
+from .captain_agent_data import (
+    WaldiezCaptainAgentData,
+)
+from .captain_agent_lib_entry import WaldiezCaptainAgentLibEntry
+
+__all__ = [
+    "WaldiezCaptainAgentData",
+    "WaldiezCaptainAgent",
+    "WaldiezCaptainAgentLibEntry",
+]

waldiez/models/agents/captain_agent/captain_agent.py

@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Waldiez captain agent model."""
+
+from typing import Literal
+
+from pydantic import Field
+from typing_extensions import Annotated
+
+from ..agent import WaldiezAgent
+from .captain_agent_data import WaldiezCaptainAgentData
+
+
+class WaldiezCaptainAgent(WaldiezAgent):
+    """Captain agent.
+
+    A `WaldiezAgent` with agent_type `captain` and
+    captain agent's related config for the agent.
+    Also see `WaldiezAgent`, `WaldiezCaptainData`, `WaldiezAgentData`
+
+    Attributes
+    ----------
+    agent_type : Literal["captain"]
+        The agent type: 'captain' for a captain agent
+    data : WaldiezCaptainAgentData
+        The captain agent's data.
+    """
+
+    agent_type: Annotated[
+        Literal["captain"],
+        Field(
+            "captain",
+            title="Agent type",
+            description="The agent type: 'captain' for a captain agent",
+            alias="agentType",
+        ),
+    ]
+    data: Annotated[
+        WaldiezCaptainAgentData,
+        Field(
+            title="Data",
+            description="The captain agent's data",
+            default_factory=WaldiezCaptainAgentData,
+        ),
+    ]

waldiez/models/agents/captain_agent/captain_agent_data.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Waldiez captain agent data."""
+
+from typing import List, Optional
+
+from pydantic import Field
+from typing_extensions import Annotated, Literal
+
+from ..agent import WaldiezAgentData
+from .captain_agent_lib_entry import WaldiezCaptainAgentLibEntry
+
+
+class WaldiezCaptainAgentData(WaldiezAgentData):
+    """Captain agent data class.
+
+    The data for a captain agent.
+    Extends `WaldiezAgentData`.
+    Extra attributes:
+    - `agent_lib`: Optional list of agent lib entries
+    - `tool_lib`:
+    - `max_round`: The maximum number of rounds in a group chat
+    - `max_turns`: The maximum number of turns for a chat
+    See the parent's docs (`WaldiezAgentData`) for the rest of the properties.
+    """
+
+    agent_lib: Annotated[
+        List[WaldiezCaptainAgentLibEntry],
+        Field(
+            default_factory=list,
+            title="Agent lib",
+            description="The agent lib",
+            alias="agentLib",
+        ),
+    ] = []
+    tool_lib: Annotated[
+        Optional[Literal["default"]],
+        Field(
+            None,
+            title="Tool lib",
+            description="Whether to use the default tool lib",
+            alias="toolLib",
+        ),
+    ] = None
+    max_round: Annotated[
+        int,
+        Field(
+            10,
+            title="Max round",
+            description="The maximum number of rounds in a group chat",
+            alias="maxRound",
+        ),
+    ] = 10
+    max_turns: Annotated[
+        int,
+        Field(
+            5,
+            title="Max turns",
+            description="The maximum number of turns for a chat",
+            alias="maxTurns",
+        ),
+    ] = 5