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

waldiez/running/standard_runner.py

@@ -2,7 +2,8 @@
 # Copyright (c) 2024 - 2025 Waldiez and contributors.
 
 # pyright: reportUnknownMemberType=false, reportAttributeAccessIssue=false
-# pylint: disable=duplicate-code
+# pyright: reportMissingTypeStubs=false, reportDeprecated=false
+# pylint: disable=duplicate-code,too-few-public-methods
 """Run a waldiez flow.
 
 The flow is first converted to an autogen flow with agents, chats, and tools.
@@ -16,12 +17,16 @@ import traceback
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Union
 
+from typing_extensions import override
+
 from waldiez.models.waldiez import Waldiez
-from waldiez.running.run_results import WaldiezRunResults
 
 from .base_runner import WaldiezBaseRunner
+from .events_mixin import EventsMixin
+from .results_mixin import WaldiezRunResults
 
 if TYPE_CHECKING:
+    from autogen.agentchat import ConversableAgent  # type: ignore
     from autogen.events import BaseEvent  # type: ignore
     from autogen.messages import BaseMessage  # type: ignore
 
@@ -59,21 +64,8 @@ class WaldiezStandardRunner(WaldiezBaseRunner):
         self._event_count = 0
         self._processed_events = 0
 
-    # pylint: disable=no-self-use
-    # noinspection PyMethodMayBeStatic
-    def print(self, *args: Any, **kwargs: Any) -> None:
-        """Print.
-
-        Parameters
-        ----------
-        *args : Any
-            Positional arguments to print.
-        **kwargs : Any
-            Keyword arguments to print.
-        """
-        WaldiezBaseRunner._print(*args, **kwargs)
-
     # pylint: disable=unused-argument
+    @override
     def _run(
         self,
         temp_dir: Path,
@@ -109,9 +101,9 @@ class WaldiezStandardRunner(WaldiezBaseRunner):
                 )
             else:
                 stream = IOStream.get_default()
-            WaldiezBaseRunner._print = stream.print
-            WaldiezBaseRunner._input = stream.input
-            WaldiezBaseRunner._send = stream.send
+            EventsMixin.set_print_function(stream.print)
+            EventsMixin.set_input_function(stream.input)
+            EventsMixin.set_send_function(stream.send)
             self.print(MESSAGES["workflow_starting"])
             self.print(self.waldiez.info.model_dump_json())
             results = loaded_module.main(
@@ -134,6 +126,7 @@ class WaldiezStandardRunner(WaldiezBaseRunner):
     def _on_event(
         self,
         event: Union["BaseEvent", "BaseMessage"],
+        agents: list["ConversableAgent"],
     ) -> bool:
         """Process an event from the workflow."""
         self._event_count += 1
@@ -143,7 +136,7 @@ class WaldiezStandardRunner(WaldiezBaseRunner):
             )
             return False
         try:
-            WaldiezBaseRunner.process_event(event)
+            EventsMixin.process_event(event, agents)
             self._processed_events += 1
         except SystemExit:  # pragma: no cover
             self.log.debug("Execution stopped by user (sync)")
@@ -160,6 +153,7 @@ class WaldiezStandardRunner(WaldiezBaseRunner):
     async def _a_on_event(
         self,
         event: Union["BaseEvent", "BaseMessage"],
+        agents: list["ConversableAgent"],
     ) -> bool:
         """Process an event from the workflow asynchronously."""
         self._event_count += 1
@@ -169,7 +163,7 @@ class WaldiezStandardRunner(WaldiezBaseRunner):
             )
             return False
         try:
-            await WaldiezBaseRunner.a_process_event(event)
+            await EventsMixin.a_process_event(event, agents)
             self._processed_events += 1
         except SystemExit:  # pragma: no cover
             self.log.debug("Execution stopped by user (async)")
@@ -184,6 +178,7 @@ class WaldiezStandardRunner(WaldiezBaseRunner):
         return not self._stop_requested.is_set()
 
     # pylint: disable=too-complex
+    @override
     async def _a_run(
         self,
         temp_dir: Path,
@@ -200,7 +195,7 @@ class WaldiezStandardRunner(WaldiezBaseRunner):
             # fmt: on
             """Execute the workflow in an async context."""
             # pylint: disable=import-outside-toplevel
-            from autogen.io import IOStream  # pyright: ignore
+            from autogen.io import IOStream
 
             from waldiez.io import StructuredIOStream
 
@@ -209,23 +204,25 @@ class WaldiezStandardRunner(WaldiezBaseRunner):
             try:
                 loaded_module = self._load_module(output_file, temp_dir)
                 if self._stop_requested.is_set():  # pragma: no cover
-                    self.log.debug(
+                    msg = (
                         "Execution stopped before AG2 "
                         "workflow event processing (async)"
                     )
+                    self.log.debug(msg)
                     return []
+                # noinspection DuplicatedCode
                 if self.structured_io:
                     stream = StructuredIOStream(
                         uploads_root=uploads_root, is_async=True
                     )
                 else:
                     stream = IOStream.get_default()
-                WaldiezBaseRunner._print = stream.print
-                WaldiezBaseRunner._input = stream.input
-                WaldiezBaseRunner._send = stream.send
+                EventsMixin.set_print_function(stream.print)
+                EventsMixin.set_input_function(stream.input)
+                EventsMixin.set_send_function(stream.send)
                 self.print(MESSAGES["workflow_starting"])
                 self.print(self.waldiez.info.model_dump_json())
-                results = await loaded_module.main(  # pyright: ignore
+                results = await loaded_module.main(
                     on_event=self._a_on_event
                 )
                 self.print(MESSAGES["workflow_finished"])