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

waldiez/exporting/models/exporter.py

@@ -0,0 +1,233 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Models exporter module."""
+
+from pathlib import Path
+from typing import Any, Optional, Union
+
+from waldiez.models import WaldiezAgent, WaldiezModel
+
+from ..core import (
+    AgentPosition,
+    ContentOrder,
+    DefaultSerializer,
+    Exporter,
+    ExporterContext,
+    ExportPosition,
+    ModelExtras,
+    get_comment,
+)
+from .processor import ModelProcessor
+
+
+class ModelsExporter(Exporter[ModelExtras]):
+    """Mdels exporter with structured extras."""
+
+    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 = False,
+        cache_seed: Optional[int] = None,
+        output_dir: Optional[Union[str, Path]] = None,
+        context: Optional[ExporterContext] = None,
+        **kwargs: Any,
+    ):
+        """Initialize the models exporter."""
+        super().__init__(context, **kwargs)
+
+        self.flow_name = flow_name
+        self.agents = agents
+        self.agent_names = agent_names
+        self.models = models
+        self.model_names = model_names
+        self.for_notebook = for_notebook
+        self.cache_seed = cache_seed
+        self.output_dir = Path(output_dir) if output_dir else None
+
+        # Initialize extras with processed model content
+        self._extras = self._create_model_extras()
+
+    @property
+    def extras(self) -> ModelExtras:
+        """Get the model extras."""
+        return self._extras
+
+    def _create_model_extras(self) -> ModelExtras:
+        """Create and populate model extras."""
+        extras = ModelExtras("models")
+
+        # Process models to generate LLM configs
+        model_processor = ModelProcessor(
+            flow_name=self.flow_name,
+            models=self.models,
+            model_names=self.model_names,
+            serializer=(
+                self.context.serializer
+                if self.context.serializer
+                else DefaultSerializer()
+            ),
+            output_dir=self.output_dir,
+        )
+
+        llm_configs_content = model_processor.process()
+
+        # Set LLM config content
+        if llm_configs_content:
+            extras.set_llm_config({"content": llm_configs_content})
+            self.add_content(
+                llm_configs_content,
+                ExportPosition.MODELS,
+                order=ContentOrder.MAIN_CONTENT,
+            )
+        # Add environment variables for API keys
+        for model in self.models:
+            if model.api_key:
+                self.add_env_var(
+                    model.api_key_env_key,
+                    model.api_key,
+                    f"API key for {self.model_names[model.id]} model",
+                )
+
+        return extras
+
+    def generate_main_content(self) -> Optional[str]:
+        """Generate the main models content (LLM configs).
+
+        Returns
+        -------
+        Optional[str]
+            The main content string, or None if no content is available.
+        """
+        # handled in extras._contribute_specific_content(...)
+        # also here for direct access
+        if self.extras.has_specific_content():
+            return self.extras.get_content()
+        return None
+
+    def _add_additional_content(self) -> None:
+        """Add model related additional content."""
+        if self.output_dir is not None:
+            # Add API key loader script if output directory is set
+            api_key_loader_script = self.get_api_key_loader_script()
+            self.add_content(
+                api_key_loader_script,
+                ExportPosition.IMPORTS,
+                order=ContentOrder.LATE_CLEANUP,
+                skip_strip=True,
+            )
+        for agent in self.agents:
+            llm_config_arg = self.get_agent_llm_config_arg(agent)
+            if llm_config_arg:
+                # Position as agent argument
+                self.add_content(
+                    llm_config_arg,
+                    ExportPosition.AGENTS,
+                    order=ContentOrder.MAIN_CONTENT,
+                    agent_id=agent.id,
+                    agent_position=AgentPosition.AS_ARGUMENT,
+                )
+
+    def get_agent_llm_config_arg(self, agent: WaldiezAgent) -> str:
+        """Get LLM config argument for agent.
+
+        Parameters
+        ----------
+        agent : WaldiezAgent
+            The agent for which to get the LLM config argument.
+
+        Returns
+        -------
+        str
+            The LLM config argument string for the agent,
+            or "llm_config=False" if no models are configured.
+        """
+        if not agent.data.model_ids:
+            return "    llm_config=False,  # pyright: ignore\n"
+
+        # Get model configs for this agent
+        model_configs: list[str] = []
+        for model_id in agent.data.model_ids:
+            model_name = self.model_names.get(model_id)
+            if model_name:
+                model_configs.append(f"{model_name}_llm_config")
+        tab = " " * 4
+        if not model_configs:
+            return f"{tab}llm_config=False, # pyright: ignore\n"
+
+        config_list = f",\n{tab}{tab}{tab}".join(model_configs)
+        llm_config = f"""{tab}llm_config=autogen.LLMConfig(
+        config_list=[
+            {config_list},
+        ]"""
+        #
+        # Add cache seed if provided
+        if self.cache_seed is not None:
+            llm_config += f",\n{tab}{tab}cache_seed={self.cache_seed},\n"
+        else:
+            llm_config += f",\n{tab}{tab}cache_seed=None,\n"
+        llm_config += f"{tab}),\n"
+
+        return llm_config
+
+    def get_api_key_loader_script(self) -> str:
+        """Get the api key loader script.
+
+        Returns
+        -------
+        str
+            The api key loader script.
+        """
+        comment = get_comment(
+            "Load model API keys",
+            for_notebook=self.config.for_notebook,
+        )
+        loader_script = f'''{comment}# NOTE:
+# This section assumes that a file named "{self.flow_name}_api_keys"
+# exists in the same directory as this file.
+# This file contains the API keys for the models used in this flow.
+# It should be .gitignored and not shared publicly.
+# If this file is not present, you can either create it manually
+# or change the way API keys are loaded in the flow.
+
+
+def load_api_key_module(flow_name: str) -> ModuleType:
+    """Load the api key module.
+
+    Parameters
+    ----------
+    flow_name : str
+        The flow name.
+
+    Returns
+    -------
+    ModuleType
+        The api keys loading module.
+    """
+    module_name = f"{{flow_name}}_api_keys"
+    if module_name in sys.modules:
+        return importlib.reload(sys.modules[module_name])
+    return importlib.import_module(module_name)
+
+__MODELS_MODULE__ = load_api_key_module("{self.flow_name}")
+
+
+def get_{self.flow_name}_model_api_key(model_name: str) -> str:
+    """Get the model api key.
+    Parameters
+    ----------
+    model_name : str
+        The model name.
+
+    Returns
+    -------
+    str
+        The model api key.
+    """
+    return __MODELS_MODULE__.get_{self.flow_name}_model_api_key(model_name)
+
+'''
+        return loader_script

waldiez/exporting/models/factory.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Factory function for creating a ModelsExporter instance."""
+
+from pathlib import Path
+from typing import Optional, Union
+
+from waldiez.models import WaldiezAgent, WaldiezModel
+
+from ..core import ExporterContext, get_default_exporter_context
+from .exporter import ModelsExporter
+
+
+def create_models_exporter(
+    # Factory function for models exporter creation
+    flow_name: str,
+    agents: list[WaldiezAgent],
+    agent_names: dict[str, str],
+    models: list[WaldiezModel],
+    model_names: dict[str, str],
+    for_notebook: bool = False,
+    cache_seed: Optional[int] = None,
+    output_dir: Optional[Union[str, Path]] = None,
+    context: Optional[ExporterContext] = None,
+) -> ModelsExporter:
+    """Create a models exporter.
+
+    Parameters
+    ----------
+    flow_name : str
+        The name of the flow.
+    agents : list[WaldiezAgent]
+        The agents that use models.
+    agent_names : dict[str, str]
+        Mapping of agent IDs to names.
+    models : list[WaldiezModel]
+        The models to export.
+    model_names : dict[str, str]
+        Mapping of model IDs to names.
+    for_notebook : bool, optional
+        Whether the export is for a notebook, by default False
+    cache_seed : Optional[int], optional
+        The cache seed if any, 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
+
+    Returns
+    -------
+    ModelsExporter
+        The created models exporter.
+    """
+    if context is None:
+        context = get_default_exporter_context()
+    return ModelsExporter(
+        flow_name=flow_name,
+        agents=agents,
+        agent_names=agent_names,
+        models=models,
+        model_names=model_names,
+        for_notebook=for_notebook,
+        cache_seed=cache_seed,
+        output_dir=output_dir,
+        context=context,
+    )

waldiez/exporting/models/processor.py

@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+# pylint: disable=too-few-public-methods,unused-argument
+"""Model exporting utilities for Waldiez."""
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+from waldiez.models import WaldiezModel
+
+from ..core.constants import FILE_HEADER
+from ..core.errors import ExporterContentError
+from ..core.extras.serializer import DefaultSerializer
+from ..core.protocols import Serializer
+
+
+@dataclass
+class ModelProcessingResult:
+    """Result from processing models."""
+
+    llm_configs_content: str = ""
+    api_keys_file: Optional[Path] = None
+    needs_api_key_loader: bool = False
+
+
+class ModelProcessor:
+    """Model processor for generating LLM configs and API key loading."""
+
+    def __init__(
+        self,
+        flow_name: str,
+        models: list[WaldiezModel],
+        model_names: dict[str, str],
+        serializer: Optional[Serializer] = None,
+        output_dir: Optional[Path] = None,
+    ):
+        self.flow_name = flow_name
+        self.models = models
+        self.model_names = model_names
+        self.serializer = serializer or DefaultSerializer()
+        self.output_dir = output_dir
+
+    def process(self) -> str:
+        """Process the flow models.
+
+        Returns
+        -------
+        str
+            The string representation of all models' LLM configs.
+        """
+        content = ""
+        for model in self.models:
+            model_name = self.model_names[model.id]
+            model_config = model.get_llm_config()
+
+            # Remove api_key if present
+            api_key = model_config.pop("api_key", None)
+            model_dict_str = self.serializer.serialize(model_config, tabs=0)
+            # and use the getter function to get it when needed
+            if api_key:  # pragma: no branch
+                extra_arg = (
+                    f'get_{self.flow_name}_model_api_key("{model_name}")'
+                )
+                # remove the \n}, from the end of the dict string
+                model_dict_str = model_dict_str.rstrip("\n},")
+                model_dict_str += f',\n    "api_key": {extra_arg}\n}}'
+            content += (
+                f"\n{model_name}_llm_config: dict[str, Any] = "
+                f"{model_dict_str}\n"
+            )
+
+        # Write API keys file if output directory provided
+        if self.output_dir:
+            self.output_dir = Path(self.output_dir)
+            self.output_dir.mkdir(parents=True, exist_ok=True)
+            self._write_api_keys()
+
+        return content
+
+    def _write_api_keys(self) -> None:
+        """Write API keys file."""
+        flow_name_upper = self.flow_name.upper()
+        api_keys_content = f'''{FILE_HEADER}
+# flake8: noqa: E501
+# pylint: disable=line-too-long
+"""API keys for the {self.flow_name} models."""
+
+import os
+
+__{flow_name_upper}_MODEL_API_KEYS__ = {{'''
+
+        for model in self.models:
+            model_name = self.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_{self.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", "")
+'''
+
+        # Write the file
+        file_name = f"{self.flow_name}_api_keys.py"
+        output_path = (
+            self.output_dir / file_name if self.output_dir else Path(file_name)
+        )
+        try:
+            with open(output_path, "w", encoding="utf-8", newline="\n") as f:
+                f.write(api_keys_content)
+        except Exception as e:  # pragma: no cover
+            raise ExporterContentError(
+                f"Failed to write API keys file: {e}"
+            ) from e

waldiez/exporting/tools/__init__.py

@@ -0,0 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Export tool."""
+
+from .exporter import ToolsExporter
+from .factory import create_tools_exporter
+
+__all__ = [
+    "ToolsExporter",
+    "create_tools_exporter",
+]

waldiez/exporting/tools/exporter.py

@@ -0,0 +1,207 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+
+"""Tools exporter."""
+
+from pathlib import Path
+from typing import Any, Optional, Union
+
+from waldiez.models import WaldiezAgent, WaldiezTool
+
+from ..core import (
+    ContentOrder,
+    Exporter,
+    ExporterContext,
+    ExportPosition,
+    ImportPosition,
+    ImportStatement,
+    ToolExtras,
+    get_comment,
+)
+from .processor import ToolProcessor
+from .registration import ToolRegistrationProcessor
+
+
+class ToolsExporter(Exporter[ToolExtras]):
+    """Rools exporter with structured extras."""
+
+    def __init__(
+        self,
+        flow_name: str,
+        agents: list[WaldiezAgent],
+        agent_names: dict[str, str],
+        tools: list[WaldiezTool],
+        tool_names: dict[str, str],
+        output_dir: Optional[Union[str, Path]] = None,
+        context: Optional[ExporterContext] = None,
+        **kwargs: Any,
+    ):
+        """Initialize the tools exporter.
+
+        Parameters
+        ----------
+        flow_name : str
+            The name of the flow.
+        agents : list[WaldiezAgent]
+            The agents that use tools.
+        agent_names : dict[str, str]
+            Mapping of agent IDs to names.
+        tools : list[WaldiezTool]
+            The tools to export.
+        tool_names : dict[str, str]
+            Mapping of tool IDs to names.
+        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
+            Additional keyword arguments.
+        """
+        super().__init__(context, **kwargs)
+
+        self.flow_name = flow_name
+        self.agents = agents
+        self.agent_names = agent_names
+        self.tools = tools
+        self.tool_names = tool_names
+        self.output_dir = Path(output_dir) if output_dir else None
+
+        # Initialize extras with processed tool content
+        self._extras = self._create_tool_extras()
+
+    @property
+    def extras(self) -> ToolExtras:
+        """Get the tool extras."""
+        return self._extras
+
+    def _create_tool_extras(self) -> ToolExtras:
+        """Create and populate tool extras."""
+        extras = ToolExtras("tools")
+
+        # Process tools to generate content
+        tool_processor = ToolProcessor(
+            flow_name=self.flow_name,
+            tools=self.tools,
+            tool_names=self.tool_names,
+            output_dir=self.output_dir,
+        )
+
+        tool_result = tool_processor.process()
+
+        # Add tool content
+        if tool_result.content:
+            extras.add_function_content(tool_result.content)
+
+        # Add imports
+        for import_stmt in tool_result.builtin_imports:
+            extras.add_import(
+                ImportStatement(import_stmt, ImportPosition.BUILTINS)
+            )
+            # Override import position for builtins
+            self.add_import(import_stmt, ImportPosition.BUILTINS)
+
+        for import_stmt in tool_result.third_party_imports:
+            extras.add_import(
+                ImportStatement(import_stmt, ImportPosition.THIRD_PARTY)
+            )
+
+        # Add environment variables
+        for environment_variable in tool_result.environment_variables:
+            self.add_env_var(
+                name=environment_variable.name,
+                value=environment_variable.value,
+                description=environment_variable.description,
+                required=environment_variable.required,
+            )
+
+        # Process tool registrations
+
+        registration_processor = ToolRegistrationProcessor(
+            agents=self.agents,
+            agent_names=self.agent_names,
+            tools=self.tools,
+            tool_names=self.tool_names,
+        )
+
+        registration_content = registration_processor.process()
+        if registration_content:
+            extras.add_registration_content(registration_content)
+
+        if self.output_dir is not None:
+            # add the tool secrets loader script
+            for tool in self.tools:
+                if tool.secrets:
+                    tool_name = self.tool_names.get(tool.id, tool.id)
+                    tool_secrets_loader_script = (
+                        self.get_tool_secrets_loader_script(tool_name=tool_name)
+                    )
+                    self.add_content(
+                        tool_secrets_loader_script,
+                        ExportPosition.TOOLS,
+                        order=ContentOrder.EARLY_SETUP,
+                        skip_strip=True,
+                    )
+        return extras
+
+    def generate_main_content(self) -> Optional[str]:
+        """Generate the main tools content."""
+        # handled as positioned content
+        return None
+
+    def _add_default_imports(self) -> None:
+        """Add default imports for tools."""
+        # Tools might need register_function import
+        if any(agent.data.tools for agent in self.agents):  # pragma: no branch
+            self.add_import("from autogen import register_function")
+
+        # Add interop import if needed
+        if any(tool.is_interop for tool in self.tools):
+            self.add_import("from autogen.interop import Interoperability")
+
+    def get_tool_secrets_loader_script(self, tool_name: str) -> str:
+        """Get the tool secrets loader script.
+
+        Parameters
+        ----------
+        tool_name : str
+            The name of the tool for which to generate the loader script.
+
+        Returns
+        -------
+        str
+            The tool secrets loader script.
+        """
+        comment = get_comment(
+            "Load tool secrets module if needed",
+            for_notebook=self.config.for_notebook,
+        )
+        loader_script = f'''{comment}# NOTE:
+# This section assumes that a file named "{self.flow_name}_{tool_name}_secrets"
+# exists in the same directory as this file.
+# This file contains the secrets for the tool used in this flow.
+# It should be .gitignored and not shared publicly.
+# If this file is not present, you can either create it manually
+# or change the way secrets are loaded in the flow.
+
+
+def load_tool_secrets_module(flow_name: str, tool_name: str) -> ModuleType:
+    """Load the tool secrets module for the given flow name and tool name.
+
+    Parameters
+    ----------
+    flow_name : str
+        The flow name.
+
+    Returns
+    -------
+    ModuleType
+        The loaded module.
+    """
+    module_name = f"{{flow_name}}_{{tool_name}}_secrets"
+    if module_name in sys.modules:
+        return importlib.reload(sys.modules[module_name])
+    return importlib.import_module(module_name)
+
+load_tool_secrets_module("{self.flow_name}", "{tool_name}")
+'''
+        return loader_script