waldiez 0.1.0__py3-none-any.whl

waldiez/models/flow/flow.py

@@ -0,0 +1,302 @@
+"""Waldiez flow model."""
+
+import uuid
+from typing import List, Optional, Tuple
+
+from pydantic import Field, model_validator
+from typing_extensions import Annotated, Literal, Self
+
+from ..agents import WaldiezAgent
+from ..chat import WaldiezChat
+from ..common import WaldiezBase, now
+from .flow_data import WaldiezFlowData
+
+
+class WaldiezFlow(WaldiezBase):
+    """Flow data class.
+
+    Attributes
+    ----------
+    id : str
+        The ID of the flow.
+    type : Literal["flow"], optional
+        The type of the "node" in a graph: "flow".
+    name : str
+        The name of the flow.
+    description : str
+        The description of the flow.
+    tags : List[str]
+        The tags of the flow.
+    requirements : List[str]
+        The requirements of the flow.
+    storage_id : str
+        The storage ID of the flow (ignored, UI related).
+    created_at : str
+        The date and time when the flow was created.
+    updated_at : str
+        The date and time when the flow was last updated.
+    data : WaldiezFlowData
+        The data of the flow. See `WaldiezFlowData`.
+    """
+
+    id: Annotated[
+        str,
+        Field(
+            description="The ID of the flow",
+            title="ID",
+            default_factory=uuid.uuid4,
+        ),
+    ]
+    type: Annotated[
+        Literal["flow"],
+        Field(
+            "flow",
+            description="The type of the 'node' in a graph",
+            title="Type",
+        ),
+    ]
+    name: Annotated[
+        str,
+        Field(
+            ...,
+            description="The name of the flow",
+            title="Name",
+        ),
+    ]
+    description: Annotated[
+        str,
+        Field(
+            ...,
+            description="The description of the flow",
+            title="Description",
+        ),
+    ]
+    tags: Annotated[
+        List[str],
+        Field(
+            description="The tags of the flow",
+            title="Tags",
+            default_factory=list,
+        ),
+    ]
+    requirements: Annotated[
+        List[str],
+        Field(
+            description="The requirements of the flow",
+            title="Requirements",
+            default_factory=list,
+        ),
+    ]
+    data: Annotated[
+        WaldiezFlowData,
+        Field(
+            ...,
+            description="The data of the flow",
+            title="Data",
+        ),
+    ]
+    storage_id: Annotated[
+        str,
+        Field(
+            uuid.uuid4(),
+            description="The storage ID of the flow (ignored, UI related)",
+            title="Storage ID",
+            alias="storageId",
+        ),
+    ]
+    created_at: Annotated[
+        str,
+        Field(
+            default_factory=now,
+            title="Created At",
+            description="The date and time when the flow was created.",
+        ),
+    ]
+    updated_at: Annotated[
+        str,
+        Field(
+            default_factory=now,
+            title="Updated At",
+            description="The date and time when the flow was last updated.",
+        ),
+    ]
+    _ordered_flow: Optional[
+        List[Tuple[WaldiezChat, WaldiezAgent, WaldiezAgent]]
+    ] = None
+
+    @property
+    def ordered_flow(
+        self,
+    ) -> List[Tuple[WaldiezChat, WaldiezAgent, WaldiezAgent]]:
+        """Get the ordered flow."""
+        if not self._ordered_flow:
+            self._ordered_flow = self._get_flow_order()
+        return self._ordered_flow
+
+    def get_agent_by_id(self, agent_id: str) -> WaldiezAgent:
+        """Get the agent by ID.
+
+        Parameters
+        ----------
+        agent_id : str
+            The ID of the agent.
+
+        Returns
+        -------
+        WaldiezAgent
+            The agent.
+
+        Raises
+        ------
+        ValueError
+            If the agent with the given ID is not found.
+        """
+        for user in self.data.agents.users:
+            if user.id == agent_id:
+                return user
+        for assistant in self.data.agents.assistants:
+            if assistant.id == agent_id:
+                return assistant
+        for manager in self.data.agents.managers:
+            if manager.id == agent_id:
+                return manager
+        for rag_user in self.data.agents.rag_users:
+            if rag_user.id == agent_id:
+                return rag_user
+        raise ValueError(f"Agent with ID {agent_id} not found.")
+
+    def _get_flow_order(
+        self,
+    ) -> List[Tuple[WaldiezChat, WaldiezAgent, WaldiezAgent]]:
+        """Get the ordered flow."""
+        # in the chats, there is the 'order' field, we use this,
+        # we only keep the ones with order >=0
+        # and sort them by this property
+        ordered_flow: List[Tuple[WaldiezChat, WaldiezAgent, WaldiezAgent]] = []
+        sorted_chats_by_order = sorted(
+            self.data.chats, key=lambda chat: chat.data.order
+        )
+        for chat in sorted_chats_by_order:
+            if chat.data.order < 0:
+                continue
+            source = self.get_agent_by_id(chat.source)
+            target = self.get_agent_by_id(chat.target)
+            ordered_flow.append((chat, source, target))
+        if not ordered_flow:
+            if len(self.data.chats) == 1:
+                chat = self.data.chats[0]
+                source = self.get_agent_by_id(chat.source)
+                target = self.get_agent_by_id(chat.target)
+                ordered_flow.append((chat, source, target))
+        return ordered_flow
+
+    def get_agent_connections(
+        self, agent_id: str, all_chats: bool = True
+    ) -> List[str]:
+        """Get the agent connections.
+
+        Parameters
+        ----------
+        agent_id : str
+            The ID of the agent.
+        all_chats : bool, optional
+            If True, get the connections from all the chats, otherwise
+            get the connections from the ordered flow (main chat flow).
+
+        Returns
+        -------
+        List[str]
+            The list of agent ids that the agent with the given ID connects to.
+        """
+        connections = []
+        if all_chats:
+            for chat in self.data.chats:
+                if chat.source == agent_id:
+                    connections.append(chat.target)
+                if chat.target == agent_id:
+                    connections.append(chat.source)
+        else:
+            for _, source, target in self.ordered_flow:
+                if source.id == agent_id:
+                    connections.append(target.id)
+                if target.id == agent_id:
+                    connections.append(source.id)
+        return connections
+
+    def get_group_chat_members(
+        self, group_manager_id: str
+    ) -> List[WaldiezAgent]:
+        """Get the group chat members.
+
+        Parameters
+        ----------
+        group_manager_id : str
+            The ID of the group manager.
+
+        Returns
+        -------
+        List[WaldiezAgent]
+            The list of group chat
+        """
+        agent = self.get_agent_by_id(group_manager_id)
+        if agent.agent_type != "manager":
+            return []
+        connections = self.get_agent_connections(group_manager_id)
+        return [self.get_agent_by_id(member_id) for member_id in connections]
+
+    def _validate_agent_connections(self) -> None:
+        for agent in self.data.agents.members:
+            if not any(
+                agent.id in (chat.source, chat.target)
+                for chat in self.data.chats
+            ):
+                raise ValueError(
+                    f"Agent {agent.id} ({agent.name}) "
+                    "does not connect to any other node."
+                )
+            # already covered above
+            # if agent.agent_type == "manager":
+            #     chat_member_ids = self.get_agent_connections(agent.id)
+            #     if not chat_member_ids:
+            #         raise ValueError(
+            #             f"Manager's {agent.id} group chat has no members."
+            #         )
+
+    @model_validator(mode="after")
+    def validate_flow(self) -> Self:
+        """Flow validation.
+
+        - unique node ids
+        - there are at least two agents
+        - all the agents connect to at least one other agent
+        - all the linked agent skills are found in the flow
+        - all the linked agent models are found in the flow
+        - all the managers have at least one member in the chat group
+        - the ordered flow (chats with position >=0) is not empty
+        - all agents' code execution config functions exist in the flow skills
+
+        Returns
+        -------
+        WaldiezFlow
+            The validated flow.
+
+        Raises
+        ------
+        ValueError
+            If the ordered flow is empty.
+            If the model IDs are not unique.
+            If the skill IDs are not unique.
+            If the agents do not connect to any other node.
+            If the manager's group chat has no members.
+        """
+        if not self.ordered_flow:
+            raise ValueError("The ordered flow is empty.")
+        model_ids = [model.id for model in self.data.models]
+        if len(model_ids) != len(set(model_ids)):
+            raise ValueError("Model IDs must be unique.")
+        skills_ids = [skill.id for skill in self.data.skills]
+        if len(skills_ids) != len(set(skills_ids)):
+            raise ValueError("Skill IDs must be unique.")
+        self.data.agents.validate_flow(model_ids, skills_ids)
+        self._validate_agent_connections()
+        return self

