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

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

@@ -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/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",

waldiez/models/model/extra_requirements.py

@@ -0,0 +1,55 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Waldiez model extra requirements."""
+
+from typing import Iterator, Set
+
+from .model import WaldiezModel
+
+
+def get_models_extra_requirements(
+    models: Iterator[WaldiezModel], autogen_version: str
+) -> Set[str]:
+    """Get the models extra requirements.
+
+    Parameters
+    ----------
+    models : List[WaldiezModel]
+        The models.
+    autogen_version : str
+        The autogen version.
+
+    Returns
+    -------
+    List[str]
+        The models extra requirements.
+    """
+    model_requirements: Set[str] = set()
+    # ref: https://github.com/ag2ai/ag2/blob/main/pyproject.toml
+    models_with_additional_requirements = [
+        "together",
+        "gemini",
+        "mistral",
+        "groq",
+        "anthropic",
+        "cohere",
+        "bedrock",  # we might add this later
+    ]
+    for model in models:
+        if model.data.api_type == "google":
+            model_requirements.add(f"pyautogen[gemini]=={autogen_version}")
+            continue
+        if model.data.api_type in models_with_additional_requirements:
+            model_requirements.add(
+                f"pyautogen[{model.data.api_type}]=={autogen_version}"
+            )
+    return model_requirements
+
+
+# for bedrock, there are some additional params needed
+# (both here and on the ui part):
+# aws_region (mandatory)
+# aws_access_key (or environment variable: AWS_ACCESS_KEY)
+# aws_secret_key (or environment variable: AWS_SECRET_KEY)
+# aws_session_token (or environment variable: AWS_SESSION_TOKEN)
+# aws_profile_name

waldiez/models/model/model.py

@@ -19,6 +19,7 @@ DEFAULT_BASE_URLS: Dict[WaldiezModelAPIType, str] = {
     "groq": "https://api.groq.com/openai/v1",
     "together": "https://api.together.xyz/v1",
     "nim": "https://integrate.api.nvidia.com/v1",
+    "cohere": "https://api.cohere.com",
 }
 
 
@@ -63,7 +64,7 @@ class WaldiezModel(WaldiezBase):
     description: Annotated[
         str,
         Field(
-            ...,
+            "Model's Description",
             title="Description",
             description="The description of the model.",
         ),
@@ -118,6 +119,7 @@ class WaldiezModel(WaldiezBase):
         - groq: 'GROQ_API_KEY',
         - together: 'TOGETHER_API_KEY',
         - nim: 'NIM_API_KEY',
+        - cohere: 'COHERE_API_KEY',
         - other: 'OPENAI_API_KEY'
         """
         env_key = "OPENAI_API_KEY"
@@ -142,6 +144,7 @@ class WaldiezModel(WaldiezBase):
             - groq: 'GROQ_API_KEY',
             - together: 'TOGETHER_API_KEY',
             - nim: 'NIM_API_KEY',
+            - cohere: 'COHERE_API_KEY',
             - other: 'OPENAI_API_KEY'
         """
         if self.data.api_key and self.data.api_key != "REPLACE_ME":
@@ -222,7 +225,7 @@ 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"):
+    if api_type in ("openai", "other", "azure", "cohere"):
         return llm_config
     if "base_url" not in llm_config or not llm_config["base_url"]:
         dict_copy = llm_config.copy()

waldiez/models/model/model_data.py

@@ -20,6 +20,7 @@ WaldiezModelAPIType = Literal[
     "groq",
     "together",
     "nim",
+    "cohere",
     "other",
 ]
 
@@ -53,7 +54,7 @@ class WaldiezModelData(WaldiezBase):
         The base url of the model, by default None.
     api_key : Optional[str]
         The api key to use with the model, by default None.
-    api_type : Literal["openai","azure","deepseek","google","anthropic","mistral","groq","together","nim","other"]
+    api_type : WaldiezModelAPIType
         The api type of the model.
     api_version : Optional[str]
         The api version of the model, by default None.

waldiez/models/waldiez.py

@@ -8,16 +8,18 @@ definitions and their optional additional skills to be used.
 """
 
 import json
-import warnings
 from dataclasses import dataclass
-from functools import cache
 from pathlib import Path
 from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
 
-from .agents import WaldiezAgent
+from .agents import (
+    WaldiezAgent,
+    get_retrievechat_extra_requirements,
+)
 from .chat import WaldiezChat
-from .flow import WaldiezFlow
-from .model import WaldiezModel
+from .common import get_autogen_version
+from .flow import WaldiezFlow, get_flow_data
+from .model import WaldiezModel, get_models_extra_requirements
 from .skill import WaldiezSkill
 
 
@@ -62,7 +64,7 @@ class Waldiez:
         Waldiez
             The Waldiez.
         """
-        flow = _get_flow(
+        flow = get_flow_data(
             data,
             flow_id=flow_id,
             name=name,
@@ -224,8 +226,9 @@ class Waldiez:
     @property
     def requirements(self) -> List[str]:
         """Get the flow requirements."""
-        autogen_version = _get_autogen_version()
+        autogen_version = get_autogen_version()
         requirements_list = filter(
+            # we use the fixed "pyautogen=={autogen_version}" below
             lambda requirement: not (
                 requirement.startswith("pyautogen")
                 or requirement.startswith("ag2")
@@ -234,30 +237,15 @@ class Waldiez:
             self.flow.requirements,
         )
         requirements = set(requirements_list)
+        requirements.add(f"pyautogen=={autogen_version}")
         if self.has_rag_agents:
-            requirements.add(f"pyautogen[retrievechat]=={autogen_version}")
-        else:
-            requirements.add(f"pyautogen=={autogen_version}")
+            rag_extras = get_retrievechat_extra_requirements(self.agents)
+            requirements.update(rag_extras)
         if self.has_multimodal_agents:
             requirements.add(f"pyautogen[lmm]=={autogen_version}")
-        # ref: https://github.com/ag2ai/ag2/blob/main/pyproject.toml
-        models_with_additional_requirements = [
-            "together",
-            "gemini",
-            "mistral",
-            "groq",
-            "anthropic",
-            "cohere",
-            "bedrock",
-        ]
-        for model in self.models:
-            if model.data.api_type == "google":
-                requirements.add(f"pyautogen[gemini]=={autogen_version}")
-                continue
-            if model.data.api_type in models_with_additional_requirements:
-                requirements.add(
-                    f"pyautogen[{model.data.api_type}]=={autogen_version}"
-                )
+        requirements.update(
+            get_models_extra_requirements(self.models, autogen_version)
+        )
         return sorted(list(requirements))
 
     def get_flow_env_vars(self) -> List[Tuple[str, str]]:
@@ -307,50 +295,3 @@ class Waldiez:
             The swarm agents and the user agent.
         """
         return self.flow.get_swarm_chat_members(initial_agent)
-
-
-def _get_flow(
-    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."""
-    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
-
-
-@cache
-def _get_autogen_version() -> str:
-    """Get the autogen version."""
-    # 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