refactor(workforce): Refactor WorkforceCallback and all related callbacks to async interface

Author: coolbeevip

camel/benchmarks/browsecomp.py

@@ -11,7 +11,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
-
+import asyncio
 import base64
 import hashlib
 import json
@@ -585,15 +585,17 @@ def process_benchmark_row(row: Dict[str, Any]) -> Dict[str, Any]:
                 input_message = QUERY_TEMPLATE.format(question=problem)
 
                 if isinstance(pipeline_template, (ChatAgent)):
-                    pipeline = pipeline_template.clone()  # type: ignore[assignment]
+                    chat_pipeline = pipeline_template.clone()
 
-                    response_text = pipeline.step(
+                    response_text = chat_pipeline.step(
                         input_message, response_format=QueryResponse
                     )
                 elif isinstance(pipeline_template, Workforce):
-                    pipeline = pipeline_template.clone()  # type: ignore[assignment]
+                    workforce_pipeline = asyncio.run(pipeline_template.clone())
                     task = Task(content=input_message, id="0")
-                    task = pipeline.process_task(task)  # type: ignore[attr-defined]
+                    task = asyncio.run(
+                        workforce_pipeline.process_task_async(task)
+                    )  # type: ignore[attr-defined]
                     if task_json_formatter:
                         formatter_in_process = task_json_formatter.clone()
                     else:
@@ -607,16 +609,16 @@ def process_benchmark_row(row: Dict[str, Any]) -> Dict[str, Any]:
 
                 elif isinstance(pipeline_template, RolePlaying):
                     # RolePlaying is different.
-                    pipeline = pipeline_template.clone(  # type: ignore[assignment]
+                    rp_pipeline = pipeline_template.clone(
                         task_prompt=input_message
                     )
 
                     n = 0
-                    input_msg = pipeline.init_chat()  # type: ignore[attr-defined]
+                    input_msg = rp_pipeline.init_chat()
                     chat_history = []
                     while n < chat_turn_limit:
                         n += 1
-                        assistant_response, user_response = pipeline.step(
+                        assistant_response, user_response = rp_pipeline.step(
                             input_msg
                         )
                         if assistant_response.terminated:  # type: ignore[union-attr]

camel/societies/workforce/workforce.py

@@ -340,6 +340,9 @@ def __init__(
         self.snapshot_interval: float = 30.0
         # Shared memory UUID tracking to prevent re-sharing duplicates
         self._shared_memory_uuids: Set[str] = set()
+        # Defer initial worker-created callbacks until an event loop is
+        # available in async context.
+        self._pending_worker_created: Deque[BaseNode] = deque(self._children)
         self._initialize_callbacks(callbacks)
 
         # Set up coordinator agent with default system message
@@ -533,10 +536,7 @@ def _initialize_callbacks(
                 "WorkforceLogger addition."
             )
 
-        for child in self._children:
-            self._notify_worker_created(child)
-
-    def _notify_worker_created(
+    async def _notify_worker_created(
         self,
         worker_node: BaseNode,
         *,
@@ -552,7 +552,19 @@ def _notify_worker_created(
             metadata=metadata,
         )
         for cb in self._callbacks:
-            cb.log_worker_created(event)
+            await cb.log_worker_created(event)
+
+    async def _flush_initial_worker_created_callbacks(self) -> None:
+        r"""Flush pending worker-created callbacks that were queued during
+        initialization before an event loop was available."""
+        if not self._pending_worker_created:
+            return
+
+        pending = list(self._pending_worker_created)
+        self._pending_worker_created.clear()
+
+        for child in pending:
+            await self._notify_worker_created(child)
 
     def _get_or_create_shared_context_utility(
         self,
@@ -1669,7 +1681,7 @@ async def _apply_recovery_strategy(
                         subtask_ids=[st.id for st in subtasks],
                     )
                     for cb in self._callbacks:
-                        cb.log_task_decomposed(task_decomposed_event)
+                        await cb.log_task_decomposed(task_decomposed_event)
                     for subtask in subtasks:
                         task_created_event = TaskCreatedEvent(
                             task_id=subtask.id,
@@ -1679,7 +1691,7 @@ async def _apply_recovery_strategy(
                             metadata=subtask.additional_info,
                         )
                         for cb in self._callbacks:
-                            cb.log_task_created(task_created_event)
+                            await cb.log_task_created(task_created_event)
 
                 # Insert subtasks at the head of the queue
                 self._pending_tasks.extendleft(reversed(subtasks))
@@ -2288,7 +2300,7 @@ async def handle_decompose_append_task(
             return [task]
 
         if reset and self._state != WorkforceState.RUNNING:
-            self.reset()
+            await self.reset()
             logger.info("Workforce reset before handling task.")
 
         # Focus on the new task
@@ -2302,7 +2314,7 @@ async def handle_decompose_append_task(
             metadata=task.additional_info,
         )
         for cb in self._callbacks:
-            cb.log_task_created(task_created_event)
+            await cb.log_task_created(task_created_event)
 
         # The agent tend to be overconfident on the whole task, so we
         # decompose the task into subtasks first
@@ -2323,7 +2335,7 @@ async def handle_decompose_append_task(
                 subtask_ids=[st.id for st in subtasks],
             )
             for cb in self._callbacks:
-                cb.log_task_decomposed(task_decomposed_event)
+                await cb.log_task_decomposed(task_decomposed_event)
             for subtask in subtasks:
                 task_created_event = TaskCreatedEvent(
                     task_id=subtask.id,
@@ -2333,7 +2345,7 @@ async def handle_decompose_append_task(
                     metadata=subtask.additional_info,
                 )
                 for cb in self._callbacks:
-                    cb.log_task_created(task_created_event)
+                    await cb.log_task_created(task_created_event)
 
         if subtasks:
             # _pending_tasks will contain both undecomposed
@@ -2361,6 +2373,9 @@ async def process_task_async(
         Returns:
             Task: The updated task.
         """
+        # Emit worker-created callbacks lazily once an event loop is present.
+        await self._flush_initial_worker_created_callbacks()
+
         # Delegate to intervention pipeline when requested to keep
         # backward-compat.
         if interactive:
@@ -2709,7 +2724,7 @@ def _start_child_node_when_paused(
             # Close the coroutine to prevent RuntimeWarning
             start_coroutine.close()
 
-    def add_single_agent_worker(
+    async def add_single_agent_worker(
         self,
         description: str,
         worker: ChatAgent,
@@ -2765,13 +2780,13 @@ def add_single_agent_worker(
         # If workforce is paused, start the worker's listening task
         self._start_child_node_when_paused(worker_node.start())
 
-        self._notify_worker_created(
+        await self._notify_worker_created(
             worker_node,
             worker_type='SingleAgentWorker',
         )
         return self
 
-    def add_role_playing_worker(
+    async def add_role_playing_worker(
         self,
         description: str,
         assistant_role_name: str,
@@ -2842,7 +2857,7 @@ def add_role_playing_worker(
         # If workforce is paused, start the worker's listening task
         self._start_child_node_when_paused(worker_node.start())
 
-        self._notify_worker_created(
+        await self._notify_worker_created(
             worker_node,
             worker_type='RolePlayingWorker',
         )
@@ -2884,7 +2899,7 @@ async def _async_reset(self) -> None:
         self._pause_event.set()
 
     @check_if_running(False)
-    def reset(self) -> None:
+    async def reset(self) -> None:
         r"""Reset the workforce and all the child nodes under it. Can only
         be called when the workforce is not running.
         """
@@ -2919,9 +2934,7 @@ def reset(self) -> None:
         if self._loop and not self._loop.is_closed():
             # If we have a loop, use it to set the event safely
             try:
-                asyncio.run_coroutine_threadsafe(
-                    self._async_reset(), self._loop
-                ).result()
+                await self._async_reset()
             except RuntimeError as e:
                 logger.warning(f"Failed to reset via existing loop: {e}")
                 # Fallback to direct event manipulation
@@ -2932,7 +2945,7 @@ def reset(self) -> None:
 
         for cb in self._callbacks:
             if isinstance(cb, WorkforceMetrics):
-                cb.reset_task_data()
+                await cb.reset_task_data()
 
     def save_workflow_memories(
         self,
@@ -3800,7 +3813,7 @@ async def _post_task(self, task: Task, assignee_id: str) -> None:
             task_id=task.id, worker_id=assignee_id
         )
         for cb in self._callbacks:
-            cb.log_task_started(task_started_event)
+            await cb.log_task_started(task_started_event)
 
         try:
             await self._channel.post_task(task, self.node_id, assignee_id)
@@ -3947,7 +3960,7 @@ async def _create_worker_node_for_task(self, task: Task) -> Worker:
 
         self._children.append(new_node)
 
-        self._notify_worker_created(
+        await self._notify_worker_created(
             new_node,
             worker_type='SingleAgentWorker',
             role=new_node_conf.role,
@@ -4085,7 +4098,7 @@ async def _post_ready_tasks(self) -> None:
                 for cb in self._callbacks:
                     # queue_time_seconds can be derived by logger if task
                     # creation time is logged
-                    cb.log_task_assigned(task_assigned_event)
+                    await cb.log_task_assigned(task_assigned_event)
 
         # Step 2: Iterate through all pending tasks and post those that are
         # ready
@@ -4243,7 +4256,7 @@ async def _post_ready_tasks(self) -> None:
                                 },
                             )
                             for cb in self._callbacks:
-                                cb.log_task_failed(task_failed_event)
+                                await cb.log_task_failed(task_failed_event)
 
                             self._completed_tasks.append(task)
                             self._cleanup_task_tracking(task.id)
@@ -4306,7 +4319,7 @@ async def _handle_failed_task(self, task: Task) -> bool:
             },
         )
         for cb in self._callbacks:
-            cb.log_task_failed(task_failed_event)
+            await cb.log_task_failed(task_failed_event)
 
         # Check for immediate halt conditions after max retries.
         if task.failure_count >= MAX_TASK_RETRIES:
@@ -4493,7 +4506,7 @@ async def _handle_completed_task(self, task: Task) -> None:
             metadata={'current_state': task.state.value},
         )
         for cb in self._callbacks:
-            cb.log_task_completed(task_completed_event)
+            await cb.log_task_completed(task_completed_event)
 
         # Find and remove the completed task from pending tasks
         tasks_list = list(self._pending_tasks)
@@ -4609,7 +4622,7 @@ async def _graceful_shutdown(self, failed_task: Task) -> None:
         # Wait for the full timeout period
         await asyncio.sleep(self.graceful_shutdown_timeout)
 
-    def get_workforce_log_tree(self) -> str:
+    async def get_workforce_log_tree(self) -> str:
         r"""Returns an ASCII tree representation of the task hierarchy and
         worker status.
         """
@@ -4619,19 +4632,19 @@ def get_workforce_log_tree(self) -> str:
         if len(metrics_cb) == 0:
             return "Metrics Callback not initialized."
         else:
-            return metrics_cb[0].get_ascii_tree_representation()
+            return await metrics_cb[0].get_ascii_tree_representation()
 
-    def get_workforce_kpis(self) -> Dict[str, Any]:
+    async def get_workforce_kpis(self) -> Dict[str, Any]:
         r"""Returns a dictionary of key performance indicators."""
         metrics_cb: List[WorkforceMetrics] = [
             cb for cb in self._callbacks if isinstance(cb, WorkforceMetrics)
         ]
         if len(metrics_cb) == 0:
             return {"error": "Metrics Callback not initialized."}
         else:
-            return metrics_cb[0].get_kpis()
+            return await metrics_cb[0].get_kpis()
 
-    def dump_workforce_logs(self, file_path: str) -> None:
+    async def dump_workforce_logs(self, file_path: str) -> None:
         r"""Dumps all collected logs to a JSON file.
 
         Args:
@@ -4643,7 +4656,7 @@ def dump_workforce_logs(self, file_path: str) -> None:
         if len(metrics_cb) == 0:
             print("Logger not initialized. Cannot dump logs.")
             return
-        metrics_cb[0].dump_to_json(file_path)
+        await metrics_cb[0].dump_to_json(file_path)
         # Use logger.info or print, consistent with existing style
         logger.info(f"Workforce logs dumped to {file_path}")
 
@@ -5119,7 +5132,7 @@ async def _listen_to_channel(self) -> None:
             logger.info("All tasks completed.")
             all_tasks_completed_event = AllTasksCompletedEvent()
             for cb in self._callbacks:
-                cb.log_all_tasks_completed(all_tasks_completed_event)
+                await cb.log_all_tasks_completed(all_tasks_completed_event)
 
         # shut down the whole workforce tree
         self.stop()
@@ -5187,7 +5200,7 @@ def stop(self) -> None:
                 f"(event-loop not yet started)."
             )
 
-    def clone(self, with_memory: bool = False) -> 'Workforce':
+    async def clone(self, with_memory: bool = False) -> 'Workforce':
         r"""Creates a new instance of Workforce with the same configuration.
 
         Args:
@@ -5219,13 +5232,13 @@ def clone(self, with_memory: bool = False) -> 'Workforce':
         for child in self._children:
             if isinstance(child, SingleAgentWorker):
                 cloned_worker = child.worker.clone(with_memory)
-                new_instance.add_single_agent_worker(
+                await new_instance.add_single_agent_worker(
                     child.description,
                     cloned_worker,
                     pool_max_size=10,
                 )
             elif isinstance(child, RolePlayingWorker):
-                new_instance.add_role_playing_worker(
+                await new_instance.add_role_playing_worker(
                     child.description,
                     child.assistant_role_name,
                     child.user_role_name,
@@ -5235,7 +5248,7 @@ def clone(self, with_memory: bool = False) -> 'Workforce':
                     child.chat_turn_limit,
                 )
             elif isinstance(child, Workforce):
-                new_instance.add_workforce(child.clone(with_memory))
+                new_instance.add_workforce(await child.clone(with_memory))
             else:
                 logger.warning(f"{type(child)} is not being cloned.")
                 continue
@@ -5470,7 +5483,7 @@ def get_children_info():
             return children_info
 
         # Add single agent worker
-        def add_single_agent_worker(
+        async def add_single_agent_worker(
             description,
             system_message=None,
             role_name="Assistant",
@@ -5534,7 +5547,9 @@ def add_single_agent_worker(
                         "message": str(e),
                     }
 
-                workforce_instance.add_single_agent_worker(description, agent)
+                await workforce_instance.add_single_agent_worker(
+                    description, agent
+                )
 
                 return {
                     "status": "success",
@@ -5545,7 +5560,7 @@ def add_single_agent_worker(
                 return {"status": "error", "message": str(e)}
 
         # Add role playing worker
-        def add_role_playing_worker(
+        async def add_role_playing_worker(
             description,
             assistant_role_name,
             user_role_name,
@@ -5602,7 +5617,7 @@ def add_role_playing_worker(
                         "message": "Cannot add workers while workforce is running",  # noqa: E501
                     }
 
-                workforce_instance.add_role_playing_worker(
+                await workforce_instance.add_role_playing_worker(
                     description=description,
                     assistant_role_name=assistant_role_name,
                     user_role_name=user_role_name,