waldiez/models/flow/flow_data.py

@@ -0,0 +1,87 @@
+"""Waldiez flow data."""
+
+from typing import Any, Dict, List
+
+from pydantic import Field
+from typing_extensions import Annotated
+
+from ..agents import WaldiezAgents
+from ..chat import WaldiezChat
+from ..common import WaldiezBase
+from ..model import WaldiezModel
+from ..skill import WaldiezSkill
+
+
+class WaldiezFlowData(WaldiezBase):
+    """Flow data class.
+
+    Attributes
+    ----------
+    nodes : List[Dict[str, Any]]
+        The nodes of the flow. We ignore this (UI-related)
+    edges : List[Dict[str, Any]]
+        The edges of the flow. We ignore this (UI-related)
+    viewport : Dict[str, Any]
+        The viewport of the flow. We ignore this (UI-related)
+    agents : WaldiezAgents
+        The agents of the flow:
+        users: List[WaldiezUserProxy]
+        assistants: List[WaldiezAssistant]
+        managers: List[WaldiezGroupManager]
+        rag_users : List[WaldiezRagUser]
+        See `WaldiezAgents` for more info.
+    models : List[WaldiezModel]
+        The models of the flow. See `WaldiezModel`.
+    skills : List[WaldiezSkill]
+        The skills of the flow. See `WaldiezSkill`.
+    chats : List[WaldiezChat]
+        The chats of the flow. See `WaldiezChat`.
+    """
+
+    # the ones below (nodes,edges, viewport) we ignore
+    # (they for graph connections, positions, etc.)
+    nodes: Annotated[
+        List[Dict[str, Any]],
+        Field(default_factory=list),
+    ]
+    edges: Annotated[
+        List[Dict[str, Any]],
+        Field(default_factory=list),
+    ]
+    viewport: Annotated[
+        Dict[str, Any],
+        Field(default_factory=dict),
+    ]
+    # these are the ones we use.
+    agents: Annotated[
+        WaldiezAgents,
+        Field(
+            description="The agents of the flow",
+            title="Agents",
+            default_factory=WaldiezAgents,
+        ),
+    ]
+    models: Annotated[
+        List[WaldiezModel],
+        Field(
+            description="The models of the flow",
+            title="Models",
+            default_factory=list,
+        ),
+    ]
+    skills: Annotated[
+        List[WaldiezSkill],
+        Field(
+            description="The skills of the flow",
+            title="Skills",
+            default_factory=list,
+        ),
+    ]
+    chats: Annotated[
+        List[WaldiezChat],
+        Field(
+            description="The chats of the flow",
+            title="Chats",
+            default_factory=list,
+        ),
+    ]

