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

waldiez/runner.py

@@ -14,7 +14,7 @@ variables specified in the waldiez file are set.
 """
 
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Union
+from typing import Any
 
 from typing_extensions import Literal
 
@@ -24,12 +24,6 @@ from .running import (
     WaldiezStandardRunner,
 )
 
-if TYPE_CHECKING:
-    from autogen.io.run_response import (  # type: ignore[import-untyped]
-        AsyncRunResponseProtocol,
-        RunResponseProtocol,
-    )
-
 
 def create_runner(
     waldiez: Waldiez,
@@ -38,6 +32,7 @@ def create_runner(
     output_path: str | Path | None = None,
     uploads_root: str | Path | None = None,
     structured_io: bool = False,
+    dot_env: str | Path | None = None,
     **kwargs: Any,
 ) -> WaldiezBaseRunner:
     """Create a Waldiez runner of the specified type.
@@ -54,6 +49,8 @@ def create_runner(
         Uploads root directory, by default None
     structured_io : bool, optional
         Use structured I/O, by default False
+    dot_env : str | Path | None, optional
+        Path to a .env file for environment variables, by default None
     **kwargs
         Additional arguments for specific runner types
 
@@ -84,6 +81,7 @@ def create_runner(
         output_path=output_path,
         uploads_root=uploads_root,
         structured_io=structured_io,
+        dot_env=dot_env,
         **kwargs,
     )
 
@@ -92,6 +90,7 @@ class WaldiezRunner(WaldiezBaseRunner):
     """Factory class for creating Waldiez runners."""
 
     # pylint: disable=super-init-not-called
+    # noinspection PyMissingConstructor
     def __init__(
         self,
         waldiez: Waldiez,
@@ -100,6 +99,7 @@ class WaldiezRunner(WaldiezBaseRunner):
         output_path: str | Path | None = None,
         uploads_root: str | Path | None = None,
         structured_io: bool = False,
+        dot_env: str | Path | None = None,
         **kwargs: Any,
     ):
         """Create a runner instance.
@@ -116,9 +116,10 @@ class WaldiezRunner(WaldiezBaseRunner):
             Uploads root directory, by default None
         structured_io : bool, optional
             Use structured I/O, by default False
+        dot_env : str | Path | None, optional
+            Path to a .env file for environment variables, by default None
         **kwargs
             Additional arguments for specific runner types
-
         """
         self._runner = create_runner(
             waldiez=waldiez,
@@ -126,6 +127,7 @@ class WaldiezRunner(WaldiezBaseRunner):
             output_path=output_path,
             uploads_root=uploads_root,
             structured_io=structured_io,
+            dot_env=dot_env,
             **kwargs,
         )
 
@@ -164,10 +166,7 @@ class WaldiezRunner(WaldiezBaseRunner):
         skip_mmd: bool,
         skip_timeline: bool,
         **kwargs: Any,
-    ) -> Union[
-        list["RunResponseProtocol"],
-        list["AsyncRunResponseProtocol"],
-    ]:
+    ) -> list[dict[str, Any]]:
         return self._runner._run(
             temp_dir=temp_dir,
             output_file=output_file,
@@ -204,14 +203,16 @@ class WaldiezRunner(WaldiezBaseRunner):
         uploads_root: Path | None,
         skip_mmd: bool,
         skip_timeline: bool,
+        dot_env: str | Path | None = None,
         **kwargs: Any,
-    ) -> Union[list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]]:
+    ) -> list[dict[str, Any]]:
         return await self._runner._a_run(
             temp_dir=temp_dir,
             output_file=output_file,
             uploads_root=uploads_root,
             skip_mmd=skip_mmd,
             skip_timeline=skip_timeline,
+            dot_env=dot_env,
             **kwargs,
         )
 
@@ -245,6 +246,7 @@ class WaldiezRunner(WaldiezBaseRunner):
         output_path: str | Path | None = None,
         uploads_root: str | Path | None = None,
         structured_io: bool = False,
