waldiez 0.5.2__py3-none-any.whl → 0.5.4__py3-none-any.whl

waldiez/models/agents/agent/agent.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0.
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
+# pylint: disable=too-many-public-methods
 """Base agent class to be inherited by all agents."""
 
 import warnings
@@ -150,6 +151,28 @@ class WaldiezAgent(WaldiezBase):
     )
     _checked_handoffs: Annotated[bool, Field(init=False, default=False)] = False
 
+    @property
+    def args_to_skip(self) -> list[str]:
+        """Get the set of arguments to skip when generating the agent string.
+
+        Returns
+        -------
+        list[str]
+            The list of arguments to skip.
+        """
+        if self.is_doc_agent:
+            return [
+                "description",
+                "human_input_mode",
+                "max_consecutive_auto_reply",
+                "default_auto_reply",
+                "code_execution_config",
+                "is_termination_msg",
+                "functions",
+                "update_agent_state_before_reply",
+            ]
+        return []
+
     @property
     def handoffs(self) -> list[WaldiezHandoff]:
         """Get the handoffs for this agent.
@@ -266,6 +289,17 @@ class WaldiezAgent(WaldiezBase):
             "rag_user_proxy",
         )
 
+    @property
+    def is_assistant(self) -> bool:
+        """Check if the agent is an assistant.
+
+        Returns
+        -------
+        bool
+            True if the agent is an assistant, False otherwise.
+        """
+        return self.agent_type == "assistant"
+
     @property
     def is_rag_user(self) -> bool:
         """Check if the agent is a RAG user.
@@ -277,6 +311,17 @@ class WaldiezAgent(WaldiezBase):
         """
         return self.agent_type in ("rag_user", "rag_user_proxy")
 
+    @property
+    def is_doc_agent(self) -> bool:
+        """Check if the agent is a doc agent.
+
+        Returns
+        -------
+        bool
+            True if the agent is a doc agent, False otherwise.
+        """
+        return self.agent_type == "doc_agent"
+
     @property
     def is_group_manager(self) -> bool:
         """Check if the agent is a group manager.
@@ -293,16 +338,9 @@ class WaldiezAgent(WaldiezBase):
         """Return the AG2 class of the agent."""
         class_name = "ConversableAgent"
         if self.is_group_member:
-            if (
-                getattr(self.data, "is_multimodal", False) is True
-            ):  # pragma: no branch
-                class_name = "MultimodalConversableAgent"
-            return class_name
-        if self.agent_type == "assistant":
-            if getattr(self.data, "is_multimodal", False) is True:
-                class_name = "MultimodalConversableAgent"
-            else:
-                class_name = "AssistantAgent"
+            return self.get_group_member_class_name()
+        if self.is_assistant:  # pragma: no branch
+            return self.get_assistant_class_name()
         if self.is_user:
             class_name = "UserProxyAgent"
         if self.is_rag_user:
@@ -313,40 +351,100 @@ class WaldiezAgent(WaldiezBase):
             class_name = "CaptainAgent"
         if self.is_group_manager:
             class_name = "GroupChatManager"
+        if self.is_doc_agent:
+            class_name = "DocAgent"
+        return class_name  # pragma: no cover
+
+    def get_group_member_class_name(self) -> str:
+        """Get the class name for a group member agent.
+
+        Returns
+        -------
+        str
+            The class name for the group member agent.
+        """
+        if (
+            getattr(self.data, "is_multimodal", False) is True
+        ):  # pragma: no branch
+            return "MultimodalConversableAgent"
+        if self.is_captain:  # pragma: no branch
+            return "CaptainAgent"
+        if self.is_reasoning:  # pragma: no branch
+            return "ReasoningAgent"
+        if self.is_doc_agent:  # pragma: no branch
+            return "DocAgent"
+        return "ConversableAgent"  # pragma: no cover
+
+    def get_assistant_class_name(self) -> str:
+        """Get the class name for an assistant agent.
+
+        Returns
+        -------
+        str
+            The class name for the assistant agent.
+        """
+        if getattr(self.data, "is_multimodal", False) is True:
+            class_name = "MultimodalConversableAgent"
+        else:
+            class_name = "AssistantAgent"
         return class_name  # pragma: no cover
 
     @property
     def ag2_imports(self) -> set[str]:
