waldiez 0.5.6__py3-none-any.whl → 0.5.8__py3-none-any.whl

waldiez/exporting/flow/file_generator.py

@@ -17,6 +17,7 @@ from .execution_generator import ExecutionGenerator
 from .utils.common import generate_header
 
 
+# noinspection PyProtocol
 class FileGenerator(ContentGenerator):
     """Generate the complete flow notebook content."""
 
@@ -34,7 +35,7 @@ class FileGenerator(ContentGenerator):
         self.context = context
         self.config = context.get_config()
 
-    # pylint: disable=too-many-locals
+    # pylint: disable=too-many-locals,unused-argument
     def generate(
         self,
         merged_result: ExportResult,
@@ -42,26 +43,26 @@ class FileGenerator(ContentGenerator):
         after_run: str,
         skip_logging: bool,
         **kwargs: Any,
-    ) -> str:
-        """Generate the complete flow notebook content.
+    ) -> str:  # pyright: ignore
+        """Generate content based on provided parameters.
 
         Parameters
         ----------
         merged_result : ExportResult
             The merged export result containing all content.
         is_async : bool
-            Whether to generate async conten
+            Whether to generate async content.
         after_run : str
             Additional content to add after the main flow execution.
         skip_logging : bool
             Whether to skip logging setup.
         **kwargs : Any
-            Additional keyword arguments for the generator.
+            Parameters to influence content generation.
 
         Returns
         -------
         str
-            The complete flow notebook content.
+            The generated content.
 
         Raises
         ------
@@ -214,7 +215,7 @@ class FileGenerator(ContentGenerator):
                 tags=self.config.tags,
                 for_notebook=self.config.for_notebook,
             )
-        header_string = "\n".join(content.content for content in from_result)
+        header_string = "\n".join(item.content for item in from_result)
         while not header_string.endswith("\n\n"):
             header_string += "\n"
         return header_string

waldiez/exporting/flow/merger.py

@@ -4,7 +4,6 @@
 """Content merger for combining multiple export results."""
 
 from dataclasses import dataclass, field
-from typing import Optional, Union
 
 from ..core import (
     ContentOrder,
@@ -34,12 +33,12 @@ class MergeStatistics:
 class ContentMerger:
     """Intelligently merges multiple ExportResult objects."""
 
-    def __init__(self, context: Optional[ExporterContext] = None):
+    def __init__(self, context: ExporterContext | None = None):
         """Initialize the content merger.
 
         Parameters
         ----------
-        context : Optional[ExporterContext], optional
+        context : ExporterContext | None, optional
             The exporter context, by default None
         """
         self.context = context or ExporterContext()
@@ -245,12 +244,13 @@ class ContentMerger:
         )
 
     # pylint: disable=no-self-use
-    def _get_order_value(self, order: Union[int, ContentOrder]) -> int:
+    # noinspection PyMethodMayBeStatic
+    def _get_order_value(self, order: int | ContentOrder) -> int:
         """Get numeric value from order for sorting.
 
         Parameters
         ----------
-        order : Union[int, ContentOrder]
+        order : int | ContentOrder
             The order value
 
         Returns
@@ -352,7 +352,8 @@ class ContentMerger:
             warnings=all_warnings,
         )
 
-    def _merge_main_content(self, results: list[ExportResult]) -> Optional[str]:
+    # noinspection PyMethodMayBeStatic
+    def _merge_main_content(self, results: list[ExportResult]) -> str | None:
         """Merge main content from results.
 
         For flow exports, main content is typically empty since all content
@@ -365,7 +366,7 @@ class ContentMerger:
 
         Returns
         -------