+        dot_env: str | Path | None = None,
         **kwargs: Any,
     ) -> "WaldiezRunner":
         """Load a waldiez flow and create a runner.
@@ -267,6 +269,8 @@ class WaldiezRunner(WaldiezBaseRunner):
             Uploads root directory, by default None
         structured_io : bool, optional
             Use structured I/O, by default False
+        dot_env : str | Path | None, optional
+            Path to a .env file for environment variables, by default None
         **kwargs
             Additional arguments for specific runner types
 
@@ -294,5 +298,6 @@ class WaldiezRunner(WaldiezBaseRunner):
             output_path=output_path,
             uploads_root=uploads_root,
             structured_io=structured_io,
+            dot_env=dot_env,
             **kwargs,
         )

waldiez/running/base_runner.py

@@ -7,6 +7,7 @@
 
 """Base runner for Waldiez workflows."""
 
+import shutil
 import sys
 import tempfile
 import threading
@@ -19,9 +20,9 @@ from typing import (
     Coroutine,
     Optional,
     Type,
-    Union,
 )
 
+from aiofiles.os import wrap
 from anyio.from_thread import start_blocking_portal
 from typing_extensions import Self
 
@@ -43,15 +44,19 @@ from .utils import (
 
 if TYPE_CHECKING:
     from autogen.events import BaseEvent  # type: ignore[import-untyped]
-    from autogen.io.run_response import (  # type: ignore[import-untyped]
-        AsyncRunResponseProtocol,
-        RunResponseProtocol,
-    )
 
 
 class WaldiezBaseRunner(WaldiezRunnerProtocol):
     """Base runner for Waldiez.
 
+    Initialization parameters:
+        - waldiez: The Waldiez flow to run.
+        - output_path: Path to save the output file.
+        - uploads_root: Root directory for uploads.
+        - structured_io: Whether to use structured I/O.
+        - skip_patch_io: Whether to skip patching I/O functions.
+        - dot_env: Path to a .env file for environment variables.
+
     Methods to override:
         - _before_run: Actions to perform before running the flow.
         - _a_before_run: Async actions to perform before running the flow.
@@ -68,6 +73,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
     _structured_io: bool
     _output_path: str | Path | None
     _uploads_root: str | Path | None
+    _dot_env_path: str | Path | None
     _skip_patch_io: bool
     _running: bool
 
@@ -78,6 +84,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         uploads_root: str | Path | None,
         structured_io: bool,
         skip_patch_io: bool = False,
+        dot_env: str | Path | None = None,
     ) -> None:
         """Initialize the Waldiez manager."""
         self._waldiez = waldiez
@@ -86,6 +93,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         WaldiezBaseRunner._output_path = output_path
         WaldiezBaseRunner._uploads_root = uploads_root
         WaldiezBaseRunner._skip_patch_io = skip_patch_io
+        WaldiezBaseRunner._dot_env_path = dot_env
         self._called_install_requirements = False
         self._exporter = WaldiezExporter(waldiez)
         self._stop_requested = threading.Event()
@@ -95,9 +103,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         self._input: (
             Callable[..., str] | Callable[..., Coroutine[Any, Any, str]]
         ) = input
-        self._last_results: Union[
-            list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]
-        ] = []
+        self._last_results: list[dict[str, Any]] = []
         self._last_exception: Exception | None = None
         self._execution_complete_event = threading.Event()
         self._execution_thread: Optional[threading.Thread] = None
@@ -186,6 +192,11 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
                 uploads_root=uploads_root,
                 structured_io=WaldiezBaseRunner._structured_io,
             )
+            if self._dot_env_path:
+                shutil.copyfile(
+                    str(self._dot_env_path),
+                    str(temp_dir / ".env"),
+                )
         return temp_dir
 
     async def _a_before_run(
@@ -203,6 +214,12 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
                 structured_io=self.structured_io,
                 force=True,
             )