-        """Return the AG2 imports of the agent."""
+        """Return the AG2 imports of the agent.
+
+        Returns
+        -------
+        set[str]
+            A set of import statements required for the agent class.
+
+        Raises
+        ------
+        ValueError
+            If the agent class is unknown and no imports are defined.
+        """
         agent_class = self.ag2_class
         imports = {"import autogen"}
-        if agent_class == "AssistantAgent":
-            imports.add("from autogen import AssistantAgent")
-        elif agent_class == "UserProxyAgent":
-            imports.add("from autogen import UserProxyAgent")
-        elif agent_class == "RetrieveUserProxyAgent":
-            imports.add(
-                "from autogen.agentchat.contrib.retrieve_user_proxy_agent "
-                "import RetrieveUserProxyAgent"
-            )
-        elif agent_class == "MultimodalConversableAgent":
-            imports.add(
-                "from autogen.agentchat.contrib.multimodal_conversable_agent "
-                "import MultimodalConversableAgent"
-            )
-        elif agent_class == "ReasoningAgent":
-            imports.add(
-                "from autogen.agents.experimental import ReasoningAgent"
-            )
-        elif agent_class == "CaptainAgent":
-            imports.add(
-                "from autogen.agentchat.contrib.captainagent "
-                "import CaptainAgent"
-            )
-        elif agent_class == "GroupChatManager":  # pragma: no branch
-            imports.add("from autogen import GroupChat")
-            imports.add("from autogen.agentchat import GroupChatManager")
-            imports.add("from autogen.agentchat.group import ContextVariables")
+        match agent_class:
+            case "AssistantAgent":
+                imports.add("from autogen import AssistantAgent")
+            case "UserProxyAgent":
+                imports.add("from autogen import UserProxyAgent")
+            case "RetrieveUserProxyAgent":
+                imports.add(
+                    "from autogen.agentchat.contrib.retrieve_user_proxy_agent "
+                    "import RetrieveUserProxyAgent"
+                )
+            case "MultimodalConversableAgent":
+                imports.add(
+                    "from "
+                    "autogen.agentchat.contrib.multimodal_conversable_agent "
+                    "import MultimodalConversableAgent"
+                )
+            case "ReasoningAgent":
+                imports.add(
+                    "from autogen.agents.experimental import ReasoningAgent"
+                )
+            case "CaptainAgent":
+                imports.add(
+                    "from autogen.agentchat.contrib.captainagent "
+                    "import CaptainAgent"
+                )
+            case "GroupChatManager":  # pragma: no branch
+                imports.add("from autogen import GroupChat")
+                imports.add("from autogen.agentchat import GroupChatManager")
+                imports.add(
+                    "from autogen.agentchat.group import ContextVariables"
+                )
+            case "DocAgent":
+                imports.add("from autogen.agents.experimental import DocAgent")
+            case "ConversableAgent":
+                imports.add("from autogen import ConversableAgent")
+            case _:  # pragma: no cover
+                raise ValueError(
+                    f"Unknown agent class: {agent_class}. "
+                    "Please implement the imports for this class."
+                )
         return imports
 
     def validate_linked_tools(

waldiez/models/agents/agent/agent_type.py

@@ -6,14 +6,15 @@ from typing_extensions import Literal
 
 # pylint: disable=line-too-long
 # fmt: off
-WaldiezAgentType = Literal["user_proxy", "assistant", "group_manager", "manager", "rag_user", "swarm", "reasoning", "captain", "user", "rag_user_proxy"]  # noqa: E501
+WaldiezAgentType = Literal["user_proxy", "assistant", "group_manager", "manager", "rag_user", "swarm", "reasoning", "captain", "user", "rag_user_proxy", "doc_agent"]  # noqa: E501
 """Possible types of a Waldiez Agent:
     - user_proxy,
     - assistant,
     - group_manager,
-    - rag_user_proxy,
+    - rag_user_proxy (deprecated: use doc_agent),
     - reasoning,
     - captain,
+    - doc_agent,
     - swarm (deprecated: do not use it),
     - user (deprecated: use user_proxy)
     - rag_user (deprecated: user rag_user_proxy)

waldiez/models/agents/agents.py

@@ -11,6 +11,7 @@ from ..common import WaldiezBase
 from .agent import WaldiezAgent
 from .assistant import WaldiezAssistant
 from .captain import WaldiezCaptainAgent
+from .doc_agent import WaldiezDocAgent
 from .group_manager import WaldiezGroupManager
 from .rag_user_proxy import WaldiezRagUserProxy
 from .reasoning import WaldiezReasoningAgent
@@ -84,6 +85,14 @@ class WaldiezAgents(WaldiezBase):
             default_factory=list,
         ),
     ] = []
+    docAgents: Annotated[
+        list[WaldiezDocAgent],
+        Field(
+            title="Document Agents.",
+            description="The Document agents in the flow.",
+            default_factory=list,
+        ),
+    ] = []
 
     @property
     def members(self) -> Iterator[WaldiezAgent]:
@@ -100,6 +109,7 @@ class WaldiezAgents(WaldiezBase):
         yield from self.reasoningAgents
         yield from self.captainAgents
         yield from self.groupManagerAgents
+        yield from self.docAgents
 
     @model_validator(mode="after")
     def validate_agents(self) -> Self:

waldiez/models/agents/doc_agent/__init__.py

@@ -0,0 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Document agent model."""
+
+from .doc_agent import WaldiezDocAgent
+from .doc_agent_data import WaldiezDocAgentData
+from .rag_query_engine import WaldiezDocAgentQueryEngine
+
+__all__ = [
+    "WaldiezDocAgent",
+    "WaldiezDocAgentData",
+    "WaldiezDocAgentQueryEngine",
+]

waldiez/models/agents/doc_agent/doc_agent.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Document agent model."""
+
+from typing import Literal
+
+from pydantic import Field
+from typing_extensions import Annotated
+
+from waldiez.models.agents.doc_agent.rag_query_engine import (
+    WaldiezDocAgentQueryEngine,
+)
+
+from ...model import WaldiezModel
+from ..agent import WaldiezAgent
+from .doc_agent_data import WaldiezDocAgentData
+
+
+class WaldiezDocAgent(WaldiezAgent):
+    """Document agent class.
+
+    The agent for handling document-related tasks.
+    Extends `WaldiezAgent`.
+    """
+
+    agent_type: Annotated[  # pyright: ignore
+        Literal["doc_agent"],
+        Field(
+            default="doc_agent",
+            title="Agent type",
+            description="The agent type: 'doc_agent' for a document agent",
+            alias="agentType",
+        ),
+    ] = "doc_agent"
+    data: Annotated[  # pyright: ignore
+        WaldiezDocAgentData,
+        Field(
+            title="Data",
+            description="The document agent's data",
+            default_factory=WaldiezDocAgentData,  # pyright: ignore
+        ),
+    ]
+
+    @property
+    def reset_collection(self) -> bool:
+        """Get whether to reset the collection for the document agent.
+
+        Returns
+        -------
+        bool
+            Whether to reset the collection for the document agent.
+        """
+        return self.data.reset_collection
+
+    def get_collection_name(self) -> str:
+        """Get the collection name for the document agent.
+
+        Returns
+        -------
+        str
+            The collection name for the document agent.
+        """
+        return self.data.get_collection_name()
+
+    def get_query_engine(self) -> WaldiezDocAgentQueryEngine:
+        """Get the query engine for the document agent.
+
+        Returns
+        -------
+        WaldiezDocAgentQueryEngine
+            The query engine for the document agent.
+        """
+        return self.data.get_query_engine()
+
+    def get_db_path(self) -> str:
+        """Get the database path for the query engine.
+
+        Returns
+        -------
+        str
+            The database path for the query engine.
+        """
+        return self.data.get_db_path()
+
+    def get_parsed_docs_path(self) -> str:
+        """Get the parsed documents path for the document agent.
+
+        Returns
+        -------
+        str
+            The parsed documents path for the document agent.
+        """
+        return self.data.get_parsed_docs_path()
+
+    def get_llm_requirements(
+        self,
+        ag2_version: str,
+        all_models: list[WaldiezModel],
+    ) -> set[str]:
+        """Get the LLM requirements for the document agent.
+
+        Parameters
+        ----------
+        ag2_version : str
+            The version of AG2 to use for the requirements.
+        all_models : list[WaldiezModel]
+            All the models in the flow.
+
+        Returns
+        -------
+        set[str]
+            The set of LLM requirements for the document agent.
+        """
+        requirements = {
+            "llama-index",
+            "llama-index-core",
+            f"ag2[rag]=={ag2_version}",
+        }
+        if not self.data.model_ids:
+            requirements.add("llama-index-llms-openai")
+        else:
+            for model_id in self.data.model_ids:
+                model = next((m for m in all_models if m.id == model_id), None)
+                if model:
+                    return model.get_llm_requirements(ag2_version=ag2_version)
+        return requirements

waldiez/models/agents/doc_agent/doc_agent_data.py

@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Document agent data model."""
+
+from pathlib import Path
+from typing import Optional, Union
+
+from platformdirs import user_data_dir
+from pydantic import Field, model_validator
+from typing_extensions import Annotated, Self
+
+from ..agent import WaldiezAgentData
+from .rag_query_engine import WaldiezDocAgentQueryEngine
+
+
+class WaldiezDocAgentData(WaldiezAgentData):
+    """Document agent data class.
+
+    The data for a document agent.
+    Extends `WaldiezAgentData`.
+    Extra attributes:
+    - `collection_name`: Optional string, the name of the collection.
+    - `reset_collection`: Optional boolean, whether to reset the collection.
+    - `parsed_docs_path`: Optional string, the path to the parsed documents.
+    - `query_engine`: Optional `RAGQueryEngine`, the query engine to use.
+    """
+
+    collection_name: Annotated[
+        Optional[str],
+        Field(
+            title="Collection Name",
+            description="The name of the collection for the document agent.",
+            default=None,
+            alias="collectionName",
+        ),
+    ] = None
+    reset_collection: Annotated[
+        bool,
+        Field(
+            title="Reset Collection",
+            description=(
+                "Whether to reset the collection for the document agent."
+            ),
+            default=False,
+            alias="resetCollection",
+        ),
+    ] = False
+    parsed_docs_path: Annotated[
+        Optional[Union[str, Path]],
+        Field(
+            title="Parsed Documents Path",
+            description=(
+                "The path to the parsed documents for the document agent."
+            ),
+            default=None,
+            alias="parsedDocsPath",
+        ),
+    ] = None
+    query_engine: Annotated[
+        Optional[WaldiezDocAgentQueryEngine],
+        Field(
+            title="Query Engine",
+            description="The query engine to use for the document agent.",
+            default=None,
+            alias="queryEngine",
+        ),
+    ] = None
+
+    @model_validator(mode="after")
+    def validate_parsed_docs_path(self) -> Self:
+        """Ensure the parsed documents path is set and is a directory.
+
+        If not set, create a default path in the user data directory.
+
+        Returns
+        -------
+        Self
+            The instance with validated `parsed_docs_path`.
+        """
+        if not self.parsed_docs_path:
+            data_dir = user_data_dir(
+                appname="waldiez",
+                appauthor="waldiez",
+            )
+            parsed_docs_path = Path(data_dir) / "parsed_docs"
+            parsed_docs_path.mkdir(parents=True, exist_ok=True)
+            self.parsed_docs_path = str(parsed_docs_path.resolve())
+        resolved = Path(self.parsed_docs_path).resolve()
+        if not resolved.is_absolute():
+            self.parsed_docs_path = str(Path.cwd() / self.parsed_docs_path)
+        if not Path(self.parsed_docs_path).is_dir():
+            Path(self.parsed_docs_path).mkdir(parents=True, exist_ok=True)
+        self.parsed_docs_path = str(Path(self.parsed_docs_path).resolve())
+        return self
+
+    def get_query_engine(self) -> WaldiezDocAgentQueryEngine:
+        """Get the query engine for the document agent.
+
+        Returns
+        -------
+        WaldiezDocAgentQueryEngine
+            The query engine for the document agent.
+        """
+        if not self.query_engine:
+            self.query_engine = WaldiezDocAgentQueryEngine()
+        return self.query_engine
+
+    def get_db_path(self) -> str:
+        """Get the database path for the query engine.
+
+        Returns
+        -------
+        str
+            The database path for the query engine.
+        """
+        return self.get_query_engine().get_db_path()
+
+    def get_collection_name(self) -> str:
+        """Get the collection name for the document agent.
+
+        Returns
+        -------
+        str
+            The collection name for the document agent.
+        """
+        return self.collection_name or "docling-parsed-docs"
+
+    def get_parsed_docs_path(self) -> str:
+        """Get the parsed documents path for the document agent.
+
+        Returns
+        -------
+        str
+            The parsed documents path for the document agent.
+        """
+        if not self.parsed_docs_path:
+            data_dir = user_data_dir(
+                appname="waldiez",
+                appauthor="waldiez",
+            )
+            parsed_docs_path = Path(data_dir) / "parsed_docs"
+            parsed_docs_path.mkdir(parents=True, exist_ok=True)
+            self.parsed_docs_path = str(parsed_docs_path.resolve())
+        resolved = Path(self.parsed_docs_path).resolve()
+        if not resolved.is_absolute():
+            self.parsed_docs_path = str(Path.cwd() / self.parsed_docs_path)
+        if not Path(self.parsed_docs_path).is_dir():
+            Path(self.parsed_docs_path).mkdir(parents=True, exist_ok=True)
+        return str(self.parsed_docs_path)

waldiez/models/agents/doc_agent/rag_query_engine.py

@@ -0,0 +1,127 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Document agent data model."""
+
+from pathlib import Path
+from typing import Optional, Union
+
+from platformdirs import user_data_dir
+from pydantic import Field, model_validator
+from typing_extensions import Annotated, Literal, Self
+
+from ...common import WaldiezBase
+
+
+class WaldiezDocAgentQueryEngine(WaldiezBase):
+    """RAG Query Engine class.
+
+    The data for a RAG query engine.
+    """
+
+    type: Annotated[
+        Optional[
+            Literal[
+                "VectorChromaQueryEngine",
+                "VectorChromaCitationQueryEngine",
+                "InMemoryQueryEngine",
+            ]
+        ],
+        "RAG Query Engine type",
+    ] = "VectorChromaQueryEngine"
+    db_path: Annotated[
+        Optional[Union[str, Path]],
+        Field(
+            title="Database Path",
+            description="The path to the database for the query engine.",
+            default=None,
+            alias="dbPath",
+        ),
+    ] = None
+    enable_query_citations: Annotated[
+        bool,
+        Field(
+            title="Enable Query Citations",
+            description=(
+                "Whether to enable query citations for the query engine."
+            ),
+            default=False,
+            alias="enableQueryCitations",
+        ),
+    ] = False
+    citation_chunk_size: Annotated[
+        int,
+        Field(
+            title="Citation Chunk Size",
+            description="The size of the citation chunks for the query engine.",
+            default=512,
+            alias="citationChunkSize",
+        ),
+    ] = 512
+
+    @model_validator(mode="after")
+    def validate_db_path(self) -> Self:
+        """Validate the db_path field.
+
+        Ensure the db_path is set and is a directory.
+
+        Returns
+        -------
+        Self
+            The instance of WaldiezDocAgentQueryEngine with validated db_path.
+        """
+        if not self.type:
+            self.type = "VectorChromaQueryEngine"
+        self.db_path = ensure_db_path(self.db_path)
+        return self
+
+    def get_db_path(self) -> str:
+        """Get the database path for the query engine.
+
+        Returns
+        -------
+        str
+            The database path for the query engine.
+        """
+        db_path = self.db_path or ensure_db_path(self.db_path)
+        return str(db_path)
+
+
+def ensure_db_path(db_path: str | Path | None) -> str:
+    """Get the database path for the query engine.
+
+    Ensure the database path is set and is a directory.
+
+    Parameters
+    ----------
+    db_path : Optional[str]
+        The database path to validate.
+
+    Returns
+    -------
+    str
+        The database path for the query engine.
+
+    Raises
+    ------
+    ValueError
+        If the database path is not set or is not a directory.
+    """
+    if not db_path:
+        data_dir = user_data_dir(
+            appname="waldiez",
+            appauthor="waldiez",
+        )
+        data_dir_path = Path(data_dir) / "rag"
+        data_dir_path.mkdir(parents=True, exist_ok=True)
+        db_path_to_resolve = data_dir_path / "chroma"
+        db_path_to_resolve.mkdir(parents=True, exist_ok=True)
+        return str(db_path_to_resolve)
+
+    resolved = Path(db_path).resolve()
+    if not resolved.is_absolute():
+        resolved = (Path.cwd() / db_path).resolve()
+    if not resolved.exists():
+        resolved.mkdir(parents=True, exist_ok=True)
+    if not resolved.is_dir():
+        raise ValueError(f"The path {resolved} is not a directory.")
+    return str(resolved)