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

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

@@ -9,12 +9,9 @@ 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", "captain"
-]
-
 
 class WaldiezAgent(WaldiezBase):
     """Waldiez Agent.
@@ -171,11 +168,11 @@ class WaldiezAgent(WaldiezBase):
         elif agent_class == "SwarmAgent":
             imports.add(
                 "from autogen import "
-                "AFTER_WORK, "
-                "ON_CONDITION, "
-                "UPDATE_SYSTEM_MESSAGE, "
+                "register_hand_off, "
+                "AfterWork, "
+                "OnCondition, "
+                "UpdateSystemMessage, "
                 "AfterWorkOption, "
-                "SwarmAgent, "
                 "SwarmResult"
             )
         elif agent_class == "ReasoningAgent":

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

@@ -22,7 +22,7 @@ class WaldiezCaptainAgent(WaldiezAgent):
     ----------
     agent_type : Literal["captain"]
         The agent type: 'captain' for a captain agent
-    data : WaldiezCaptainData
+    data : WaldiezCaptainAgentData
         The captain agent's data.
     """
 

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

@@ -8,6 +8,7 @@ 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,
@@ -17,6 +18,7 @@ __all__ = [
     "WaldiezBase",
     "now",
     "check_function",
+    "gather_code_imports",
     "get_autogen_version",
     "get_function",
     "generate_function",

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/model/__init__.py

@@ -3,12 +3,13 @@
 """Waldiez model."""
 
 from .extra_requirements import get_models_extra_requirements
-from .model import DEFAULT_BASE_URLS, WaldiezModel
+from .model import DEFAULT_BASE_URLS, MODEL_NEEDS_BASE_URL, WaldiezModel
 from .model_data import WaldiezModelAPIType, WaldiezModelData, WaldiezModelPrice
 
 __all__ = [
     "get_models_extra_requirements",
     "DEFAULT_BASE_URLS",
+    "MODEL_NEEDS_BASE_URL",
     "WaldiezModel",
     "WaldiezModelData",
     "WaldiezModelPrice",

waldiez/models/model/extra_requirements.py

@@ -36,6 +36,8 @@ def get_models_extra_requirements(
         "bedrock",  # we might add this later
     ]
     for model in models:
+        for requirement in model.requirements:
+            model_requirements.add(requirement)
         if model.data.api_type == "google":
             model_requirements.add(f"pyautogen[gemini]=={autogen_version}")
             continue

waldiez/models/model/model.py

@@ -23,6 +23,22 @@ DEFAULT_BASE_URLS: Dict[WaldiezModelAPIType, str] = {
 }
 
 
+# we can omit the base_url for these models
+MODEL_NEEDS_BASE_URL: Dict[WaldiezModelAPIType, bool] = {
+    "openai": False,
+    "azure": False,
+    "google": False,
+    "anthropic": False,
+    "cohere": False,
+    "other": False,  # falls back to openai
+    "deepseek": True,
+    "mistral": True,
+    "groq": True,
+    "together": True,
+    "nim": True,
+}
+
+
 class WaldiezModel(WaldiezBase):
     """Waldiez Model class.
 
@@ -225,10 +241,13 @@ def set_default_base_url(
     Dict[str, Any]
         The llm config dictionary with the default base url set.
     """
-    if api_type in ("openai", "other", "azure", "cohere"):
-        return llm_config
+    dict_copy = llm_config.copy()
     if "base_url" not in llm_config or not llm_config["base_url"]:
-        dict_copy = llm_config.copy()
-        dict_copy["base_url"] = DEFAULT_BASE_URLS.get(api_type, "")
-        return dict_copy
-    return llm_config
+        if MODEL_NEEDS_BASE_URL.get(api_type, True):
+            dict_copy["base_url"] = DEFAULT_BASE_URLS.get(api_type, "")
+    if (
+        not llm_config.get("base_url", "")
+        and MODEL_NEEDS_BASE_URL.get(api_type, True) is False
+    ):
+        dict_copy.pop("base_url", None)
+    return dict_copy

waldiez/models/model/model_data.py

@@ -23,6 +23,7 @@ WaldiezModelAPIType = Literal[
     "cohere",
     "other",
 ]
+"""Possible API types for the model."""
 
 
 class WaldiezModelPrice(WaldiezBase):

waldiez/models/skill/__init__.py

@@ -2,11 +2,15 @@
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Waldiez Skill related models."""
 
+from .extra_requirements import get_skills_extra_requirements
 from .skill import SHARED_SKILL_NAME, WaldiezSkill
 from .skill_data import WaldiezSkillData
+from .skill_type import WaldiezSkillType
 
 __all__ = [
     "SHARED_SKILL_NAME",
     "WaldiezSkill",
     "WaldiezSkillData",
+    "WaldiezSkillType",
+    "get_skills_extra_requirements",
 ]

waldiez/models/skill/extra_requirements.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Waldiez skill extra requirements."""
+
+from typing import Iterator, Set
+
+from .skill import WaldiezSkill
+
+
+def get_skills_extra_requirements(
+    skills: Iterator[WaldiezSkill],
+    autogen_version: str,
+) -> Set[str]:
+    """Get the skills extra requirements.
+
+    Parameters
+    ----------
+    skills : List[WaldiezSkill]
+        The skills.
+    autogen_version : str
+        The ag2 version.
+    Returns
+    -------
+    List[str]
+        The skills extra requirements.
+    """
+    skill_requirements: Set[str] = set()
+    for skill in skills:
+        if skill.skill_type == "langchain":
+            skill_requirements.add(
+                f"pyautogen[interop-langchain]=={autogen_version}"
+            )
+        if skill.skill_type == "crewai":
+            skill_requirements.add(
+                f"pyautogen[interop-crewai]=={autogen_version}"
+            )
+        for requirement in skill.requirements:
+            skill_requirements.add(requirement)
+    return skill_requirements