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

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/agents/group_manager/speakers.py

@@ -16,8 +16,11 @@ WaldiezGroupManagerSpeakersSelectionMethod = Literal[
     "round_robin",
     "custom",
 ]
+"""Possible methods for the speaker selection."""
 WaldiezGroupManagerSpeakersSelectionMode = Literal["repeat", "transition"]
+"""Possible selection modes: repeat, transition."""
 WaldiezGroupManagerSpeakersTransitionsType = Literal["allowed", "disallowed"]
+"""Possible transitions types: allowed, disallowed."""
 
 CUSTOM_SPEAKER_SELECTION = "custom_speaker_selection"
 CUSTOM_SPEAKER_SELECTION_ARGS = ["last_speaker", "groupchat"]

waldiez/models/agents/rag_user/retrieve_config.py

@@ -13,8 +13,11 @@ from ...common import WaldiezBase, check_function, generate_function
 from .vector_db_config import WaldiezRagUserVectorDbConfig
 
 WaldiezRagUserTask = Literal["code", "qa", "default"]
+"""Possible tasks for the retrieve chat."""
 WaldiezRagUserVectorDb = Literal["chroma", "pgvector", "mongodb", "qdrant"]
+"""Possible vector dbs for the retrieve chat."""
 WaldiezRagUserChunkMode = Literal["multi_lines", "one_line"]
+"""Possible chunk modes for the retrieve chat."""
 WaldiezRagUserModels: Dict[WaldiezRagUserVectorDb, str] = {
     "chroma": "all-MiniLM-L6-v2",
     "mongodb": "all-MiniLM-L6-v2",

waldiez/models/agents/reasoning/reasoning_agent_reason_config.py

@@ -8,6 +8,7 @@ from typing_extensions import Annotated, Literal
 from ...common import WaldiezBase
 
 ReasoningConfigMethod = Literal["beam_search", "mcts", "lats", "dfs"]
+"""Possible reasoning methods."""
 
 
 class WaldiezReasoningAgentReasonConfig(WaldiezBase):

waldiez/models/agents/swarm_agent/after_work.py

@@ -17,16 +17,18 @@ from typing_extensions import Annotated, Literal, Self
 from ...common import WaldiezBase, check_function, generate_function
 
 WaldiezSwarmAfterWorkRecipientType = Literal["agent", "option", "callable"]
+"""The possible AfterWork recipient types."""
 WaldiezSwarmAfterWorkOption = Literal[
     "TERMINATE", "REVERT_TO_USER", "STAY", "SWARM_MANAGER"
 ]
+"""The possible AfterWork options."""
 
 
 CUSTOM_AFTER_WORK = "custom_after_work"
 CUSTOM_AFTER_WORK_ARGS = ["last_speaker", "messages", "groupchat"]
 CUSTOM_AFTER_WORK_TYPES = (
-    ["SwarmAgent", "List[Dict[str, Any]]", "GroupChat"],
-    "Union[AfterWorkOption, SwarmAgent, str]",
+    ["ConversableAgent", "List[Dict[str, Any]]", "GroupChat"],
+    "Union[AfterWorkOption, ConversableAgent, str]",
 )
 
 
@@ -45,15 +47,15 @@ class WaldiezSwarmAfterWork(WaldiezBase):
     recipient_type : WaldiezSwarmAfterWorkRecipientType
         The type of recipient.
         Can be 'agent', 'option', or 'callable'.
-        If 'agent', the recipient is a SwarmAgent.
+        If 'agent', the recipient is a Swarm Agent.
         If 'option', the recipient is an AfterWorkOption :
             ('TERMINATE', 'REVERT_TO_USER', 'STAY', 'SWARM_MANAGER').
         If 'callable', it should have the signature:
         def custom_after_work(
-            last_speaker: SwarmAgent,
+            last_speaker: ConversableAgent,
             messages: List[dict],
             groupchat: GroupChat,
-        ) -> Union[AfterWorkOption, SwarmAgent, str]:
+        ) -> Union[AfterWorkOption, ConversableAgent, str]:
 
     """
 
@@ -79,15 +81,15 @@ class WaldiezSwarmAfterWork(WaldiezBase):
             description=(
                 "The type of recipient. "
                 "Can be 'agent', 'option', or 'callable'. "
-                "If 'agent', the recipient is a SwarmAgent.  "
+                "If 'agent', the recipient is a Swarm Agent.  "
                 "If 'option', the recipient is an AfterWorkOption :"
                 "    ('TERMINATE', 'REVERT_TO_USER', 'STAY', 'SWARM_MANAGER'). "
                 "If 'callable', it should have the signature: "
                 "def custom_after_work("
-                "    last_speaker: SwarmAgent,"
+                "    last_speaker: ConversableAgent,"
                 "    messages: List[Dict[str, Any]],"
                 "    groupchat: GroupChat,"
-                ") -> Union[AfterWorkOption, SwarmAgent, str]:"
+                ") -> Union[AfterWorkOption, ConversableAgent, str]:"
             ),
         ),
     ]
@@ -117,13 +119,13 @@ class WaldiezSwarmAfterWork(WaldiezBase):
             The recipient string and the function content if applicable.
         """
         if self.recipient_type == "option":
