feat: implemented extra log level for LLM token stream

.logging_config.yaml

@@ -3,6 +3,7 @@ version: 1
 custom_levels:
   OBSERVATION: 25
   ACTION: 26
+  LLM: 9
 
 formatters:
   # Console output
@@ -26,7 +27,7 @@ handlers:
     stream: ext://sys.stdout
   ui:
     class: zmq.log.handlers.PUBHandler
-    level: DEBUG
+    level: LLM
     formatter: json_experiment
 
 # Level of external libraries
@@ -36,5 +37,5 @@ root:
 
 loggers:
   control_backend:
-    level: DEBUG
+    level: LLM
     handlers: [ui]

src/control_backend/agents/llm/llm_agent.py

@@ -125,7 +125,7 @@ class LLMAgent(BaseAgent):
                 full_message += token
                 current_chunk += token
 
-                self.logger.info(
+                self.logger.llm(
                     "Received token: %s",
                     full_message,
                     extra={"reference": message_id},  # Used in the UI to update old logs

src/control_backend/logging/setup_logging.py

@@ -4,6 +4,7 @@ import os
 
 import yaml
 import zmq
+from zmq.log.handlers import PUBHandler
 
 from control_backend.core.config import settings
 
@@ -51,15 +52,27 @@ def setup_logging(path: str = ".logging_config.yaml") -> None:
                 logging.warning(f"Could not load logging configuration: {e}")
                 config = {}
 
-            if "custom_levels" in config:
+            custom_levels = config.get("custom_levels", {}) or {}
-                for level_name, level_num in config["custom_levels"].items():
+            for level_name, level_num in custom_levels.items():
-                    add_logging_level(level_name, level_num)
+                add_logging_level(level_name, level_num)
 
             if config.get("handlers") is not None and config.get("handlers").get("ui"):
                 pub_socket = zmq.Context.instance().socket(zmq.PUB)
                 pub_socket.connect(settings.zmq_settings.internal_pub_address)
                 config["handlers"]["ui"]["interface_or_socket"] = pub_socket
+
             logging.config.dictConfig(config)
 
+            # Patch ZMQ PUBHandler to know about custom levels
+            if custom_levels:
+                for logger_name in ("control_backend",):
+                    logger = logging.getLogger(logger_name)
+                    for handler in logger.handlers:
+                        if isinstance(handler, PUBHandler):
+                            # Use the INFO formatter as the default template
+                            default_fmt = handler.formatters[logging.INFO]
+                            for level_num in custom_levels.values():
+                                handler.setFormatter(default_fmt, level=level_num)
+
     else:
         logging.warning("Logging config file not found. Using default logging configuration.")

test/unit/agents/llm/test_llm_agent.py

@@ -49,6 +49,9 @@ async def test_llm_processing_success(mock_httpx_client, mock_settings):
     agent = LLMAgent("llm_agent")
     agent.send = AsyncMock()  # Mock the send method to verify replies
 
+    mock_logger = MagicMock()
+    agent.logger = mock_logger
+
     # Simulate receiving a message from BDI
     prompt = LLMPromptMessage(text="Hi", norms=[], goals=[])
     msg = InternalMessage(