waldiez 0.4.7__py3-none-any.whl → 0.4.8__py3-none-any.whl

waldiez/exporting/agent/extras/handoffs/target.py

@@ -0,0 +1,189 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+# pylint: disable=too-few-public-methods
+"""Transition target processor for Waldiez agents."""
+
+from dataclasses import dataclass, field
+from typing import Callable, Set
+
+from waldiez.exporting.chats.utils.nested import get_nested_chat_queue
+from waldiez.models import (
+    WaldiezAgent,
+    WaldiezAgentTarget,
+    WaldiezChat,
+    WaldiezGroupOrNestedTarget,
+    WaldiezRandomAgentTarget,
+    WaldiezTransitionTarget,
+)
+
+
+@dataclass
+class TargetResult:
+    """Result from processing a transition target.
+
+    Attributes
+    ----------
+    content : str
+        Content to be used for the target processing.
+    before_content : str
+        Content to be placed before the target processing.
+    extra_imports : Set[str]
+        Additional imports required for the target processing.
+    """
+
+    content: str = ""
+    before_content: str = ""
+    extra_imports: Set[str] = field(default_factory=set)  # pyright: ignore
+
+
+class TransitionTargetProcessor:
+    """Processor for transition targets."""
+
+    def __init__(
+        self,
+        agent: WaldiezAgent,
+        agent_names: dict[str, str],
+        chat_names: dict[str, str],
+        all_chats: list[WaldiezChat],
+        serializer: Callable[..., str],
+    ) -> None:
+        """Initialize the processor with agent names and context.
+
+        Parameters
+        ----------
+        agent_names : dict[str, str]
+            Mapping of agent IDs to their names.
+        """
+        self.agent_names = agent_names
+        self.chat_names = chat_names
+        self.agent = agent
+        self.agent_name = agent_names[agent.id]
+        self.all_chats = all_chats
+        self.serializer = serializer
+
+    def process(self, target: WaldiezTransitionTarget) -> TargetResult:
+        """Process transition target based on its type.
+
+        Parameters
+        ----------
+        target : WaldiezTransitionTarget
+            The transition target to process.
+
+        Returns
+        -------
+        TargetResult
+            The processed result containing the target string,
+            before_content code, and extra imports.
+
+        Raises
+        ------
+        ValueError
+            If the target type is unknown.
+        """
+        result = TargetResult()
+        where = "autogen.agentchat.group"
+        if target.target_type == "GroupManagerTarget":
+            where += ".targets.group_manager_target"
+        what = target.target_type
+        result.extra_imports.add(f"from {where} import {what}")
+
+        processors: dict[str, Callable[[WaldiezTransitionTarget], str]] = {
+            "AgentTarget": self._process_agent_target,
+            "RandomAgentTarget": self._process_random_agent_target,
+            "GroupChatTarget": self._process_group_chat_target,
+            "NestedChatTarget": self._process_nested_chat_target,
+            "AskUserTarget": self._process_simple_target,
+            "GroupManagerTarget": self._process_simple_target,
+            "RevertToUserTarget": self._process_simple_target,
+            "StayTarget": self._process_simple_target,
+            "TerminateTarget": self._process_simple_target,
+        }
+
+        processor = processors.get(target.target_type)
+        if not processor:
+            raise ValueError(f"Unknown target type: {target.target_type}")
+
+        result.content = processor(target)
+
+        # Special handling for nested chat targets
+        if target.target_type == "NestedChatTarget":
+            nested_result = self._process_nested_chat_target_full(target)
+            result.before_content = nested_result.before_content
+
+        return result
+
+    def _process_agent_target(self, target: WaldiezTransitionTarget) -> str:
+        """Process agent target."""
+        if not isinstance(target, WaldiezAgentTarget):
+            raise ValueError(
+                "Expected WaldiezAgentTarget for agent target processing."
+            )
+        agent_name = self.agent_names[target.value[0]]
+        return f"AgentTarget({agent_name})"
+
+    def _process_random_agent_target(
+        self, target: WaldiezTransitionTarget
+    ) -> str:
+        """Process random agent target."""
+        if not isinstance(target, WaldiezRandomAgentTarget):
+            raise ValueError(
+                "Expected WaldiezRandomAgentTarget"
+                " for random agent target processing."
+            )
+        agent_vars = [self.agent_names[agent_id] for agent_id in target.value]
+        agents_str = ", ".join(agent_vars)
+        return f"RandomAgentTarget([{agents_str}])"
+
+    def _process_group_chat_target(
+        self, target: WaldiezTransitionTarget
+    ) -> str:
+        """Process group chat target."""
+        if not isinstance(target, WaldiezGroupOrNestedTarget):
+            raise ValueError(
+                "Expected WaldiezGroupOrNestedTarget for group chat target "
+                "processing."
+            )
+        chat_name = self.chat_names[target.value[0]]
+        return f"GroupChatTarget({chat_name})"
+
+    def _process_nested_chat_target(self, _: WaldiezTransitionTarget) -> str:
+        """Process nested chat target (simple string only)."""
+        chat_name = f"{self.agent_name}_handoff_nested_chat_queue"
+        target_arg = f'nested_chat_config={{"chat_queue": {chat_name}}}'
+        return f"NestedChatTarget({target_arg})"
+
+    def _process_nested_chat_target_full(
+        self, _: WaldiezTransitionTarget
+    ) -> TargetResult:
+        """Process nested chat target with full configuration."""
+        result = TargetResult()
+
+        chat_queue, extra_methods = get_nested_chat_queue(
+            nested_chat=self.agent.data.nested_chats[0],
+            agent=self.agent,
+            agent_names=self.agent_names,
+            chat_names=self.chat_names,
+            all_chats=self.all_chats,
+            serializer=self.serializer,
+        )
+
+        chat_name = f"{self.agent_name}_handoff_nested_chat_queue"
+
+        if extra_methods:
+            result.before_content += "\n".join(extra_methods) + "\n"
+
+        result.before_content += (
+            f"{chat_name}: list[dict[str, Any]] = {chat_queue}\n\n"
+        )
+        return result
+
+    # pylint: disable=no-self-use
+    def _process_simple_target(self, target: WaldiezTransitionTarget) -> str:
+        """Process simple targets that don't need parameters.
+
+        Parameters
+        ----------
+        target : WaldiezTransitionTarget
+            The target to process.
+        """
+        return f"{target.target_type}()"