-        Optional[str]
+        str | None
             The merged main content, or None
         """
         main_contents: list[str] = []

waldiez/exporting/flow/orchestrator.py

@@ -22,6 +22,7 @@ from .merger import ContentMerger
 from .utils import (
     generate_header,
     get_after_run_content,
+    get_common_env_var_setup,
     get_np_no_nep50_handle,
     get_sqlite_out,
     get_start_logging,
@@ -31,6 +32,7 @@ from .utils import (
 
 
 # pylint: disable=no-self-use,too-many-instance-attributes
+# noinspection PyMethodMayBeStatic
 class ExportOrchestrator:
     """Coordinates the export process."""
 
@@ -168,10 +170,15 @@ class ExportOrchestrator:
             position=ExportPosition.TOP,  # befoe everything
             order=ContentOrder.EARLY_SETUP,
         )
+        merged_result.add_content(
+            get_common_env_var_setup(),
+            position=ExportPosition.IMPORTS,  # after imports (need os)
+            order=ContentOrder.CLEANUP,
+        )
         merged_result.add_content(
             get_np_no_nep50_handle(),
             position=ExportPosition.IMPORTS,  # after imports (need np)
-            order=ContentOrder.CLEANUP,
+            order=ContentOrder.CLEANUP.value + 1,
         )
         if not self.should_skip_logging():
             merged_result.add_content(
@@ -180,7 +187,7 @@ class ExportOrchestrator:
                     for_notebook=self.config.for_notebook,
                 ),
                 position=ExportPosition.IMPORTS,  # after imports, before models
-                order=ContentOrder.CLEANUP.value + 1,  # after imports
+                order=ContentOrder.CLEANUP.value + 2,  # after imports and np
                 skip_strip=True,  # keep newlines
             )
         # merged_result.add_content
@@ -266,6 +273,7 @@ class ExportOrchestrator:
     def _get_chats_exporter(self) -> ChatsExporter:
         """Get or create chats exporter."""
         if self._chats_exporter is None:
+            # noinspection PyTypeChecker
             self._chats_exporter = create_chats_exporter(
                 waldiez=self.waldiez,
                 all_agents=self.agents,
@@ -325,6 +333,7 @@ class ExportOrchestrator:
         agent_arguments: dict[str, list[str]],
     ) -> AgentExporter:
         """Create an exporter for a specific agent."""
+        # noinspection PyTypeChecker
         return create_agent_exporter(
             agent=agent,
             agent_names=self.agent_names,
@@ -343,6 +352,7 @@ class ExportOrchestrator:
             ),
         )
 
+    # noinspection PyMethodMayBeStatic
     def _extract_agent_arguments_from_result(
         self, result: ExportResult
     ) -> dict[str, list[str]]:
@@ -416,14 +426,18 @@ class ExportOrchestrator:
     def should_skip_logging(self) -> bool:
         """Determine if logging should be skipped.
 
+        In case there is any compatibility issue with logging
+        we can skip logging entirely (until handled/fixed).
+
+        e.g:
+        return self.waldiez.has_doc_agents
+        or:
+        if self.waldiez.flow.is_group_chat:
+            return any(tool.is_predefined for tool in self.waldiez.tools)
+
         Returns
         -------
         bool
             True if logging should be skipped, False otherwise.
         """
-        return self.waldiez.has_doc_agents
-        # if not self.waldiez.tools:
-        #     return False
-        # if self.waldiez.flow.is_group_chat:
-        #     return any(tool.is_predefined for tool in self.waldiez.tools)
-        # return False
+        return False

waldiez/exporting/flow/utils/__init__.py

@@ -5,6 +5,7 @@
 from .common import (
     generate_header,
     get_after_run_content,
+    get_common_env_var_setup,
     get_np_no_nep50_handle,
 )
 from .importing import get_sorted_imports, get_the_imports_string
@@ -12,6 +13,7 @@ from .logging import get_sqlite_out, get_start_logging, get_stop_logging
 
 __all__ = [
     "generate_header",
+    "get_common_env_var_setup",
     "get_np_no_nep50_handle",
     "get_after_run_content",
     "get_the_imports_string",

waldiez/exporting/flow/utils/common.py

@@ -102,29 +102,21 @@ Tags: {", ".join(tags)}
 """'''
 
 
-def main_doc_string(is_async: bool) -> str:
+def main_doc_string() -> str:
     """Generate the docstring for the main function.
 
-    Parameters
-    ----------
-    is_async : bool
-        Whether the main function is asynchronous.
-
     Returns
     -------
     str
         The docstring for the main function.
     """
-    return_type_hint = (
-        "AsyncRunResponseProtocol" if is_async else "RunResponseProtocol"
-    )
+    return_type_hint = "list[dict[str, Any]]"
     return f'''"""Start chatting.
 
     Returns
     -------
     {return_type_hint}
-        The result of the chat session, which can be a single ChatResult,
-        a list of ChatResults, or a dictionary mapping integers to ChatResults.
+        The result of the chat session.
 
     Raises
     ------
@@ -183,6 +175,23 @@ def get_after_run_content(
     return content
 
 
+def get_common_env_var_setup() -> str:
+    """Get common environment variable setup for Waldiez flows.
+
+    Returns
+    -------
+    str
+        The content to set up common environment variables.
+    """
+    content = """
+# Common environment variable setup for Waldiez flows
+load_dotenv(override=True)
+os.environ["AUTOGEN_USE_DOCKER"] = "0"
+os.environ["ANONYMIZED_TELEMETRY"] = "False"
+"""
+    return content
+
+
 def get_np_no_nep50_handle() -> str:
     """Handle the "module numpy has no attribute _no_pep50_warning" error.
 

