refactor(workforce): Enhance worker initialization and callback handling in async context

Author: coolbeevip

camel/societies/workforce/workforce.py

@@ -339,6 +339,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
@@ -532,9 +535,6 @@ def _initialize_callbacks(
                 "WorkforceLogger addition."
             )
 
-        for child in self._children:
-            asyncio.run(self._notify_worker_created(child))
-
     async def _notify_worker_created(
         self,
         worker_node: BaseNode,
@@ -553,6 +553,18 @@ async def _notify_worker_created(
         for cb in self._callbacks:
             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,
         session_id: Optional[str] = None,
@@ -2258,6 +2270,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:

test/workforce/test_workforce_callbacks.py

@@ -18,6 +18,8 @@
 from camel.agents import ChatAgent
 from camel.messages import BaseMessage
 from camel.models import ModelFactory
+from camel.societies.workforce import SingleAgentWorker
+from camel.societies.workforce.base import BaseNode
 from camel.societies.workforce.events import (
     AllTasksCompletedEvent,
     TaskAssignedEvent,
@@ -142,6 +144,80 @@ def _build_stub_agent() -> ChatAgent:
     return ChatAgent(model=model)
 
 
+def _build_persona_agent(role_name: str, content: str) -> ChatAgent:
+    """Construct a stub-backed ChatAgent with a system persona."""
+    return ChatAgent(
+        system_message=BaseMessage.make_assistant_message(
+            role_name=role_name,
+            content=content,
+        ),
+        model=ModelFactory.create(
+            model_platform=ModelPlatformType.OPENAI,
+            model_type=ModelType.STUB,
+        ),
+    )
+
+
+def _build_worker_specs() -> list[tuple[str, str, ChatAgent]]:
+    """Build the standard trio of workers used across tests."""
+    return [
+        (
+            "A researcher who can search online for information.",
+            "SearchWork",
+            _build_persona_agent(
+                "Research Specialist",
+                "You are a research specialist who excels at finding and "
+                "gathering information from the web.",
+            ),
+        ),
+        (
+            "An analyst who can process research findings.",
+            "AnalystWorker",
+            _build_persona_agent(
+                "Business Analyst",
+                "You are an expert business analyst. Your job is "
+                "to analyze research findings, identify key insights, "
+                "opportunities, and challenges.",
+            ),
+        ),
+        (
+            "A writer who can create a final report from the analysis.",
+            "WriterWorker",
+            _build_persona_agent(
+                "Report Writer",
+                "You are a professional report writer. You take "
+                "analytical insights and synthesize them into a clear, "
+                "concise, and well-structured final report.",
+            ),
+        ),
+    ]
+
+
+async def _assert_metrics_callbacks(
+    workforce: Workforce, cb: _MetricsCallback
+):
+    """Verify metrics callback toggles across reset cycles."""
+    assert not cb.dump_to_json_called
+    assert not cb.get_ascii_tree_called
+    assert not cb.get_kpis_called
+    await workforce.dump_workforce_logs("foo.log")
+    assert cb.dump_to_json_called
+
+    await workforce.reset()
+    assert not cb.dump_to_json_called
+    assert not cb.get_ascii_tree_called
+    assert not cb.get_kpis_called
+    await workforce.get_workforce_kpis()
+    assert cb.get_kpis_called
+
+    await workforce.reset()
+    assert not cb.dump_to_json_called
+    assert not cb.get_ascii_tree_called
+    assert not cb.get_kpis_called
+    await workforce.get_workforce_log_tree()
+    assert cb.get_ascii_tree_called
+
+
 @pytest.mark.asyncio
 async def test_workforce_callback_registration_and_metrics_handling():
     """Verify default logger addition and metrics-callback skip logic.
@@ -176,165 +252,48 @@ async def test_workforce_callback_registration_and_metrics_handling():
         Workforce("CB Test - Invalid", callbacks=[object()])
 
 
-def assert_event_sequence(events: list[str], min_worker_count: int):
-    """
-    Validate that the given event sequence follows the expected logical order.
-    This version is flexible to handle:
-    - Task retries and dynamic worker creation
-    - Cases where tasks are not decomposed (e.g., when using stub models)
-    """
-    idx = 0
-    n = len(events)
-
-    # 1. Expect at least min_worker_count WorkerCreatedEvent events first
-    initial_worker_count = 0
-    while idx < n and events[idx] == "WorkerCreatedEvent":
-        initial_worker_count += 1
-        idx += 1
-    assert initial_worker_count >= min_worker_count, (
-        f"Expected at least {min_worker_count} initial "
-        f"WorkerCreatedEvents, got {initial_worker_count}"
-    )
-
-    # 2. Expect one main TaskCreatedEvent
-    assert idx < n and events[idx] == "TaskCreatedEvent", (
-        f"Event {idx} should be TaskCreatedEvent, got "
-        f"{events[idx] if idx < n else 'END'}"
-    )
-    idx += 1
-
-    # 3. TaskDecomposedEvent may or may not be present
-    # (depends on coordinator behavior)
-    # If the coordinator can't parse stub responses, it may skip
-    # decomposition
-    has_decomposition = idx < n and events[idx] == "TaskDecomposedEvent"
-    if has_decomposition:
-        idx += 1
-
-    # 4. Count all event types in the remaining events
-    all_events = events[idx:]
-    task_assigned_count = all_events.count("TaskAssignedEvent")
-    task_started_count = all_events.count("TaskStartedEvent")
-    task_completed_count = all_events.count("TaskCompletedEvent")
-    all_tasks_completed_count = all_events.count("AllTasksCompletedEvent")
-
-    # 5. Validate basic invariants
-    # At minimum, the main task should be assigned and processed
-    assert (
-        task_assigned_count >= 1
-    ), f"Expected at least 1 TaskAssignedEvent, got {task_assigned_count}"
-    assert (
-        task_started_count >= 1
-    ), f"Expected at least 1 TaskStartedEvent, got {task_started_count}"
-    assert (
-        task_completed_count >= 1
-    ), f"Expected at least 1 TaskCompletedEvent, got {task_completed_count}"
-
-    # 6. Expect exactly one AllTasksCompletedEvent at the end
-    assert all_tasks_completed_count == 1, (
-        f"Expected exactly 1 AllTasksCompletedEvent, got "
-        f"{all_tasks_completed_count}"
-    )
-    assert (
-        events[-1] == "AllTasksCompletedEvent"
-    ), "Last event should be AllTasksCompletedEvent"
-
-    # 7. All events should be of expected types
-    allowed_events = {
-        "WorkerCreatedEvent",
-        "WorkerDeletedEvent",
-        "TaskCreatedEvent",
-        "TaskDecomposedEvent",
-        "TaskAssignedEvent",
-        "TaskStartedEvent",
-        "TaskCompletedEvent",
-        "TaskFailedEvent",
-        "AllTasksCompletedEvent",
-    }
-    for i, e in enumerate(events):
-        assert e in allowed_events, f"Unexpected event type at {i}: {e}"
-
-
 @pytest.mark.asyncio
-async def test_workforce_emits_expected_event_sequence():
-    # Use STUB model to avoid real API calls and ensure fast,
-    # deterministic execution
-    search_agent = ChatAgent(
-        system_message=BaseMessage.make_assistant_message(
-            role_name="Research Specialist",
-            content="You are a research specialist who excels at finding and "
-            "gathering information from the web.",
-        ),
-        model=ModelFactory.create(
-            model_platform=ModelPlatformType.OPENAI,
-            model_type=ModelType.STUB,
-        ),
-    )
-
-    analyst_agent = ChatAgent(
-        system_message=BaseMessage.make_assistant_message(
-            role_name="Business Analyst",
-            content="You are an expert business analyst. Your job is "
-            "to analyze research findings, identify key insights, "
-            "opportunities, and challenges.",
-        ),
-        model=ModelFactory.create(
-            model_platform=ModelPlatformType.OPENAI,
-            model_type=ModelType.STUB,
-        ),
-    )
-
-    writer_agent = ChatAgent(
-        system_message=BaseMessage.make_assistant_message(
-            role_name="Report Writer",
-            content="You are a professional report writer. You take "
-            "analytical insights and synthesize them into a clear, "
-            "concise, and well-structured final report.",
-        ),
-        model=ModelFactory.create(
-            model_platform=ModelPlatformType.OPENAI,
-            model_type=ModelType.STUB,
-        ),
-    )
-
+@pytest.mark.parametrize(
+    "preconfigure_children",
+    [False, True],
+    ids=["add_workers_at_runtime", "preconfigure_children"],
+)
+async def test_workforce_emits_expected_events_for_worker_init_modes(
+    preconfigure_children: bool,
+):
+    """Validate event ordering for both worker setup paths."""
     cb = _MetricsCallback()
-
-    # Use STUB models for coordinator and task agents to avoid real API calls
-    coordinator_agent = ChatAgent(
-        model=ModelFactory.create(
-            model_platform=ModelPlatformType.OPENAI,
-            model_type=ModelType.STUB,
+    coordinator_agent = _build_stub_agent()
+    task_agent = _build_stub_agent()
+    worker_specs = _build_worker_specs()
+
+    if preconfigure_children:
+        children: list[BaseNode] = [
+            SingleAgentWorker(description=child_desc, worker=agent)
+            for _, child_desc, agent in worker_specs
+        ]
+        workforce = Workforce(
+            'Business Analysis Team',
+            graceful_shutdown_timeout=30.0,
+            callbacks=[cb],
+            coordinator_agent=coordinator_agent,
+            task_agent=task_agent,
+            children=children,
         )
-    )
-    task_agent = ChatAgent(
-        model=ModelFactory.create(
-            model_platform=ModelPlatformType.OPENAI,
-            model_type=ModelType.STUB,
+    else:
+        workforce = Workforce(
+            'Business Analysis Team',
+            graceful_shutdown_timeout=30.0,
+            callbacks=[cb],
+            coordinator_agent=coordinator_agent,
+            task_agent=task_agent,
         )
-    )
-
-    workforce = Workforce(
-        'Business Analysis Team',
-        graceful_shutdown_timeout=30.0,
-        callbacks=[cb],
-        coordinator_agent=coordinator_agent,
-        task_agent=task_agent,
-    )
+        for add_desc, _, agent in worker_specs:
+            await workforce.add_single_agent_worker(
+                add_desc,
+                worker=agent,
+            )
 
-    await workforce.add_single_agent_worker(
-        "A researcher who can search online for information.",
-        worker=search_agent,
-    )
-    await workforce.add_single_agent_worker(
-        "An analyst who can process research findings.",
-        worker=analyst_agent,
-    )
-    await workforce.add_single_agent_worker(
-        "A writer who can create a final report from the analysis.",
-        worker=writer_agent,
-    )
-
-    # Use a simpler task to ensure fast and deterministic execution
     human_task = Task(
         content=(
             "Create a simple report about electric scooters. "
@@ -348,27 +307,18 @@ async def test_workforce_emits_expected_event_sequence():
 
     await workforce.process_task_async(human_task)
 
-    # test that the event sequence is as expected
+    expected_events = [
+        "WorkerCreatedEvent",
+        "WorkerCreatedEvent",
+        "WorkerCreatedEvent",
+        "TaskCreatedEvent",
+        "WorkerCreatedEvent",
+        "TaskAssignedEvent",
+        "TaskStartedEvent",
+        "TaskCompletedEvent",
+        "AllTasksCompletedEvent",
+    ]
     actual_events = [e.__class__.__name__ for e in cb.events]
-    assert_event_sequence(actual_events, min_worker_count=3)
+    assert actual_events == expected_events
 
-    # test that metrics callback methods work as expected
-    assert not cb.dump_to_json_called
-    assert not cb.get_ascii_tree_called
-    assert not cb.get_kpis_called
-    await workforce.dump_workforce_logs("foo.log")
-    assert cb.dump_to_json_called
-
-    await workforce.reset()
-    assert not cb.dump_to_json_called
-    assert not cb.get_ascii_tree_called
-    assert not cb.get_kpis_called
-    await workforce.get_workforce_kpis()
-    assert cb.get_kpis_called
-
-    await workforce.reset()
-    assert not cb.dump_to_json_called
-    assert not cb.get_ascii_tree_called
-    assert not cb.get_kpis_called
-    await workforce.get_workforce_log_tree()
-    assert cb.get_ascii_tree_called
+    await _assert_metrics_callbacks(workforce, cb)