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

waldiez/exporting/models/models_exporter.py

@@ -0,0 +1,192 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Export models (llm_configs)."""
+
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple, Union
+
+from waldiez.models import WaldiezAgent, WaldiezModel
+
+from ..base import (
+    AgentPosition,
+    AgentPositions,
+    BaseExporter,
+    ExporterMixin,
+    ExporterReturnType,
+    ExportPosition,
+    ImportPosition,
+)
+from .utils import export_models, get_agent_llm_config_arg
+
+
+class ModelsExporter(BaseExporter, ExporterMixin):
+    """Models exporter."""
+
+    _exported_string: Optional[str]
+
+    def __init__(
+        self,
+        flow_name: str,
+        agents: List[WaldiezAgent],
+        agent_names: Dict[str, str],
+        models: List[WaldiezModel],
+        model_names: Dict[str, str],
+        for_notebook: bool,
+        output_dir: Optional[Union[str, Path]] = None,
+    ) -> None:
+        """Initialize the models exporter.
+
+        Parameters
+        ----------
+        agents : List[WaldiezAgent]
+            The agents.
+        agent_names : Dict[str, str]
+            The agent names.
+        models : List[WaldiezModel]
+            The models.
+        model_names : Dict[str, str]
+            The model names.
+        output_dir : Optional[Union[str, Path]], optional
+            The output directory if any, by default None
+        """
+        self.for_notebook = for_notebook
+        self.flow_name = flow_name
+        self.agents = agents
+        self.agent_names = agent_names
+        self.models = models
+        self.model_names = model_names
+        if output_dir is not None and not isinstance(output_dir, Path):
+            output_dir = Path(output_dir)
+        self.output_dir = output_dir
+        self._exported_string = None
+
+    def get_imports(self) -> Optional[List[Tuple[str, ImportPosition]]]:
+        """Generate the imports string.
+
+        Returns
+        -------
+        Optional[Tuple[str, ImportPosition]]
+            The exported imports and the position of the imports.
+        """
+        if not self.output_dir:
+            return None
+        file_path = self.output_dir / f"{self.flow_name}_api_keys.py"
+        if not file_path.exists():
+            # might be because the models are not exported yet
+            if not self._exported_string:
+                self.generate()
+            # if still not exported, return None
+            if not file_path.exists():  # pragma: no cover
+                return None
+        import_string = f"from {self.flow_name}_api_keys import (" + "\n"
+        import_string += f"    get_{self.flow_name}_model_api_key," + "\n"
+        import_string += ")\n"
+        return [(import_string, ImportPosition.LOCAL)]
+
+    def get_after_export(
+        self,
+    ) -> Optional[List[Tuple[str, Union[ExportPosition, AgentPosition]]]]:
+        # fmt: off
+        """Generate the after export strings.
+
+        The arguments for the agent's initialization.
+        example generated args:
+        >>> agent1 = ConversableAgent(
+        >>>     ...
+        >>>     llm_config=False,
+        >>>     ...
+        >>> )
+
+        >>> agent2 = ConversableAgent(
+        >>>     ...
+        >>>     llm_config={
+        >>>         "config_list": [
+        >>>             model1_llm_config,
+        >>>             model2_llm_config,
+        >>>         ],
+        >>>     },
+        >>>     ...
+        >>> )
+
+        where `model1_llm_config` and `model2_llm_config`
+        are the exported models using `self.generate()`
+
+        Returns
+        -------
+        Optional[List[Tuple[str, Union[ExportPosition, AgentPosition]]]]
+            The exported after export strings and their positions.
+        """
+        # fmt: on
+        agent_llm_config_args: List[
+            Tuple[str, Union[ExportPosition, AgentPosition]]
+        ] = []
+        for agent in self.agents:
+            agent_llm_config_args.append(
+                (
+                    get_agent_llm_config_arg(
+                        agent,
+                        all_models=self.models,
+                        model_names=self.model_names,
+                    ),
+                    AgentPosition(
+                        agent=agent, position=AgentPositions.AS_ARGUMENT
+                    ),
+                )
+            )
+        return agent_llm_config_args
+
+    def generate(self) -> str:
+        """Export the models.
+
+        Returns
+        -------
+        str
+            The exported models.
+        """
+        if not self._exported_string:  # pragma: no cover
+            self._exported_string = export_models(
+                flow_name=self.flow_name,
+                all_models=self.models,
+                model_names=self.model_names,
+                output_dir=self.output_dir,
+                serializer=self.serializer,
+            )
+        return self._exported_string
+
+    def get_environment_variables(self) -> Optional[List[Tuple[str, str]]]:
+        """Get the environment variables to set.
+
+        Returns
+        -------
+        Optional[List[Tuple[str, str]]
+            The environment variables to set.
+        """
+        env_vars = []
+        for model in self.models:
+            if model.api_key:
+                env_vars.append((model.api_key_env_key, model.api_key))
+        return env_vars
+
+    def export(self) -> ExporterReturnType:
+        """Export the models.
+
+        Returns
+        -------
+        ExporterReturnType
+            The exported models,
+            the imports,
+            the before export strings,
+            the after export strings,
+            and the environment variables.
+        """
+        exported_string = self.generate()
+        imports = self.get_imports()
+        after_export = self.get_after_export()
+        result: ExporterReturnType = {
+            "content": exported_string,
+            "imports": imports,
+            "before_export": None,
+            "after_export": after_export,
+            "environment_variables": None,
+        }
+        return result

waldiez/exporting/models/utils.py

@@ -0,0 +1,166 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Model/LLM related string generation functions.
+
+Functions
+---------
+export_models
+    Get the string representations of the LLM configs.
+"""
+
+from pathlib import Path
+from typing import Callable, Dict, List, Optional
+
+from waldiez.models import WaldiezAgent, WaldiezModel
+
+
+def export_models(
+    flow_name: str,
+    all_models: List[WaldiezModel],
+    model_names: Dict[str, str],
+    serializer: Callable[..., str],
+    output_dir: Optional[Path] = None,
+) -> str:
+    """Get the string representations of all the models in the flow.
+
+    Parameters
+    ----------
+    flow_name : str
+        The name of the flow.
+    all_models : List[WaldiezModel]
+        All the models in the flow.
+    model_names : Dict[str, str]
+        A mapping of model ids to model names.
+    serializer : Callable[..., str]
+        The serializer function.
+    output_dir : Optional[Path]
+        The output directory to write the api keys.
+
+    Returns
+    -------
+    str
+        The models' llm config string.
+    """
+    content = ""
+    for model in all_models:
+        model_name = model_names[model.id]
+        model_config = model.get_llm_config()
+        model_config["api_key"] = (
+            f'get_{flow_name}_model_api_key("{model_name}")'
+        )
+        model_dict_str = serializer(model_config, tabs=0)
+        model_dict_str = model_dict_str.replace(
+            f'"get_{flow_name}_model_api_key("{model_name}")"',
+            f'get_{flow_name}_model_api_key("{model_name}")',
+        )
+        content += "\n" + f"{model_name}_llm_config = {model_dict_str}" + "\n"
+    if output_dir:
+        write_api_keys(flow_name, all_models, model_names, output_dir)
+    return content
+
+
+def get_agent_llm_config_arg(
+    agent: WaldiezAgent,
+    all_models: List[WaldiezModel],
+    model_names: Dict[str, str],
+    tabs: int = 1,
+) -> str:
+    """Get the string representation of the agent's llm config argument.
+
+    Parameters
+    ----------
+    agent : WaldiezAgent
+        The agent.
+    all_models : List[WaldiezModel]
+        All the models in the flow.
+    model_names : Dict[str, str]
+        A mapping of model ids to model names.
+    tabs : int, optional
+        The number of tabs for indentation, by default 1.
+
+    Returns
+    -------
+    str
+        The agent's llm config argument to use.
+    """
+    tab = "    " * tabs if tabs > 0 else ""
+    if not agent.data.model_ids:
+        return f"{tab}llm_config=False," + "\n"
+    content = f"{tab}llm_config=" + "{\n"
+    content += f'{tab}    "config_list": ['
+    got_at_least_one_model = False
+    for model_id in agent.data.model_ids:
+        model = next((m for m in all_models if m.id == model_id), None)
+        if model is not None:
+            model_name = model_names[model_id]
+            content += "\n" + f"{tab}        {model_name}_llm_config,"
+            got_at_least_one_model = True
+    if not got_at_least_one_model:  # pragma: no cover
+        return f"{tab}llm_config=False," + "\n"
+    content += "\n" + f"{tab}    ]," + "\n"
+    content += tab + "},\n"
+    return content
+
+
+def write_api_keys(
+    flow_name: str,
+    all_models: List[WaldiezModel],
+    model_names: Dict[str, str],
+    output_dir: Path,
+) -> None:
+    """Write the api keys to a separate file.
+
+    Parameters
+    ----------
+    flow_name : str
+        The name of the flow.
+    all_models : List[WaldiezModel]
+        All the models in the flow.
+    model_names : Dict[str, str]
+        A mapping of model ids to model names.
+    output_dir : Path
+        The output directory to write the api keys.
+    """
+    flow_name_upper = flow_name.upper()
+    api_keys_content = f'''
+"""API keys for the {flow_name} models."""
+
+import os
+
+__{flow_name_upper}_MODEL_API_KEYS__ = {{'''
+    for model in all_models:
+        model_name = model_names[model.id]
+        key_env = model.api_key_env_key
+        api_keys_content += (
+            "\n" + f'    "{model_name}": '
+            f'{{"key": "{model.api_key}", "env_key": "{key_env}"}},'
+        )
+    api_keys_content += "\n}\n"
+    api_keys_content += f'''
+
+def get_{flow_name}_model_api_key(model_name: str) -> str:
+    """Get the api key for the model.
+
+    Parameters
+    ----------
+    model_name : str
+        The name of the model.
+
+    Returns
+    -------
+    str
+        The api key for the model.
+    """
+    entry = __{flow_name_upper}_MODEL_API_KEYS__.get(model_name, {{}})
+    if not entry:
+        return ""
+    env_key = entry.get("env_key", "")
+    if env_key:
+        from_env = os.environ.get(env_key, "")
+        if from_env:
+            return from_env
+    return entry.get("key", "")
+'''
+    file_name = f"{flow_name}_api_keys.py"
+    with open(output_dir / file_name, "w", encoding="utf-8", newline="\n") as f:
+        f.write(api_keys_content)

waldiez/exporting/skills/__init__.py

@@ -1,163 +1,9 @@
-"""Skills/tools related string generation functions.
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Export skill."""
 
-Functions
----------
-get_agent_skill_registration
-    Get an agent's skill registration string.
-export_skills
-    Get the skills content and secrets.
-"""
+from .skills_exporter import SkillsExporter
 
-from pathlib import Path
-from typing import Dict, List, Optional, Set, Tuple, Union
-
-from waldiez.models import WaldiezSkill
-
-from ..utils import get_escaped_string
-
-
-def get_agent_skill_registration(
-    caller_name: str,
-    executor_name: str,
-    skill_name: str,
-    skill_description: str,
-) -> str:
-    """Get the agent skill string and secrets.
-
-    Parameters
-    ----------
-    caller_name : str
-        The name of the caller (agent).
-    executor_name : str
-        The name of the executor (agent).
-    skill_name : str
-        The name of the skill.
-    skill_description : str
-        The skill description.
-
-    Returns
-    -------
-    str
-        The agent skill string.
-
-    Example
-    -------
-    ```python
-    >>> get_agent_skill_registration(
-    ...     caller_name="agent1",
-    ...     executor_name="agent2",
-    ...     skill_name="skill1",
-    ...     skill_description="A skill that does something.",
-    ... )
-    register_function(
-        skill1,
-        caller=agent1,
-        executor=agent2,
-        name="skill1",
-        description="A skill that does something.",
-    )
-    ```
-    """
-    skill_description = get_escaped_string(skill_description)
-    content = f"""register_function(
-    {skill_name},
-    caller={caller_name},
-    executor={executor_name},
-    name="{skill_name}",
-    description="{skill_description}",
-)"""
-    return content
-
-
-def _write_skill_secrets(
-    skill: WaldiezSkill,
-    skill_name: str,
-    output_dir: Path,
-) -> None:
-    """Write the skill secrets to a file.
-
-    Parameters
-    ----------
-    skill : WaldiezSkill
-        The skill.
-    skill_name : str
-        The name of the skill.
-    output_dir : Path
-        The output directory to save the secrets to.
-    """
-    if not skill.secrets:
-        return
-    secrets_file = output_dir / f"{skill_name}_secrets.py"
-    with secrets_file.open("w", encoding="utf-8", newline="\n") as f:
-        f.write('"""Secrets for the skill."""\n')
-        f.write("from os import environ\n\n")
-        for key, value in skill.secrets.items():
-            f.write(f'environ["{key}"] = "{value}"\n')
-
-
-def export_skills(
-    skills: List[WaldiezSkill],
-    skill_names: Dict[str, str],
-    output_dir: Optional[Union[str, Path]] = None,
-) -> Tuple[Set[str], Set[Tuple[str, str]]]:
-    """Get the skills' contents and secrets.
-
-    If `output_dir` is provided, the contents are saved to that directory.
-
-    Parameters
-    ----------
-    skills : List[WaldiezSkill]
-        The skills.
-    skill_names : Dict[str, str]
-        The skill names.
-    output_dir : Optional[Union[str, Path]]
-        The output directory to save the skills to.
-
-    Returns
-    -------
-    Tuple[Set[str], Set[Tuple[str, str]]]
-        - The skill imports to use in the main file.
-        - The skill secrets to set as environment variables.
-
-    Example
-    -------
-    ```python
-    >>> from waldiez.models import WaldiezSkill, WaldiezSkillData
-    >>> skill1 = WaldiezSkill(
-    ...     id="ws-1",
-    ...     name="skill1",
-    ...     description="A skill that does something.",
-    ...     tags=["skill", "skill1"],
-    ...     requirements=[],
-    ...     data=WaldiezSkillData(
-    ...         content="def skill1():\\n    pass",
-    ...         secrets={"API_KEY": "1234567890"},
-    ... )
-    >>> skill_names = {"ws-1": "skill1"}
-    >>> export_skills([skill1], skill_names, None)
-    ({'from skill1 import skill1'}, {('API_KEY', '1234567890')})
-    ```
-    """
-    skill_imports: Set[str] = set()
-    skill_secrets: Set[Tuple[str, str]] = set()
-    for skill in skills:
-        skill_name = skill_names[skill.id]
-        skill_secrets.update(skill.secrets.items())
-        if not output_dir:
-            skill_imports.add(f"from {skill_name} import {skill_name}")
-            continue
-        if not isinstance(output_dir, Path):
-            output_dir = Path(output_dir)
-        if not skill.secrets:
-            skill_imports.add(f"from {skill_name} import {skill_name}")
-        else:
-            # have the secrets before the skill
-            skill_imports.add(
-                f"import {skill_name}_secrets  # noqa\n"
-                f"from {skill_name} import {skill_name}"
-            )
-        _write_skill_secrets(skill, skill_name, output_dir)
-        skill_file = output_dir / f"{skill_name}.py"
-        with skill_file.open("w", encoding="utf-8", newline="\n") as f:
-            f.write(skill.content)
-    return skill_imports, skill_secrets
+__all__ = [
+    "SkillsExporter",
+]

waldiez/exporting/skills/skills_exporter.py

@@ -0,0 +1,169 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Skills/tools related string generation functions.
+
+Functions
+---------
+get_agent_skill_registration
+    Get an agent's skill registration string.
+export_skills
+    Get the skills content and secrets.
+"""
+
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple, Union
+
+from waldiez.models import WaldiezAgent, WaldiezSkill
+
+from ..base import (
+    AgentPosition,
+    AgentPositions,
+    BaseExporter,
+    ExporterMixin,
+    ExporterReturnType,
+    ExportPosition,
+    ImportPosition,
+)
+from .utils import export_skills, get_agent_skill_registrations
+
+
+class SkillsExporter(BaseExporter, ExporterMixin):
+    """Skill exporter."""
+
+    def __init__(
+        self,
+        flow_name: str,
+        agents: List[WaldiezAgent],
+        agent_names: Dict[str, str],
+        skills: List[WaldiezSkill],
+        skill_names: Dict[str, str],
+        output_dir: Optional[Union[str, Path]] = None,
+    ) -> None:
+        """Initialize the skill exporter.
+
+        Parameters
+        ----------
+        flow_name : str
+            The name of the flow.
+        agents : List[WaldiezAgent]
+            The agents.
+        agent_names : Dict[str, str]
+            The agent names.
+        skills : List[WaldiezSkill]
+            The skills.
+        skill_names : Dict[str, str]
+            The skill names.
+        output_dir : Optional[Union[str, Path]], optional
+            The output directory if any, by default None
+        """
+        self.flow_name = flow_name
+        self.agents = agents
+        self.agent_names = agent_names
+        self.skills = skills
+        self.skill_names = skill_names
+        self.output_dir = output_dir
+        self.skill_imports, self.skill_secrets, self.skills_contents = (
+            export_skills(
+                flow_name=flow_name,
+                skills=skills,
+                skill_names=skill_names,
+                output_dir=output_dir,
+            )
+        )
+
+    def get_environment_variables(self) -> List[Tuple[str, str]]:
+        """Get the environment variables to set.
+
+        Returns
+        -------
+        List[Tuple[str, str]]
+            The environment variables to set.
+        """
+        return self.skill_secrets
+
+    def get_imports(self) -> List[Tuple[str, ImportPosition]]:
+        """ "Generate the imports string.
+
+        Returns
+        -------
+        Tuple[str, int]
+            The exported imports and the position of the imports.
+        """
+        if not self.skill_imports:
+            return []
+        imports: List[Tuple[str, ImportPosition]] = []
+        for skill_import in self.skill_imports:
+            if (skill_import, ImportPosition.LOCAL) not in imports:
+                imports.append((skill_import, ImportPosition.LOCAL))
+        return imports
+
+    def get_before_export(
+        self,
+    ) -> Optional[List[Tuple[str, Union[ExportPosition, AgentPosition]]]]:
+        """Generate the content before the main export.
+
+        Returns
+        -------
+        Optional[List[Tuple[str, Union[ExportPosition, AgentPosition]]]]
+            The exported content before the main export and its position.
+        """
+
+    def generate(self) -> Optional[str]:
+        """Generate the main export.
+
+        Returns
+        -------
+        Optional[str]
+            The exported content.
+        """
+        return self.skills_contents
+
+    def get_after_export(
+        self,
+    ) -> Optional[List[Tuple[str, Union[ExportPosition, AgentPosition]]]]:
+        """Generate the content after the main export.
+
+        Returns
+        -------
+        Optional[List[Tuple[str, Union[ExportPosition, AgentPosition]]]]
+            The exported content after the main export and its position.
+        """
+        agent_registrations: List[
+            Tuple[str, Union[ExportPosition, AgentPosition]]
+        ] = []
+        for agent in self.agents:
+            agent_registration = get_agent_skill_registrations(
+                agent=agent,
+                agent_names=self.agent_names,
+                all_skills=self.skills,
+                skill_names=self.skill_names,
+                string_escape=self.string_escape,
+            )
+            if agent_registration:
+                # after all agents since we use the executor
+                # (it might not yet be defined)
+                position = AgentPosition(None, AgentPositions.AFTER_ALL, 1)
+                agent_registrations.append((agent_registration, position))
+        return agent_registrations
+
+    def export(self) -> ExporterReturnType:
+        """Export the skills.
+
+        Returns
+        -------
+        ExporterReturnType
+            The exported skills content, the imports,
+            the before export strings, the after export strings,
+            and the environment variables.
+        """
+        imports = self.get_imports()
+        after_export = self.get_after_export()
+        environment_variables = self.get_environment_variables()
+        result: ExporterReturnType = {
+            "content": self.generate(),
+            "imports": imports,
+            "before_export": None,
+            "after_export": after_export,
+            "environment_variables": environment_variables,
+        }
+        return result