+            if self._dot_env_path:
+                wrapped = wrap(shutil.copyfile)
+                await wrapped(
+                    str(self._dot_env_path),
+                    str(temp_dir / ".env"),
+                )
         return temp_dir
 
     def _run(
@@ -212,7 +229,8 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         uploads_root: Path | None,
         skip_mmd: bool,
         skip_timeline: bool,
-    ) -> Union[list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]]:
+        **kwargs: Any,
+    ) -> list[dict[str, Any]]:
         """Run the Waldiez flow."""
         raise NotImplementedError(
             "The _run method must be implemented in the subclass."
@@ -225,7 +243,8 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         uploads_root: Path | None,
         skip_mmd: bool,
         skip_timeline: bool,
-    ) -> Union[list["AsyncRunResponseProtocol"], list["RunResponseProtocol"]]:
+        **kwargs: Any,
+    ) -> list[dict[str, Any]]:
         """Run the Waldiez flow asynchronously."""
         raise NotImplementedError(
             "The _a_run method must be implemented in the subclass."
@@ -262,8 +281,6 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             The path to the output file.
         uploads_root : Path | None
             The root path for uploads, if any.
-        structured_io : bool
-            Whether to use structured IO instead of the default 'input/print'.
         skip_mmd : bool
             Whether to skip generating the mermaid diagram.
         """
@@ -273,9 +290,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
 
     def _after_run(
         self,
-        results: Union[
-            list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]
-        ],
+        results: list[dict[str, Any]],
         output_file: Path,
         uploads_root: Path | None,
         temp_dir: Path,
@@ -300,9 +315,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
 
     async def _a_after_run(
         self,
-        results: Union[
-            list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]
-        ],
+        results: list[dict[str, Any]],
         output_file: Path,
         uploads_root: Path | None,
         temp_dir: Path,
@@ -367,6 +380,8 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             for req in self.waldiez.requirements
             if req not in sys.modules and "waldiez" not in req
         }
+        if "python-dotenv" not in extra_requirements:
+            extra_requirements.add("python-dotenv")
         return extra_requirements
 
     def install_requirements(self) -> None:
@@ -444,7 +459,9 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         structured_io: bool | None = None,
         skip_mmd: bool = False,
         skip_timeline: bool = False,
-    ) -> Union[list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]]:
+        dot_env: str | Path | None = None,
+        **kwargs: Any,
+    ) -> list[dict[str, Any]]:
         """Run the Waldiez flow in blocking mode.
 
         Parameters
@@ -460,18 +477,25 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             Whether to skip generating the mermaid diagram, by default False.
         skip_timeline : bool
             Whether to skip generating the timeline JSON.
+        dot_env : str | Path | None
+            The path to the .env file, if any.
+        **kwargs : Any
+            Additional keyword arguments for the run method.
 
         Returns
         -------
-        Union[list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]]
-            The result of the run, which can be a list of RunResponseProtocol
-            or a list of AsyncRunResponseProtocol.
+        list[dict[str, Any]]
+            The result of the run.
 
         Raises
         ------
         RuntimeError
             If the runner is already running.
         """
+        if dot_env is not None:
+            resolved = Path(dot_env).resolve()
+            if resolved.is_file():
+                WaldiezBaseRunner._dot_env_path = resolved
         if structured_io is not None:
             WaldiezBaseRunner._structured_io = structured_io
         if self.is_running():
@@ -496,9 +520,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         self.install_requirements()
         refresh_environment()
         WaldiezBaseRunner._running = True
-        results: Union[
-            list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]
-        ] = []
+        results: list[dict[str, Any]] = []
         old_env_vars = set_env_vars(self.waldiez.get_flow_env_vars())
         try:
             with chdir(to=temp_dir):
@@ -526,6 +548,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             sys.path.pop(0)
         return results
 
+    # noinspection DuplicatedCode
     async def a_run(
         self,
         output_path: str | Path | None = None,
@@ -533,7 +556,9 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         structured_io: bool | None = None,
         skip_mmd: bool = False,
         skip_timeline: bool = False,
-    ) -> Union[list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]]:
+        dot_env: str | Path | None = None,
+        **kwargs: Any,
+    ) -> list[dict[str, Any]]:
         """Run the Waldiez flow asynchronously.
 
         Parameters
