waldiez 0.2.1__py3-none-any.whl → 0.3.0__py3-none-any.whl

waldiez/exporting/flow/flow_exporter.py

@@ -0,0 +1,503 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Flow Exporter class.
+
+- We gather all the exports (models, skills, agents, chats).
+
+- We first add all the imports from the above exports.
+- If we have skills, we include their imports.
+    (their files were generated when exporting the skills).
+- Then, we write the all model configs.
+- Next, we write the agent definitions
+    (using the `llm_config=...` argument from the model exports).
+- If additional (nested_chats) are defined,
+    we write their registrations after all agents are defined.
+- Next, we write the chat definitions
+    (using the agent names from the agent exports).
+- If exporting to py,
+    we add the `run` function and the `def main()` to call the run function.
+- If the flow is async, the `run` function is async.
+"""
+
+# flake8: noqa E501
+# pylint: disable=line-too-long
+
+from functools import partial
+from pathlib import Path
+from typing import List, Optional, Tuple, Union
+
+from waldiez.models import Waldiez, WaldiezAgent
+
+from ..agent import AgentExporter
+from ..base import (
+    AgentPosition,
+    AgentPositions,
+    BaseExporter,
+    ExporterMixin,
+    ExporterReturnType,
+    ExportPosition,
+    ExportPositions,
+    ImportPosition,
+)
+from ..chats import ChatsExporter
+from ..models import ModelsExporter
+from ..skills import SkillsExporter
+from .utils import (
+    add_after_agent_content,
+    add_after_chat_content,
+    add_before_agent_content,
+    ensure_unique_names,
+    gather_agent_outputs,
+    gather_imports,
+    get_def_main,
+    get_ipynb_content_start,
+    get_logging_start_string,
+    get_logging_stop_string,
+    get_py_content_start,
+    get_sqlite_out,
+    get_sqlite_out_call,
+    get_the_imports_string,
+)
+
+
+# pylint: disable=too-many-instance-attributes
+class FlowExporter(BaseExporter, ExporterMixin):
+    """Flow exporter."""
+
+    def __init__(
+        self,
+        waldiez: Waldiez,
+        for_notebook: bool,
+        output_dir: Optional[Union[str, Path]] = None,
+    ) -> None:
+        """Initialize the flow exporter."""
+        self.waldiez = waldiez
+        self.for_notebook = for_notebook
+        if output_dir is not None and not isinstance(output_dir, Path):
+            output_dir = Path(output_dir).resolve()
+        self.output_dir = output_dir
+        self.initialize()
+
+    def initialize(
+        self,
+    ) -> None:
+        """Get all the names in the flow.
+
+        We need to make sure that no duplicate names are used,
+        and that the names can be used as python variables.
+        """
+        unique_names = ensure_unique_names(
+            self.waldiez,
+            self.get_valid_instance_name,
+        )
+        self.flow_name = unique_names["flow_name"]
+        self.agents = unique_names["agents"]
+        self.models = unique_names["models"]
+        self.skills = unique_names["skills"]
+        self.chats = unique_names["chats"]
+        self.agent_names = unique_names["agent_names"]
+        self.model_names = unique_names["model_names"]
+        self.skill_names = unique_names["skill_names"]
+        self.chat_names = unique_names["chat_names"]
+
+    def export_flow(self) -> ExporterReturnType:
+        """Export the flow.
+
+        Returns
+        -------
+        ExporterReturnType
+            The exported flow.
+        """
+        models_output = self.export_models()
+        skills_output = self.export_skills()
+        chats_output = self.export_chats()
+        env_vars = self.gather_environment_variables(
+            model_env_vars=models_output["environment_variables"],
+            skill_env_vars=skills_output["environment_variables"],
+            chat_env_vars=chats_output["environment_variables"],
+        )
+        before_export = self.gather_exports(
+            model_export=models_output["before_export"],
+            skill_export=skills_output["before_export"],
+            chat_export=chats_output["before_export"],
+        )
+        after_export = self.gather_exports(
+            model_export=models_output["after_export"],
+            skill_export=skills_output["after_export"],
+            chat_export=chats_output["after_export"],
+        )
+        # agents last (to make sure we have any needed arguments)
+        # like `llm_config=...` from the models
+        agents_output = self.export_agents(
+            before_export=before_export,
+            after_export=after_export,
+        )
+        imports = gather_imports(
+            model_imports=models_output["imports"],
+            skill_imports=skills_output["imports"],
+            chat_imports=chats_output["imports"],
+            agent_imports=agents_output["imports"],
+        )
+        if agents_output["environment_variables"]:
+            env_vars.extend(agents_output["environment_variables"])
+        if agents_output["before_export"]:
+            before_export.extend(agents_output["before_export"])
+        if agents_output["after_export"]:
+            after_export.extend(agents_output["after_export"])
+        all_imports = (
+            get_the_imports_string(imports, is_async=self.waldiez.is_async),
+            ImportPosition.LOCAL,
+        )
+        before_chats_export = chats_output["before_export"] or []
+        content_before_chats = [
+            x[0]
+            for x in before_chats_export
+            if isinstance(x[1], ExportPosition)
+            and x[1].position == ExportPositions.CHATS
+        ]
+        before_chats = "\n".join(content_before_chats)
+        content = self.merge_exports(
+            imports=all_imports,
+            models_output=models_output["content"] or "",
+            skills_output=skills_output["content"] or "",
+            agents_content=agents_output["content"] or "",
+            chats_content=chats_output["content"] or "",
+            before_chats=before_chats,
+        )
+        return {
+            "content": content,
+            "imports": [all_imports],
+            "after_export": after_export,
+            "before_export": before_export,
+            "environment_variables": env_vars,
+        }
+
+    def merge_exports(
+        self,
+        imports: Tuple[str, ImportPosition],
+        models_output: str,
+        skills_output: str,
+        agents_content: str,
+        chats_content: str,
+        before_chats: str,
+    ) -> str:
+        """Merge all the export contents.
+
+        Parameters
+        ----------
+        imports : Tuple[str, ImportPosition]
+            The imports.
+        models_output : str
+            The models output.
+        skills_output : str
+            The skills output.
+        agents_content : str
+            The agents content.
+        chats_content : str
+            The chats content.
+        before_chats : str
+        Returns
+        -------
+        str
+            The merged export contents.
+        """
+        content = (
+            get_py_content_start(self.waldiez)
+            if not self.for_notebook
+            else get_ipynb_content_start(self.waldiez, comment=self.comment)
+        )
+        content += self.get_comment("imports", self.for_notebook) + "\n"
+        content += imports[0] + "\n"
+        content += self.get_comment("logging", self.for_notebook) + "\n"
+        content += get_logging_start_string(tabs=0) + "\n\n"
+        if models_output:
+            content += self.get_comment("models", self.for_notebook) + "\n"
+            content += models_output + "\n"
+        if skills_output:
+            content += self.get_comment("skills", self.for_notebook) + "\n"
+            content += skills_output + "\n"
+        if agents_content:
+            content += self.get_comment("agents", self.for_notebook) + "\n"
+            content += agents_content + "\n"
+        if before_chats:
+            content += before_chats + "\n"
+        content += get_sqlite_out() + "\n"
+        content += self.get_comment("run", self.for_notebook) + "\n"
+        if self.for_notebook is False:
+            content += get_def_main(
+                chats_content, is_async=self.waldiez.is_async
+            )
+        else:
+            content += "\n" + chats_content + "\n"
+            content += get_logging_stop_string(tabs=0) + "\n"
+            content += get_sqlite_out_call(tabs=0) + "\n"
+        content = content.replace("\n\n\n\n", "\n\n\n")
+        return content
+
+    @staticmethod
+    def gather_environment_variables(
+        model_env_vars: Optional[List[Tuple[str, str]]],
+        skill_env_vars: Optional[List[Tuple[str, str]]],
+        chat_env_vars: Optional[List[Tuple[str, str]]],
+    ) -> List[Tuple[str, str]]:
+        """
+        Gather all the environment variables.
+
+        Parameters
+        ----------
+        model_env_vars : Optional[List[Tuple[str, str]]]
+            The model environment variables.
+        skill_env_vars : Optional[List[Tuple[str, str]]]
+            The skill environment variables.
+        chat_env_vars : Optional[List[Tuple[str, str]]]
+            The chat environment variables.
+
+        Returns
+        -------
+        List[Tuple[str, str]]
+            The gathered environment variables.
+        """
+        all_env_vars: List[Tuple[str, str]] = []
+        if model_env_vars:
+            all_env_vars.extend(model_env_vars)
+        if skill_env_vars:
+            all_env_vars.extend(skill_env_vars)
+        if chat_env_vars:
+            all_env_vars.extend(chat_env_vars)
+        return all_env_vars
+
+    @staticmethod
+    def gather_exports(
+        model_export: Optional[
+            List[Tuple[str, Union[ExportPosition, AgentPosition]]]
+        ],
+        skill_export: Optional[
+            List[Tuple[str, Union[ExportPosition, AgentPosition]]]
+        ],
+        chat_export: Optional[
+            List[Tuple[str, Union[ExportPosition, AgentPosition]]]
+        ],
+    ) -> List[Tuple[str, Union[ExportPosition, AgentPosition]]]:
+        """Gather all (but agents) the before or after exports.
+
+        Parameters
+        ----------
+        model_export : Optional[List[Tuple[str, Union[ExportPosition, AgentPosition]]]]
+            The model exports.
+        skill_export : Optional[List[Tuple[str, Union[ExportPosition, AgentPosition]]]]
+            The skill exports.
+        chat_export : Optional[List[Tuple[str, Union[ExportPosition, AgentPosition]]]]
+            The chat exports.
+
+        Returns
+        -------
+        List[Tuple[str, Union[ExportPosition, AgentPosition]]]
+            The gathered exports.
+        """
+        all_exports: List[Tuple[str, Union[ExportPosition, AgentPosition]]] = []
+        if model_export:
+            all_exports.extend(model_export)
+        if skill_export:
+            all_exports.extend(skill_export)
+        if chat_export:
+            all_exports.extend(chat_export)
+        return all_exports
+
+    def export_models(self) -> ExporterReturnType:
+        """Export the models.
+
+        Returns
+        -------
+        str
+            The exported models.
+        """
+        exporter = ModelsExporter(
+            flow_name=self.flow_name,
+            agents=self.agents,
+            agent_names=self.agent_names,
+            models=self.models,
+            model_names=self.model_names,
+            for_notebook=self.for_notebook,
+            output_dir=self.output_dir,
+        )
+        return exporter.export()
+
+    def export_skills(self) -> ExporterReturnType:
+        """Export the skills.
+
+        Returns
+        -------
+        str
+            The exported skills.
+        """
+        exporter = SkillsExporter(
+            flow_name=self.flow_name,
+            agents=self.agents,
+            agent_names=self.agent_names,
+            skills=self.skills,
+            skill_names=self.skill_names,
+            output_dir=self.output_dir,
+        )
+        return exporter.export()
+
+    @staticmethod
+    def gather_agent_arguments(
+        before_export: List[Tuple[str, Union[ExportPosition, AgentPosition]]],
+        after_export: List[Tuple[str, Union[ExportPosition, AgentPosition]]],
+    ) -> List[Tuple[str, AgentPosition]]:
+        """Gather the agent arguments.
+
+        Parameters
+        ----------
+        before_export : List[Tuple[str, Union[ExportPosition, AgentPosition]]]
+            The before export.
+        after_export : List[Tuple[str, Union[ExportPosition, AgentPosition]]]
+            The after export.
+
+        Returns
+        -------
+        List[Tuple[str, AgentPosition]]
+            The gathered agent arguments.
+        """
+        exported_with_agent_arg: List[Tuple[str, AgentPosition]] = []
+        for before in before_export:
+            position = before[1]
+            if (
+                isinstance(position, AgentPosition)
+                and position.position == AgentPositions.AS_ARGUMENT
+            ):
+                exported_with_agent_arg.append((before[0], position))
+        for after in after_export:
+            position = after[1]
+            if (
+                isinstance(position, AgentPosition)
+                and position.position == AgentPositions.AS_ARGUMENT
+            ):
+                exported_with_agent_arg.append((after[0], position))
+        return exported_with_agent_arg
+
+    def export_agents(
+        self,
+        before_export: List[Tuple[str, Union[ExportPosition, AgentPosition]]],
+        after_export: List[Tuple[str, Union[ExportPosition, AgentPosition]]],
+    ) -> ExporterReturnType:
+        """Export the agents.
+
+        Parameters
+        ----------
+        before_export : List[Tuple[str, Union[ExportPosition, AgentPosition]]]
+            The before export.
+        after_export : List[Tuple[str, Union[ExportPosition, AgentPosition]]]
+            The after export.
+        Returns
+        -------
+        str
+            The exported agents.
+        """
+        agent_outputs = []
+        for agent in self.agents:
+            exported_with_agent_arg = self.gather_agent_arguments(
+                before_export, after_export
+            )
+            arguments_resolver = partial(
+                self.agent_arguments_resolver,
+                exported_with_agent_arg,
+            )
+            group_chat_members = self.waldiez.get_group_chat_members(agent)
+            exporter = AgentExporter(
+                agent=agent,
+                agent_names=self.agent_names,
+                model_names=self.model_names,
+                skill_names=self.skill_names,
+                chats=(self.chats, self.chat_names),
+                is_async=self.waldiez.is_async,
+                for_notebook=self.for_notebook,
+                output_dir=self.output_dir,
+                group_chat_members=group_chat_members,
+                arguments_resolver=arguments_resolver,
+            )
+            agent_output = exporter.export()
+            agent_content = agent_output["content"] or ""
+            after_agent_export = agent_output["after_export"]
+            if after_agent_export:
+                after_export.extend(after_agent_export)
+            before_agent_export = agent_output["before_export"]
+            if before_agent_export:
+                before_export.extend(before_agent_export)
+            if agent_content:
+                agent_content = add_before_agent_content(
+                    agent_content,
+                    before_export,
+                    agent,
+                )
+                agent_content = add_after_agent_content(
+                    agent_content,
+                    after_export,
+                    agent,
+                )
+            agent_output["content"] = agent_content
+            agent_outputs.append(agent_output)
+        return gather_agent_outputs(
+            before_export=before_export,
+            after_export=after_export,
+            agent_outputs=agent_outputs,
+        )
+
+    @staticmethod
+    def agent_arguments_resolver(
+        additional_exports: List[Tuple[str, AgentPosition]], agent: WaldiezAgent
+    ) -> List[str]:
+        """Resolve the arguments for the agent.
+
+        Parameters
+        ----------
+        additional_exports : List[Tuple[str, AgentPosition]]
+            The additional exports.
+        agent : WaldiezAgent
+            The agent.
+
+        Returns
+        -------
+        List[str]
+            The arguments for the agent.
+        """
+        return [x[0] for x in additional_exports if x[1].agent == agent]
+
+    def export_chats(self) -> ExporterReturnType:
+        """Export the chats.
+
+        Returns
+        -------
+        str
+            The exported chats.
+        """
+        exporter = ChatsExporter(
+            get_swarm_members=self.waldiez.get_swarm_members,
+            all_agents=self.agents,
+            agent_names=self.agent_names,
+            all_chats=self.chats,
+            chat_names=self.chat_names,
+            main_chats=self.waldiez.chats,
+            for_notebook=self.for_notebook,
+            is_async=self.waldiez.is_async,
+        )
+        output = exporter.export()
+        chat_contents = output["content"] or ""
+        after_chats = output["after_export"]
+        if chat_contents and after_chats:
+            chat_contents = add_after_chat_content(
+                chat_contents,
+                after_chats,
+            )
+        output["content"] = chat_contents
+        return output
+
+    def export(self) -> ExporterReturnType:
+        """Export the flow.
+
+        Returns
+        -------
+        SubExporterReturnType
+            The exported flow.
+        """
+        return self.export_flow()

waldiez/exporting/flow/utils/__init__.py

@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Utility functions for exporting waldiez to different formats."""
+
+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 get_ipynb_content_start, get_py_content_start
+from .flow_names import ensure_unique_names
+from .importing_utils import (
+    gather_imports,
+    get_standard_imports,
+    get_the_imports_string,
+)
+from .logging_utils import (
+    get_logging_start_string,
+    get_logging_stop_string,
+    get_sqlite_out,
+    get_sqlite_out_call,
+)
+
+__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_def_main",
+    "get_py_content_start",
+    "get_ipynb_content_start",
+    "get_logging_start_string",
+    "get_logging_stop_string",
+    "get_sqlite_out",
+    "get_sqlite_out_call",
+    "get_standard_imports",
+    "get_the_imports_string",
+]