waldiez/exporting/agent/extras/rag/__init__.py

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Agent exporter rag extras."""
+
+from .vector_db_extras import VectorDBExtras, get_vector_db_extras
+
+__all__ = [
+    "get_vector_db_extras",
+    "VectorDBExtras",
+]

waldiez/exporting/agent/{utils/rag_user/chroma_utils.py → extras/rag/chroma_extras.py}

@@ -3,24 +3,24 @@
 """Get chroma db related imports and content."""
 
 from pathlib import Path
-from typing import Set, Tuple
+from typing import Set
 
-from waldiez.models import WaldiezRagUser
+from waldiez.models import WaldiezRagUserProxy
 
 
-def _get_chroma_client_string(agent: WaldiezRagUser) -> Tuple[str, str]:
+def _get_chroma_client_string(agent: WaldiezRagUserProxy) -> tuple[str, str]:
     """Get the ChromaVectorDB client string.
 
     Parameters
     ----------
-    agent : WaldiezRagUser
+    agent : WaldiezRagUserProxy
         The agent.
     agent_name : str
         The agent's name.
 
     Returns
     -------
-    Tuple[str, str]
+    tuple[str, str]
         The 'client' and what to import.
     """
     # TODO: also check `connection_url` (chroma in client-server mode)
@@ -46,20 +46,20 @@ def _get_chroma_client_string(agent: WaldiezRagUser) -> Tuple[str, str]:
 
 
 def _get_chroma_embedding_function_string(
-    agent: WaldiezRagUser, agent_name: str
-) -> Tuple[str, str, str]:
+    agent: WaldiezRagUserProxy, agent_name: str
+) -> tuple[str, str, str]:
     """Get the ChromaVectorDB embedding function string.
 
     Parameters
     ----------
-    agent : WaldiezRagUser
+    agent : WaldiezRagUserProxy
         The agent.
     agent_name : str
         The agent's name.
 
     Returns
     -------
-    Tuple[str, str, str]
+    tuple[str, str, str]
         The 'embedding_function', the import and the custom embedding function.
     """
     to_import = ""
@@ -68,8 +68,9 @@ def _get_chroma_embedding_function_string(
     vector_db_model = agent.retrieve_config.db_config.model
     if not agent.retrieve_config.use_custom_embedding:
         to_import = (
-            "from chromadb.utils.embedding_functions "
-            "import SentenceTransformerEmbeddingFunction"
+            "from chromadb.utils.embedding_functions."
+            "sentence_transformer_embedding_function import "
+            "SentenceTransformerEmbeddingFunction"
         )
         embedding_function_arg = "SentenceTransformerEmbeddingFunction("
         embedding_function_arg += f'model_name="{vector_db_model}")'
@@ -84,20 +85,20 @@ def _get_chroma_embedding_function_string(
 
 
 def get_chroma_db_args(
-    agent: WaldiezRagUser, agent_name: str
-) -> Tuple[str, Set[str], str, str]:
+    agent: WaldiezRagUserProxy, agent_name: str
+) -> tuple[str, Set[str], str, str]:
     """Get the 'kwargs to use for ChromaVectorDB.
 
     Parameters
     ----------