waldiez/exporting/flow/utils/importing.py

@@ -100,7 +100,8 @@ def sort_imports(
             third_party_imports.append(import_string)
         elif position == ImportPosition.LOCAL:  # pragma: no branch
             local_imports.append(import_string)
-
+    if "from dotenv import load_dotenv" not in third_party_imports:
+        third_party_imports.append("from dotenv import load_dotenv")
     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)

waldiez/exporting/flow/utils/logging.py

@@ -49,7 +49,7 @@ def get_start_logging(is_async: bool, for_notebook: bool) -> str:
         "Start logging.",
         for_notebook=for_notebook,
     )
-    if is_async is False:
+    if not is_async:
         return f'''
 {tab}{comment}
 {tab}def start_logging() -> None:
@@ -83,6 +83,7 @@ def get_start_logging(is_async: bool, for_notebook: bool) -> str:
 
 
 # pylint: disable=differing-param-doc,differing-type-doc
+# noinspection PyUnresolvedReferences
 def get_sync_sqlite_out() -> str:
     r"""Get the sqlite to csv and json conversion code string.
 
@@ -165,6 +166,7 @@ def get_sync_sqlite_out() -> str:
 
 
 # pylint: disable=differing-param-doc,differing-type-doc,line-too-long
+# noinspection PyUnresolvedReferences
 def get_async_sqlite_out() -> str:
     r"""Get the sqlite to csv and json conversion code string.
 
@@ -325,6 +327,7 @@ def get_sqlite_out_call(tabs: int, is_async: bool) -> str:
     return content
 
 
+# noinspection PyUnresolvedReferences
 def get_stop_logging(is_async: bool, tabs: int = 0) -> str:
     r"""Get the function to stop logging and gather logs.
 
@@ -343,7 +346,7 @@ def get_stop_logging(is_async: bool, tabs: int = 0) -> str:
     Example
     -------
     ```python
-    >>> get_logging_stop_string()
+    >>> get_stop_logging()
     def stop_logging() -> None:
         \"\"\"Stop logging.\"\"\"
         runtime_logging.stop()

waldiez/exporting/models/exporter.py

@@ -186,7 +186,8 @@ class ModelsExporter(Exporter[ModelExtras]):
             for_notebook=self.config.for_notebook,
         )
         loader_script = f'''{comment}# NOTE:
-# This section assumes that a file named "{self.flow_name}_api_keys"
+# This section assumes that a file named:
+# "{self.flow_name}_api_keys.py"
 # exists in the same directory as this file.
 # This file contains the API keys for the models used in this flow.
 # It should be .gitignored and not shared publicly.

waldiez/exporting/models/factory.py

@@ -3,7 +3,6 @@
 """Factory function for creating a ModelsExporter instance."""
 
 from pathlib import Path
-from typing import Optional, Union
 
 from waldiez.models import WaldiezAgent, WaldiezModel
 
@@ -19,9 +18,9 @@ def create_models_exporter(
     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,
 ) -> ModelsExporter:
     """Create a models exporter.
 
@@ -39,11 +38,11 @@ def create_models_exporter(
         Mapping of model IDs to names.
     for_notebook : bool, optional
         Whether the export is for a notebook, by default False
-    cache_seed : Optional[int], optional
+    cache_seed : int | None, optional
         The cache seed if any, by default None
-    output_dir : Optional[Union[str, Path]], optional
+    output_dir : str | Path | None, optional
         Output directory for generated files, by default None
-    context : Optional[ExporterContext], optional
+    context : ExporterContext | None, optional
         Exporter context with dependencies, by default None
 
     Returns

waldiez/exporting/tools/exporter.py

@@ -4,7 +4,7 @@
 """Tools exporter."""
 
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 from waldiez.models import WaldiezAgent, WaldiezTool
 
@@ -32,8 +32,8 @@ class ToolsExporter(Exporter[ToolExtras]):
         agent_names: dict[str, str],
         tools: list[WaldiezTool],
         tool_names: dict[str, str],
-        output_dir: Optional[Union[str, Path]] = None,
-        context: Optional[ExporterContext] = None,
+        output_dir: str | Path | None = None,
+        context: ExporterContext | None = None,
         **kwargs: Any,
     ):
         """Initialize the tools exporter.
