waldiez 0.1.0__py3-none-any.whl

waldiez/exporting/chats/nested.py

@@ -0,0 +1,264 @@
+"""Nested chats exporting."""
+
+from typing import Dict, List, Optional, Tuple
+
+from waldiez.models import (
+    WaldiezAgent,
+    WaldiezAgentNestedChat,
+    WaldiezAgentNestedChatMessage,
+    WaldiezChat,
+)
+
+from ..utils import get_escaped_string, get_object_string
+
+
+def get_nested_chat_trigger_agent_names(
+    all_chats: List[WaldiezChat],
+    nested_chat: WaldiezAgentNestedChat,
+    agent_names: Dict[str, str],
+) -> str:
+    """Get the trigger agent names for the nested chat.
+
+    Parameters
+    ----------
+    all_chats : List[WaldiezChat]
+        All the chats in the flow.
+    nested_chat : WaldiezAgentNestedChat
+        The nested chat.
+    agent_names : Dict[str, str]
+        A mapping of agent id to agent name.
+
+    Returns
+    -------
+    str
+        The trigger agent names.
+    """
+    trigger_agent_ids: List[str] = []
+    for message in nested_chat.triggered_by:
+        waldiez_chat = next(chat for chat in all_chats if chat.id == message.id)
+        if message.is_reply:
+            trigger_agent_ids.append(waldiez_chat.target)
+        else:
+            trigger_agent_ids.append(waldiez_chat.source)
+    agents = [agent_names[agent_id] for agent_id in trigger_agent_ids]
+    trigger_string = f'{[", ".join(agents)]}'
+    return trigger_string.replace("'", '"')
+
+
+def get_nested_chat_message_string(
+    waldiez_chat: WaldiezChat,
+    message: WaldiezAgentNestedChatMessage,
+    agent: WaldiezAgent,
+    agent_names: Dict[str, str],
+    chat_names: Dict[str, str],
+) -> Tuple[str, Optional[str]]:
+    """Get the nested chat message string.
+
+    Parameters
+    ----------
+    waldiez_chat : WaldiezChat
+        The chat.
+    message : WaldiezAgentNestedChatMessage
+        The message.
+    agent : WaldiezAgent
+        The agent.
+    agent_names : Dict[str, str]
+        A mapping of agent id to agent name.
+    chat_names : Dict[str, str]
+        A mapping of chat id to chat name.
+
+    Returns
+    -------
+    Tuple[str, Optional[str]]
+        The message string and the method name if the message is a method.
+    """
+    sender_name: Optional[str] = None
+    sender_id = waldiez_chat.target if message.is_reply else waldiez_chat.source
+    recipient_id = (
+        waldiez_chat.source if message.is_reply else waldiez_chat.target
+    )
+    if sender_id != agent.id:
+        sender_name = agent_names[sender_id]
+    recipient_name = agent_names[recipient_id]
+    chat_dict = waldiez_chat.get_chat_args()
+    chat_dict["recipient"] = recipient_name
+    if sender_name:
+        chat_dict["sender"] = sender_name
+    message_value, message_source = get_chat_nested_string(
+        chat=waldiez_chat, is_reply=message.is_reply, chat_names=chat_names
+    )
+    chat_dict["message"] = message_value
+    message_dict_str = get_object_string(chat_dict, tabs=1)
+    if message_source:
+        # it's not a string, its the name of the function
+        message_dict_str = message_dict_str.replace(
+            f': "{message_value}"', f": {message_value}"
+        )
+    if sender_name:
+        message_dict_str = message_dict_str.replace(
+            f': "{sender_name}"', f": {sender_name}"
+        )
+    if recipient_name:
+        message_dict_str = message_dict_str.replace(
+            f': "{recipient_name}"', f": {recipient_name}"
+        )
+    return message_dict_str, message_source
+
+
+def get_nested_chat_queue(
+    nested_chat: WaldiezAgentNestedChat,
+    agent: WaldiezAgent,
+    agent_names: Dict[str, str],
+    chat_names: Dict[str, str],
+    all_chats: List[WaldiezChat],
+) -> Tuple[str, List[str]]:
+    """Get the nested chat queue.
+
+    Parameters
+    ----------
+    nested_chat : WaldiezAgentNestedChat
+        The nested chat.
+    agent : WaldiezAgent
+        The agent.
+    agent_names : Dict[str, str]
+        A mapping of agent id to agent name.
+    chat_names : Dict[str, str]
+        A mapping of chat id to chat name.
+    all_chats : List[WaldiezChat]
+        All the chats in the flow.
+
+    Returns
+    -------
+    Tuple[str, List[str]]
+        The nested chat queue and the methods to include.
+    """
+    message_methods_to_include = []
+    chat_messages_str = "[\n"
+    for message in nested_chat.messages:
+        waldiez_chat = next(chat for chat in all_chats if chat.id == message.id)
+        message_str, message_source = get_nested_chat_message_string(
+            waldiez_chat=waldiez_chat,
+            message=message,
+            agent=agent,
+            agent_names=agent_names,
+            chat_names=chat_names,
+        )
+        if message_source:
+            message_methods_to_include.append(message_source)
+        chat_messages_str += f"    {message_str}," + "\n"
+    chat_messages_str += "]"
+    if chat_messages_str == "[\n]":
+        return "", message_methods_to_include
+    return chat_messages_str, message_methods_to_include
+
+
+def get_chat_nested_string(
+    chat: WaldiezChat,
+    is_reply: bool,
+    chat_names: Dict[str, str],
+) -> Tuple[str, Optional[str]]:
+    """Get the nested chat message.
+
+    Parameters
+    ----------
+    chat : WaldiezChat
+        The chat.
+    is_reply : bool
+        Whether to use the nested chat's reply message or not.
+    chat_names : Dict[str, str]
+        A mapping of chat id to chat name
+
+    Returns
+    -------
+    Tuple[str, Optional[str]]
+        If the message is a string, the message content and None.
+        If the message is a method, the method name and the method content.
+        If the message is None, 'None' and None.
+    """
+    message = (
+        chat.data.nested_chat.reply
+        if is_reply
+        else chat.data.nested_chat.message
+    )
+    if not message or message.type == "none" or message.content is None:
+        return "None", None
+    if message.type == "string":
+        return get_escaped_string(message.content), None
+    chat_name = chat_names[chat.id]
+    method_args = "recipient, messages, sender, config"
+    function_name = "nested_chat_reply" if is_reply else "nested_chat_message"
+    new_function_name = f"{function_name}_{chat_name}"
+    function_def = f"\ndef {new_function_name}({method_args}):"
+    attribute_name = "reply_content" if is_reply else "message_content"
+    function_content = getattr(chat.data.nested_chat, attribute_name)
+    return new_function_name, function_def + "\n" + function_content + "\n"
+
+
+def export_nested_chat(
+    agent: WaldiezAgent,
+    all_chats: List[WaldiezChat],
+    chat_names: Dict[str, str],
+    agent_names: Dict[str, str],
+) -> str:
+    """Get the nested chat string.
+
+    Parameters
+    ----------
+    agent : WaldiezAgent
+        The agent.
+    all_chats : List[WaldiezChat]
+        All the chats in the flow.
+    chat_names : Dict[str, str]
+        The chat names.
+    agent_names : Dict[str, str]
+        The agent names.
+
+    Returns
+    -------
+    str
+        The nested chat string.
+    """
+    if not agent.data.nested_chats:
+        return ""
+    content = ""
+    extra_contents = []
+    agent_name = agent_names[agent.id]
+    use_suffix = len(agent.data.nested_chats) > 1
+    for index, entry in enumerate(agent.data.nested_chats):
+        trigger_names = get_nested_chat_trigger_agent_names(
+            all_chats=all_chats, nested_chat=entry, agent_names=agent_names
+        )
+        chat_queue, extra_methods = get_nested_chat_queue(
+            nested_chat=entry,
+            agent=agent,
+            agent_names=agent_names,
+            chat_names=chat_names,
+            all_chats=all_chats,
+        )
+        if not chat_queue:
+            continue
+        extra_contents.extend(extra_methods)
+        var_name = (
+            f"{agent_name}_chat_queue_{index}"
+            if use_suffix
+            else f"{agent_name}_chat_queue"
+        )
+        content += f"\n{var_name} = {chat_queue}" + "\n"
+        content += f"""\n
+{agent_name}.register_nested_chats(
+    trigger={trigger_names},
+    chat_queue={var_name},
+)\n
+"""
+    functions_string = "\n".join(sorted(extra_contents))
+    if functions_string:
+        functions_string = functions_string + "\n"
+    content = f"{functions_string}{content}"
+    return (
+        content.replace('"None"', "None")
+        .replace("'None'", "None")
+        .replace('"False"', "False")
+        .replace("'False'", "False")
+        .replace("'True'", "True")
+        .replace('"True"', "True")
+    )

