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

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"
@@ -127,24 +200,39 @@ class ExecutionGenerator:
         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

@@ -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,8 +2,6 @@
 # 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 = [
@@ -14,6 +12,7 @@ BUILTIN_IMPORTS = [
     "import os",
     "import sqlite3",
     "import sys",
+    "import traceback",
     "from dataclasses import asdict",
     "from pprint import pprint",
     "from types import ModuleType",
@@ -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",
 ]
 

waldiez/exporting/flow/utils/logging.py

@@ -126,7 +126,7 @@ def get_sync_sqlite_out() -> str:
     content += "        csv_writer.writeheader()\n"
     content += "        csv_writer.writerows(data)\n"
     content += '    json_file = csv_file.replace(".csv", ".json")\n'
-    content += '    with open(json_file, "w", encoding="utf-8") as file:\n'
+    content += '    with open(json_file, "w", encoding="utf-8", newline="\\n") as file:\n'
     content += "        json.dump(data, file, indent=4, ensure_ascii=False)\n"
     content += "\n"
     return content
@@ -172,7 +172,7 @@ def get_async_sqlite_out() -> str:
     content += "        await csv_writer.writeheader()\n"
     content += "        await csv_writer.writerows(data)\n"
     content += '    json_file = csv_file.replace(".csv", ".json")\n'
-    content += '    async with aiofiles.open(json_file, "w", encoding="utf-8") as file:\n'
+    content += '    async with aiofiles.open(json_file, "w", encoding="utf-8", newline="\\n") as file:\n'
     content += "        await file.write(json.dumps(data, indent=4, ensure_ascii=False))\n"
     content += "\n"
     # fmt: on

waldiez/exporting/models/exporter.py

@@ -3,7 +3,7 @@
 """Models exporter module."""
 
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import Any
 
 from waldiez.models import WaldiezAgent, WaldiezModel
 
@@ -31,9 +31,9 @@ class ModelsExporter(Exporter[ModelExtras]):
         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,
+        cache_seed: int | None = None,
+        output_dir: str | Path | None = None,
+        context: ExporterContext | None = None,
         **kwargs: Any,
     ):
         """Initialize the models exporter."""
@@ -94,7 +94,7 @@ class ModelsExporter(Exporter[ModelExtras]):
 
         return extras
 
-    def generate_main_content(self) -> Optional[str]:
+    def generate_main_content(self) -> str | None:
         """Generate the main models content (LLM configs).
 
         Returns
@@ -146,7 +146,7 @@ class ModelsExporter(Exporter[ModelExtras]):
             or "llm_config=False" if no models are configured.
         """
         if not agent.data.model_ids:
-            return "    llm_config=False,  # pyright: ignore\n"
+            return "    llm_config=False,\n"
 
         # Get model configs for this agent
         model_configs: list[str] = []
@@ -154,9 +154,9 @@ class ModelsExporter(Exporter[ModelExtras]):
             model_name = self.model_names.get(model_id)
             if model_name:
                 model_configs.append(f"{model_name}_llm_config")
-        tab = " " * 4
+        tab: str = " " * 4
         if not model_configs:
-            return f"{tab}llm_config=False, # pyright: ignore\n"
+            return f"{tab}llm_config=False,\n"
 
         config_list = f",\n{tab}{tab}{tab}".join(model_configs)
         llm_config = f"""{tab}llm_config=autogen.LLMConfig(

waldiez/exporting/models/processor.py

@@ -5,7 +5,6 @@
 
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Optional
 
 from waldiez.models import WaldiezModel
 
@@ -20,7 +19,7 @@ class ModelProcessingResult:
     """Result from processing models."""
 
     llm_configs_content: str = ""
-    api_keys_file: Optional[Path] = None
+    api_keys_file: Path | None = None
     needs_api_key_loader: bool = False
 
 
@@ -32,8 +31,8 @@ class ModelProcessor:
         flow_name: str,
         models: list[WaldiezModel],
         model_names: dict[str, str],
-        serializer: Optional[Serializer] = None,
-        output_dir: Optional[Path] = None,
+        serializer: Serializer | None = None,
+        output_dir: Path | None = None,
     ):
         self.flow_name = flow_name
         self.models = models
@@ -93,8 +92,9 @@ __{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"
             api_keys_content += (
-                "\n" + f'    "{model_name}": '
+                f'    "{model_name}": '
                 f'{{"key": "{model.api_key}", "env_key": "{key_env}"}},'
             )
 

waldiez/exporting/tools/exporter.py

@@ -4,7 +4,7 @@
 """Tools exporter."""
 
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any
 
 from waldiez.models import WaldiezAgent, WaldiezTool
 
@@ -148,7 +148,7 @@ class ToolsExporter(Exporter[ToolExtras]):
                     )
         return extras
 
-    def generate_main_content(self) -> Optional[str]:
+    def generate_main_content(self) -> str | None:
         """Generate the main tools content."""
         # handled as positioned content
         return None