update WorkforceLogger with log_message (#3485)

Co-authored-by: Tao Sun <[email protected]>
Co-authored-by: JINO ROHIT <[email protected]>

Author: 3 people

camel/societies/workforce/events.py

@@ -22,6 +22,7 @@
 class WorkforceEventBase(BaseModel):
     model_config = ConfigDict(frozen=True, extra='forbid')
     event_type: Literal[
+        "log",
         "task_decomposed",
         "task_created",
         "task_assigned",
@@ -39,6 +40,25 @@ class WorkforceEventBase(BaseModel):
     )
 
 
+class LogEvent(WorkforceEventBase):
+    event_type: Literal["log"] = "log"
+    message: str
+    level: Literal["debug", "info", "warning", "error", "critical"]
+    color: (
+        Literal[
+            "red",
+            "green",
+            "yellow",
+            "blue",
+            "cyan",
+            "magenta",
+            "gray",
+            "black",
+        ]
+        | None
+    ) = None
+
+
 class WorkerCreatedEvent(WorkforceEventBase):
     event_type: Literal["worker_created"] = "worker_created"
     worker_id: str
@@ -109,6 +129,7 @@ class QueueStatusEvent(WorkforceEventBase):
 
 
 WorkforceEvent = Union[
+    LogEvent,
     TaskDecomposedEvent,
     TaskCreatedEvent,
     TaskAssignedEvent,

camel/societies/workforce/workforce.py

@@ -97,6 +97,7 @@
 
 from .events import (
     AllTasksCompletedEvent,
+    LogEvent,
     TaskAssignedEvent,
     TaskCompletedEvent,
     TaskCreatedEvent,
@@ -3880,10 +3881,17 @@ async def _post_task(self, task: Task, assignee_id: str) -> None:
             logger.error(
                 f"Failed to post task {task.id} to {assignee_id}: {e}"
             )
-            print(
-                f"{Fore.RED}Failed to post task {task.id} to {assignee_id}: "
-                f"{e}{Fore.RESET}"
-            )
+            for cb in self._callbacks:
+                cb.log_message(
+                    LogEvent(
+                        message=(
+                            f"Failed to post task {task.id} to {assignee_id}: "
+                            f"{e}"
+                        ),
+                        level="error",
+                        color="red",
+                    )
+                )
 
     async def _post_dependency(self, dependency: Task) -> None:
         await self._channel.post_dependency(dependency, self.node_id)
@@ -4010,7 +4018,12 @@ async def _create_worker_node_for_task(self, task: Task) -> Worker:
         )
         new_node.set_channel(self._channel)
 
-        print(f"{Fore.CYAN}{new_node} created.{Fore.RESET}")
+        for cb in self._callbacks:
+            cb.log_message(
+                LogEvent(
+                    message=f"{new_node} created.", level="info", color="cyan"
+                )
+            )
 
         self._children.append(new_node)
 
@@ -4657,10 +4670,17 @@ async def _handle_completed_task(self, task: Task) -> None:
                 tasks_list.pop(i)
                 self._pending_tasks = deque(tasks_list)
                 found_and_removed = True
-                print(
-                    f"{Fore.GREEN}✅ Task {task.id} completed and removed "
-                    f"from queue.{Fore.RESET}"
-                )
+                for cb in self._callbacks:
+                    cb.log_message(
+                        LogEvent(
+                            message=(
+                                f"✅ Task {task.id} completed and removed "
+                                f"from queue."
+                            ),
+                            level="info",
+                            color="green",
+                        )
+                    )
                 break
 
         if not found_and_removed:

camel/societies/workforce/workforce_callback.py

@@ -13,10 +13,14 @@
 # ========= Copyright 2023-2025 @ CAMEL-AI.org. All Rights Reserved. =========
 from __future__ import annotations
 
+import typing
 from abc import ABC, abstractmethod
 
+from colorama import Fore
+
 from .events import (
     AllTasksCompletedEvent,
+    LogEvent,
     TaskAssignedEvent,
     TaskCompletedEvent,
     TaskCreatedEvent,
@@ -34,6 +38,30 @@ class WorkforceCallback(ABC):
     Implementations should persist or stream events as appropriate.
     """
 
+    __COLOR_MAP: typing.ClassVar = {
+        "yellow": Fore.YELLOW,
+        "red": Fore.RED,
+        "green": Fore.GREEN,
+        "cyan": Fore.CYAN,
+        "magenta": Fore.MAGENTA,
+        "gray": Fore.LIGHTBLACK_EX,
+        "black": Fore.BLACK,
+    }
+
+    def _get_color_message(self, event: LogEvent) -> str:
+        r"""Gets a colored message for a log event."""
+        if event.color is None or event.color not in self.__COLOR_MAP:
+            return event.message
+        color = self.__COLOR_MAP.get(event.color)
+        return f"{color}{event.message}{Fore.RESET}"
+
+    @abstractmethod
+    def log_message(
+        self,
+        event: LogEvent,
+    ) -> None:
+        pass
+
     @abstractmethod
     def log_task_created(
         self,

camel/societies/workforce/workforce_logger.py

@@ -18,6 +18,7 @@
 from camel.logger import get_logger
 from camel.societies.workforce.events import (
     AllTasksCompletedEvent,
+    LogEvent,
     QueueStatusEvent,
     TaskAssignedEvent,
     TaskCompletedEvent,
@@ -50,6 +51,11 @@ def __init__(self, workforce_id: str):
         self._worker_information: Dict[str, Dict[str, Any]] = {}
         self._initial_worker_logs: List[Dict[str, Any]] = []
 
+    def log_message(self, event: LogEvent) -> None:
+        r"""Logs a message to the console with color."""
+        colored_message = self._get_color_message(event)
+        print(colored_message)
+
     def _log_event(self, event_type: str, **kwargs: Any) -> None:
         r"""Internal method to create and store a log entry.
 

examples/workforce/workforce_callbacks_example.py

@@ -27,6 +27,7 @@
 from camel.models import ModelFactory
 from camel.societies.workforce.events import (
     AllTasksCompletedEvent,
+    LogEvent,
     TaskAssignedEvent,
     TaskCompletedEvent,
     TaskCreatedEvent,
@@ -38,6 +39,7 @@
 )
 from camel.societies.workforce.workforce import Workforce
 from camel.societies.workforce.workforce_callback import WorkforceCallback
+from camel.societies.workforce.workforce_logger import WorkforceLogger
 from camel.types import ModelPlatformType, ModelType
 
 logger = get_logger(__name__)
@@ -46,6 +48,12 @@
 class PrintCallback(WorkforceCallback):
     r"""Simple callback printing events to logs to observe ordering."""
 
+    def log_message(self, event: LogEvent) -> None:
+        print(
+            f"[PrintCallback] {event.message} level={event.level}, "
+            f"color={event.color}"
+        )
+
     def log_task_created(self, event: TaskCreatedEvent) -> None:
         print(
             f"[PrintCallback] task_created: id={event.task_id}, "
@@ -115,7 +123,9 @@ def build_student_agent() -> ChatAgent:
 
 
 async def run_demo() -> None:
-    callbacks = [PrintCallback()]
+    logger_cb = WorkforceLogger('demo-logger')
+    print_cb = PrintCallback()
+    callbacks = [logger_cb, print_cb]
 
     workforce = Workforce(
         "Workforce Callbacks Demo",