waldiez 0.1.0__py3-none-any.whl

waldiez/models/agents/rag_user/vector_db_config.py

@@ -0,0 +1,162 @@
+"""The vector db config for the RAG user agent."""
+
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+from pydantic import ConfigDict, Field, model_validator
+from pydantic.alias_generators import to_camel
+from typing_extensions import Annotated, Self
+
+from ...common import WaldiezBase
+
+
+# pylint: disable=line-too-long
+class WaldiezRagUserVectorDbConfig(WaldiezBase):
+    """The config for the vector db.
+
+    Attributes
+    ----------
+    model : str
+        The model to use for the vector db embeddings.
+    use_memory : bool
+        Whether to use memory for the vector db (if `qdrant` is used).
+    use_local_storage : bool
+        Whether to use local storage for the db (if `qdrant` or `chroma` is used).
+    local_storage_path : Optional[str]
+        The path to the local storage for the vector db (if `qdrant` or `chroma` is used).
+    connection_url : Optional[str]
+        The connection url for the vector db.
+    wait_until_index_ready : Optional[float]
+        Blocking call to wait until the database indexes are ready (if `mongodb` is used).
+        None, the default, means no wait.
+    wait_until_document_ready : Optional[float]
+        Blocking call to wait until the database documents are ready (if `mongodb` is used).
+        None, the default, means no wait.
+    metadata : Optional[Dict[str, Any]]
+        The metadata to use for the vector db.
+        Example: {"hnsw:space": "ip", "hnsw:construction_ef": 30, "hnsw:M": 32}
+
+    Functions
+    ---------
+    validate_vector_db_config()
+        Validate the vector db config.
+    """
+
+    model_config = ConfigDict(
+        extra="forbid",
+        alias_generator=to_camel,
+        populate_by_name=True,
+        frozen=False,
+        protected_namespaces=(),  # we use "model" as a field
+    )
+
+    model: Annotated[
+        Optional[str],
+        Field(
+            None,
+            title="Model",
+            description="The model to use for the vector db embeddings.",
+        ),
+    ]
+    use_memory: Annotated[
+        bool,
+        Field(
+            True,
+            title="Use Memory",
+            description=(
+                "Whether to use memory for the vector db (if qdrant is used)."
+            ),
+        ),
+    ]
+    use_local_storage: Annotated[
+        bool,
+        Field(
+            False,
+            title="Use Local Storage",
+            description=(
+                "Whether to use local storage for the vector db "
+                "(if qdrant or chroma is used)."
+            ),
+        ),
+    ]
+    local_storage_path: Annotated[
+        Optional[str],
+        Field(
+            None,
+            title="Local Storage Path",
+            description=(
+                "The path to the local storage for the vector db "
+                "(if qdrant is used)."
+            ),
+        ),
+    ]
+    connection_url: Annotated[
+        Optional[str],
+        Field(
+            None,
+            title="Connection URL",
+            description="The connection url for the vector db.",
+        ),
+    ]
+    wait_until_index_ready: Annotated[
+        Optional[float],
+        Field(
+            None,
+            title="Wait Until Index Ready",
+            description=(
+                "The time to wait/block until the database indexes are ready. "
+                "None, the default, means no wait."
+            ),
+        ),
+    ]
+    wait_until_document_ready: Annotated[
+        Optional[float],
+        Field(
+            None,
+            title="Wait Until Document Ready",
+            description=(
+                "The time to wait/block until the database documents "
+                "are ready. None, the default, means no wait."
+            ),
+        ),
+    ]
+    metadata: Annotated[
+        Optional[Dict[str, Any]],
+        Field(
+            None,
+            title="Metadata",
+            description=(
+                "The metadata to use for the vector db. Example: "
+                '"{"hnsw:space": "ip", '
+                '"hnsw:construction_ef": 30, '
+                '"hnsw:M": 32}'
+            ),
+        ),
+    ]
+
+    @model_validator(mode="after")
+    def validate_vector_db_config(self) -> Self:
+        """Validate the vector db config.
+
+        if local storage is used, make sure the path is provided,
+        and make it absolute if not already.
+
+        Returns
+        -------
+        WaldiezRagUserVectorDbConfig
+            The vector db config.
+
+        Raises
+        ------
+        ValueError
+            If the validation fails.
+        """
+        if self.use_local_storage:
+            if self.local_storage_path is None:
+                raise ValueError(
+                    "The local storage path must be provided if local storage is used."
+                )
+            as_path = Path(self.local_storage_path)
+            if not as_path.is_absolute():
+                self.local_storage_path = str(as_path.resolve())
+        return self

waldiez/models/agents/user_proxy/__init__.py

@@ -0,0 +1,6 @@
+"""User proxy agent model."""
+
+from .user_proxy import WaldiezUserProxy
+from .user_proxy_data import WaldiezUserProxyData
+
+__all__ = ["WaldiezUserProxy", "WaldiezUserProxyData"]

waldiez/models/agents/user_proxy/user_proxy.py

@@ -0,0 +1,41 @@
+"""User proxy agent model."""
+
+from pydantic import Field
+from typing_extensions import Annotated, Literal
+
+from ..agent import WaldiezAgent
+from .user_proxy_data import WaldiezUserProxyData
+
+
+class WaldiezUserProxy(WaldiezAgent):
+    """User proxy agent model.
+
+    A `WaldiezAgent` with agent_type `user` and
+    default `human_input_mode`: `"ALWAYS"`
+    See `WaldiezAgent`,`WaldiezUserProxyData`,`WaldiezAgentData` for more info.
+
+    Attributes
+    ----------
+    agent_type : Literal["user"]
+        The agent type: 'user' for a user proxy agent
+    data : WaldiezUserProxyData
+        The user proxy agent's data
+    """
+
+    agent_type: Annotated[
+        Literal["user"],
+        Field(
+            "user",
+            title="Agent type",
+            description="The agent type in a graph: 'user'",
+            alias="agentType",
+        ),
+    ]
+    data: Annotated[
+        WaldiezUserProxyData,
+        Field(
+            title="Data",
+            description="The user proxy agent's data",
+            default_factory=WaldiezUserProxyData,
+        ),
+    ]