waldiez/exporting/flow/__init__.py

@@ -0,0 +1,5 @@
+"""Export the entire flow to string."""
+
+from .flow import export_flow
+
+__all__ = ["export_flow"]

waldiez/exporting/flow/def_main.py

@@ -0,0 +1,37 @@
+"""Get the main function (if exporting to python)."""
+
+from ..utils import (
+    get_logging_start_string,
+    get_logging_stop_string,
+    get_sqlite_to_csv_call_string,
+)
+
+
+def get_def_main(waldiez_chats: str) -> str:
+    """Get the main function.
+
+    When exporting to python, waldiez_chats string will be the
+    content of the main function. It contains either a
+    `{sender.initiate_chat(recipient, ...)}` (if there is only one chat)
+    or `initiate_chats([..])`, with the list of chats to initiate.
+
+    Parameters
+    ----------
+    waldiez_chats : str
+        The content of the main function.
+
+    Returns
+    -------
+    str
+        The main function.
+    """
+    content = """def main():
+    # type: () -> Union[ChatResult, List[ChatResult]]
+    \"\"\"Start chatting.\"\"\"
+"""
+    content += get_logging_start_string(1)
+    content += f"    results = {waldiez_chats}" + "\n"
+    content += get_logging_stop_string(1) + "\n"
+    content += get_sqlite_to_csv_call_string(1) + "\n"
+    content += "    return results\n"
+    return content