-    agent : WaldiezRagUser
+    agent : WaldiezRagUserProxy
         The agent.
     agent_name : str
         The agent's name.
 
     Returns
     -------
-    Tuple[str, Set[str], str]
+    tuple[str, Set[str], str]
 
         - The 'kwargs' string.
         - What to import.

waldiez/exporting/agent/{utils/rag_user/mongo_utils.py → extras/rag/mongo_extras.py}

@@ -2,26 +2,26 @@
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Get mongodb related content and imports."""
 
-from typing import Set, Tuple
+from typing import Set
 
-from waldiez.models import WaldiezRagUser
+from waldiez.models import WaldiezRagUserProxy
 
 
 def _get_mongodb_embedding_function_string(
-    agent: WaldiezRagUser, agent_name: str
-) -> Tuple[str, str, str]:
+    agent: WaldiezRagUserProxy, agent_name: str
+) -> tuple[str, str, str]:
     """Get the MongoDBAtlasVectorDB embedding function string.
 
     Parameters
     ----------
-    agent : WaldiezRagUser
+    agent : WaldiezRagUserProxy
         The agent.
     agent_name : str
         The agent's name.
 
     Returns
     -------
-    Tuple[str, str, str]
+    tuple[str, str, str]
         The 'embedding_function', the import and the custom_embedding_function.
     """
     to_import = ""
@@ -45,20 +45,20 @@ def _get_mongodb_embedding_function_string(
 
 
 def get_mongodb_db_args(
-    agent: WaldiezRagUser, agent_name: str
-) -> Tuple[str, Set[str], str]:
+    agent: WaldiezRagUserProxy, agent_name: str
+) -> tuple[str, Set[str], str]:
     """Get the kwargs to use for MongoDBAtlasVectorDB.
 
     Parameters
     ----------
-    agent : WaldiezRagUser
+    agent : WaldiezRagUserProxy
         The agent.
     agent_name : str
         The agent's name.
 
     Returns
     -------
-    Tuple[str, Set[str], str]
+    tuple[str, Set[str], str]
         The kwargs to use, what to import and the custom_embedding_function.
     """
     embedding_function_arg, to_import_embedding, embedding_function_body = (

waldiez/exporting/agent/{utils/rag_user/pgvector_utils.py → extras/rag/pgvector_extras.py}

@@ -2,22 +2,22 @@
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Get pgvector related content and imports."""
 
-from typing import Set, Tuple
+from typing import Set
 
-from waldiez.models import WaldiezRagUser
+from waldiez.models import WaldiezRagUserProxy
 
 
-def _get_pgvector_client_string(agent: WaldiezRagUser) -> Tuple[str, str]:
+def _get_pgvector_client_string(agent: WaldiezRagUserProxy) -> tuple[str, str]:
     """Get the PGVectorDB client string.
 
     Parameters
     ----------
-    agent : WaldiezRagUser
+    agent : WaldiezRagUserProxy
         The agent.
 
     Returns
     -------
-    Tuple[str, str]
+    tuple[str, str]
         The 'client' and what to import.
     """
     to_import = "import psycopg"
@@ -28,20 +28,20 @@ def _get_pgvector_client_string(agent: WaldiezRagUser) -> Tuple[str, str]:
 
 
 def _get_pgvector_embedding_function_string(
-    agent: WaldiezRagUser, agent_name: str
-) -> Tuple[str, str, str]:
+    agent: WaldiezRagUserProxy, agent_name: str
+) -> tuple[str, str, str]:
     """Get the PGVectorDB embedding function string.
 
     Parameters
     ----------
-    agent : WaldiezRagUser
+    agent : WaldiezRagUserProxy
         The agent.
     agent_name : str
         The agent's name.
 
     Returns
     -------
-    Tuple[str, str, str]
+    tuple[str, str, str]
         The 'embedding_function', the import and the custom_embedding_function.
     """
     to_import = ""
