waldiez 0.4.7__py3-none-any.whl → 0.4.9__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,33 +46,32 @@ 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 = ""
     embedding_function_arg = ""
     embedding_function_content = ""
-    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}")'
+        embedding_function_arg = f"{agent_name}_embedding_function"
     else:
         embedding_function_content, embedding_function_arg = (
             agent.retrieve_config.get_custom_embedding_function(
@@ -84,20 +83,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.
@@ -123,21 +122,35 @@ def get_chroma_db_args(
     # manually initializing the collection before running the flow,
     # might be a workaround.
     content_before = f"{agent_name}_client = {client_str}" + "\n"
+    vector_db_model = agent.retrieve_config.db_config.model
+    if not embedding_function_body:
+        content_before += (
+            f"{agent_name}_embedding_function = "
+            "SentenceTransformerEmbeddingFunction(\n"
+            f'    model_name="{vector_db_model}",' + "\n"
+            ")\n"
+        )
     collection_name = agent.retrieve_config.collection_name
     get_or_create = agent.retrieve_config.get_or_create
     if collection_name:
         if get_or_create:
             content_before += (
-                f"{agent_name}_client.get_or_create_collection("
-                f'"{collection_name}")' + "\n"
+                f"{agent_name}_client.get_or_create_collection" + "(\n"
+                f'    "{collection_name}",' + "\n"
+                f"    embedding_function={embedding_function_arg}," + "\n"
+                ")" + "\n"
             )
         else:
             content_before += (
                 "try:\n"
-                f'    {agent_name}_client.get_collection("{collection_name}")'
-                + "\n"
+                f"    {agent_name}_client.get_collection(" + "\n"
+                f'        "{collection_name}",' + "\n"
+                f"        embedding_function={embedding_function_arg}," + "\n"
+                "    )\n"
                 "except ValueError:\n"
-                f"    {agent_name}_client.create_collection("
-                f'"{collection_name}")' + "\n"
+                f"    {agent_name}_client.create_collection(" + "\n"
+                f'        "{collection_name}",' + "\n"
+                f"        embedding_function={embedding_function_arg}," + "\n"
+                "    )\n"
             )
     return kwarg_string, to_import, embedding_function_body, content_before

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)