waldiez/exporting/flow/flow.py

@@ -0,0 +1,185 @@
+"""Export the entire flow to string."""
+
+from pathlib import Path
+from typing import Dict, List, Optional, Set, Tuple
+
+from waldiez.models import (
+    Waldiez,
+    WaldiezAgent,
+    WaldiezChat,
+    WaldiezModel,
+    WaldiezSkill,
+)
+
+from ..agents import export_agent
+from ..chats import export_chats, export_nested_chat
+from ..models import export_models
+from ..skills import export_skills
+from ..utils import (
+    get_comment,
+    get_imports_string,
+    get_logging_start_string,
+    get_logging_stop_string,
+    get_pylint_ignore_comment,
+    get_sqlite_to_csv_call_string,
+    get_sqlite_to_csv_string,
+)
+from .def_main import get_def_main
+
+
+# pylint: disable=too-many-locals
+def export_flow(
+    waldiez: Waldiez,
+    agents: Tuple[List[WaldiezAgent], Dict[str, str]],
+    chats: Tuple[List[WaldiezChat], Dict[str, str]],
+    models: Tuple[List[WaldiezModel], Dict[str, str]],
+    skills: Tuple[List[WaldiezSkill], Dict[str, str]],
+    output_dir: Optional[Path],
+    notebook: bool,
+) -> str:
+    """Export the entire flow to a string.
+
+    It contains the required imports, the model and skill definitions,
+    the agent definitions, the links between agents, models and skills,
+    the agents' nested chats, the chat definitions, and the actual call
+    to start the chat(s).
+
+    Parameters
+    ----------
+    waldiez : Waldiez
+        The Waldiez instance.
+    agents : Tuple[List[WaldiezAgent], Dict[str, str]]
+        The agents and their names.
+    chats : Tuple[List[WaldiezChat], Dict[str, str]]
+        The chats and their names.
+    models : Tuple[List[WaldiezModel], Dict[str, str]]
+        The models and their names.
+    skills : Tuple[List[WaldiezSkill], Dict[str, str]]
+        The skills and their names.
+    output_dir : Optional[Path]
+        The output directory.
+    notebook : bool
+        Whether the export is for a jupyter notebook or a python script.
+
+    Returns
+    -------
+    str
+        The flow string.
+    """
+    all_agents, agent_names = agents
+    all_models, model_names = models
+    all_skills, skill_names = skills
+    all_chats, chat_names = chats
+    agent_strings = ""
+    # we need to add `skipped_agent_strings` after the other agents are defined
+    # for example, a group_manager needs the group members to have been defined
+    skipped_agent_strings = ""
+    nested_chats_strings = ""
+    builtin_imports: Set[str] = {
+        "import csv",
+        "import os",
+        "import sqlite3",
+    }
+    other_imports: Set[str] = {
+        "from autogen import Agent",
+        "from autogen import ConversableAgent",
+        "from autogen import ChatResult",
+        "from autogen import runtime_logging",
+    }
+    skill_imports, _ = export_skills(
+        skills=all_skills,
+        skill_names=skill_names,
+        output_dir=output_dir,
+    )
+    if len(waldiez.chats) > 1:
+        other_imports.add("from autogen import initiate_chats")
+    for agent in all_agents:
+        agent_string, after_agent, agent_imports = export_agent(
+            agent=agent,
+            agent_names=agent_names,
+            model_names=model_names,
+            skill_names=skill_names,
+            all_models=all_models,
+            all_skills=all_skills,
+            group_chat_members=waldiez.flow.get_group_chat_members(agent.id),
+        )
+        other_imports.update(agent_imports)
+        if after_agent:
+            skipped_agent_strings += after_agent
+        if agent.agent_type == "manager":
+            skipped_agent_strings += agent_string
+        else:
+            agent_strings += agent_string
+        agent_nested_chats_string = export_nested_chat(
+            agent=agent,
+            agent_names=agent_names,
+            all_chats=all_chats,
+            chat_names=chat_names,
+        )
+        if agent_nested_chats_string:
+            nested_chats_strings += "\n" + agent_nested_chats_string
+    agent_strings += skipped_agent_strings
+    models_string = export_models(
+        all_models=all_models,
+        model_names=model_names,
+        notebook=notebook,
+    )
+    all_imports_string = get_imports_string(
+        imports=other_imports,
+        builtin_imports=builtin_imports,
+        skill_imports=skill_imports,
+    )
+    return _combine_strings(
+        waldiez=waldiez,
+        imports_string=all_imports_string,
+        agents_string=agent_strings,
+        nested_chats_string=nested_chats_strings,
+        models_string=models_string,
+        agent_names=agent_names,
+        chat_names=chat_names,
+        notebook=notebook,
+    )
+
+
+# pylint: disable=too-many-arguments
+def _combine_strings(
+    waldiez: Waldiez,
+    imports_string: str,
+    agents_string: str,
+    nested_chats_string: str,
+    models_string: str,
+    agent_names: Dict[str, str],
+    chat_names: Dict[str, str],
+    notebook: bool,
+) -> str:
+    content = get_pylint_ignore_comment(notebook)
+    content += imports_string
+    # content += get_comment("logging", notebook) + "\n"
+    # content += get_logging_start_string(tabs=0) + "\n\n"
+    content += models_string
+    content += get_comment("agents", notebook) + "\n"
+    content += agents_string
+    if nested_chats_string:
+        content += get_comment("nested", notebook) + "\n"
+    content += nested_chats_string
+    chats_content, additional_methods = export_chats(
+        main_chats=waldiez.chats,
+        agent_names=agent_names,
+        chat_names=chat_names,
+        tabs=0 if notebook else 1,
+    )
+    if additional_methods:
+        while not content.endswith("\n\n"):  # pragma: no cover
+            content += "\n"
+        content += "\n" + additional_methods + "\n"
+    content += get_sqlite_to_csv_string()
+    content += get_comment("run", notebook) + "\n"
+    if not notebook:
+        content += get_def_main(chats_content)
+    else:
+        content += get_logging_start_string(tabs=0)
+        content += "\n" + chats_content + "\n"
+        content += get_logging_stop_string(tabs=0) + "\n"
+        content += get_sqlite_to_csv_call_string(tabs=0) + "\n"
+    content = content.replace("\n\n\n\n", "\n\n\n")
+    return content