@@ -65,20 +65,20 @@ def _get_pgvector_embedding_function_string(
 
 
 def get_pgvector_db_args(
-    agent: WaldiezRagUser, agent_name: str
-) -> Tuple[str, Set[str], str]:
+    agent: WaldiezRagUserProxy, agent_name: str
+) -> tuple[str, Set[str], str]:
     """Get the kwargs to use for PGVectorDB.
 
     Parameters
     ----------
-    agent : WaldiezRagUser
+    agent : WaldiezRagUserProxy
         The agent.
     agent_name : str
         The agent's name.
 
     Returns
     -------
-    Tuple[str, Set[str], str]
+    tuple[tuple[str,str], Set[str], str]
         The kwargs to use, what to import and the custom_embedding_function.
     """
     client_str, to_import_client = _get_pgvector_client_string(agent)

waldiez/exporting/agent/{utils/rag_user/qdrant_utils.py → extras/rag/qdrant_extras.py}

@@ -3,24 +3,24 @@
 """Get qdrant db related imports and content."""
 
 from pathlib import Path
-from typing import Set, Tuple
+from typing import Set
 
-from waldiez.models import WaldiezRagUser
+from waldiez.models import WaldiezRagUserProxy
 
 
-def _get_qdrant_client_string(agent: WaldiezRagUser) -> Tuple[str, str]:
+def _get_qdrant_client_string(agent: WaldiezRagUserProxy) -> tuple[str, str]:
     """Get the QdrantVectorDB client string.
 
     Parameters
     ----------
-    agent : WaldiezRagUser
+    agent : WaldiezRagUserProxy
         The agent.
     agent_name : str
         The agent's name.
 
     Returns
     -------
-    Tuple[str, str, str]
+    tuple[str, str, str]
         The 'client' argument, and the module to import.
     """
     to_import: str = "from qdrant_client import QdrantClient"
@@ -44,20 +44,20 @@ def _get_qdrant_client_string(agent: WaldiezRagUser) -> Tuple[str, str]:
 
 
 def _get_qdrant_embedding_function_string(
-    agent: WaldiezRagUser, agent_name: str
-) -> Tuple[str, str, str]:
+    agent: WaldiezRagUserProxy, agent_name: str
+) -> tuple[str, str, str]:
     """Get the QdrantVectorDB embedding function string.
 
     Parameters
     ----------
-    agent : WaldiezRagUser
+    agent : WaldiezRagUserProxy
         The agent.
     agent_name : str
         The agent's name.
 
     Returns
     -------
-    Tuple[str, str, str]
+    tuple[str, str, str]
         The 'embedding_function', the module to import
         and the custom_embedding_function if used.
     """
@@ -83,20 +83,20 @@ def _get_qdrant_embedding_function_string(
 
 
 def get_qdrant_db_args(
-    agent: WaldiezRagUser, agent_name: str
-) -> Tuple[str, Set[str], str]:
+    agent: WaldiezRagUserProxy, agent_name: str
+) -> tuple[str, Set[str], str]:
     """Get the kwargs to use for QdrantVectorDB.
 
     Parameters
     ----------
-    agent : WaldiezRagUser
+    agent : WaldiezRagUserProxy
         The agent.
     agent_name : str
         The agent's name.
 
     Returns
     -------
-    Tuple[str, Set[str], str]
+    tuple[str, Set[str], str]
         The kwargs to use, the imports and the embedding function body if used.
     """
     client_str, to_import_client = _get_qdrant_client_string(agent)

waldiez/exporting/agent/{utils/rag_user/vector_db.py → extras/rag/vector_db_extras.py}

@@ -4,51 +4,30 @@
 # pylint: disable=line-too-long
 """Vector DB exporting utils for RAG user agents."""
 
-from typing import Any, Set, Tuple
+from dataclasses import dataclass, field
+from typing import Any
 
-from waldiez.models import WaldiezRagUser
+from waldiez.models import WaldiezRagUserProxy
 
