waldiez 0.2.2__py3-none-any.whl → 0.3.1__py3-none-any.whl

waldiez/exporting/agent/utils/agent_imports.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Get the imports needed for the agent."""
+
+from typing import Set
+
+
+def get_agent_imports(agent_class: str) -> Set[str]:
+    """Get the imports needed for the agent.
+
+    Parameters
+    ----------
+    agent_class : str
+        The agent class name.
+
+    Returns
+    -------
+    Set[str]
+        The imports needed for the agent.
+    """
+    imports = set(["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 == "GroupChatManager":
+        imports.add("from autogen import GroupChatManager")
+    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 == "SwarmAgent":
+        imports.add(
+            "from autogen import "
+            "AFTER_WORK, "
+            "ON_CONDITION, "
+            "UPDATE_SYSTEM_MESSAGE, "
+            "AfterWorkOption, "
+            "SwarmAgent, "
+            "SwarmResult"
+        )
+    else:
+        imports.add("from autogen import ConversableAgent")
+    return imports

waldiez/exporting/{agents → agent/utils}/code_execution.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Code execution related functions for exporting agents."""
 
 from typing import Dict, Tuple
@@ -36,15 +38,15 @@ def get_agent_code_execution_config(
         if use_docker
         else "LocalCommandLineCodeExecutor"
     )
-    executor_content = f"{agent_name}_executor = {executor_class_name}(\n"
+    executor_content = f"{agent_name}_executor = {executor_class_name}(" + "\n"
     if agent.data.code_execution_config.work_dir:
         wok_dir = agent.data.code_execution_config.work_dir.replace(
             '"', '\\"'
         ).replace("\n", "\\n")
-        executor_content += f'    work_dir="{wok_dir}",\n'
+        executor_content += f'    work_dir="{wok_dir}",' + "\n"
     if agent.data.code_execution_config.timeout:
         executor_content += (
-            f"    timeout={agent.data.code_execution_config.timeout},\n"
+            f"    timeout={agent.data.code_execution_config.timeout}," + "\n"
         )
     if use_docker is False and agent.data.code_execution_config.functions:
         function_names = []
