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

waldiez/exporting/flow/orchestrator.py

@@ -0,0 +1,411 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Flow export orchestrator."""
+
+from typing import Any, Callable
+
+from waldiez.models import Waldiez, WaldiezAgent
+
+from ..agent import AgentExporter, create_agent_exporter
+from ..chats import ChatsExporter, create_chats_exporter
+from ..core import (
+    AgentPosition,
+    ContentOrder,
+    ExporterContext,
+    ExportPosition,
+    ExportResult,
+    ImportPosition,
+)
+from ..models import ModelsExporter, create_models_exporter
+from ..tools import ToolsExporter, create_tools_exporter
+from .merger import ContentMerger
+from .utils import (
+    generate_header,
+    get_after_run_content,
+    get_np_no_nep50_handle,
+    get_sqlite_out,
+    get_start_logging,
+    get_stop_logging,
+    get_the_imports_string,
+)
+
+
+# pylint: disable=no-self-use,too-many-instance-attributes
+class ExportOrchestrator:
+    """Coordinates the export process."""
+
+    def __init__(
+        self,
+        waldiez: Waldiez,
+        context: ExporterContext,
+    ) -> None:
+        """Initialize the export orchestrator.
+
+        Parameters
+        ----------
+        waldiez : Waldiez
+            The Waldiez instance containing the flow to export.
+        context : ExporterContext
+            The exporter context containing dependencies and configuration.
+        """
+        self.waldiez = waldiez
+        self.context = context
+        self.config = context.get_config()
+        self._tools_exporter: ToolsExporter | None = None
+        self._models_exporter: ModelsExporter | None = None
+        self._chats_exporter: ChatsExporter | None = None
+        self.logger = context.get_logger()
+        self._initialize()
+
+    def _initialize(self) -> None:
+        """Initialize the orchestrator with necessary configurations."""
+        unique_names = self.waldiez.flow.unique_names
+        self.flow_name = unique_names["flow_name"]
+        self.agents = unique_names["agents"]
+        self.models = unique_names["models"]
+        self.tools = unique_names["tools"]
+        self.chats = unique_names["chats"]
+        self.agent_names = unique_names["agent_names"]
+        self.model_names = unique_names["model_names"]
+        self.tool_names = unique_names["tool_names"]
+        self.chat_names = unique_names["chat_names"]
+
+    def orchestrate(self) -> ExportResult:
+        """Orchestrate the export process.
+
+        Returns
+        -------
+        ExportResult
+            The result of the export process,
+            containing the generated script and any additional metadata.
+        """
+        results: list[ExportResult] = []
+        agent_arguments: dict[str, list[str]] = {}
+
+        # 1. Tools first (needed by agents)
+        if self.waldiez.tools:
+            self.logger.info("Exporting tools ...")
+            tools_result = self._get_tools_exporter().export()
+            # Extract tool arguments for agents
+            tool_arguments = self._extract_agent_arguments_from_result(
+                tools_result
+            )
+            # Merge tool arguments into agent arguments
+            self._merge_agent_arguments(
+                source=tool_arguments,
+                target=agent_arguments,
+            )
+            results.append(tools_result)
+            self.logger.debug("Exported %s", tools_result)
+
+        # 2. Models second (needed by agents)
+        if self.waldiez.models:
+            self.logger.info("Exporting models ...")
+            models_result = self._get_models_exporter().export()
+            # Extract model arguments for agents
+            model_arguments = self._extract_agent_arguments_from_result(
+                models_result
+            )
+            # Merge model arguments into agent arguments
+            self._merge_agent_arguments(
+                source=model_arguments,
+                target=agent_arguments,
+            )
+            results.append(models_result)
+            self.logger.debug("Exported %s", models_result)
+        # 3. Chats third (agents might need agent chat registrations)
+        # we always have at least one chat (already validated in Waldiez init)
+        self.logger.info("Exporting chats ...")
+        chats_result = self._get_chats_exporter().export()
+        self.logger.debug("Exported %s", chats_result)
+        # Extract chat arguments for agents
+        chat_arguments = self._extract_agent_arguments_from_result(chats_result)
+        # Merge chat arguments into agent arguments
+        self._merge_agent_arguments(
+            source=chat_arguments,
+            target=agent_arguments,
+        )
+        results.append(chats_result)
+
+        # 4. Agents last
+        # we always have at least one agent (already validated in Waldiez init)
+        agent_results = self._export_all_agents(agent_arguments)
+        results.extend(agent_results)
+
+        # 5. Merge everything
+        merger = ContentMerger(self.context)
+        merged_result = merger.merge_results(results)
+        # Check for issues
+        stats = merger.get_merge_statistics()
+        if stats.conflicts_found:
+            self.logger.info(
+                "Resolved %d merge conflicts", len(stats.conflicts_found)
+            )
+        self.logger.debug("Merged result: %s", merged_result)
+        return self._finalize_export(merged_result)
+
+    def _finalize_export(self, merged_result: ExportResult) -> ExportResult:
+        """Finalize the export result with additional content.
+
+        Parameters
+        ----------
+        merged_result : ExportResult
+            The merged export result containing all content.
+
+        Returns
+        -------
+        ExportResult
+            The finalized export result with additional content.
+        """
+        merged_result.add_content(
+            generate_header(
+                name=self.waldiez.name,
+                description=self.waldiez.description,
+                requirements=self.waldiez.requirements,
+                tags=self.waldiez.tags,
+                for_notebook=self.config.for_notebook,
+            ),
+            position=ExportPosition.TOP,  # befoe everything
+            order=ContentOrder.EARLY_SETUP,
+        )
+        merged_result.add_content(
+            get_np_no_nep50_handle(),
+            position=ExportPosition.IMPORTS,  # after imports (need np)
+            order=ContentOrder.CLEANUP,
+        )
+        merged_result.add_content(
+            get_start_logging(
+                is_async=self.waldiez.is_async,
+                for_notebook=self.config.for_notebook,
+            ),
+            position=ExportPosition.IMPORTS,  # after imports, before models
+            order=ContentOrder.CLEANUP.value + 1,  # after imports
+            skip_strip=True,  # keep newlines
+        )
+        merged_result.add_content(
+            get_sqlite_out(is_async=self.waldiez.is_async),
+            position=ExportPosition.AGENTS,
+            order=ContentOrder.LATE_CLEANUP.value + 1,  # after all agents
+        )
+        merged_result.add_content(
+            get_stop_logging(is_async=self.waldiez.is_async),
+            position=ExportPosition.AGENTS,
+            order=ContentOrder.LATE_CLEANUP.value
+            + 2,  # before def main (chats)
+        )
+        all_imports: list[tuple[str, ImportPosition]] = [
+            (item.statement, item.position)
+            for item in merged_result.get_sorted_imports()
+        ]
+        import_string = get_the_imports_string(
+            all_imports=all_imports,
+            is_async=self.waldiez.is_async,
+        )
+        if self.config.for_notebook:
+            cell_ignore_comment = (
+                "# pyright: "
+                "reportUnusedImport=false,"
+                "reportMissingTypeStubs=false"
+                "\n"
+            )
+            import_string = cell_ignore_comment + import_string
+        merged_result.add_content(
+            import_string,
+            position=ExportPosition.IMPORTS,  # imports section
+            order=ContentOrder.EARLY_SETUP,  # top position
+        )
+        return merged_result
+
+    def get_after_run_content(self) -> str:
+        """Get the content to be executed after the main flow run.
+
+        Returns
+        -------
+        str
+            The content to be executed after the main flow run.
+        """
+        return get_after_run_content(
+            waldiez=self.waldiez,
+            agent_names=self.agent_names,
+            tabs=1,
+        )
+
+    def _get_tools_exporter(self) -> ToolsExporter:
+        """Get or create tools exporter."""
+        if self._tools_exporter is None:
+            self._tools_exporter = create_tools_exporter(
+                flow_name=self.flow_name,
+                agents=self.agents,
+                agent_names=self.agent_names,
+                tools=self.tools,
+                tool_names=self.tool_names,
+                output_dir=self.config.output_directory,
+                context=self.context,
+            )
+        return self._tools_exporter
+
+    def _get_models_exporter(self) -> ModelsExporter:
+        """Get or create models exporter."""
+        if self._models_exporter is None:
+            self._models_exporter = create_models_exporter(
+                flow_name=self.flow_name,
+                agents=self.agents,
+                agent_names=self.agent_names,
+                models=self.models,
+                model_names=self.model_names,
+                for_notebook=self.config.for_notebook,
+                cache_seed=self.config.cache_seed,
+                output_dir=self.config.output_directory,
+                context=self.context,
+            )
+        return self._models_exporter
+
+    def _get_chats_exporter(self) -> ChatsExporter:
+        """Get or create chats exporter."""
+        if self._chats_exporter is None:
+            self._chats_exporter = create_chats_exporter(
+                waldiez=self.waldiez,
+                all_agents=self.agents,
+                agent_names=self.agent_names,
+                models=self.models,
+                model_names=self.model_names,
+                tools=self.tools,
+                tool_names=self.tool_names,
+                all_chats=self.chats,
+                main_chats=self.waldiez.initial_chats,
+                chat_names=self.chat_names,
+                output_dir=self.config.output_directory,
+                context=self.context,
+            )
+        return self._chats_exporter
+
+    # pylint: disable=no-self-use
+    def _create_agent_arguments_resolver(
+        self, exported_arguments: dict[str, Any]
+    ) -> Callable[[WaldiezAgent], list[str]]:
+        """Create an arguments resolver function for a specific agent.
+
+        Parameters
+        ----------
+        exported_arguments : dict[str, Any]
+            A dictionary containing exported arguments for agents,
+            models, and tools.
+
+        Returns
+        -------
+        Callable[[WaldiezAgent], list[str]]
+            A function that takes a WaldiezAgent and returns a list of
+            arguments to be used for that agent.
+        """
+
+        def arguments_resolver(target_agent: WaldiezAgent) -> list[str]:
+            """Resolve arguments for the target agent.
+
+            Parameters
+            ----------
+            target_agent : WaldiezAgent
+                The agent for which to resolve arguments.
+
+            Returns
+            -------
+            list[str]
+                A list of arguments to be used for the agent.
+            """
+            # Extract arguments from exported_arguments
+            return exported_arguments.get(target_agent.id, [])
+
+        return arguments_resolver
+
+    def _create_agent_exporter(
+        self,
+        agent: WaldiezAgent,
+        agent_arguments: dict[str, list[str]],
+    ) -> AgentExporter:
+        """Create an exporter for a specific agent."""
+        return create_agent_exporter(
+            agent=agent,
+            agent_names=self.agent_names,
+            models=(self.models, self.model_names),
+            chats=(self.chats, self.chat_names),
+            initial_chats=self.waldiez.initial_chats,
+            tool_names=self.tool_names,
+            cache_seed=self.waldiez.cache_seed,
+            is_async=self.waldiez.is_async,
+            for_notebook=self.config.for_notebook,
+            output_dir=self.config.output_directory,
+            context=self.context,
+            group_chat_members=self.waldiez.get_group_chat_members(agent),
+            arguments_resolver=self._create_agent_arguments_resolver(
+                exported_arguments=agent_arguments,
+            ),
+        )
+
+    def _extract_agent_arguments_from_result(
+        self, result: ExportResult
+    ) -> dict[str, list[str]]:
+        """Extract agent-specific arguments from an export result."""
+        agent_args: dict[str, list[str]] = {}
+
+        for content in result.positioned_content:
+            # Look for content positioned as agent arguments
+            if (
+                content.position == ExportPosition.AGENTS
+                and content.agent_position == AgentPosition.AS_ARGUMENT
+                and content.agent_id
+            ):
+                if content.agent_id not in agent_args:
+                    agent_args[content.agent_id] = []
+
+                # The content itself is the argument string
+                strip_content = content.content.strip()
+                if strip_content:
+                    # Add the content as an argument for the agent
+                    agent_args[content.agent_id].append(f"    {strip_content}")
+
+        return agent_args
+
+    def _merge_agent_arguments(
+        self, target: dict[str, list[str]], source: dict[str, list[str]]
+    ) -> None:
+        """Merge agent arguments from source into target."""
+        for agent_id, args in source.items():
+            if agent_id not in target:
+                target[agent_id] = []
+            target[agent_id].extend(args)
+
+    def _gather_agent_arguments(self) -> dict[str, Any]:
+        """Gather exported arguments from all agents, models, and tools.
+
+        Returns
+        -------
+        dict[str, Any]
+            A dictionary containing exported arguments for agents,
+            models, and tools.
+        """
+        exported_arguments: dict[str, Any] = {}
+        return exported_arguments
+
+    def _export_all_agents(
+        self, exported_arguments: dict[str, Any]
+    ) -> list[ExportResult]:
+        """Export all agents in dependency order.
+
+        Parameters
+        ----------
+        exported_arguments : dict[str, Any]
+            A dictionary containing exported arguments for agents,
+            models, and tools.
+
+        Returns
+        -------
+        list[ExportResult]
+            A list of export results for each agent.
+        """
+        results: list[ExportResult] = []
+        for agent in self.waldiez.agents:
+            agent_exporter = self._create_agent_exporter(
+                agent, exported_arguments
+            )
+            agent_result = agent_exporter.export()
+            results.append(agent_result)
+        return results

waldiez/exporting/flow/utils/__init__.py

@@ -1,49 +1,22 @@
 # SPDX-License-Identifier: Apache-2.0.
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
-"""Utility functions for exporting waldiez to different formats."""
+"""Flow exporter utils."""
 
-from .agent_utils import (
-    add_after_agent_content,
-    add_after_all_agents_content,
-    add_before_agent_content,
-    add_before_all_agents_content,
-    gather_agent_outputs,
-)
-from .chat_utils import add_after_chat_content, add_before_chat_content
-from .def_main import get_def_main
-from .flow_content import (
+from .common import (
+    generate_header,
     get_after_run_content,
-    get_ipynb_content_start,
     get_np_no_nep50_handle,
-    get_py_content_start,
-)
-from .flow_names import ensure_unique_names
-from .importing_utils import gather_imports, get_the_imports_string
-from .logging_utils import (
-    get_sqlite_out,
-    get_sqlite_out_call,
-    get_start_logging,
-    get_stop_logging,
 )
+from .importing import get_sorted_imports, get_the_imports_string
+from .logging import get_sqlite_out, get_start_logging, get_stop_logging
 
 __all__ = [
-    "add_after_agent_content",
-    "add_after_all_agents_content",
-    "add_before_agent_content",
-    "add_before_all_agents_content",
-    "add_after_chat_content",
-    "add_before_chat_content",
-    "ensure_unique_names",
-    "gather_agent_outputs",
-    "gather_imports",
-    "get_after_run_content",
-    "get_def_main",
+    "generate_header",
     "get_np_no_nep50_handle",
-    "get_py_content_start",
-    "get_ipynb_content_start",
+    "get_after_run_content",
+    "get_the_imports_string",
+    "get_sorted_imports",
     "get_start_logging",
     "get_stop_logging",
     "get_sqlite_out",
-    "get_sqlite_out_call",
-    "get_the_imports_string",
 ]

waldiez/exporting/flow/utils/common.py

@@ -0,0 +1,206 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+# pylint: disable=line-too-long
+# flake8: noqa: E501
+"""Common utils for the final generatio."""
+
+from waldiez.models import Waldiez
+
+from ...core import FILE_HEADER
+from .linting import (
+    get_flake8_ignore_comment,
+    get_mypy_ignore_comment,
+    get_pylint_ignore_comment,
+    get_pyright_ignore_comment,
+)
+
+RETURN_TYPE_HINT = "Union[ChatResult, list[ChatResult], dict[int, ChatResult]]"
+GENERATED_WITH = "🧩 generated with ❤️ by Waldiez."
+
+
+def generate_header(
+    name: str,
+    description: str,
+    requirements: list[str],
+    tags: list[str],
+    for_notebook: bool,
+) -> str:
+    """Generate the header for the script or notebook.
+
+    Parameters
+    ----------
+    name : str
+        The name of the flow.
+    description : str
+        A brief description of the flow.
+    requirements : list[str]
+        A list of requirements for the flow.
+    tags : list[str]
+        A list of tags associated with the flow.
+    for_notebook : bool
+        Whether the header is for a notebook or a script.
+
+    Returns
+    -------
+    str
+        The header content.
+    """
+    if not for_notebook:
+        return _get_py_header(
+            name=name,
+            description=description,
+            requirements=requirements,
+            tags=tags,
+        )
+    return _get_ipynb_heeader(
+        name=name,
+        description=description,
+        requirements=requirements,
+        tags=tags,
+    )
+
+
+def _get_ipynb_heeader(
+    name: str,
+    description: str,
+    requirements: list[str],
+    tags: list[str],
+) -> str:
+    requirements_str = " ".join(requirements)
+    tags_str = ", ".join(tags)
+    content = "# %% [markdown]\n"
+    content += f"## Name: {name}\n\n"
+    content += f"### Description: {description}\n\n"
+    content += f"### Tags: {tags_str}\n\n"
+    content += f"####{GENERATED_WITH}\n\n"
+    content += "#### Requirements\n\n# %%\n"
+    content += "import sys  # pyright: ignore\n"
+    # fmt: off
+    content += "# # " + f"!{{sys.executable}} -m pip install -q {requirements_str}" + "\n"
+    # fmt: on
+    return content
+
+
+def _get_py_header(
+    name: str,
+    description: str,
+    requirements: list[str],
+    tags: list[str],
+) -> str:
+    return f'''#!/usr/bin/env python
+{FILE_HEADER}
+{get_flake8_ignore_comment()}
+{get_pylint_ignore_comment()}
+{get_mypy_ignore_comment([])}
+{get_pyright_ignore_comment()}
+"""{name}.
+
+{description}
+
+Requirements: {", ".join(requirements)}
+Tags: {", ".join(tags)}
+{GENERATED_WITH}
+"""'''
+
+
+def main_doc_string() -> str:
+    """Generate the docstring for the main function.
+
+    Returns
+    -------
+    str
+        The docstring for the main function.
+    """
+    return f'''"""Start chatting.
+
+    Returns
+    -------
+    {RETURN_TYPE_HINT}
+        The result of the chat session, which can be a single ChatResult,
+        a list of ChatResults, or a dictionary mapping integers to ChatResults.
+    """'''
+
+
+def get_after_run_content(
+    waldiez: Waldiez,
+    agent_names: dict[str, str],
+    tabs: int,
+) -> str:
+    """Get content to add after the flow is run.
+
+    Parameters
+    ----------
+    waldiez : Waldiez
+        The waldiez object.
+    agent_names : dict[str, str]
+        The dictionary of agent names and their corresponding ids
+    tabs : int
+        The number of tabs to add before the content.
+
+    Returns
+    -------
+    str
+        The content to add after the flow is run.
+    """
+    # if the flow has a reasoning agents, we add
+    # agent.visualize_tree() for each agent
+    content = ""
+    tab = "    "
+    space = tab * tabs
+    for agent in waldiez.agents:
+        if agent.is_reasoning:
+            agent_name = agent_names[agent.id]
+            content += f"""
+{space}# pylint: disable=broad-except,too-many-try-statements
+{space}try:
+{space}{tab}{agent_name}.visualize_tree()
+{space}{tab}if os.path.exists("tree_of_thoughts.png"):
+{space}{tab}{tab}new_name = "{agent_name}_tree_of_thoughts.png"
+{space}{tab}{tab}os.rename("tree_of_thoughts.png", new_name)
+{space}except BaseException:
+{space}{tab}pass
+{space}# save the tree to json
+{space}try:
+{space}{tab}data = {agent_name}._root.to_dict()  # pylint: disable=protected-access  # pyright: ignore
+{space}{tab}with open("{agent_name}_reasoning_tree.json", "w", encoding="utf-8") as f:
+{space}{tab}{tab}json.dump(data, f)
+{space}except BaseException:
+{space}{tab}pass
+"""
+    return content
+
+
+def get_np_no_nep50_handle() -> str:
+    """Handle the "module numpy has no attribute _no_pep50_warning" error.
+
+    Returns
+    -------
+    str
+        The content to handle the error.
+    """
+    content = '''
+#
+# let's try to avoid:
+# module 'numpy' has no attribute '_no_nep50_warning'"
+# ref: https://github.com/numpy/numpy/blob/v2.2.2/doc/source/release/2.2.0-notes.rst#nep-50-promotion-state-option-removed
+os.environ["NEP50_DEPRECATION_WARNING"] = "0"
+os.environ["NEP50_DISABLE_WARNING"] = "1"
+os.environ["NPY_PROMOTION_STATE"] = "weak"
+if not hasattr(np, "_no_pep50_warning"):
+
+    import contextlib
+    from typing import Generator
+
+    @contextlib.contextmanager
+    def _np_no_nep50_warning() -> Generator[None, None, None]:
+        """Dummy function to avoid the warning.
+
+        Yields
+        ------
+        None
+            Nothing.
+        """
+        yield
+    setattr(np, "_no_pep50_warning", _np_no_nep50_warning)  # noqa
+'''
+    return content