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

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

waldiez/exporting/agent/extras/rag_user_proxy_agent_extras.py

@@ -0,0 +1,245 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""RAG user proxy agent configuration processor."""
+
+from typing import Union
+
+from waldiez.models import (
+    WaldiezAgent,
+    WaldiezRagUserProxy,
+    WaldiezRagUserProxyModels,
+    WaldiezRagUserProxyRetrieveConfig,
+)
+
+from ...core import (
+    DefaultPathResolver,
+    DefaultSerializer,
+    ImportPosition,
+    ImportStatement,
+    InstanceArgument,
+    PathResolver,
+    Serializer,
+)
+from ...core.extras.agent_extras import RAGUserExtras
+from .rag import VectorDBExtras, get_vector_db_extras
+
+
+# pylint: disable=too-few-public-methods
+class RagUserProxyAgentProcessor:
+    """Processor for RAG user proxy agent configuration."""
+
+    _before_agent: str = ""
+
+    def __init__(
+        self,
+        agent: WaldiezAgent,
+        agent_name: str,
+        model_names: dict[str, str],
+        path_resolver: PathResolver | None = None,
+        serializer: Serializer | None = None,
+    ):
+        """Initialize the processor.
+
+        Parameters
+        ----------
+        agent : WaldiezAgent
+            The Waldiez RAG user proxy agent to process.
+        model_names : dict[str, str]
+            A mapping from model id to model name.
+        path_resolver : PathResolver | None
+            Optional path resolver for resolving paths.
+            Defaults to DefaultPathResolver if not provided.
+        serializer : Serializer | None
+            Optional serializer for the RAG configuration.
+            Defaults to DefaultSerializer if not provided.
+        """
+        self.agent = agent
+        self.agent_name = agent_name
+        self.model_names = model_names
+        self.path_resolver = path_resolver or DefaultPathResolver()
+        self.serializer = serializer or DefaultSerializer()
+
+    def process(self) -> RAGUserExtras:
+        """Process RAG user proxy agent configuration.
+
+        Returns
+        -------
+        RAGConfig
+            The processed result containing extra arguments, before content,
+            imports, and environment variables.
+        """
+        result = RAGUserExtras(self.agent.id)
+        if not self.agent.is_rag_user or not isinstance(
+            self.agent, WaldiezRagUserProxy
+        ):
+            return result
+        # Get the extra args
+        vector_db_extras = get_vector_db_extras(
+            agent=self.agent, agent_name=self.agent_name
+        )
+        before_agent, arg_value = self._get_retrieve_config(vector_db_extras)
+        if before_agent:
+            self._before_agent += before_agent
+        if arg_value:
+            retrieve_arg = InstanceArgument(
+                instance_id=self.agent.id,
+                name="retrieve_config",
+                value=arg_value,
+                tabs=1,
+            )
+            result.add_arg(retrieve_arg)
+        for import_statement in vector_db_extras.imports:
+            result.add_import(
+                ImportStatement(
+                    statement=import_statement,
+                    position=ImportPosition.THIRD_PARTY,
+                )
+            )
+        if self._before_agent:
+            result.prepend_before_agent(self._before_agent)
+        return result
+
+    def _get_retrieve_config(
+        self, vector_db_extras: VectorDBExtras
+    ) -> tuple[str, str]:
+        """Get the retrieve config argument for the agent.
+
+        Returns
+        -------
+        InstanceArgument | None
+            The retrieve config argument if applicable, otherwise None.
+        """
+        if not isinstance(self.agent, WaldiezRagUserProxy):
+            return "", ""
+        retrieve_config: WaldiezRagUserProxyRetrieveConfig = (
+            self.agent.retrieve_config
+        )
+        if not retrieve_config:
+            return "", ""
+        args_dict = self._get_args_dict()
+        if not args_dict:
+            return "", ""
+        before_agent = vector_db_extras.before_arg
+        if self.agent.retrieve_config.use_custom_token_count:
+            function_content, token_count_arg_name = (
+                retrieve_config.get_custom_token_count_function(
+                    name_suffix=self.agent_name
+                )
+            )
+            args_dict["custom_token_count_function"] = token_count_arg_name
+            before_agent += "\n" + function_content + "\n"
+        if self.agent.retrieve_config.use_custom_text_split:
+            function_content, text_split_arg_name = (
+                retrieve_config.get_custom_text_split_function(
+                    name_suffix=self.agent_name
+                )
+            )
+            args_dict["custom_text_split_function"] = text_split_arg_name
+            before_agent += "\n" + function_content + "\n"
+        if before_agent.strip() and not before_agent.endswith("\n"):
+            before_agent += "\n"
+        return before_agent, self._get_args_string(
+            args_dict,
+            vector_db_extras,
+            before_agent,
+        )
+
+    def _get_args_string(
+        self,
+        args_dict: dict[str, Union[str, list[str]]],
+        vector_db_extras: VectorDBExtras,
+        before_agent: str,
+    ) -> str:
+        if not isinstance(self.agent, WaldiezRagUserProxy):
+            return ""
+        args_content = self.serializer.serialize(args_dict)
+        # get the last line (where the dict ends)
+        args_parts = args_content.split("\n")
+        before_vector_db = args_parts[:-1]
+        closing_arg = args_parts[-1]
+        args_content = "\n".join(before_vector_db)
+        # add the vector_db arg
+        args_content += (
+            ",\n"
+            + f'        "vector_db": {vector_db_extras.vector_db_arg},'
+            + "\n"
+        )
+        # we should not need to include the client, but let's do it
+        # to avoid later issues (with telemetry or other client settings)
+        # https://github.com/ag2ai/ag2/blob/main/autogen/agentchat/\
+        #   contrib/retrieve_user_proxy_agent.py#L265-L266
+        if (
+            f"{self.agent_name}_client" in before_agent
+            and self.agent.retrieve_config.vector_db == "chroma"
+        ):
+            args_content += (
+                f'        "client": {self.agent_name}_client,' + "\n"
+            )
+        args_content += closing_arg
+        return args_content
+
+    def _get_model_arg(
+        self,
+    ) -> str:
+        if not isinstance(self.agent, WaldiezRagUserProxy):
+            return ""
+        if self.agent.data.model_ids:
+            model_name = self.model_names[self.agent.data.model_ids[0]]
+            new_model_name = f"{model_name}"
+            return f"{new_model_name}"
+        if self.agent.retrieve_config.model in self.model_names:
+            selected_model = self.model_names[self.agent.retrieve_config.model]
+            new_model_name = f"{selected_model}"
+            return f"{new_model_name}"
+        return WaldiezRagUserProxyModels[self.agent.retrieve_config.vector_db]
+
+    def _get_args_dict(self) -> dict[str, Union[str, list[str]]]:
+        if not isinstance(self.agent, WaldiezRagUserProxy):
+            return {}
+        model_arg = self._get_model_arg()
+        args_dict: dict[str, Union[str, list[str]]] = {
+            "task": self.agent.retrieve_config.task,
+            "model": model_arg,
+        }
+        optional_args = [
+            "chunk_token_size",
+            "context_max_tokens",
+            "customized_prompt",
+            "customized_answer_prefix",
+        ]
+        for arg in optional_args:
+            arg_value = getattr(self.agent.retrieve_config, arg)
+            if arg_value is not None:
+                args_dict[arg] = arg_value
+                args_dict[arg] = getattr(self.agent.retrieve_config, arg)
+        docs_path: Union[str, list[str]] = []
+        if self.agent.retrieve_config.docs_path:
+            doc_paths = (
+                self.agent.retrieve_config.docs_path
+                if isinstance(self.agent.retrieve_config.docs_path, list)
+                else [self.agent.retrieve_config.docs_path]
+            )
+            docs_path = [
+                item
+                for item in [
+                    self.path_resolver.resolve(path) for path in doc_paths
+                ]
+                if item
+            ]
+            args_dict["docs_path"] = docs_path
+        if docs_path:
+            args_dict["docs_path"] = docs_path
+        non_optional_args = [
+            "new_docs",
+            "update_context",
+            "get_or_create",
+            "overwrite",
+            "recursive",
+            "chunk_mode",
+            "must_break_at_empty_line",
+            "collection_name",
+            "distance_threshold",
+        ]
+        for arg in non_optional_args:
+            args_dict[arg] = getattr(self.agent.retrieve_config, arg)
+        return args_dict

waldiez/exporting/agent/extras/reasoning_agent_extras.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+# pylint: disable=too-few-public-methods,no-self-use
+"""Reasoning agent configuration processor."""
+
+from typing import Optional
+
+from waldiez.models import WaldiezAgent, WaldiezReasoningAgent
+
+from ...core import (
+    CodeExecutionConfig,
+    DefaultSerializer,
+    InstanceArgument,
+    Serializer,
+    SystemMessageConfig,
+    TerminationConfig,
+)
+from ...core.extras.agent_extras import ReasoningExtras
+
+
+class ReasoningAgentProcessor:
+    """Processor for reasoning agent configuration."""
+
+    def __init__(self, agent: WaldiezAgent, serializer: Serializer | None):
+        """Initialize the processor with the agent and serializer.
+
+        Parameters
+        ----------
+        agent : WaldiezAgent
+            The Waldiez agent to process.
+        serializer : Serializer | None
+            Optional serializer for the reasoning configuration.
+            Defaults to DefaultSerializer if not provided.
+        """
+        self.agent = agent
+        self.serializer = serializer or DefaultSerializer()
+
+    def process(
+        self,
+        code_execution_config: Optional[CodeExecutionConfig] = None,
+        termination_config: Optional[TerminationConfig] = None,
+        system_message_config: Optional[SystemMessageConfig] = None,
+    ) -> ReasoningExtras:
+        """Process reasoning agent configuration.
+
+        Parameters
+        ----------
+        code_execution_config : CodeExecutionConfig, optional
+            Configuration for code execution, if applicable.
+        termination_config : TerminationConfig, optional
+            Configuration for termination, if applicable.
+        system_message_config : SystemMessageConfig, optional
+            Configuration for system messages, if applicable.
+
+        Returns
+        -------
+        ReasoningExtras
+            The processed result containing extra arguments, before content,
+            imports, and environment variables.
+        """
+        result = ReasoningExtras(
+            instance_id=self.agent.id,
+            code_execution_config=code_execution_config,
+            termination_config=termination_config,
+            system_message_config=system_message_config,
+        )
+        if not self.agent.is_reasoning or not isinstance(
+            self.agent, WaldiezReasoningAgent
+        ):  # pragma: no cover
+            return result
+
+        reasoning_config = self.agent.get_reasoning_config()
+        serialized = self.serializer.serialize(reasoning_config)
+        reason_arg = InstanceArgument(
+            instance_id=self.agent.id,
+            name="reason_config",
+            value=serialized,
+            tabs=1,
+        )
+        result.add_arg(reason_arg)
+        verbose_arg = InstanceArgument(
+            instance_id=self.agent.id,
+            name="verbose",
+            value=self.agent.data.verbose,
+            tabs=1,
+        )
+        result.add_arg(verbose_arg)
+        return result

waldiez/exporting/agent/factory.py

@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+# pylint: disable=too-few-public-methods,
+# pylint: disable=too-many-arguments,too-many-positional-arguments
+"""Factory for creating agent exporter."""
+
+from pathlib import Path
+from typing import Any, Callable, Optional, Union
+
+from waldiez.models import (
+    WaldiezAgent,
+    WaldiezAgentConnection,
+    WaldiezChat,
+    WaldiezModel,
+)
+
+from ..core import (
+    ExporterContext,
+    get_default_exporter_context,
+)
+from .exporter import AgentExporter
+
+
+def create_agent_exporter(
+    agent: WaldiezAgent,
+    agent_names: dict[str, str],
+    models: tuple[list[WaldiezModel], dict[str, str]],
+    chats: tuple[list[WaldiezChat], dict[str, str]],
+    tool_names: dict[str, str],
+    initial_chats: list[WaldiezAgentConnection],
+    is_async: bool = False,
+    for_notebook: bool = False,
+    cache_seed: Optional[int] = None,
+    group_chat_members: Optional[list[WaldiezAgent]] = None,
+    arguments_resolver: Optional[Callable[[WaldiezAgent], list[str]]] = None,
+    output_dir: Optional[Union[str, Path]] = None,
+    context: Optional[ExporterContext] = None,
+    **kwargs: Any,
+) -> AgentExporter:
+    """Create an agent exporter.
+
+    Parameters
+    ----------
+    agent : WaldiezAgent
+        The agent to export.
+    agent_names : dict[str, str]
+        Mapping of agent IDs to names.
+    models : tuple[list[WaldiezModel], dict[str, str]]
+        All models and model names mapping.
+    chats : tuple[list[WaldiezChat], dict[str, str]]
+        All chats and chat names mapping.
+    tool_names : dict[str, str]
+        Mapping of tool IDs to names.
+    is_async : bool, optional
+        Whether the flow is async, by default False
+    for_notebook : bool, optional
+        Whether exporting for notebook, by default False
+    cache_seed : Optional[int], optional
+        Cache seed if any, by default None
+    initial_chats : Optional[list[WaldiezAgentConnection]], optional
+        Initial chats for group managers, by default None
+    group_chat_members : Optional[list[WaldiezAgent]], optional
+        Group chat members if group manager, by default None
+    arguments_resolver : Optional[Callable], optional
+        Function to resolve additional arguments, by default None
+    output_dir : Optional[Union[str, Path]], optional
+        Output directory for generated files, by default None
+    context : Optional[ExporterContext], optional
+        Exporter context with dependencies, by default None
+    **kwargs : Any
+        Additional keyword arguments.
+
+    Returns
+    -------
+    AgentExporter
+        The created agent exporter.
+    """
+    if context is None:
+        context = get_default_exporter_context()
+    return AgentExporter(
+        agent=agent,
+        agent_names=agent_names,
+        models=models,
+        chats=chats,
+        tool_names=tool_names,
+        is_async=is_async,
+        for_notebook=for_notebook,
+        cache_seed=cache_seed,
+        initial_chats=initial_chats,
+        group_chat_members=group_chat_members,
+        arguments_resolver=arguments_resolver,
+        output_dir=output_dir,
+        context=context,
+        **kwargs,
+    )