waldiez 0.3.10__py3-none-any.whl → 0.3.12__py3-none-any.whl

waldiez/exporting/flow/utils/importing_utils.py

@@ -28,6 +28,7 @@ TYPING_IMPORTS = [
 ]
 COMMON_AUTOGEN_IMPORTS = [
     "from autogen import Agent",
+    "from autogen import Cache",
     "from autogen import ConversableAgent",
     "from autogen import ChatResult",
     "from autogen import GroupChat",

waldiez/exporting/models/models_exporter.py

@@ -32,6 +32,7 @@ class ModelsExporter(BaseExporter, ExporterMixin):
         models: List[WaldiezModel],
         model_names: Dict[str, str],
         for_notebook: bool,
+        cache_seed: Optional[int],
         output_dir: Optional[Union[str, Path]] = None,
     ) -> None:
         """Initialize the models exporter.
@@ -46,6 +47,10 @@ class ModelsExporter(BaseExporter, ExporterMixin):
             The models.
         model_names : Dict[str, str]
             The model names.
+        for_notebook : bool
+            Whether the export is for a notebook or not.
+        cache_seed : Optional[int]
+            The cache seed if any, by default None
         output_dir : Optional[Union[str, Path]], optional
             The output directory if any, by default None
         """
@@ -57,6 +62,7 @@ class ModelsExporter(BaseExporter, ExporterMixin):
         self.model_names = model_names
         if output_dir is not None and not isinstance(output_dir, Path):
             output_dir = Path(output_dir)
+        self.cache_seed = cache_seed
         self.output_dir = output_dir
         self._exported_string = None
 
@@ -127,6 +133,7 @@ class ModelsExporter(BaseExporter, ExporterMixin):
                         agent,
                         all_models=self.models,
                         model_names=self.model_names,