@@ -549,18 +574,25 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             Whether to skip generating the mermaid diagram, by default False.
         skip_timeline : bool
             Whether to skip generating the timeline JSON, by default False.
+        dot_env : str | Path | None
+            The path to the .env file, if any.
+        **kwargs : Any
+            Additional keyword arguments for the run method.
 
         Returns
         -------
-        Union[list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]]
-            The result of the run, which can be a list of RunResponseProtocol
-            or a list of AsyncRunResponseProtocol.
+        list[dict[str, Any]]
+            The result of the run.
 
         Raises
         ------
         RuntimeError
             If the runner is already running.
         """
+        if dot_env is not None:
+            resolved = Path(dot_env).resolve()
+            if resolved.is_file():
+                WaldiezBaseRunner._dot_env_path = resolved
         if structured_io is not None:
             WaldiezBaseRunner._structured_io = structured_io
         if self.is_running():
@@ -576,9 +608,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         await self.a_install_requirements()
         refresh_environment()
         WaldiezBaseRunner._running = True
-        results: Union[
-            list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]
-        ] = []
+        results: list[dict[str, Any]] = []
         old_env_vars = set_env_vars(self.waldiez.get_flow_env_vars())
         try:
             async with a_chdir(to=temp_dir):
@@ -612,6 +642,8 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         structured_io: bool | None = None,
         skip_mmd: bool = False,
         skip_timeline: bool = False,
+        dot_env: str | Path | None = None,
+        **kwargs: Any,
     ) -> None:
         """Start running the Waldiez flow in a non-blocking way.
 
@@ -627,12 +659,20 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             Whether to skip generating the mermaid diagram, by default False.
         skip_timeline : bool
             Whether to skip generating the timeline JSON, by default False.
+        dot_env : str | Path | None
+            The path to the .env file, if any.
+        **kwargs : Any
+            Additional keyword arguments for the start method.
 
         Raises
         ------
         RuntimeError
             If the runner is already running.
         """
+        if dot_env is not None:
+            resolved = Path(dot_env).resolve()
+            if resolved.is_file():
+                WaldiezBaseRunner._dot_env_path = resolved
         if structured_io is not None:
             WaldiezBaseRunner._structured_io = structured_io
         if self.is_running():
@@ -656,6 +696,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             skip_timeline=skip_timeline,
         )
 
+    # noinspection DuplicatedCode
     async def a_start(
         self,
         output_path: str | Path | None,
@@ -663,6 +704,8 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         structured_io: bool | None = None,
         skip_mmd: bool = False,
         skip_timeline: bool = False,
+        dot_env: str | Path | None = None,
+        **kwargs: Any,
     ) -> None:
         """Asynchronously start running the Waldiez flow in a non-blocking way.
 
@@ -678,12 +721,20 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             Whether to skip generating the mermaid diagram, by default False.
         skip_timeline : bool = False
             Whether to skip generating the timeline JSON, by default False.
+        dot_env : str | Path | None = None
+            The path to the .env file, if any.
+        **kwargs : Any
+            Additional keyword arguments for the start method.
 
         Raises
         ------
         RuntimeError
             If the runner is already running.
         """
+        if dot_env is not None:
+            resolved = Path(dot_env).resolve()
+            if resolved.is_file():
+                WaldiezBaseRunner._dot_env_path = resolved
         if structured_io is not None:
             WaldiezBaseRunner._structured_io = structured_io
         if self.is_running():
@@ -709,9 +760,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
 
     def after_run(
         self,
-        results: Union[
-            list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]
-        ],
+        results: list[dict[str, Any]],
         output_file: Path,
         uploads_root: Path | None,
         temp_dir: Path,
@@ -722,7 +771,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
 
         Parameters
         ----------
-        results : Union[ChatResult, list[ChatResult], dict[int, ChatResult]]
+        results : list[dict[str, Any]]
             The results of the flow run.
         output_file : Path
             The path to the output file.
@@ -746,9 +795,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
 
     async def a_after_run(
         self,
-        results: Union[
-            list["RunResponseProtocol"], list["AsyncRunResponseProtocol"]
-        ],
+        results: list[dict[str, Any]],
         output_file: Path,
         uploads_root: Path | None,
         temp_dir: Path,
@@ -759,7 +806,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
 
         Parameters
         ----------
-        results : Union[ChatResult, list[ChatResult], dict[int, ChatResult]]
+        results : list[dict[str, Any]]
             The results of the flow run.
         output_file : Path
             The path to the output file.
@@ -828,6 +875,11 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         """Check if the runner is using structured IO."""
         return WaldiezBaseRunner._structured_io
 
+    @property
+    def dot_env_path(self) -> str | Path | None:
+        """Get the path to the .env file, if any."""
+        return WaldiezBaseRunner._dot_env_path
+
     @property
     def output_path(self) -> str | Path | None:
         """Get the output path for the runner."""
@@ -854,6 +906,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         output_path: str | Path | None = None,
         uploads_root: str | Path | None = None,
         structured_io: bool = False,
+        dot_env: str | Path | None = None,
     ) -> "WaldiezBaseRunner":
         """Load a waldiez flow from a file and create a runner.
 