waldiez/models/model/__init__.py

@@ -0,0 +1,11 @@
+"""Waldiez model."""
+
+from .model import WaldiezModel
+from .model_data import WaldiezModelAPIType, WaldiezModelData, WaldiezModelPrice
+
+__all__ = [
+    "WaldiezModel",
+    "WaldiezModelData",
+    "WaldiezModelPrice",
+    "WaldiezModelAPIType",
+]

waldiez/models/model/model.py

@@ -0,0 +1,169 @@
+"""Waldiez model(llm) model."""
+
+import os
+from typing import Any, Dict, List, Optional
+
+from pydantic import Field
+from typing_extensions import Annotated, Literal
+
+from ..common import WaldiezBase, now
+from .model_data import WaldiezModelData
+
+
+class WaldiezModel(WaldiezBase):
+    """Waldiez Model class.
+
+    Attributes
+    ----------
+    id : str
+        The ID of the model.
+    name : str
+        The name of the model.
+    description : str
+        The description of the model.
+    tags : List[str]
+        The tags of the model.
+    requirements : List[str]
+        The requirements of the model.
+    created_at : str
+        The date and time when the model was created.
+    updated_at : str
+        The date and time when the model was last updated.
+    data : WaldiezModelData
+        The data of the model.
+        See `waldiez.models.model.WaldiezModelData` for more information.
+    """
+
+    id: Annotated[
+        str, Field(..., title="ID", description="The ID of the model.")
+    ]
+    type: Annotated[
+        Literal["model"],
+        Field(
+            default="model",
+            title="Type",
+            description="The type of the 'node' in a graph.",
+        ),
+    ]
+    name: Annotated[
+        str, Field(..., title="Name", description="The name of the model.")
+    ]
+    description: Annotated[
+        str,
+        Field(
+            ...,
+            title="Description",
+            description="The description of the model.",
+        ),
+    ]
+    tags: Annotated[
+        List[str],
+        Field(
+            default_factory=list,
+            title="Tags",
+            description="The tags of the model.",
+        ),
+    ]
+    requirements: Annotated[
+        List[str],
+        Field(
+            default_factory=list,
+            title="Requirements",
+            description="The requirements of the model.",
+        ),
+    ]
+    created_at: Annotated[
+        str,
+        Field(
+            default_factory=now,
+            title="Create At",
+            description="The date and time when the model was created.",
+        ),
+    ]
+    updated_at: Annotated[
+        str,
+        Field(
+            default_factory=now,
+            title="Updated At",
+            description="The date and time when the model was last updated.",
+        ),
+    ]
+    data: Annotated[
+        WaldiezModelData,
+        Field(..., title="Data", description="The data of the model."),
+    ]
+
+    @property
+    def api_key(self) -> str:
+        """Get the model's api key.
+
+        Either from the model's data or from the environment variables:
+
+            - openai: 'OPENAI_API_KEY',
+            - azure: 'AZURE_API_KEY',
+            - google: 'GOOGLE_GEMINI_API_KEY',
+            - anthropic: 'ANTHROPIC_API_KEY',
+            - mistral: 'MISTRAL_API_KEY',
+            - groq: 'GROQ_API_KEY',
+            - together: 'TOGETHER_API_KEY',
+            - nim: 'NIM_API_KEY',
+            - other: 'OPENAI_API_KEY'
+        """
+        if self.data.api_key:
+            return self.data.api_key
+        env_key = "OPENAI_API_KEY"
+        if self.data.api_type == "google":
+            env_key = "GOOGLE_GEMINI_API_KEY"
+        elif self.data.api_type not in ["openai", "other"]:
+            env_key = f"{self.data.api_type.upper()}_API_KEY"
+        api_key = os.environ.get(env_key, "")
+        return api_key
+
+    @property
+    def price(self) -> Optional[List[float]]:
+        """Get the model's price."""
+        if self.data.price is None:
+            return None
+        if isinstance(
+            self.data.price.prompt_price_per_1k, float
+        ) and isinstance(self.data.price.completion_token_price_per_1k, float):
+            return [
+                self.data.price.prompt_price_per_1k,
+                self.data.price.completion_token_price_per_1k,
+            ]
+        return None
+
+    def get_llm_config(self, skip_price: bool = False) -> Dict[str, Any]:
+        """Get the model's llm config.
+
+        Parameters
+        ----------
+        skip_price : bool, optional
+            Whether to skip the price, by default False
+
+        Returns
+        -------
+        Dict[str, Any]
+            The model's llm config dictionary.
+        """
+        _llm_config: Dict[str, Any] = {}
+        _llm_config["model"] = self.name
+        for attr, atr_type in [
+            ("base_url", str),
+            ("max_tokens", int),
+            ("temperature", float),
+            ("top_p", float),
+            ("api_version", str),
+            ("default_headers", dict),
+        ]:
+            value = getattr(self.data, attr)
+            if value and isinstance(value, atr_type):
+                _llm_config[attr] = value
+        if self.data.api_type not in ["nim", "other"]:
+            _llm_config["api_type"] = self.data.api_type
+        other_attrs = ["api_key"] if skip_price else ["api_key", "price"]
+        for attr in other_attrs:
+            value = getattr(self, attr)
+            if value:
+                _llm_config[attr] = value
+        return _llm_config