-            return f"AFTER_WORK(AfterWorkOption.{self.recipient})", ""
+            return f"AfterWork(AfterWorkOption.{self.recipient})", ""
         if self.recipient_type == "agent":
             # the the recipient is passed as the agent name
             # (and not its id), care should be taken to ensure
             # the all the agents in the flow have unique names
             agent_instance = agent_names.get(self.recipient, self.recipient)
-            return f"AFTER_WORK({agent_instance})", ""
+            return f"AfterWork({agent_instance})", ""
 
         function_name = CUSTOM_AFTER_WORK
         if name_prefix:
@@ -131,7 +133,7 @@ class WaldiezSwarmAfterWork(WaldiezBase):
         if name_suffix:
             function_name = f"{function_name}_{name_suffix}"
         return (
-            f"AFTER_WORK({function_name})",
+            f"AfterWork({function_name})",
             generate_function(
                 function_name=function_name,
                 function_args=CUSTOM_AFTER_WORK_ARGS,

waldiez/models/agents/swarm_agent/on_condition.py

@@ -12,6 +12,7 @@ from .on_condition_available import WaldiezSwarmOnConditionAvailable
 from .on_condition_target import WaldiezSwarmOnConditionTarget
 
 WaldiezSwarmOnConditionTargetType = Literal["agent", "nested_chat"]
+"""Possible types for the target of the OnCondition."""
 
 
 class WaldiezSwarmOnCondition(WaldiezBase):
@@ -30,7 +31,7 @@ class WaldiezSwarmOnCondition(WaldiezBase):
         The condition for transitioning to the target agent
 
     available: str, optional
-        Optional condition to determine if this ON_CONDITION is available.
+        Optional condition to determine if this OnCondition is available.
         Can be a Callable or a string.  If a string, it will look up the
         value of the context variable with that name, which should be a bool.
 
@@ -73,7 +74,7 @@ class WaldiezSwarmOnCondition(WaldiezBase):
             default_factory=WaldiezSwarmOnConditionAvailable,
             title="Available",
             description=(
-                "Optional condition to determine if this ON_CONDITION "
+                "Optional condition to determine if this OnCondition "
                 "is available."
             ),
         ),

waldiez/models/agents/swarm_agent/on_condition_available.py

@@ -12,6 +12,7 @@ from ...common import WaldiezBase, check_function, generate_function
 WaldiezSwarmOnConditionAvailableCheckType = Literal[
     "string", "callable", "none"
 ]
+"""Possible types for the `available` check."""
 
 CUSTOM_ON_CONDITION_AVAILABLE = "custom_on_condition_available"
 CUSTOM_ON_CONDITION_AVAILABLE_ARGS = ["agent", "message"]

waldiez/models/agents/swarm_agent/swarm_agent_data.py

@@ -29,7 +29,7 @@ class WaldiezSwarmAgentData(WaldiezAgentData):
         A list of functions (skill ids) to register with the agent.
 
     update_agent_state_before_reply : List[str]
-        A list of functions, including `UPDATE_SYSTEM_MESSAGE`,
+        A list of functions, including `UpdateSystemMessage`,
         called to update the agent's state before it replies. Each function
         is called when the agent is selected and before it speaks.
 
@@ -67,10 +67,10 @@ class WaldiezSwarmAgentData(WaldiezAgentData):
             title="Update Agent State Before Reply",
             alias="updateAgentStateBeforeReply",
             description=(
-                "A list of functions, including UPDATE_SYSTEM_MESSAGEs,"
+                "A list of functions, including UpdateSystemMessage,"
                 "called to update the agent's state before it replies. "
                 " Each function is called when the agent is selected "
-                "and before it speaks. If not an UPDATE_SYSTEM_MESSAGE, "
+                "and before it speaks. If not an UpdateSystemMessage, "
                 "it should be a skill id."
             ),
             default_factory=list,

waldiez/models/agents/swarm_agent/update_system_message.py

@@ -10,6 +10,7 @@ from typing_extensions import Annotated, Literal, Self
 from ...common import WaldiezBase, check_function, generate_function
 
 WaldiezSwarmUpdateFunctionType = Literal["string", "callable"]
+"""Possible types for the update function."""
 
 CUSTOM_UPDATE_SYSTEM_MESSAGE = "custom_update_system_message"
 CUSTOM_UPDATE_SYSTEM_MESSAGE_ARGS = ["agent", "messages"]

waldiez/models/chat/chat_message.py

@@ -12,6 +12,7 @@ from ..common import WaldiezBase, check_function, update_dict
 WaldiezChatMessageType = Literal[
     "string", "method", "rag_message_generator", "none"
 ]
+"""Possible types for the message."""
 
 CALLABLE_MESSAGE = "callable_message"
 CALLABLE_MESSAGE_ARGS = ["sender", "recipient", "context"]

waldiez/models/chat/chat_summary.py

@@ -20,6 +20,7 @@ WaldiezChatSummaryMethod = Literal[
     "reflection_with_llm",
     "last_msg",
 ]
+"""Possible methods for the LLM summary."""
 
 
 class WaldiezChatSummary(WaldiezBase):

waldiez/models/common/__init__.py

@@ -2,11 +2,13 @@
 # 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
 from .method_utils import (
     check_function,
+    gather_code_imports,
     generate_function,
     get_function,
     parse_code_string,
@@ -16,6 +18,8 @@ __all__ = [
     "WaldiezBase",
     "now",
     "check_function",
+    "gather_code_imports",
+    "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

@@ -65,6 +65,6 @@ class WaldiezBase(BaseModel):
             by_alias = True
         return super().model_dump_json(by_alias=by_alias, **kwargs)
 
-    def __hash__(self) -> int:
+    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/common/method_utils.py

@@ -3,6 +3,10 @@
 """Function related utilities."""
 
 import ast
+import importlib.util
+import sys
+import sysconfig
+from pathlib import Path
 from typing import List, Optional, Tuple
 
 import parso
@@ -12,6 +16,35 @@ import parso.tree
 MAX_VAR_NAME_LENGTH = 64
 
 
+def is_standard_library(module_name: str) -> bool:
+    """Check if the module is part of the standard library.
+
+    Parameters
+    ----------
+    module_name : str
+        The module name.
+
+    Returns
+    -------
+    bool
+        True if the module is part of the standard library.
+    """
+    if module_name in sys.builtin_module_names:
+        return True
+    try:
+        spec = importlib.util.find_spec(module_name)
+    except BaseException:  # pylint: disable=broad-except
+        return False
+    if spec is None or not spec.origin:
+        return False
+    if "site-packages" in spec.origin:
+        return False
+    if spec.origin.startswith(sys.prefix) or spec.origin == "frozen":
+        return True
+    stdlib_path = str(Path(sysconfig.get_path("stdlib")).resolve())
+    return spec.origin.startswith(stdlib_path)
+
+
 def parse_code_string(
     code_string: str,
 ) -> Tuple[Optional[str], Optional[ast.Module]]:
@@ -38,6 +71,71 @@ def parse_code_string(
     return None, tree
 
 
+def gather_code_imports(
+    code_string: str,
+    is_interop: bool,
+) -> Tuple[List[str], List[str]]:
+    """Gather the imports from the code string.
+
+    Parameters
+    ----------
+    code_string : str
+        The code string.
+    is_interop : bool
+        If True, make sure the interoperability import is present.
+
+    Returns
+    -------
+    Tuple[List[str], List[str]]
+        The standard library imports and the third party imports.
+    """
+    standard_lib_imports: List[str] = []
+    third_party_imports: List[str] = []
+    tree = parso.parse(code_string)  # type: ignore
+    for node in tree.iter_imports():
+        if node.type == "import_name":
+            full_import_statement = node.get_code().strip()
+            module_name = (
+                node.get_code().replace("import", "").strip().split(" ")[0]
+            )
+            if not module_name:
+                continue
+            if is_standard_library(module_name):
+                standard_lib_imports.append(full_import_statement)
+            else:
+                third_party_imports.append(full_import_statement)
+        elif node.type == "import_from":
+            full_import_statement = node.get_code().strip()
+            module_name = (
+                node.get_code().replace("from", "").strip().split(" ")[0]
+            )
+            if not module_name:
+                continue
+            if is_standard_library(module_name):
+                standard_lib_imports.append(full_import_statement)
+            else:
+                third_party_imports.append(full_import_statement)
+    if is_interop and (
+        "from autogen.interop import Interoperability"
+        not in third_party_imports
+    ):
+        third_party_imports.append(
+            "from autogen.interop import Interoperability"
+        )
+    # sorted_standard_lib_imports =  # first import x, then from a import b
+    sorted_standard_lib_imports = sorted(
+        [stmt for stmt in standard_lib_imports if stmt.startswith("import ")]
+    ) + sorted(
+        [stmt for stmt in standard_lib_imports if stmt.startswith("from ")]
+    )
+    sorted_third_party_imports = sorted(
+        [stmt for stmt in third_party_imports if stmt.startswith("import ")]
+    ) + sorted(
+        [stmt for stmt in third_party_imports if stmt.startswith("from ")]
+    )
+    return sorted_standard_lib_imports, sorted_third_party_imports
+
+
 def check_function(
     code_string: str,
     function_name: str,

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",
 ]

waldiez/models/flow/utils.py

@@ -4,7 +4,7 @@
 
 import uuid
 from datetime import datetime, timezone
-from typing import List
+from typing import Any, Dict, List, Optional
 
 from ..agents import (
     WaldiezAgent,
@@ -170,3 +170,63 @@ def merge_nested_chat_messages(
             chat_ids_added.append(chat.id)
     nested_chat.messages.sort(key=lambda x: chat_ids_added.index(x.id))
     return [nested_chat]
+
+
+def get_flow_data(
+    data: Dict[str, Any],
+    flow_id: Optional[str] = None,
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+    tags: Optional[List[str]] = None,
+    requirements: Optional[List[str]] = None,
+) -> Dict[str, Any]:
+    """Get the flow from the passed data dict.
+
+    Parameters
+    ----------
+    data : Dict[str, Any]
+        The data dict.
+    flow_id : Optional[str], optional
+        The flow ID, by default None.
+    name : Optional[str], optional
+        The flow name, by default None.
+    description : Optional[str], optional
+        The flow description, by default None.
+    tags : Optional[List[str]], optional
+        The flow tags, by default None.
+    requirements : Optional[List[str]], optional
+        The flow requirements, by default None.
+
+    Returns
+    -------
+    Dict[str, Any]
+        The flow data.
+
+    Raises
+    ------
+    ValueError
+        If the flow type is not "flow".
+    """
+    item_type = data.get("type", "flow")
+    if item_type != "flow":
+        # empty flow (from exported model/skill ?)
+        raise ValueError(f"Invalid flow type: {item_type}")
+    from_args = {
+        "id": flow_id,
+        "name": name,
+        "description": description,
+        "tags": tags,
+        "requirements": requirements,
+    }
+    for key, value in from_args.items():
+        if value:
+            data[key] = value
+    if "name" not in data:
+        data["name"] = "Waldiez Flow"
+    if "description" not in data:
+        data["description"] = "Waldiez Flow description"
+    if "tags" not in data:
+        data["tags"] = []
+    if "requirements" not in data:
+        data["requirements"] = []
+    return data

waldiez/models/model/__init__.py

@@ -2,10 +2,12 @@
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Waldiez model."""
 
+from .extra_requirements import get_models_extra_requirements
 from .model import DEFAULT_BASE_URLS, WaldiezModel
 from .model_data import WaldiezModelAPIType, WaldiezModelData, WaldiezModelPrice
 
 __all__ = [
+    "get_models_extra_requirements",
     "DEFAULT_BASE_URLS",
     "WaldiezModel",
     "WaldiezModelData",