waldiez/exporting/models/__init__.py

@@ -0,0 +1,193 @@
+"""Model/LLM related string generation functions.
+
+Functions
+---------
+export_models
+    Get the string representations of the LLM configs.
+"""
+
+from typing import Dict, List
+
+from waldiez.models import WaldiezModel
+
+from ..utils import get_comment, get_object_string
+
+
+def export_models(
+    all_models: List[WaldiezModel],
+    model_names: Dict[str, str],
+    notebook: bool,
+) -> str:
+    """Get the string representations of the LLM configs.
+
+    Parameters
+    ----------
+    all_models : List[WaldiezModel]
+        The models.
+    model_names : Dict[str, str]
+        A mapping of model ids to model names.
+    notebook : bool
+        Whether to export the string for a jupyter notebook.
+
+    Returns
+    -------
+    str
+        The string representation of the models.
+
+    Example
+    -------
+    ```python
+    >>> from waldiez.models import WaldiezModel, WaldiezModelData
+    >>> model = WaldiezModel(
+    ...     id="wm-1",
+    ...     name="llama3.1"  ,
+    ...     description="A model for llamas :P.",
+    ...     tags=["llama", "llama3.1"],
+    ...     requirements=[],
+    ...     data=WaldiezModelData(
+    ...         base_url="https://example.com/v1",
+    ...         api_key="1234567890",
+    ...         api_type="openai",
+    ...         temperature=0.5,
+    ...         price={
+    ...             "prompt_price_per_1k": 0.0001,
+    ...             "completion_token_price_per_1k": 0.0002,
+    ...         },
+    ...     ),
+    ... )
+    >>> model_names = {"wm-1": "llama3_1"}
+    >>> export_models([model], model_names, True)
+
+    # # Models
+    llama3_1_llm_config = {
+        "config_list": [{
+            "model": "llama3.1",
+            "base_url": "https://example.com/v1",
+            "api_key": "1234567890",
+            "api_type": "openai",
+            "temperature": 0.5,
+            "price": [0.0001, 0.0002],
+        }]
+    }
+    ```
+    """
+    content = get_comment("models", notebook) + "\n"
+    if len(all_models) == 1:
+        only_model = all_models[0]
+        model_name = model_names[only_model.id]
+        llm_config = only_model.get_llm_config()
+        model_dict_str = get_object_string(llm_config, tabs=2)
+        content += f"{model_name}_llm_config = " + "{\n"
+        content += '    "config_list": [\n'
+        content += f"        {model_dict_str}\n"
+        content += "    ]\n"
+        content += "}\n"
+    else:
+        for model in all_models:
+            model_name = model_names[model.id]
+            llm_config = model.get_llm_config()
+            model_dict_str = get_object_string(llm_config, tabs=2)
+            content += f"{model_name}_llm_config = " + "{\n"
+            content += '    "config_list": [\n'
+            content += f"        {model_dict_str}\n"
+            content += "    ]\n"
+            content += "}\n"
+    return content
+
+
+def export_agent_models(
+    agent_model_ids: List[str],
+    all_models: List[WaldiezModel],
+    agent_name: str,
+) -> str:
+    """Get the string representations of the agent's registered models.
+
+    Parameters
+    ----------
+    agent_model_ids : List[str]
+        The model ids registered to the agent.
+    all_models : List[WaldiezModel]
+        All the models in the flow.
+    agent_name : str
+        The name of the agent.
+
+    Returns
+    -------
+    str
+        The agent's llm config string.
+
+    Example
+    -------
+    ```python
+    >>> from waldiez.models import WaldiezModel, WaldiezModelData
+    >>> model1 = WaldiezModel(
+    ...     id="wm-1",
+    ...     name="llama3.1"  ,
+    ...     description="A model for llamas :P.",
+    ...     tags=["llama", "llama3.1"],
+    ...     requirements=[],
+    ...     data=WaldiezModelData(
+    ...         base_url="https://example.com/v1",
+    ...         api_key="1234567890",
+    ...         api_type="openai",
+    ...         temperature=0.5,
+    ...         price={
+    ...             "prompt_price_per_1k": 0.0001,
+    ...             "completion_token_price_per_1k": 0.0002,
+    ...         },
+    ...     ),
+    ... )
+    >>> model2 = WaldiezModel(
+    ...     id="wm-2",
+    ...     name="llama3.2"  ,
+    ...     description="A model for llamas :P.",
+    ...     tags=["llama", "llama3.2"],
+    ...     requirements=[],
+    ...     data=WaldiezModelData(
+    ...         base_url="https://example.com/v1",
+    ...        api_key="1234567890",
+    ...         api_type="openai",
+    ...         temperature=0.5,
+    ...         price={
+    ...             "prompt_price_per_1k": 0.0001,
+    ...             "completion_token_price_per_1k": 0.0002,
+    ...         },
+    ...     ),
+    ... )
+    >>> all_models = [model1, model2]
+    >>> agent_model_ids = ["wm-1", "wm-2"]
+    >>> export_agent_models(agent_model_ids, all_models, "llama_agent")
+
+    llama_agent_llm_config = {
+        "config_list": [
+            {
+                "model": "llama3.1",
+                "base_url": "https://example.com/v1",
+                "api_key": "1234567890",
+                "api_type": "openai",
+                "temperature": 0.5,
+                "price": [0.0001, 0.0002],
+            },
+            {
+                "model": "llama3.2",
+                "base_url": "https://example.com/v1",
+                "api_key": "1234567890",
+                "api_type": "openai",
+                "temperature": 0.5,
+                "price": [0.0001, 0.0002],
+            },
+        ]
+    }
+    ```
+    """
+    content = f"{agent_name}_llm_config = " + "{\n"
+    content += '    "config_list": [\n'
+    for model_id in agent_model_ids:
+        model = next((m for m in all_models if m.id == model_id), None)
+        if model is not None:
+            llm_config = model.get_llm_config()
+            model_dict_str = get_object_string(llm_config, tabs=2)
+            content += f"        {model_dict_str},\n"
+    content += "    ]\n"
+    content += "}\n"
+    return content