Models

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 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

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