@@ -876,6 +929,8 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
         structured_io : bool, optional
             Whether to use structured IO instead of the default 'input/print',
             by default False.
+        dot_env : str | Path | None, optional
+            The path to the .env file, if any, by default None.
 
         Returns
         -------
@@ -894,6 +949,7 @@ class WaldiezBaseRunner(WaldiezRunnerProtocol):
             output_path=output_path,
             uploads_root=uploads_root,
             structured_io=structured_io,
+            dot_env=dot_env,
         )
 
     def __enter__(self) -> Self:

waldiez/running/environment.py

@@ -135,6 +135,7 @@ def reload_autogen() -> None:  # noqa: C901  # pragma: no cover
 
     except Exception as e:
         # If reload fails, at least try to re-import autogen
+        # noinspection PyBroadException
         try:
             import autogen  # type: ignore  # noqa: F401
 
@@ -174,6 +175,7 @@ def reload_chroma_if_needed() -> None:  # pragma: no cover
         return
 
 
+# noinspection DuplicatedCode
 def try_handle_the_np_thing() -> None:
     """Try to handle the numpy deprecation warning."""
     # we might get:

waldiez/running/patch_io_stream.py

@@ -95,6 +95,7 @@ def get_printer() -> Callable[..., None]:  # noqa: C901
         The printer function that handles Unicode encoding errors gracefully.
     """
     try:
+        # noinspection PyUnresolvedReferences
         from autogen.io import IOStream  # pyright: ignore
 
         printer = IOStream.get_default().print
@@ -102,6 +103,7 @@ def get_printer() -> Callable[..., None]:  # noqa: C901
         # Fallback to standard print if autogen is not available
         printer = print
 
+    # noinspection PyBroadException
     def safe_printer(*args: Any, **kwargs: Any) -> None:  # noqa: C901
         """Safe printer that handles Unicode encoding errors.
 

waldiez/running/post_run.py

@@ -15,6 +15,7 @@ from .patch_io_stream import get_printer
 from .timeline_processor import TimelineProcessor
 
 
+# noinspection PyUnusedLocal
 def after_run(
     temp_dir: Path,
     output_file: Optional[Union[str, Path]],
@@ -71,6 +72,7 @@ def after_run(
     shutil.rmtree(temp_dir)
 
 
+# noinspection PyBroadException
 def _make_mermaid_diagram(
     temp_dir: Path,
     output_file: Optional[Union[str, Path]],
@@ -93,6 +95,7 @@ def _make_mermaid_diagram(
                 pass
 
 
+# noinspection PyBroadException
 def _make_timeline_json(
     temp_dir: Path,
 ) -> None:

waldiez/running/pre_run.py

@@ -13,6 +13,7 @@ from .environment import in_virtualenv, is_root
 from .utils import strip_ansi
 
 
+# noinspection PyUnresolvedReferences
 def install_requirements(
     extra_requirements: set[str],
     upgrade: bool = False,