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

waldiez/exporting/flow/utils/importing.py

@@ -0,0 +1,373 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Get the standard imports for the flow exporter."""
+
+from typing import Optional
+
+from waldiez.exporting.core import ImportPosition
+
+BUILTIN_IMPORTS = [
+    "import csv",
+    "import importlib",
+    "import json",
+    "import os",
+    "import sqlite3",
+    "import sys",
+    "from dataclasses import asdict",
+    "from pprint import pprint",
+    "from types import ModuleType",
+]
+TYPING_IMPORT_NAMES = [
+    "Annotated",
+    "Any",
+    "Callable",
+    "Dict",
+    "List",
+    "Optional",
+    "Set",
+    "Tuple",
+    "Union",
+]
+# pylint: disable=inconsistent-quotes
+TYPING_IMPORTS = [f"from typing import {', '.join(TYPING_IMPORT_NAMES)}"]
+COMMON_AUTOGEN_IMPORTS = [
+    "from autogen import Agent",
+    "from autogen import Cache",
+    "from autogen import ConversableAgent",
+    "from autogen import ChatResult",
+    "from autogen import GroupChat",
+    "from autogen import runtime_logging",
+]
+
+
+def get_sorted_imports(collected_imports: list[str]) -> list[str]:
+    """Get the sorted imports.
+
+    Parameters
+    ----------
+    collected_imports : list[str]
+        The collected imports.
+
+    Returns
+    -------
+    list[str]
+        The sorted imports.
+    """
+    # Remove duplicates while preserving order by converting to dict
+    unique_imports = list(dict.fromkeys(collected_imports))
+
+    sorted_imports = sorted(
+        [imp for imp in unique_imports if imp.startswith("import ")]
+    ) + sorted([imp for imp in unique_imports if imp.startswith("from ")])
+    return sorted_imports
+
+
+def sort_imports(
+    all_imports: list[tuple[str, ImportPosition]],
+) -> tuple[list[str], list[str], list[str], list[str], bool]:
+    """Sort the imports.
+
+    Parameters
+    ----------
+    all_imports : list[tuple[str, ImportPosition]]
+        All the imports.
+
+    Returns
+    -------
+    tuple[list[str], list[str], list[str], list[str], bool]
+        The sorted imports and a flag if we got `import autogen`.
+    """
+    builtin_imports: list[str] = BUILTIN_IMPORTS.copy() + TYPING_IMPORTS.copy()
+    third_party_imports: list[str] = []
+    local_imports: list[str] = []
+    autogen_imports: list[str] = COMMON_AUTOGEN_IMPORTS.copy()
+    got_import_autogen = False
+
+    for import_string, position in all_imports:
+        if "import autogen" in import_string:
+            got_import_autogen = True
+            continue
+        if import_string.startswith("from autogen"):
+            autogen_imports.append(import_string)
+            continue
+        if position == ImportPosition.BUILTINS:
+            builtin_imports.append(import_string)
+        elif position == ImportPosition.THIRD_PARTY:
+            third_party_imports.append(import_string)
+        elif position == ImportPosition.LOCAL:  # pragma: no branch
+            local_imports.append(import_string)
+
+    autogen_imports = clean_and_group_autogen_imports(autogen_imports)
+    third_party_imports = ensure_np_import(third_party_imports)
+    sorted_builtins = get_sorted_imports(builtin_imports)
+    sorted_third_party = get_sorted_imports(third_party_imports)
+    sorted_locals = get_sorted_imports(local_imports)
+
+    return (
+        sorted_builtins,
+        sorted(autogen_imports),
+        sorted_third_party,
+        sorted_locals,
+        got_import_autogen,
+    )
+
+
+def clean_and_group_autogen_imports(autogen_imports: list[str]) -> list[str]:
+    """Cleanup and group autogen imports.
+
+    Parameters
+    ----------
+    autogen_imports : list[str]
+        List of autogen import statements
+
+    Returns
+    -------
+    list[str]
+        Cleaned and grouped autogen imports
+    """
+    # Group imports by module path
+    import_groups: dict[str, set[str]] = {}
+    direct_imports: set[str] = set()
+
+    for imp in autogen_imports:
+        imp = imp.strip()
+        if not imp:
+            continue
+
+        if imp.startswith("import autogen"):
+            direct_imports.add(imp)
+            continue
+
+        # Parse "from autogen.module import items"
+        if imp.startswith("from autogen"):  # pragma: no branch
+            parts = imp.split(" import ")
+            if len(parts) == 2:  # pragma: no branch
+                module_path = parts[0]  # "from autogen.module"
+                items = parts[1].strip()
+
+                if module_path not in import_groups:
+                    import_groups[module_path] = set()
+
+                # Handle multiple imports in one line
+                for item in items.split(","):
+                    import_groups[module_path].add(item.strip())
+
+    # Build cleaned import list
+    cleaned_imports: list[str] = []
+
+    # Add direct imports first
+    cleaned_imports.extend(sorted(direct_imports))
+
+    # Add grouped imports, sorted by module path
+    for module_path in sorted(import_groups.keys()):
+        sorted_items = sorted(import_groups[module_path])
+        items_str = ", ".join(sorted_items)
+        import_statement = f"{module_path} import {items_str}"
+        cleaned_imports.append(import_statement)
+
+    return cleaned_imports
+
+
+def get_the_imports_string(
+    all_imports: list[tuple[str, ImportPosition]],
+    is_async: bool,
+) -> str:
+    """Get the final imports string.
+
+    Parameters
+    ----------
+    all_imports : list[tuple[str, ImportPosition]]
+        All the imports.
+    is_async : bool
+        If the flow is async.
+
+    Returns
+    -------
+    str
+        The final imports string.
+    """
+    (
+        builtin_imports,
+        autogen_imports,
+        third_party_imports,
+        local_imports,
+        got_import_autogen,
+    ) = sort_imports(all_imports)
+
+    # Get the final imports string.
+    # Making sure that there are two lines
+    # after each import section
+    # (builtin, third party, local)
+    final_string = "\n".join(builtin_imports) + "\n"
+    while not final_string.endswith("\n\n"):
+        final_string += "\n"
+
+    if is_async:
+        final_string += (
+            "\nimport aiofiles"
+            "\nimport aiosqlite"
+            "\nimport anyio"
+            "\nimport nest_asyncio"
+            "\nfrom aiocsv import AsyncDictWriter\n"
+        )
+
+    if got_import_autogen:
+        final_string += "\nimport autogen  # type: ignore\n"
+
+    if autogen_imports:  # pragma: no branch
+        final_string += "\n".join(autogen_imports) + "\n"
+
+    if third_party_imports:  # pragma: no branch
+        final_string += "\n".join(third_party_imports) + "\n"
+
+    while not final_string.endswith("\n\n"):
+        final_string += "\n"
+
+    if local_imports:
+        final_string += "\n".join(local_imports) + "\n"
+
+    while not final_string.endswith("\n\n"):
+        final_string += "\n"
+
+    if is_async:
+        final_string += (
+            "# pylint: disable=broad-exception-caught\n"
+            "try:\n"
+            "    nest_asyncio.apply()  # pyright: ignore\n"
+            "except BaseException:\n"
+            "    pass  # maybe on uvloop?\n"
+        )
+
+    return final_string.replace("\n\n\n", "\n\n")  # avoid too many newlines
+
+
+def ensure_np_import(third_party_imports: list[str]) -> list[str]:
+    """Ensure numpy is imported.
+
+    Parameters
+    ----------
+    third_party_imports : list[str]
+        The third party imports.
+
+    Returns
+    -------
+    list[str]
+        The third party imports with numpy.
+    """
+    if (
+        not third_party_imports
+        or "import numpy as np" not in third_party_imports
+    ):
+        third_party_imports.append("import numpy as np")
+    return third_party_imports
+
+
+def gather_imports(
+    model_imports: Optional[list[tuple[str, ImportPosition]]] = None,
+    tool_imports: Optional[list[tuple[str, ImportPosition]]] = None,
+    chat_imports: Optional[list[tuple[str, ImportPosition]]] = None,
+    agent_imports: Optional[list[tuple[str, ImportPosition]]] = None,
+) -> list[tuple[str, ImportPosition]]:
+    """Gather all the imports.
+
+    Parameters
+    ----------
+    model_imports : tuple[str, ImportPosition]
+        The model imports.
+    tool_imports : tuple[str, ImportPosition]
+        The tool imports.
+    chat_imports : tuple[str, ImportPosition]
+        The chat imports.
+    agent_imports : tuple[str, ImportPosition]
+        The agent imports.
+
+    Returns
+    -------
+    tuple[str, ImportPosition]
+        The gathered imports.
+    """
+    all_imports: list[tuple[str, ImportPosition]] = []
+    for import_statement in BUILTIN_IMPORTS:
+        all_imports.append(
+            (
+                import_statement,
+                ImportPosition.BUILTINS,
+            )
+        )
+    if model_imports:
+        all_imports.extend(model_imports)
+    if tool_imports:
+        all_imports.extend(tool_imports)
+    if chat_imports:
+        all_imports.extend(chat_imports)
+    if agent_imports:
+        all_imports.extend(agent_imports)
+    all_imports = deduplicate_imports(all_imports)
+    one_line_typing_imports: list[tuple[str, ImportPosition]] = []
+    final_imports: list[tuple[str, ImportPosition]] = []
+    for import_tuple in all_imports:
+        if import_tuple[0].startswith("from typing import "):
+            one_line_typing_imports.append(import_tuple)
+        else:
+            final_imports.append(import_tuple)
+    final_typing_imports = get_the_typing_imports(one_line_typing_imports)
+    final_imports.insert(1, (final_typing_imports, ImportPosition.BUILTINS))
+    return final_imports
+
+
+def get_the_typing_imports(
+    one_line_typing_imports: list[tuple[str, ImportPosition]],
+) -> str:
+    """Get the typing imports as a single line.
+
+    Parameters
+    ----------
+    one_line_typing_imports : list[tuple[str, ImportPosition]]
+        The one line typing imports.
+
+    Returns
+    -------
+    str
+        The final typing imports string.
+    """
+    # merge the ones found in the imports
+    # with the default ones
+    typing_imports: set[str] = set(TYPING_IMPORT_NAMES)
+    for import_tuple in one_line_typing_imports:
+        if import_tuple[0].startswith(
+            "from typing import "
+        ):  # pragma: no branch
+            # extract the names from the import statement
+            names = import_tuple[0].split("from typing import ")[1]
+            for name in names.split(","):
+                typing_imports.add(name.strip())
+    # create the final import statement
+    final_typing_imports = (
+        f"from typing import {', '.join(sorted(typing_imports))}"
+    )
+    return final_typing_imports.strip()
+
+
+def deduplicate_imports(
+    imports: list[tuple[str, ImportPosition]],
+) -> list[tuple[str, ImportPosition]]:
+    """Deduplicate imports while preserving order.
+
+    Parameters
+    ----------
+    imports : list[tuple[str, ImportPosition]]
+        The imports to deduplicate.
+
+    Returns
+    -------
+    list[tuple[str, ImportPosition]]
+        The deduplicated imports.
+    """
+    # Remove duplicates while preserving order
+    seen: set[tuple[str, ImportPosition]] = set()
+    deduplicated_imports: list[tuple[str, ImportPosition]] = []
+    for import_tuple in imports:
+        if import_tuple not in seen:
+            seen.add(import_tuple)
+            deduplicated_imports.append(import_tuple)
+    return deduplicated_imports

waldiez/exporting/flow/utils/linting.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+"""Linting comments to include in the generated code."""
+
+import math
+
+
+def split_linter_comment(
+    prefix: str,
+    rules: list[str],
+    max_lines: int = 3,
+) -> str:
+    """Split linter comment.
+
+    Parameters
+    ----------
+    prefix : str
+        The prefix for the comment, e.g., "# pylint: disable=" or "# pyright: ".
+    rules : list[str]
+        The list of linter rules to include in the comment.
+    max_lines : int, optional
+        The maximum number of lines to split the comment into, by default 3.
+
+    Returns
+    -------
+    str
+        The formatted comment string with the rules split into lines.
+    """
+    # Calculate minimum number of rules per line to not exceed max_lines
+    rules_per_line = max(1, math.ceil(len(rules) / max_lines))
+    # pylint: disable=inconsistent-quotes
+    lines = [
+        f"{prefix}{','.join(rules[i : i + rules_per_line])}"
+        for i in range(0, len(rules), rules_per_line)
+    ]
+    return "\n".join(lines) + "\n"
+
+
+PYLINT_RULES = [
+    "line-too-long",
+    "unknown-option-value",
+    "unused-argument",
+    "unused-import",
+    "unused-variable",
+    "invalid-name",
+    "import-error",
+    "import-outside-toplevel",
+    "inconsistent-quotes",
+    "missing-function-docstring",
+    "missing-param-doc",
+    "missing-return-doc",
+    "ungrouped-imports",
+    "unnecessary-lambda-assignment",
+    "too-many-arguments",
+    "too-many-locals",
+    "too-many-try-statements",
+    "broad-exception-caught",
+]
+PYRIGHT_RULES = [
+    "reportUnusedImport",
+    "reportMissingTypeStubs",
+    "reportUnknownArgumentType",
+    "reportUnknownMemberType",
+    "reportUnknownLambdaType",
+    "reportUnnecessaryIsInstance",
+    "reportUnknownVariableType",
+]
+
+MYPY_RULES = [
+    "import-untyped",
+    "no-redef",
+    "unused-ignore",
+    "import-not-found",
+]
+
+FLAKE8_RULES = ["E501"]
+
+
+def get_flake8_ignore_comment(rules: list[str] | None = None) -> str:
+    """Get the flake8 ignore comment string.
+
+    Parameters
+    ----------
+    rules : Optional[list[str]], optional
+        The flake8 rules to ignore, by default None.
+
+    Returns
+    -------
+    str
+        The flake8 ignore comment string.
+
+    Example
+    -------
+    ```python
+    >>> get_flake8_ignore_comment(["E501", "F401"])
+
+    # flake8: noqa: E501, F401
+    ```
+    """
+    if not rules:
+        rules = FLAKE8_RULES
+    prefix = "# flake8: noqa: "
+    output = ", ".join(rules)
+    return prefix + output + "\n"
+
+
+def get_pylint_ignore_comment(rules: list[str] | None = None) -> str:
+    """Get the pylint ignore comment string.
+
+    Parameters
+    ----------
+    rules : Optional[list[str]], optional
+        The pylint rules to ignore, by default None.
+
+    Returns
+    -------
+    str
+        The pylint ignore comment string.
+
+    Example
+    -------
+    ```python
+    >>> get_pylint_ignore_comment(True, ["invalid-name", "line-too-long"])
+
+    # pylint: disable=invalid-name, line-too-long
+    ```
+    """
+    if not rules:
+        rules = PYLINT_RULES
+    prefix = "# pylint: disable="
+    return split_linter_comment(
+        prefix,
+        rules,
+        max_lines=3,
+    )
+
+
+def get_pyright_ignore_comment(rules: list[str] | None = None) -> str:
+    """Get the pyright ignore comment string.
+
+    Parameters
+    ----------
+    rules : Optional[list[str]], optional
+        The pyright rules to ignore, by default None.
+
+    Returns
+    -------
+    str
+        The pyright ignore comment string.
+
+    Example
+    -------
+    ```python
+    >>> get_pyright_ignore_comment(
+    ...     True,
+    ...     ["reportUnusedImport", "reportMissingTypeStubs"]
+    ... )
+
+    # pyright: reportUnusedImport=false, reportMissingTypeStubs=false
+    ```
+    """
+    if not rules:
+        rules = PYRIGHT_RULES
+    prefix = "# pyright: "
+    output = split_linter_comment(
+        prefix,
+        [f"{rule}=false" for rule in rules],
+        max_lines=3,
+    )
+    return output
+
+
+def get_mypy_ignore_comment(rules: list[str] | None = None) -> str:
+    """Get the mypy ignore comment string.
+
+    Parameters
+    ----------
+    rules : Optional[list[str]], optional
+        The mypy rules to ignore, by default None.
+
+    Returns
+    -------
+    str
+        The mypy ignore comment string.
+
+    Example
+    -------
+    ```python
+    >>> get_mypy_ignore_comment(["import-untyped", "no-redef"])
+
+    # mypy: disable-error-code="import-untyped,no-redef"
+    ```
+    """
+    if rules is None:
+        rules = MYPY_RULES
+    if not rules:
+        return "# type: ignore\n"
+    prefix = "# mypy: disable-error-code="
+    content = ", ".join(rules)
+    return prefix + f'"{content}"\n'

waldiez/exporting/flow/utils/{logging_utils.py → logging.py}

@@ -15,16 +15,18 @@ get_sqlite_to_csv_call_string
     Get the string to call the sqlite to csv conversion.
 """
 