waldiez/models/model/model_data.py

@@ -0,0 +1,86 @@
+"""Waldiez Model Data."""
+
+from typing import Dict, Optional
+
+from pydantic import Field
+from typing_extensions import Annotated, Literal
+
+from ..common import WaldiezBase
+
+WaldiezModelAPIType = Literal[
+    "openai",
+    "azure",
+    "google",
+    "anthropic",
+    "mistral",
+    "groq",
+    "together",
+    "nim",
+    "other",
+]
+
+
+class WaldiezModelPrice(WaldiezBase):
+    """Model Price.
+
+    Attributes
+    ----------
+    prompt_price_per_1k : float
+        The prompt price per 1k tokens.
+    completion_token_price_per_1k : float
+        The completion token price per 1k tokens.
+    """
+
+    prompt_price_per_1k: Annotated[
+        Optional[float], Field(None, alias="promptPricePer1k")
+    ]
+    completion_token_price_per_1k: Annotated[
+        Optional[float], Field(None, alias="completionTokenPricePer1k")
+    ]
+
+
+# pylint: disable=line-too-long
+class WaldiezModelData(WaldiezBase):
+    """Waldiez Model Data.
+
+    Attributes
+    ----------
+    base_url : Optional[str]
+        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","google","anthropic","mistral","groq","together","nim","other"]
+        The api type of the model.
+    api_version : Optional[str]
+        The api version of the model, by default None.
+    temperature : Optional[float]
+        The temperature of the model, by default None.
+    top_p : Optional[float]
+        The top p of the model, by default None.
+    max_tokens : Optional[int]
+        The max tokens of the model, by default None.
+    default_headers : Dict[str, str]
+        The default headers of the model.
+    price : Optional[WaldiezModelPrice]
+        The price of the model, by default None.
+    """
+
+    base_url: Annotated[
+        Optional[str],
+        Field(
+            None,
+            title="Base URL",
+            description="The base url of the model",
+            alias="baseUrl",
+        ),
+    ]
+    api_key: Annotated[Optional[str], Field(None, alias="apiKey")]
+    api_type: Annotated[WaldiezModelAPIType, Field("other", alias="apiType")]
+    api_version: Annotated[Optional[str], Field(None, alias="apiVersion")]
+    temperature: Annotated[Optional[float], Field(None, alias="temperature")]
+    top_p: Annotated[Optional[float], Field(None, alias="topP")]
+    max_tokens: Annotated[Optional[int], Field(None, alias="maxTokens")]
+    default_headers: Annotated[
+        Dict[str, str], Field(alias="defaultHeaders", default_factory=dict)
+    ]
+    price: Annotated[Optional[WaldiezModelPrice], Field(None)]

waldiez/models/skill/__init__.py

@@ -0,0 +1,9 @@
+"""Waldiez Skill related models."""
+
+from .skill import WaldiezSkill
+from .skill_data import WaldiezSkillData
+
+__all__ = [
+    "WaldiezSkill",
+    "WaldiezSkillData",
+]