@@ -50,9 +50,9 @@ class ToolsExporter(Exporter[ToolExtras]):
             The tools to export.
         tool_names : dict[str, str]
             Mapping of tool IDs to names.
-        output_dir : Optional[Union[str, Path]], optional
+        output_dir : str | Path | None
             Output directory for generated files, by default None
-        context : Optional[ExporterContext], optional
+        context : ExporterContext | None
             Exporter context with dependencies, by default None
         **kwargs
             Additional keyword arguments.
@@ -176,7 +176,8 @@ class ToolsExporter(Exporter[ToolExtras]):
             for_notebook=self.config.for_notebook,
         )
         loader_script = f'''{comment}# NOTE:
-# This section assumes that a file named "{self.flow_name}_{tool_name}_secrets"
+# This section assumes that a file named:
+# "{self.flow_name}_{tool_name}_secrets.py"
 # exists in the same directory as this file.
 # This file contains the secrets for the tool used in this flow.
 # It should be .gitignored and not shared publicly.

waldiez/exporting/tools/factory.py

@@ -3,7 +3,6 @@
 """Factory function for creating a ToolsExporter instance."""
 
 from pathlib import Path
-from typing import Optional, Union
 
 from waldiez.models import WaldiezAgent, WaldiezTool
 
@@ -17,8 +16,8 @@ def create_tools_exporter(
     agent_names: dict[str, str],
     tools: list[WaldiezTool],
     tool_names: dict[str, str],
-    output_dir: Optional[Union[str, Path]] = None,
-    context: Optional[ExporterContext] = None,
+    output_dir: str | Path | None = None,
+    context: ExporterContext | None = None,
 ) -> ToolsExporter:
     """Create a tools exporter.
 
@@ -34,9 +33,9 @@ def create_tools_exporter(
         The tools to export.
     tool_names : dict[str, str]
         Mapping of tool IDs to names.
-    output_dir : Optional[Union[str, Path]], optional
+    output_dir : str | Path | None, optional
         Output directory for generated files, by default None
-    context : Optional[ExporterContext], optional
+    context : ExporterContext | None, optional
         Exporter context with dependencies, by default None
 
     Returns

waldiez/exporting/tools/processor.py

@@ -5,7 +5,6 @@
 
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Optional
 
 from waldiez.models import WaldiezTool
 
@@ -34,7 +33,7 @@ class ToolProcessor:
         flow_name: str,
         tools: list[WaldiezTool],
         tool_names: dict[str, str],
-        output_dir: Optional[Path] = None,
+        output_dir: Path | None = None,
     ):
         """Initialize the tool processor.
 
@@ -46,7 +45,7 @@ class ToolProcessor:
             The tools to process.
         tool_names : dict[str, str]
             Mapping of tool IDs to names.
-        output_dir : Optional[Path], optional
+        output_dir : Path | None, optional
             Output directory for generated files, by default None
         """
         self.flow_name = flow_name
@@ -193,7 +192,12 @@ ag2_{tool_name} = ag2_{tool_name}_interop.convert_tool(
                 )
                 f.write("import os\n\n")
                 for key, value in tool.secrets.items():
-                    f.write(f'os.environ["{key}"] = "{value}"\n')
+                    # f.write(f'os.environ["{key}"] = "{value}"\n')
+                    # check first if the key already exists in os.environ
+                    f.write(
+                        f'os.environ["{key}"] = '
+                        f'os.environ.get("{key}", "{value}")\n'
+                    )
         except Exception as exc:  # pragma: no cover
             raise ExporterContentError(
                 f"Failed to write secrets file for tool '{tool_name}': {exc}"
@@ -217,6 +221,7 @@ ag2_{tool_name} = ag2_{tool_name}_interop.convert_tool(
         )
 
     # pylint: disable=no-self-use
+    # noinspection PyMethodMayBeStatic
     def _sort_imports(self, imports: list[str]) -> list[str]:
         """Sort imports: 'import' statements first, then 'from' statements.
 

waldiez/exporting/tools/registration.py

@@ -99,6 +99,7 @@ class ToolRegistrationProcessor:
         return "\n".join(registrations) if registrations else ""
 
     # pylint: disable=no-self-use
+    # noinspection PyMethodMayBeStatic
     def _create_registration(
         self,
         caller_name: str,

waldiez/io/_ws.py

@@ -93,6 +93,7 @@ class WebSocketsAdapter:
             The received message, decoded as a string.
         """
         # pylint: disable=too-many-try-statements
+        # noinspection PyBroadException
         try:
             response = await asyncio.wait_for(
                 self.websocket.recv(), timeout=timeout
@@ -143,6 +144,7 @@ class StarletteAdapter:
         str
             The received message, decoded as a string.
         """
+        # noinspection PyBroadException
         try:
             return await asyncio.wait_for(
                 self.websocket.receive_text(), timeout=timeout

waldiez/io/models/content/audio.py

@@ -11,6 +11,7 @@ from pydantic import BaseModel
 from .base import AudioContent
 
 
+# noinspection PyUnusedLocal
 class AudioMediaContent(BaseModel):
     """Audio media content."""
 

waldiez/io/models/content/file.py

@@ -11,6 +11,7 @@ from pydantic import BaseModel
 from .base import FileContent
 
 
+# noinspection PyUnusedLocal
 class FileMediaContent(BaseModel):
     """File media content."""
 

waldiez/io/models/content/image.py

@@ -12,6 +12,7 @@ from ...utils import get_image
 from .base import ImageContent
 
 
+# noinspection DuplicatedCode
 class ImageMediaContent(BaseModel):
     """Image media content."""
 

waldiez/io/models/content/text.py

@@ -9,6 +9,7 @@ from pydantic import BaseModel
 from typing_extensions import Literal
 
 
+# noinspection PyUnusedLocal
 class TextMediaContent(BaseModel):
     """Text media content."""
 

waldiez/io/models/content/video.py

@@ -11,6 +11,7 @@ from typing_extensions import Literal
 from .base import VideoContent
 
 
+# noinspection PyUnusedLocal
 class VideoMediaContent(BaseModel):
     """Video media content."""
 

waldiez/io/models/user_input.py

@@ -186,6 +186,7 @@ class UserInputData(BaseModel):
 
         return None  # pragma: no cover
 
+    # noinspection PyNestedDecorators
     @field_validator("content", mode="before")
     @classmethod
     def validate_content(cls, v: Any) -> MediaContent:  # noqa: C901,D102

waldiez/io/models/user_response.py

@@ -21,6 +21,7 @@ class UserResponse(StructuredBase):
     type: MessageType = "input_response"
     data: UserInputData | list[UserInputData] | str
 
+    # noinspection PyNestedDecorators
     @field_validator("data", mode="before")
     @classmethod
     def validate_data(

waldiez/io/mqtt.py

@@ -33,6 +33,7 @@ except ImportError as error:  # pragma: no cover
         "MQTT client not installed. Please install paho-mqtt with `pip install paho-mqtt`."
     ) from error
 
+from autogen.events import BaseEvent  # type: ignore
 from autogen.io import IOStream  # type: ignore
 from autogen.messages import BaseMessage  # type: ignore
 
@@ -54,6 +55,7 @@ MQTT_MAX_RECONNECT_COUNT = 12
 MQTT_MAX_RECONNECT_DELAY = 60
 
 
+# noinspection PyUnusedLocal,PyBroadException
 class MqttIOStream(IOStream):
     """MQTT I/O stream."""
 
@@ -299,6 +301,7 @@ class MqttIOStream(IOStream):
                 reconnect_count += 1
             LOG.info("Reconnect failed after %s attempts.", reconnect_count)
 
+    # noinspection PyUnresolvedReferences
     def _on_message(
         self, client: mqtt.Client, userdata: Any, msg: mqtt.MQTTMessage
     ) -> None:
@@ -476,13 +479,13 @@ class MqttIOStream(IOStream):
             )
         self._print(payload)
 
-    def send(self, message: BaseMessage) -> None:
+    def send(self, message: BaseEvent | BaseMessage) -> None:
         """Send a structured message to MQTT.
 
         Parameters
         ----------
-        message : BaseMessage
-            The message to send.
+        message : BaseEvent | BaseMessage
+            The message or event to send.
         """
         try:
             message_dump = message.model_dump(mode="json")