+                        cache_seed=self.cache_seed,
                     ),
                     AgentPosition(
                         agent=agent, position=AgentPositions.AS_ARGUMENT

waldiez/exporting/models/utils.py

@@ -63,6 +63,7 @@ def get_agent_llm_config_arg(
     agent: WaldiezAgent,
     all_models: List[WaldiezModel],
     model_names: Dict[str, str],
+    cache_seed: Optional[int],
     tabs: int = 1,
 ) -> str:
     """Get the string representation of the agent's llm config argument.
@@ -75,6 +76,8 @@ def get_agent_llm_config_arg(
         All the models in the flow.
     model_names : Dict[str, str]
         A mapping of model ids to model names.
+    cache_seed : Optional[int]
+        The cache seed.
     tabs : int, optional
         The number of tabs for indentation, by default 1.
 
@@ -98,6 +101,7 @@ def get_agent_llm_config_arg(
     if not got_at_least_one_model:  # pragma: no cover
         return f"{tab}llm_config=False," + "\n"
     content += "\n" + f"{tab}    ]," + "\n"
+    content += f'{tab}    "cache_seed": {cache_seed},' + "\n"
     content += tab + "},\n"
     return content
 

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/agent.py

@@ -2,7 +2,7 @@
 # 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
@@ -12,7 +12,7 @@ from .agent_data import WaldiezAgentData
 from .code_execution import WaldiezAgentCodeExecutionConfig
 
 WaldiezAgentType = Literal[
-    "user", "assistant", "manager", "rag_user", "swarm", "reasoning"
+    "user", "assistant", "manager", "rag_user", "swarm", "reasoning", "captain"
 ]
 
 
@@ -25,9 +25,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 +63,7 @@ class WaldiezAgent(WaldiezBase):
         ),
     ]
     agent_type: Annotated[
-        Literal[
-            "user", "assistant", "manager", "rag_user", "swarm", "reasoning"
-        ],
+        WaldiezAgentType,
         Field(
             ...,
             title="Agent type",
@@ -129,6 +125,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 "
+                "AFTER_WORK, "
+                "ON_CONDITION, "
+                "UPDATE_SYSTEM_MESSAGE, "
+                "AfterWorkOption, "
+                "SwarmAgent, "
+                "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/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:
@@ -115,8 +125,8 @@ class WaldiezAgents(WaldiezBase):
             If the agents are invalid.
         """
         all_agent_ids = [agent.id for agent in self.members]
-        if len(all_agent_ids) < 2:
-            raise ValueError("At least two agents are required.")
+        if len(all_agent_ids) < 1:
+            raise ValueError("At least one agent is required.")
         if len(all_agent_ids) != len(set(all_agent_ids)):
             raise ValueError("Agent IDs must be unique.")
         return 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 : WaldiezCaptainData
+        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

waldiez/models/agents/captain_agent/captain_agent_lib_entry.py

@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Waldiez captain agent lib entry."""
+
+from pydantic import Field
+from typing_extensions import Annotated
+
+from ...common import WaldiezBase
+
+
+class WaldiezCaptainAgentLibEntry(WaldiezBase):
+    """Captain agent lib entry."""
+
+    name: Annotated[
+        str,
+        Field(
+            ...,
+            title="Name",
+            description="The name of the agent",
+        ),
+    ]
+    description: Annotated[
+        str,
+        Field(
+            ...,
+            title="Description",
+            description="The description of the agent",
+        ),
+    ]
+    system_message: Annotated[
+        str,
+        Field(
+            ...,
+            title="System message",
+            description="The system message",
+            alias="systemMessage",
+        ),
+    ]

waldiez/models/agents/extra_requirements.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Extra requirements for agents."""
+
+# pylint: disable=line-too-long
+import platform
+from typing import Iterator, List, Set
+
+from .agent import WaldiezAgent
+from .rag_user import WaldiezRagUser
+
+
+def get_retrievechat_extra_requirements(
+    agents: Iterator[WaldiezAgent],
+) -> Set[str]:
+    """Get the retrievechat extra requirements.
+
+    Parameters
+    ----------
+    agents : List[WaldiezAgent]
+        The flow agents.
+
+    Returns
+    -------
+    Set[str]
+        The retrievechat extra requirements.
+    """
+    # https://github.com/ag2ai/ag2/blob/main/pyproject.toml
+    # with chromadb relaxed
+    # to avoid conflicts with other extras and (later) allow py3.13
+    rag_requirements: Set[str] = {
+        "protobuf==4.25.3",
+        "chromadb>=0.5.23",
+        "sentence_transformers",
+        "pypdf",
+        "ipython",
+        "beautifulsoup4",
+        "markdownify",
+    }
+    for agent in agents:
+        if agent.agent_type == "rag_user" and isinstance(agent, WaldiezRagUser):
+            # if not chroma, get the relevant db requirements
+            db_type = agent.data.retrieve_config.vector_db
+            if db_type == "pgvector":
+                rag_requirements.update(
+                    [
+                        "pgvector>=0.2.5",
+                        "psycopg[binary]>=3.2.4",
+                    ]
+                )
+            elif db_type == "mongodb":
+                rag_requirements.add("pymongo>=4.11")
+            elif db_type == "qdrant":
+                rag_requirements.update(["qdrant_client[fastembed]"])
+    return rag_requirements
+
+
+def get_captain_agent_extra_requirements() -> List[str]:
+    """Get the captain agent extra requirements.
+
+    Returns
+    -------
+    List[str]
+        The captain agent extra requirements.
+    """
+    # https://github.com/ag2ai/ag2/blob/main/autogen/agentchat/contrib/captainagent/tools/requirements.txt  # noqa: E501
+    tool_requirements = [
+        "markdownify",
+        "arxiv",
+        "pymupdf",
+        "wikipedia-api",
+        "easyocr",
+        "python-pptx",
+        "openai-whisper",
+        "pandas",
+        "scipy",
+        # "sentence-transformers", also in agent_requirements
+    ]
+    agent_requirements = [
+        "chromadb",
+        "sentence-transformers",
+        "huggingface-hub",
+    ]
+    if platform.system() == "Linux":
+        agent_requirements.append("pysqlite3-binary")
+    # on windows and OSX, installing pysqlite3-binary seem to fail in some cases
+    # we can handle/install if needed in waldiez.utils.pysqlite3_checker
+    return tool_requirements + agent_requirements

waldiez/models/common/__init__.py

@@ -2,6 +2,7 @@
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Common utils for all models."""
 
+from .ag2_version import get_autogen_version
 from .base import WaldiezBase
 from .date_utils import now
 from .dict_utils import update_dict
@@ -16,6 +17,7 @@ __all__ = [
     "WaldiezBase",
     "now",
     "check_function",
+    "get_autogen_version",
     "get_function",
     "generate_function",
     "parse_code_string",

waldiez/models/common/ag2_version.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Get the autogen version."""
+
+import warnings
+from functools import cache
+
+
+@cache
+def get_autogen_version() -> str:
+    """Get the autogen version.
+
+    Returns
+    -------
+    str
+        The autogen version.
+
+    Raises
+    ------
+    ValueError
+        If pyautogen is not installed.
+    """
+    # pylint: disable=import-outside-toplevel
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        try:
+            from autogen.version import __version__ as ag2  # type: ignore
+        except ImportError as error:  # pragma: no cover
+            raise ValueError("pyautogen is not installed.") from error
+    return ag2

waldiez/models/common/base.py

@@ -64,3 +64,7 @@ class WaldiezBase(BaseModel):
         if not isinstance(by_alias, bool):
             by_alias = True
         return super().model_dump_json(by_alias=by_alias, **kwargs)
+
+    def __hash__(self) -> int:  # pragma: no cover
+        """Return the hash of the object."""
+        return id(self)

waldiez/models/common/date_utils.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Date utilities."""
 
 from datetime import datetime, timezone

waldiez/models/common/dict_utils.py

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

waldiez/models/flow/__init__.py

@@ -4,8 +4,10 @@
 
 from .flow import WaldiezFlow
 from .flow_data import WaldiezFlowData
+from .utils import get_flow_data
 
 __all__ = [
+    "get_flow_data",
     "WaldiezFlow",
     "WaldiezFlowData",
 ]