waldiez 0.5.10__py3-none-any.whl → 0.6.1__py3-none-any.whl

waldiez/exporting/core/types.py

@@ -5,7 +5,7 @@
 
 from dataclasses import dataclass, field, fields
 from pathlib import Path
-from typing import Any, Optional, TypeVar
+from typing import Any, TypeVar
 
 from .enums import (
     ImportPosition,
@@ -25,7 +25,7 @@ class _NoExtrasType:
         return "<NoExtras>"
 
 
-NoExtras = _NoExtrasType()
+NoExtras = _NoExtrasType()  # pylint: disable=invalid-name
 
 
 # Core Data Structures
@@ -35,7 +35,7 @@ class ImportStatement:
 
     statement: str
     position: ImportPosition = ImportPosition.THIRD_PARTY
-    metadata: Optional[dict[str, Any]] = None
+    metadata: dict[str, Any] | None = None
 
     def __hash__(self) -> int:
         """Hash based on the import statement.
@@ -89,7 +89,7 @@ class EnvironmentVariable:
 
     name: str
     value: str
-    description: Optional[str] = None
+    description: str | None = None
     required: bool = True
 
     def __post_init__(self) -> None:
@@ -151,7 +151,7 @@ class InstanceArgument:
     with_new_line_if_empty: bool = False
     skip_if_empty_string: bool = True
     skip_trailing_comma: bool = False
-    comment: Optional[str] = None
+    comment: str | None = None
 
     def has_content(self) -> bool:
         """Check if the instance argument has content.
@@ -266,13 +266,13 @@ class ExportConfig:
     description: str = (
         "Make AG2 Agents Collaborate: Drag, Drop, and Orchestrate with Waldiez"
     )
-    requirements: list[str] = field(default_factory=list[str])
-    tags: list[str] = field(default_factory=list[str])
+    requirements: list[str] = field(default_factory=list)
+    tags: list[str] = field(default_factory=list)
     output_extension: str = "py"
     is_async: bool = False
-    output_directory: Optional[str | Path] = None
-    uploads_root: Optional[Path] = None
-    cache_seed: Optional[int] = None
+    output_directory: str | Path | None = None
+    uploads_root: Path | None = None
+    cache_seed: int | None = None
 
     @property
     def for_notebook(self) -> bool:
@@ -314,6 +314,7 @@ class ExportConfig:
         if for_notebook:
             output_extension = "ipynb"
         cache_seed = kwargs.pop("cache_seed", None)
+        # noinspection PyUnreachableCode
         if cache_seed is not None and not isinstance(cache_seed, int):
             cache_seed = None
         return cls(

waldiez/exporting/core/utils/llm_config.py

@@ -64,7 +64,7 @@ def _get_agent_llm_config_arg_as_arg(
     tab = " " * tab_leng * tabs if tabs > 0 else ""
     # tab = "    " * tabs if tabs > 0 else ""
     if not agent.data.model_ids:
-        return f"{tab}llm_config=False,  # pyright: ignore" + "\n"
+        return f"{tab}llm_config=False," + "\n"
     content = f"{tab}llm_config=autogen.LLMConfig(" + "\n"
     content += f"{tab}    config_list=["
     got_at_least_one_model = False
@@ -77,7 +77,7 @@ def _get_agent_llm_config_arg_as_arg(
             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,  # pyright: ignore" + "\n"
+        return f"{tab}llm_config=False," + "\n"
     content += "\n" + f"{tab}    ]," + "\n"
     content += f"{tab}    cache_seed={cache_seed}," + "\n"
     if temperature is not None:
@@ -97,7 +97,7 @@ def _get_agent_llm_config_arg_as_dict(
 ) -> str:
     tab = " " * tab_leng * tabs if tabs > 0 else ""
     if not agent.data.model_ids:
-        return f'{tab}"llm_config": False' + "\n"
+        return f'{tab}"llm_config": False,' + "\n"
     content = f'{tab}"llm_config": autogen.LLMConfig(' + "\n"
     content += f"{tab}    config_list=["
     got_at_least_one_model = False
@@ -110,7 +110,7 @@ def _get_agent_llm_config_arg_as_dict(
             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"
+        return f'{tab}"llm_config": False,' + "\n"
     content += "\n" + f"{tab}    ]," + "\n"
     content += f"{tab}    cache_seed={cache_seed}," + "\n"
     if temperature is not None:

waldiez/exporting/core/validation.py

@@ -3,7 +3,6 @@
 """Validation types and results for Waldiez exporting core."""
 
 from dataclasses import dataclass, field
-from typing import Optional
 
 
 # Validation Types
@@ -13,8 +12,8 @@ class ValidationError:
 
     message: str
     severity: str = "error"  # error, warning, info
-    location: Optional[str] = None
-    suggestion: Optional[str] = None
+    location: str | None = None
+    suggestion: str | None = None
 
 
 @dataclass
@@ -23,17 +22,17 @@ class ValidationResult:
 
     is_valid: bool
     errors: list[ValidationError] = field(
-        default_factory=list[ValidationError],
+        default_factory=list,
     )
     warnings: list[ValidationError] = field(
-        default_factory=list[ValidationError],
+        default_factory=list,
     )
 
     def add_error(
         self,
         message: str,
-        location: Optional[str] = None,
-        suggestion: Optional[str] = None,
+        location: str | None = None,
+        suggestion: str | None = None,
     ) -> None:
         """Add a validation error.
 
@@ -41,9 +40,9 @@ class ValidationResult:
         ----------
         message : str
             The error message to add.
-        location : Optional[str], optional
+        location : str | None, optional
             The location in the code where the error occurred, by default None
-        suggestion : Optional[str], optional
+        suggestion : str | None, optional
             A suggestion for fixing the error, by default None
         """
         self.errors.append(
@@ -51,14 +50,14 @@ class ValidationResult:
         )
         self.is_valid = False
 
-    def add_warning(self, message: str, location: Optional[str] = None) -> None:
+    def add_warning(self, message: str, location: str | None = None) -> None:
         """Add a validation warning.
 
         Parameters
         ----------
         message : str
             The warning message to add.
-        location : Optional[str], optional
+        location : str | None, optional
             The location in the code where the warning occurred, by default None
         """
         self.warnings.append(ValidationError(message, "warning", location))

waldiez/exporting/flow/execution_generator.py

@@ -67,6 +67,79 @@ class ExecutionGenerator:
             + "\n"
         )
 
+    @staticmethod
+    def generate_store_results(is_async: bool) -> str:
+        """Generate the part that writes the results to results.json.
+
+        Parameters
+        ----------
+        is_async : bool
+            Whether the flow is async or not.
+
+        Returns
+        -------
+        str
+            The part that generates the code to store the results.
+        """
+        content: str = "async " if is_async else ""
+        tab = "    "
+        content += (
+            "def store_results(result_dicts: list[dict[str, Any]]) -> None:\n"
+        )
+        content += f'{tab}"""Store the results to results.json.\n'
+        content += f"{tab}Parameters\n"
+        content += f"{tab}----------\n"
+        content += f"{tab}result_dicts : list[dict[str, Any]]\n"
+        content += f"{tab}{tab}The list of the results.\n"
+        content += f'{tab}"""\n'
+        if is_async:
+            content += f'{tab}async with aiofiles.open("results.json", "w", encoding="utf-8", newline="\\n") as file:\n'
+            content += f"{tab}{tab}await file.write(json.dumps({{'results': result_dicts}}, indent=4, ensure_ascii=False))\n"
+        else:
+            content += f'{tab}with open("results.json", "w", encoding="utf-8", newline="\\n") as file:\n'
+            content += f"{tab}{tab}file.write(json.dumps({{'results': result_dicts}}, indent=4, ensure_ascii=False))\n"
+        return content
+
+    @staticmethod
+    def generate_store_error(is_async: bool) -> str:
+        """Generate the part that writes an error to error.json.
+
+        Parameters
+        ----------
+        is_async : bool
+            Whether the flow is async or not.
+
+        Returns
+        -------
+        str
+            The content for writing the error to file.
+        """
+        content = "\nasync " if is_async else "\n"
+        content += '''def store_error(exc: BaseException | None = None) -> None:
+    """Store the error in error.json.
+
+    Parameters
+    ----------
+    exc : BaseException | None
+        The exception we got if any.
+    """
+    reason = "Event handler stopped processing" if not exc else traceback.format_exc()
+    try:'''
+        if is_async:
+            content += """
+        async with aiofiles.open("error.json", "w", encoding="utf-8", newline="\\n") as file:
+            await file.write(json.dumps({"error": reason}))"""
+        else:
+            content += """
+        with open("error.json", "w", encoding="utf-8", newline="\\n") as file:
+            file.write(json.dumps({"error": reason}))"""
+        content += """
+    except BaseException: # pylint: disable=broad-exception-caught
+        pass
+"""
+
+        return content
+
     @staticmethod
     def generate_main_function(
         content: str,
@@ -108,12 +181,12 @@ class ExecutionGenerator:
         flow_content += f"{comment}\n"
         if is_async:
             flow_content += "async "
-        on_event_arg = "on_event: Optional[Callable[[BaseEvent], bool]] = None"
+        on_event_arg = "on_event: Callable[[BaseEvent, list[ConversableAgent]], bool] | None = None"
         if is_async:
             on_event_arg = (
-                "on_event: Optional["
-                "Callable[[BaseEvent], Coroutine[None, None, bool]]"
-                "] = None"
+                "on_event: "
+                "Callable[[BaseEvent, list[ConversableAgent]], Coroutine[None, None, bool]]"
+                " | None = None"
             )
         return_type_hint = "list[dict[str, Any]]"
         flow_content += f"def main({on_event_arg}) -> {return_type_hint}:\n"
@@ -125,25 +198,41 @@ class ExecutionGenerator:
         flow_content += "    result_dicts: list[dict[str, Any]] = []\n"
         space = "    "
         if cache_seed is not None:
+            # noinspection SqlDialectInspection
             flow_content += (
-                f"    with Cache.disk(cache_seed={cache_seed}"
-                ") as cache:  # pyright: ignore\n"
+                f"    with Cache.disk(cache_seed={cache_seed}) as cache:\n"
             )
             space = f"{space}    "
         flow_content += f"{content}" + "\n"
         if not skip_logging:
-            if is_async:
-                flow_content += f"{space}await stop_logging()"
-            else:
-                flow_content += f"{space}stop_logging()"
+            flow_content += ExecutionGenerator._get_stop_logging_call(
+                space, is_async
+            )
         flow_content += "\n"
         if after_run:
             flow_content += after_run + "\n"
         if cache_seed is not None:
             space = space[4:]
+        flow_content += ExecutionGenerator._get_store_results_call(
+            space, is_async
+        )
         flow_content += f"{space}return result_dicts\n"
         return flow_content
 
+    @staticmethod
+    def _get_stop_logging_call(space: str, is_async: bool) -> str:
+        """Get stop logging call."""
+        if is_async:
+            return f"{space}await stop_logging()"
+        return f"{space}stop_logging()"
+
+    @staticmethod
+    def _get_store_results_call(space: str, is_async: bool) -> str:
+        """Get store results call."""
+        if is_async:
+            return f"{space}await store_results(result_dicts)\n"
+        return f"{space}store_results(result_dicts)\n"
+
     @staticmethod
     def generate_call_main_function(is_async: bool, for_notebook: bool) -> str:
         """Generate the call_main function for the flow script.

waldiez/exporting/flow/exporter.py

@@ -3,7 +3,7 @@
 """Flow exporter."""
 
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any
 
 from waldiez.models import Waldiez
 
@@ -21,7 +21,7 @@ class FlowExporter(Exporter[FlowExtras]):
         waldiez: Waldiez,
         output_dir: Path | None,
         for_notebook: bool,
-        context: Optional[ExporterContext] = None,
+        context: ExporterContext | None = None,
         **kwargs: Any,
     ) -> None:
         """Initialize the chats exporter.

waldiez/exporting/flow/factory.py

@@ -3,7 +3,7 @@
 """Factory function for creating a FlowExporter instance."""
 
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any
 
 from waldiez.logger import WaldiezLogger
 from waldiez.models import Waldiez
@@ -22,7 +22,7 @@ def create_flow_exporter(
     output_dir: Path | None,
     uploads_root: Path | None,
     for_notebook: bool,
-    context: Optional[ExporterContext] = None,
+    context: ExporterContext | None = None,
     **kwargs: Any,
 ) -> FlowExporter:
     """Create a flow exporter.

waldiez/exporting/flow/file_generator.py

@@ -43,7 +43,7 @@ class FileGenerator(ContentGenerator):
         after_run: str,
         skip_logging: bool,
         **kwargs: Any,
-    ) -> str:  # pyright: ignore
+    ) -> str:
         """Generate content based on provided parameters.
 
         Parameters
@@ -172,7 +172,9 @@ class FileGenerator(ContentGenerator):
         )
         execution_gen = ExecutionGenerator()
         chat_contents = "\n".join(chat.content for chat in chats_content)
-        main = execution_gen.generate_main_function(
+        before_main = execution_gen.generate_store_error(is_async) + "\n\n"
+        before_main += execution_gen.generate_store_results(is_async) + "\n\n"
+        main = before_main + execution_gen.generate_main_function(
             content=chat_contents,
             is_async=is_async,
             for_notebook=for_notebook,

waldiez/exporting/flow/merger.py

@@ -27,7 +27,7 @@ class MergeStatistics:
     total_content_items: int = 0
     total_env_vars: int = 0
     deduplicated_env_vars: int = 0
-    conflicts_found: list[str] = field(default_factory=list[str])
+    conflicts_found: list[str] = field(default_factory=list)
 
 
 class ContentMerger:
@@ -156,10 +156,11 @@ class ContentMerger:
                 elif imp.position.value == existing.position.value:
                     # Same position - check for conflicts
                     if imp.metadata != existing.metadata:
-                        conflicts.append(
+                        conflict = (
                             f"Import '{key}' has conflicting metadata: "
                             f"{existing.metadata} vs {imp.metadata}"
                         )
+                        conflicts.append(conflict)
                     # Keep existing (first wins for same priority)
 
         if conflicts:
@@ -292,10 +293,11 @@ class ContentMerger:
 
                     # Check for value conflicts
                     if existing.value != env_var.value:
-                        conflicts.append(
+                        conflict = (
                             f"Environment variable '{key}' has conflicting "
                             f"values: '{existing.value}' vs '{env_var.value}'"
                         )
+                        conflicts.append(conflict)
                     # Keep first occurrence
                     # (tools/models take precedence over agents)
 

waldiez/exporting/flow/orchestrator.py

@@ -4,7 +4,7 @@
 
 from typing import Any, Callable
 
-from waldiez.models import Waldiez, WaldiezAgent
+from waldiez.models import Waldiez, WaldiezAgent, WaldiezGroupManager
 
 from ..agent import AgentExporter, create_agent_exporter
 from ..chats import ChatsExporter, create_chats_exporter
@@ -167,7 +167,7 @@ class ExportOrchestrator:
                 tags=self.waldiez.tags,
                 for_notebook=self.config.for_notebook,
             ),
-            position=ExportPosition.TOP,  # befoe everything
+            position=ExportPosition.TOP,  # before everything
             order=ContentOrder.EARLY_SETUP,
         )
         merged_result.add_content(
@@ -224,6 +224,12 @@ class ExportOrchestrator:
             position=ExportPosition.IMPORTS,  # imports section
             order=ContentOrder.EARLY_SETUP,  # top position
         )
+        known_agents_string = self._get_the_known_agents_string()
+        merged_result.add_content(
+            known_agents_string,
+            position=ExportPosition.AGENTS,
+            order=ContentOrder.POST_CONTENT.value + 100,
+        )
         return merged_result
 
     def get_after_run_content(self) -> str:
@@ -240,6 +246,70 @@ class ExportOrchestrator:
             tabs=1,
         )
 
+    def _get_the_known_agents_string(self) -> str:
+        """Get the _get_known_agents method."""
+        content: str = """
+
+def _check_for_extra_agents(agent: ConversableAgent) -> list[ConversableAgent]:
+    _extra_agents: list[ConversableAgent] = []
+    _agent_cls_name = agent.__class__.__name__
+    if _agent_cls_name == "CaptainAgent":
+        _assistant_agent = getattr(agent, "assistant", None)
+        if _assistant_agent and _assistant_agent not in _extra_agents:
+            _extra_agents.append(_assistant_agent)
+        _executor_agent = getattr(agent, "executor", None)
+        if _executor_agent and _executor_agent not in _extra_agents:
+            _extra_agents.append(_executor_agent)
+    return _extra_agents
+
+
+def _check_for_group_members(agent: ConversableAgent) -> list[ConversableAgent]:
+    _extra_agents: list[ConversableAgent] = []
+    _group_chat = getattr(agent, "_groupchat", None)
+    if _group_chat:
+        _chat_agents = getattr(_group_chat, "agents", [])
+        if isinstance(_chat_agents, list):
+            for _group_member in _chat_agents:
+                if _group_member not in _extra_agents:
+                    _extra_agents.append(_group_member)
+    _manager = getattr(agent, "_group_manager", None)
+    if _manager:
+        if _manager not in _extra_agents:
+            _extra_agents.append(_manager)
+        for _group_member in _check_for_group_members(_manager):
+            if _group_member not in _extra_agents:
+                _extra_agents.append(_group_member)
+    return _extra_agents
+
+
+def _get_known_agents() -> list[ConversableAgent]:
+    _known_agents: list[ConversableAgent] = []"""
+        # group_manager: WaldiezGroupManager | None = None
+        for agent in self.waldiez.agents:
+            if (
+                isinstance(agent, WaldiezGroupManager)
+                and self.waldiez.is_group_pattern_based
+            ):
+                # not defined in agents
+                # we'll get it from "._group_manager"
+                # if found in the rest of the agents.
+                continue
+            agent_name = self.agent_names.get(agent.id)
+            if agent_name:
+                content += f"""
+    if {agent_name} not in _known_agents:
+        _known_agents.append({agent_name})
+    _known_agents.append({agent_name})
+    for _group_member in _check_for_group_members({agent_name}):
+        if _group_member not in _known_agents:
+            _known_agents.append(_group_member)
+    for _extra_agent in _check_for_extra_agents({agent_name}):
+        if _extra_agent not in _known_agents:
+            _known_agents.append(_extra_agent)
+"""
+        content += "    return _known_agents\n\n"
+        return content
+
     def _get_tools_exporter(self) -> ToolsExporter:
         """Get or create tools exporter."""
         if self._tools_exporter is None:

waldiez/exporting/flow/utils/common.py

@@ -2,7 +2,7 @@
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 # pylint: disable=line-too-long
 # flake8: noqa: E501
-"""Common utils for the final generatio."""
+"""Common utils for the final flow generation."""
 
 from waldiez.models import Waldiez
 
@@ -51,7 +51,7 @@ def generate_header(
             requirements=requirements,
             tags=tags,
         )
-    return _get_ipynb_heeader(
+    return _get_ipynb_header(
         name=name,
         description=description,
         requirements=requirements,
@@ -59,7 +59,7 @@ def generate_header(
     )
 
 
-def _get_ipynb_heeader(
+def _get_ipynb_header(
     name: str,
     description: str,
     requirements: list[str],
@@ -73,7 +73,7 @@ def _get_ipynb_heeader(
     content += f"### Tags: {tags_str}\n\n"
     content += f"####{GENERATED_WITH}\n\n"
     content += "#### Requirements\n\n# %%\n"
-    content += "import sys  # pyright: ignore\n"
+    content += "import sys\n"
     # fmt: off
     content += "# # " + f"!{{sys.executable}} -m pip install -q {requirements_str}" + "\n"
     # fmt: on
@@ -166,8 +166,8 @@ def get_after_run_content(
 {space}# save the tree to json
 {space}# pylint: disable=protected-access
 {space}try:
-{space}{tab}data = {agent_name}._root.to_dict()  # pyright: ignore
-{space}{tab}with open("{agent_name}_reasoning_tree.json", "w", encoding="utf-8") as f:
+{space}{tab}data = {agent_name}._root.to_dict()  # pyright: ignore[reportPrivateUsage]
+{space}{tab}with open("{agent_name}_reasoning_tree.json", "w", encoding="utf-8", newline="\\n") as f:
 {space}{tab}{tab}json.dump(data, f)
 {space}except BaseException:
 {space}{tab}pass

waldiez/exporting/flow/utils/importing.py

@@ -2,17 +2,17 @@
 # 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 asyncio",
     "import csv",
     "import importlib",
     "import json",
     "import os",
     "import sqlite3",
     "import sys",
+    "import traceback",
     "from dataclasses import asdict",
     "from pprint import pprint",
     "from types import ModuleType",
@@ -208,7 +208,6 @@ def get_the_imports_string(
         final_string += "\n"
 
     if is_async:
-        builtin_imports.insert(0, "import asyncio")
         final_string += (
             "\nimport aiofiles"
             "\nimport aiosqlite"
@@ -239,7 +238,7 @@ def get_the_imports_string(
         final_string += (
             "# pylint: disable=broad-exception-caught\n"
             "try:\n"
-            "    nest_asyncio.apply()  # pyright: ignore\n"
+            "    nest_asyncio.apply()\n"
             "except BaseException:\n"
             "    pass  # maybe on uvloop?\n"
         )
@@ -269,10 +268,10 @@ def ensure_np_import(third_party_imports: list[str]) -> list[str]:
 
 
 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,
+    model_imports: list[tuple[str, ImportPosition]] | None = None,
+    tool_imports: list[tuple[str, ImportPosition]] | None = None,
+    chat_imports: list[tuple[str, ImportPosition]] | None = None,
+    agent_imports: list[tuple[str, ImportPosition]] | None = None,
 ) -> list[tuple[str, ImportPosition]]:
     """Gather all the imports.
 

waldiez/exporting/flow/utils/linting.py

@@ -37,32 +37,48 @@ def split_linter_comment(
 
 
 PYLINT_RULES = [
-    "line-too-long",
-    "unknown-option-value",
-    "unused-argument",
-    "unused-import",
-    "unused-variable",
+    "broad-exception-caught",
+    "f-string-without-interpolation",
     "invalid-name",
     "import-error",
     "import-outside-toplevel",
     "inconsistent-quotes",
+    "line-too-long",
     "missing-function-docstring",
     "missing-param-doc",
     "missing-return-doc",
-    "ungrouped-imports",
-    "unnecessary-lambda-assignment",
+    "no-member",
+    "pointless-string-statement",
+    "too-complex",
     "too-many-arguments",
     "too-many-locals",
     "too-many-try-statements",
-    "broad-exception-caught",
+    "ungrouped-imports",
+    "unnecessary-lambda-assignment",
+    "unknown-option-value",
+    "unused-argument",
+    "unused-import",
+    "unused-variable",
 ]
+# including basedpyright specific (not pyright only) rules
 PYRIGHT_RULES = [
-    "reportUnusedImport",
+    "reportArgumentType",
+    "reportAttributeAccessIssue",
+    "reportCallInDefaultInitializer",
+    "reportDeprecated",
+    "reportDuplicateImport",
     "reportMissingTypeStubs",
+    "reportOperatorIssue",
+    "reportOptionalMemberAccess",
+    "reportPossiblyUnboundVariable",
+    "reportUnreachable",
+    "reportUnusedImport",
     "reportUnknownArgumentType",
     "reportUnknownMemberType",
     "reportUnknownLambdaType",
     "reportUnnecessaryIsInstance",
+    "reportUnusedParameter",
+    "reportUnusedVariable",
     "reportUnknownVariableType",
 ]