waldiez/models/agents/user_proxy/user_proxy_data.py

@@ -0,0 +1,30 @@
+"""User proxy agent data module."""
+
+from pydantic import Field
+from typing_extensions import Annotated, Literal
+
+from ..agent import WaldiezAgentData
+
+
+class WaldiezUserProxyData(WaldiezAgentData):
+    """User proxy agent data class.
+
+    The data for an agent with `human_input_mode`
+    set to `"ALWAYS"` as default.
+    See the parent's docs (`WaldiezAgentData`) for the rest of the properties.
+
+    Attributes
+    ----------
+    human_input_mode : Literal["ALWAYS", "NEVER", "TERMINATE"]
+        The human input mode, Defaults to `ALWAYS`
+    """
+
+    human_input_mode: Annotated[
+        Literal["ALWAYS", "NEVER", "TERMINATE"],
+        Field(
+            "ALWAYS",
+            title="Human input mode",
+            description="The human input mode, Defaults to `ALWAYS`",
+            alias="humanInputMode",
+        ),
+    ]

waldiez/models/chat/__init__.py

@@ -0,0 +1,22 @@
+"""Waldiez chat related models."""
+
+from .chat import WaldiezChat
+from .chat_data import WaldiezChatData
+from .chat_message import (
+    WaldiezChatMessage,
+    WaldiezChatMessageType,
+    validate_message_dict,
+)
+from .chat_nested import WaldiezChatNested
+from .chat_summary import WaldiezChatSummary, WaldiezChatSummaryMethod
+
+__all__ = [
+    "WaldiezChat",
+    "WaldiezChatData",
+    "WaldiezChatMessage",
+    "WaldiezChatMessageType",
+    "WaldiezChatNested",
+    "WaldiezChatSummary",
+    "WaldiezChatSummaryMethod",
+    "validate_message_dict",
+]

waldiez/models/chat/chat.py

@@ -0,0 +1,129 @@
+"""Waldiez chat model."""
+
+from typing import Any, Dict, Optional
+
+from pydantic import Field
+from typing_extensions import Annotated
+
+from ..agents import WaldiezAgent, WaldiezRagUser
+from ..common import WaldiezBase
+from .chat_data import WaldiezChatData
+from .chat_message import WaldiezChatMessage
+from .chat_nested import WaldiezChatNested
+
+
+class WaldiezChat(WaldiezBase):
+    """Chat class.
+
+    Attributes
+    ----------
+    id : str
+        The chat ID.
+    data : WaldiezChatData
+        The chat data.
+        See `waldiez.models.chat.WaldiezChatData` for more information.
+    name : str
+        The chat name.
+    source : str
+        The chat source.
+    target : str
+        The chat target.
+    nested_chat : WaldiezChatNested
+        The nested chat message/reply if any.
+    message : WaldiezChatMessage
+        The chat message.
+    message_content : Optional[str]
+        The chat message content if any. If method, the method's body.
+
+    Functions
+    ---------
+    get_chat_args()
+        Get the chat arguments to use in autogen.
+    """
+
+    id: Annotated[
+        str,
+        Field(
+            ...,
+            title="ID",
+            description="The chat ID.",
+        ),
+    ]
+    data: Annotated[
+        WaldiezChatData,
+        Field(
+            ...,
+            title="Data",
+            description="The chat data.",
+        ),
+    ]
+
+    @property
+    def name(self) -> str:
+        """Get the name."""
+        return self.data.name
+
+    @property
+    def source(self) -> str:
+        """Get the source."""
+        if self.data.real_source:
+            return self.data.real_source
+        return self.data.source
+
+    @property
+    def target(self) -> str:
+        """Get the target."""
+        if self.data.real_target:
+            return self.data.real_target
+        return self.data.target
+
+    @property
+    def nested_chat(self) -> WaldiezChatNested:
+        """Get the nested chat."""
+        return self.data.nested_chat
+
+    @property
+    def message(self) -> WaldiezChatMessage:
+        """Get the message."""
+        if isinstance(
+            self.data.message, str
+        ):  # pragma: no cover (just for the lint)
+            return WaldiezChatMessage(
+                type="string",
+                use_carryover=False,
+                content=self.data.message,
+                context={},
+            )
+        return self.data.message
+
+    @property
+    def message_content(self) -> Optional[str]:
+        """Get the message content."""
+        return self.data.message_content
+
+    def get_chat_args(
+        self,
+        sender: Optional[WaldiezAgent] = None,
+    ) -> Dict[str, Any]:
+        """Get the chat arguments to use in autogen.
+
+        Parameters
+        ----------
+        sender : WaldiezAgent
+            The sender agent, to check if it's a RAG user.
+        Returns
+        -------
+        dict
+            The chat arguments.
+        """
+        args_dict = self.data.get_chat_args()
+        if (
+            isinstance(sender, WaldiezRagUser)
+            and sender.agent_type == "rag_user"
+            and self.message.type == "rag_message_generator"
+        ):
+            # check for n_results in agent data, to add in context
+            n_results = sender.data.retrieve_config.n_results
+            if isinstance(n_results, int) and n_results > 0:
+                args_dict["n_results"] = n_results
+        return args_dict