diff --git a/.logging_config.yaml b/.logging_config.yaml index a244917..4af5d56 100644 --- a/.logging_config.yaml +++ b/.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] diff --git a/src/control_backend/agents/llm/llm_agent.py b/src/control_backend/agents/llm/llm_agent.py index 0263b30..55099e2 100644 --- a/src/control_backend/agents/llm/llm_agent.py +++ b/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 diff --git a/src/control_backend/logging/setup_logging.py b/src/control_backend/logging/setup_logging.py index fe40738..05ae85a 100644 --- a/src/control_backend/logging/setup_logging.py +++ b/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: - for level_name, level_num in config["custom_levels"].items(): - add_logging_level(level_name, level_num) + custom_levels = config.get("custom_levels", {}) or {} + for level_name, level_num in custom_levels.items(): + 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.") diff --git a/test/unit/agents/llm/test_llm_agent.py b/test/unit/agents/llm/test_llm_agent.py index e2b6460..62c189e 100644 --- a/test/unit/agents/llm/test_llm_agent.py +++ b/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(