@@ -53,15 +55,11 @@ def get_agent_code_execution_config(
             function_names.append(skill_name)
         if function_names:
             # pylint: disable=inconsistent-quotes
+            function_names_string = ", ".join(function_names)
             executor_content += (
-                f"    functions=[{', '.join(function_names)}],\n"
+                f"    functions=[{function_names_string}]," + "\n"
             )
     executor_content += ")\n\n"
-    # if (
-    #     executor_content
-    #     == f"{agent_name}_executor = {executor_class_name}(\n)\n\n"
-    # ):
-    #     # empty executor?
-    #     return "", "False", ""
     executor_arg = f'{{"executor": {agent_name}_executor}}'
-    return executor_content, executor_arg, executor_class_name
+    the_import = f"from autogen.coding import {executor_class_name}"
+    return executor_content, executor_arg, the_import

waldiez/exporting/{agents → agent/utils}/group_manager.py

@@ -1,16 +1,17 @@
-"""Export group manger and group chat to string."""
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Export group manager and group chat to string."""
 
-from typing import Dict, List, Optional, Tuple
+from typing import Callable, Dict, List, Optional, Tuple
 
 from waldiez.models import WaldiezAgent, WaldiezGroupManager
 
-from ..utils import get_method_string, get_object_string
-
 
 def get_group_manager_extras(
     agent: WaldiezAgent,
     group_chat_members: List[WaldiezAgent],
     agent_names: Dict[str, str],
+    serializer: Callable[..., str],
 ) -> Tuple[str, str]:
     """Get the group manager extra string and custom selection method if any.
 
@@ -22,6 +23,8 @@ def get_group_manager_extras(
         The group members.
     agent_names : Dict[str, str]
         The agent names.
+    serializer : Callable[..., str]
+        The serializer function.
 
     Returns
     -------
@@ -34,12 +37,17 @@ def get_group_manager_extras(
     custom_speaker_selection: Optional[str] = None
     if agent.agent_type == "manager" and isinstance(agent, WaldiezGroupManager):
         group_chat_string, group_chat_name, custom_speaker_selection = (
-            _get_group_manager_extras(agent, group_chat_members, agent_names)
+            _get_group_manager_extras(
+                agent=agent,
+                group_members=group_chat_members,
+                agent_names=agent_names,
+                serializer=serializer,
+            )
         )
         if group_chat_name:
-            group_chat_arg = f"\n    groupchat={group_chat_name},"
+            group_chat_arg = "\n" + f"    groupchat={group_chat_name},"
         if custom_speaker_selection:
-            before_agent_string += f"{custom_speaker_selection}\n\n"
+            before_agent_string += f"{custom_speaker_selection}" + "\n"
         if group_chat_string:
             before_agent_string += group_chat_string
     return before_agent_string, group_chat_arg
@@ -49,6 +57,7 @@ def _get_group_manager_extras(
     agent: WaldiezGroupManager,
     group_members: List[WaldiezAgent],
     agent_names: Dict[str, str],
+    serializer: Callable[..., str],
 ) -> Tuple[str, str, Optional[str]]:
     """Get the group manager extra string and custom selection method if any.
 
@@ -60,6 +69,8 @@ def _get_group_manager_extras(
         The group members.
     agent_names : Dict[str, str]
         The agent names.
+    serializer : Callable[..., str]
+        The serializer function.
 
     Returns
     -------
@@ -88,25 +99,19 @@ def _get_group_manager_extras(
         group_chat_string += f"    max_round={agent.data.max_round}," + "\n"
     if agent.data.admin_name:
         group_chat_string += f'    admin_name="{agent.data.admin_name}",' + "\n"
-    extra_group_chat_string, method_name_and_content = (
-        _get_group_chat_speakers_string(agent, agent_names)
+    extra_group_chat_string, custom_selection_method = (
+        _get_group_chat_speakers_string(agent, agent_names, serializer)
     )
-    custom_selection_method: Optional[str] = None
     group_chat_string += extra_group_chat_string
     group_chat_string += ")\n\n"
-    if method_name_and_content:
-        method_name, method_content = method_name_and_content
-        custom_selection_method = get_method_string(
-            "custom_speaker_selection",
-            method_name,
-            method_content,
-        )
     return group_chat_string, group_chat_name, custom_selection_method
 
 
 def _get_group_chat_speakers_string(
-    agent: WaldiezGroupManager, agent_names: Dict[str, str]
-) -> Tuple[str, Optional[Tuple[str, str]]]:
+    agent: WaldiezGroupManager,
+    agent_names: Dict[str, str],
+    serializer: Callable[..., str],
+) -> Tuple[str, Optional[str]]:
     """Get the group chat speakers string.
 
     Parameters
@@ -115,16 +120,18 @@ def _get_group_chat_speakers_string(
         The agent.
     agent_names : Dict[str, str]
         The agent names.
+    serializer : Callable[..., str]
+        The serializer function.
 
     Returns
     -------
     str
         The group chat speakers string.
-    Optional[Tuple[str, str]]
-        The custom selection method name and content if any.
+    Optional[str]
+        The custom custom for speaker selection if any.
     """
     speakers_string = ""
-    method_name_and_content: Optional[Tuple[str, str]] = None
+    function_content: Optional[str] = None
     if agent.data.speakers.max_retries_for_selecting is not None:
         speakers_string += (
             "    max_retries_for_selecting_speaker="
@@ -139,12 +146,14 @@ def _get_group_chat_speakers_string(
         )
     else:
         agent_name = agent_names[agent.id]
-        method_name = f"custom_speaker_selection_method_{agent_name}"
-        method_name_and_content = (
-            method_name,
-            agent.data.speakers.custom_method_string or "",
+        function_content, function_name = (
+            agent.data.speakers.get_custom_method_function(
+                name_suffix=agent_name
+            )
+        )
+        speakers_string += (
+            f"    speaker_selection_method={function_name}," + "\n"
         )
-        speakers_string += f"    speaker_selection_method={method_name}," "\n"
     # selection_mode == "repeat":
     if agent.data.speakers.selection_mode == "repeat":
         speakers_string += _get_speakers_selection_repeat_string(
@@ -156,10 +165,12 @@ def _get_group_chat_speakers_string(
         and agent.data.speakers.allowed_or_disallowed_transitions
     ):
         speakers_string += _get_speakers_selection_transition_string(
-            agent, agent_names
+            agent=agent,
+            agent_names=agent_names,
+            serializer=serializer,
         )
     speakers_string = speakers_string.replace('"None"', "None")
-    return speakers_string, method_name_and_content
+    return speakers_string, function_content
 
 
 def _get_speakers_selection_repeat_string(
@@ -168,9 +179,8 @@ def _get_speakers_selection_repeat_string(
     speakers_string = ""
     if isinstance(agent.data.speakers.allow_repeat, bool):
         speakers_string += (
-            "    allow_repeat_speaker="
-            f"{agent.data.speakers.allow_repeat},"
-            "\n"
+            f"    allow_repeat_speaker={agent.data.speakers.allow_repeat},"
+            + "\n"
         )
     elif isinstance(agent.data.speakers.allow_repeat, list):
         # get the names of the agents
@@ -178,12 +188,14 @@ def _get_speakers_selection_repeat_string(
             agent_names[agent_id]
             for agent_id in agent.data.speakers.allow_repeat
         )
-        speakers_string += f"    allow_repeat=[{allow_repeat}]," "\n"
+        speakers_string += f"    allow_repeat=[{allow_repeat}]," + "\n"
     return speakers_string
 
 
 def _get_speakers_selection_transition_string(
-    agent: WaldiezGroupManager, agent_names: Dict[str, str]
+    agent: WaldiezGroupManager,
+    agent_names: Dict[str, str],
+    serializer: Callable[..., str],
 ) -> str:
     speakers_string = ""
     allowed_or_disallowed_speaker_transitions = {}
@@ -194,7 +206,7 @@ def _get_speakers_selection_transition_string(
         allowed_or_disallowed_speaker_transitions[agent_names[agent_id]] = [
             agent_names[transition] for transition in transitions
         ]
-    transitions_string = get_object_string(
+    transitions_string = serializer(
         allowed_or_disallowed_speaker_transitions, 1
     )
     transitions_string = transitions_string.replace('"', "").replace("'", "")

waldiez/exporting/{agents → agent/utils}/rag_user/__init__.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """RAG User Agent related string generation."""
 
 from .rag_user import get_rag_user_extras, get_rag_user_retrieve_config_str

waldiez/exporting/{agents → agent/utils}/rag_user/chroma_utils.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Get chroma db related imports and content."""
 
 from pathlib import Path
@@ -22,19 +24,21 @@ def _get_chroma_client_string(agent: WaldiezRagUser) -> Tuple[str, str]:
         The 'client' and what to import.
     """
     # TODO: also check `connection_url` (chroma in client-server mode)
-    to_import = "chromadb"
+    to_import = "import chromadb"
     client_str = "chromadb."
     if (
         agent.retrieve_config.db_config.use_local_storage
         and agent.retrieve_config.db_config.local_storage_path is not None
     ):
-        # on windows, we get:
+        # on windows, we might get:
         # SyntaxError: (unicode error) 'unicodeescape' codec can't decode bytes
         # in position 2-3: truncated \UXXXXXXXX escape
         local_path = Path(agent.retrieve_config.db_config.local_storage_path)
         client_str += (
-            f'PersistentClient(path=r"{local_path}", '
-            "settings=Settings(anonymized_telemetry=False))"
+            "PersistentClient(\n"
+            f'    path=r"{local_path}",' + "\n"
+            "    settings=Settings(anonymized_telemetry=False),\n"
+            ")"
         )
     else:
         client_str += "Client(Settings(anonymized_telemetry=False))"
@@ -60,7 +64,7 @@ def _get_chroma_embedding_function_string(
     """
     to_import = ""
     embedding_function_arg = ""
-    embedding_function_body = ""
+    embedding_function_content = ""
     vector_db_model = agent.retrieve_config.db_config.model
     if not agent.retrieve_config.use_custom_embedding:
         to_import = (
@@ -70,13 +74,13 @@ def _get_chroma_embedding_function_string(
         embedding_function_arg = "SentenceTransformerEmbeddingFunction("
         embedding_function_arg += f'model_name="{vector_db_model}")'
     else:
-        embedding_function_arg = f"custom_embedding_function_{agent_name}"
-        embedding_function_body = (
-            f"\ndef custom_embedding_function_{agent_name}():\n"
-            f"{agent.retrieve_config.embedding_function_string}\n"
+        embedding_function_content, embedding_function_arg = (
+            agent.retrieve_config.get_custom_embedding_function(
+                name_suffix=agent_name
+            )
         )
-
-    return embedding_function_arg, to_import, embedding_function_body
+        embedding_function_content = "\n" + embedding_function_content
+    return embedding_function_arg, to_import, embedding_function_content
 
 
 def get_chroma_db_args(
@@ -108,8 +112,8 @@ def get_chroma_db_args(
     if to_import_embedding:
         to_import.add(to_import_embedding)
     kwarg_string = (
-        f"            client={agent_name}_client,\n"
-        f"            embedding_function={embedding_function_arg},\n"
+        f"            client={agent_name}_client," + "\n"
+        f"            embedding_function={embedding_function_arg}," + "\n"
     )
     # The RAG example:
     # https://ag2ai.github.io/ag2/docs/notebooks/agentchat_groupchat_RAG/
@@ -118,21 +122,22 @@ def get_chroma_db_args(
     # https://github.com/microsoft/autogen/issues/3551#issuecomment-2366930994
     # manually initializing the collection before running the flow,
     # might be a workaround.
-    content_before = f"{agent_name}_client = {client_str}\n"
+    content_before = f"{agent_name}_client = {client_str}" + "\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'"{collection_name}")' + "\n"
             )
         else:
             content_before += (
                 "try:\n"
-                f'    {agent_name}_client.get_collection("{collection_name}")\n'
+                f'    {agent_name}_client.get_collection("{collection_name}")'
+                + "\n"
                 "except ValueError:\n"
                 f"    {agent_name}_client.create_collection("
-                f'"{collection_name}")\n'
+                f'"{collection_name}")' + "\n"
             )
     return kwarg_string, to_import, embedding_function_body, content_before

waldiez/exporting/{agents → agent/utils}/rag_user/mongo_utils.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Get mongodb related content and imports."""
 
 from typing import Set, Tuple
@@ -24,7 +26,7 @@ def _get_mongodb_embedding_function_string(
     """
     to_import = ""
     embedding_function_arg = ""
-    embedding_function_body = ""
+    embedding_function_content = ""
     if not agent.retrieve_config.use_custom_embedding:
         to_import = "from sentence_transformers import SentenceTransformer"
         embedding_function_arg = (
@@ -33,12 +35,13 @@ def _get_mongodb_embedding_function_string(
             ").encode"
         )
     else:
-        embedding_function_arg = f"custom_embedding_function_{agent_name}"
-        embedding_function_body = (
-            f"\ndef custom_embedding_function_{agent_name}():\n"
-            f"{agent.retrieve_config.embedding_function_string}\n"
+        embedding_function_content, embedding_function_arg = (
+            agent.retrieve_config.get_custom_embedding_function(
+                name_suffix=agent_name
+            )
         )
-    return embedding_function_arg, to_import, embedding_function_body
+        embedding_function_content = "\n" + embedding_function_content
+    return embedding_function_arg, to_import, embedding_function_content
 
 
 def get_mongodb_db_args(
@@ -67,17 +70,18 @@ def get_mongodb_db_args(
     tab = " " * 12
     db_config = agent.retrieve_config.db_config
     kwarg_string = (
-        f'{tab}connection_string="{db_config.connection_url}",\n'
-        f"{tab}embedding_function={embedding_function_arg},\n"
+        f'{tab}connection_string="{db_config.connection_url}",' + "\n"
+        f"{tab}embedding_function={embedding_function_arg}," + "\n"
     )
     wait_until_document_ready = db_config.wait_until_document_ready
     wait_until_index_ready = db_config.wait_until_index_ready
     if wait_until_document_ready is not None:
         kwarg_string += (
-            f"{tab}wait_until_document_ready={wait_until_document_ready},\n"
+            f"{tab}wait_until_document_ready={wait_until_document_ready},"
+            + "\n"
         )
     if wait_until_index_ready is not None:
         kwarg_string += (
-            f"{tab}wait_until_index_ready={wait_until_index_ready},\n"
+            f"{tab}wait_until_index_ready={wait_until_index_ready}," + "\n"
         )
     return kwarg_string, to_import, embedding_function_body

waldiez/exporting/{agents → agent/utils}/rag_user/pgvector_utils.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Get pgvector related content and imports."""
 
 from typing import Set, Tuple
@@ -18,7 +20,7 @@ def _get_pgvector_client_string(agent: WaldiezRagUser) -> Tuple[str, str]:
     Tuple[str, str]
         The 'client' and what to import.
     """
-    to_import = "psycopg"
+    to_import = "import psycopg"
     client_str = "psycopg."
     connection_url = agent.retrieve_config.db_config.connection_url
     client_str += f'connect("{connection_url}")'
@@ -44,20 +46,22 @@ def _get_pgvector_embedding_function_string(
     """
     to_import = ""
     embedding_function_arg = ""
-    embedding_function_body = ""
+    embedding_function_content = ""
     if agent.retrieve_config.use_custom_embedding:
         embedding_function_arg = f"custom_embedding_function_{agent_name}"
-        embedding_function_body = (
-            f"\ndef custom_embedding_function_{agent_name}():\n"
-            f"{agent.retrieve_config.embedding_function_string}\n"
+        embedding_function_content, embedding_function_arg = (
+            agent.retrieve_config.get_custom_embedding_function(
+                name_suffix=agent_name
+            )
         )
+        embedding_function_content = "\n" + embedding_function_content
     else:
         to_import = "from sentence_transformers import SentenceTransformer"
         embedding_function_arg = "SentenceTransformer("
         embedding_function_arg += (
             f'"{agent.retrieve_config.db_config.model}").encode'
         )
-    return embedding_function_arg, to_import, embedding_function_body
+    return embedding_function_arg, to_import, embedding_function_content
 
 
 def get_pgvector_db_args(
@@ -87,7 +91,7 @@ def get_pgvector_db_args(
         else {to_import_client}
     )
     kwarg_str = (
-        f"            client={client_str},\n"
-        f"            embedding_function={embedding_function_arg},\n"
+        f"            client={client_str}," + "\n"
+        f"            embedding_function={embedding_function_arg}," + "\n"
     )
     return kwarg_str, to_import, embedding_function_body

waldiez/exporting/{agents → agent/utils}/rag_user/qdrant_utils.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Get qdrant db related imports and content."""
 
 from pathlib import Path
@@ -61,7 +63,7 @@ def _get_qdrant_embedding_function_string(
     """
     to_import = ""
     embedding_function_arg = ""
-    embedding_function_body = ""
+    embedding_function_content = ""
     vector_db_model = agent.retrieve_config.db_config.model
     if not agent.retrieve_config.use_custom_embedding:
         to_import = (
@@ -71,12 +73,13 @@ def _get_qdrant_embedding_function_string(
         embedding_function_arg = "FastEmbedEmbeddingFunction("
         embedding_function_arg += f'model_name="{vector_db_model}")'
     else:
-        embedding_function_arg = f"custom_embedding_function_{agent_name}"
-        embedding_function_body = (
-            f"\ndef custom_embedding_function_{agent_name}():\n"
-            f"{agent.retrieve_config.embedding_function_string}\n"
+        embedding_function_content, embedding_function_arg = (
+            agent.retrieve_config.get_custom_embedding_function(
+                name_suffix=agent_name
+            )
         )
-    return embedding_function_arg, to_import, embedding_function_body
+        embedding_function_content = "\n" + embedding_function_content
+    return embedding_function_arg, to_import, embedding_function_content
 
 
 def get_qdrant_db_args(
@@ -106,7 +109,7 @@ def get_qdrant_db_args(
         else {to_import_client}
     )
     kwarg_string = (
-        f"            client={client_str},\n"
-        f"            embedding_function={embedding_function_arg},\n"
+        f"            client={client_str}," + "\n"
+        f"            embedding_function={embedding_function_arg}," + "\n"
     )
     return kwarg_string, to_import, embedding_function_body