+from ...core import get_comment
 
-def get_start_logging(is_async: bool, tabs: int = 0) -> str:
+
+def get_start_logging(is_async: bool, for_notebook: bool) -> str:
     r"""Get the logging start call string.
 
     Parameters
     ----------
     is_async : bool
         Whether to use async mode.
-    tabs : int, optional
-        The number of tabs to use for indentation, by default 0
+    for_notebook : bool
+        Whether the logging is for a notebook or a script.
 
     Returns
     -------
@@ -42,17 +44,26 @@ def get_start_logging(is_async: bool, tabs: int = 0) -> str:
             config={"dbname": "flow.db"},
         )
     """
-    tab = "    " * tabs
+    tab = ""
+    comment = get_comment(
+        "Start logging.",
+        for_notebook=for_notebook,
+    )
     if is_async is False:
         return f'''
+{tab}{comment}
 {tab}def start_logging() -> None:
 {tab}    """Start logging."""
 {tab}    runtime_logging.start(
 {tab}        logger_type="sqlite",
 {tab}        config={{"dbname": "flow.db"}},
 {tab}    )
+{tab}
+{tab}
+{tab}start_logging()
 '''
     return f'''
+{tab}{comment}
 {tab}def start_logging() -> None:
 {tab}    """Start logging."""
 {tab}    # pylint: disable=import-outside-toplevel
@@ -65,6 +76,9 @@ def get_start_logging(is_async: bool, tabs: int = 0) -> str:
 {tab}            "sqlite",
 {tab}            {{"dbname": "flow.db"}},
 {tab}        )
+{tab}
+{tab}
+{tab}start_logging()
 '''
 
 