-from .chroma_utils import get_chroma_db_args
-from .mongo_utils import get_mongodb_db_args
-from .pgvector_utils import get_pgvector_db_args
-from .qdrant_utils import get_qdrant_db_args
+from .chroma_extras import get_chroma_db_args
+from .mongo_extras import get_mongodb_db_args
+from .pgvector_extras import get_pgvector_db_args
+from .qdrant_extras import get_qdrant_db_args
 
 
-def _get_metadata_arg(
-    agent: WaldiezRagUser,
-) -> str:
-    """Get the metadata arg.
-
-    Parameters
-    ----------
-    agent : WaldiezRagUser
-        The agent.
+@dataclass
+class VectorDBExtras:
+    """Vector DB exporting extras for RAG user agents."""
 
-    Returns
-    -------
-    str
-        The metadata arg.
-    """
-    metadata_arg = ""
-    if agent.retrieve_config.db_config.metadata:
-        tab = "    "
-        indent = tab * 3
-        metadata_arg += f"{indent}metadata={{" + "\n"
-        for key, value in agent.retrieve_config.db_config.metadata.items():
-            value_string: Any = f'"{value}"'
-            if str(value).isdigit():
-                value_string = int(value)
-            elif str(value).replace(".", "").isdigit():
-                value_string = float(value)
-            metadata_arg += f'{indent}    "{key}": {value_string},' + "\n"
-        metadata_arg += f"{indent}}}," + "\n"
-    return metadata_arg
+    before_arg: str
+    vector_db_arg: str
+    imports: set[str] = field(default_factory=set[str])
 
 
-def get_rag_user_vector_db_string(
-    agent: WaldiezRagUser,
+def get_vector_db_extras(
+    agent: WaldiezRagUserProxy,
     agent_name: str,
-) -> Tuple[str, str, Set[str]]:
+) -> VectorDBExtras:
     """Get the RAG user vector db string.
 
     The vector db can be one of the following:
@@ -62,20 +41,20 @@ def get_rag_user_vector_db_string(
 
     Parameters
     ----------
-    agent : WaldiezRagUser
+    agent : WaldiezRagUserProxy
         The agent.
     agent_name : str
         The agent's name.
 
     Returns
     -------
-    Tuple[str, str, Set[str]]
+    tuple[str, str, Set[str]]
         The content before the arg if any, the arg and the related imports.
     """
     before = ""
-    imports: Set[str] = set()
+    imports: set[str] = set()
     ef_body: str = ""
-    db_imports: Set[str] = set()
+    db_imports: set[str] = set()
     kwarg_string = ""
     content_before = ""
     vdb_class = "ChromaVectorDB"
@@ -97,14 +76,14 @@ def get_rag_user_vector_db_string(
     if agent.retrieve_config.vector_db == "mongodb":
         vdb_class = "MongoDBAtlasVectorDB"
         imports.add(
-            "from autogen.agentchat.contrib.vectordb.mongo import MongoDBAtlasVectorDB"
+            "from autogen.agentchat.contrib.vectordb.mongodb import MongoDBAtlasVectorDB"
         )
         kwarg_string, db_imports, ef_body = get_mongodb_db_args(
             agent, agent_name
         )
     if agent.retrieve_config.vector_db == "pgvector":
         imports.add(
-            "from autogen.agentchat.contrib.vectordb.pgvector import PGVectorDB"
+            "from autogen.agentchat.contrib.vectordb.pgvectordb import PGVectorDB"
         )
         vdb_class = "PGVectorDB"
         kwarg_string, db_imports, ef_body = get_pgvector_db_args(
@@ -114,9 +93,43 @@ def get_rag_user_vector_db_string(
         before += "\n" + f"{content_before}"
     if ef_body:
         before += "\n" + f"{ef_body}" + "\n"
-    if db_imports:
-        imports.update(db_imports)
+    imports.update(db_imports)
     kwarg_string += _get_metadata_arg(agent)
     vdb_arg = f"{vdb_class}(" + "\n"
     vdb_arg += kwarg_string + "        )"
-    return before, vdb_arg, imports
+    return VectorDBExtras(
+        before_arg=before,
+        vector_db_arg=vdb_arg,
+        imports=imports,
+    )
+
+
+def _get_metadata_arg(
+    agent: WaldiezRagUserProxy,
+) -> str:
+    """Get the metadata arg.
+
+    Parameters
+    ----------
+    agent : WaldiezRagUserProxy
+        The agent.
+
+    Returns
+    -------
+    str
+        The metadata arg.
+    """
+    metadata_arg = ""
+    if agent.retrieve_config.db_config.metadata:
+        tab = "    "
+        indent = tab * 3
+        metadata_arg += f"{indent}metadata={{" + "\n"
+        for key, value in agent.retrieve_config.db_config.metadata.items():
+            value_string: Any = f'"{value}"'
+            if str(value).isdigit():
+                value_string = int(value)
+            elif str(value).replace(".", "").isdigit():
+                value_string = float(value)
+            metadata_arg += f'{indent}    "{key}": {value_string},' + "\n"
+        metadata_arg += f"{indent}}}," + "\n"
+    return metadata_arg