@@ -146,7 +160,7 @@ def get_sync_sqlite_out() -> str:
     content += '    json_file = csv_file.replace(".csv", ".json")\n'
     content += '    with open(json_file, "w", encoding="utf-8") as file:\n'
     content += "        json.dump(data, file, indent=4, ensure_ascii=False)\n"
-    content += "\n\n"
+    content += "\n"
     return content
 
 
@@ -231,7 +245,7 @@ def get_async_sqlite_out() -> str:
     content += '    json_file = csv_file.replace(".csv", ".json")\n'
     content += '    async with aiofiles.open(json_file, "w", encoding="utf-8") as file:\n'
     content += "        await file.write(json.dumps(data, indent=4, ensure_ascii=False))\n"
-    content += "\n\n"
+    content += "\n"
     return content
 
 
@@ -312,15 +326,15 @@ def get_sqlite_out_call(tabs: int, is_async: bool) -> str:
     return content
 
 
-def get_stop_logging(tabs: int, is_async: bool) -> str:
+def get_stop_logging(is_async: bool, tabs: int = 0) -> str:
     r"""Get the function to stop logging and gather logs.
 
     Parameters
     ----------
-    tabs : int
-        The number of tabs to use for indentation
     is_async : bool
         Whether to use async mode
+    tabs : int, optional
+        The number of tabs to use for indentation, by default 0
 
     Returns
     -------

waldiez/exporting/models/__init__.py

@@ -2,8 +2,10 @@
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 """Models (llm_configs) exporter."""
 
-from .models_exporter import ModelsExporter
+from .exporter import ModelsExporter
+from .factory import create_models_exporter
 
 __all__ = [
     "ModelsExporter",
+    "create_models_exporter",
 ]