Merge pull request stacklok#129 from stacklok/fix-logging

fix: Logging module always producing JSON

Author: aponcedeleonch

docs/development.md

@@ -51,7 +51,7 @@ codegate/
 │       ├── cli.py           # Command-line interface
 │       ├── config.py        # Configuration management
 │       ├── exceptions.py    # Shared exceptions
-│       ├── logging.py       # Logging setup
+│       ├── codegate_logging.py       # Logging setup
 │       ├── prompts.py       # Prompts management
 │       ├── server.py        # Main server implementation
 │       └── providers/       # External service providers

docs/logging.md

@@ -1,6 +1,6 @@
 # Logging System
 
-The logging system in Codegate (`logging.py`) provides a flexible and structured logging solution with support for both JSON and text formats.
+The logging system in Codegate (`codegate_logging.py`) provides a flexible and structured logging solution with support for both JSON and text formats.
 
 ## Log Routing
 
@@ -18,9 +18,9 @@ When using JSON format (default), log entries include:
 ```json
 {
   "timestamp": "YYYY-MM-DDThh:mm:ss.mmmZ",
-  "level": "LOG_LEVEL",
+  "log_level": "LOG_LEVEL",
   "module": "MODULE_NAME",
-  "message": "Log message",
+  "event": "Log message",
   "extra": {
     // Additional fields as you desire
   }
@@ -49,9 +49,9 @@ YYYY-MM-DDThh:mm:ss.mmmZ - LEVEL - NAME - MESSAGE
 ### Basic Logging
 
 ```python
-import logging
+import structlog
 
-logger = logging.getLogger(__name__)
+logger = structlog.get_logger(__name__)
 
 # Different log levels
 logger.info("This is an info message")

poetry.lock

@@ -16,7 +16,7 @@ weaviate = ">=0.1.2"
 weaviate-client = ">=4.9.3"
 torch = ">=2.5.1"
 transformers = ">=4.46.3"
-
+structlog = ">=24.4.0"
 litellm = "^1.52.16"
 llama_cpp_python = ">=0.3.2"
 

pyproject.toml

@@ -5,6 +5,7 @@
 from typing import Dict, Optional
 
 import click
+import structlog
 
 from codegate.codegate_logging import LogFormat, LogLevel, setup_logging
 from codegate.config import Config, ConfigurationError
@@ -140,7 +141,8 @@ def serve(
             cli_provider_urls=cli_provider_urls,
         )
 
-        logger = setup_logging(cfg.log_level, cfg.log_format)
+        setup_logging(cfg.log_level, cfg.log_format)
+        logger = structlog.get_logger("codegate")
         logger.info(
             "Starting server",
             extra={

src/codegate/cli.py

@@ -1,9 +1,9 @@
-import datetime
-import json
 import logging
 import sys
 from enum import Enum
-from typing import Any, Optional
+from typing import Optional
+
+import structlog
 
 
 class LogLevel(str, Enum):
@@ -46,123 +46,6 @@ def _missing_(cls, value: str) -> Optional["LogFormat"]:
             )
 
 
-class JSONFormatter(logging.Formatter):
-    """Custom formatter that outputs log records as JSON."""
-
-    def __init__(self) -> None:
-        """Initialize the JSON formatter."""
-        super().__init__()
-        self.default_time_format = "%Y-%m-%dT%H:%M:%S"
-        self.default_msec_format = "%s.%03dZ"
-
-    def format(self, record: logging.LogRecord) -> str:
-        """Format the log record as a JSON string.
-
-        Args:
-            record: The log record to format
-
-        Returns:
-            str: JSON formatted log entry
-        """
-        # Create the base log entry
-        log_entry: dict[str, Any] = {
-            "timestamp": self.formatTime(record, self.default_time_format),
-            "level": record.levelname,
-            "module": record.module,
-            "message": record.getMessage(),
-            "extra": {},
-        }
-
-        # Add extra fields from the record
-        extra_attrs = {}
-        for key, value in record.__dict__.items():
-            if key not in {
-                "args",
-                "asctime",
-                "created",
-                "exc_info",
-                "exc_text",
-                "filename",
-                "funcName",
-                "levelname",
-                "levelno",
-                "lineno",
-                "module",
-                "msecs",
-                "msg",
-                "name",
-                "pathname",
-                "process",
-                "processName",
-                "relativeCreated",
-                "stack_info",
-                "thread",
-                "threadName",
-                "extra",
-            }:
-                extra_attrs[key] = value
-
-        # Handle the explicit extra parameter if present
-        if hasattr(record, "extra"):
-            try:
-                if isinstance(record.extra, dict):
-                    extra_attrs.update(record.extra)
-            except Exception:
-                extra_attrs["unserializable_extra"] = str(record.extra)
-
-        # Add all extra attributes to the log entry
-        if extra_attrs:
-            try:
-                json.dumps(extra_attrs)  # Test if serializable
-                log_entry["extra"] = extra_attrs
-            except (TypeError, ValueError):
-                # If serialization fails, convert values to strings
-                serializable_extra = {}
-                for key, value in extra_attrs.items():
-                    try:
-                        json.dumps({key: value})  # Test individual value
-                        serializable_extra[key] = value
-                    except (TypeError, ValueError):
-                        serializable_extra[key] = str(value)
-                log_entry["extra"] = serializable_extra
-
-        # Handle exception info if present
-        if record.exc_info:
-            log_entry["extra"]["exception"] = self.formatException(record.exc_info)
-
-        # Handle stack info if present
-        if record.stack_info:
-            log_entry["extra"]["stack_info"] = self.formatStack(record.stack_info)
-
-        return json.dumps(log_entry)
-
-
-class TextFormatter(logging.Formatter):
-    """Standard text formatter with consistent timestamp format."""
-
-    def __init__(self) -> None:
-        """Initialize the text formatter."""
-        super().__init__(
-            fmt="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
-            datefmt="%Y-%m-%dT%H:%M:%S.%03dZ",
-        )
-
-    def formatTime(  # noqa: N802
-        self, record: logging.LogRecord, datefmt: Optional[str] = None
-    ) -> str:
-        """Format the time with millisecond precision.
-
-        Args:
-            record: The log record
-            datefmt: The date format string (ignored as we use a fixed format)
-
-        Returns:
-            str: Formatted timestamp
-        """
-        ct = datetime.datetime.fromtimestamp(record.created, datetime.UTC)
-        return ct.strftime(self.datefmt)
-
-
 def setup_logging(
     log_level: Optional[LogLevel] = None, log_format: Optional[LogFormat] = None
 ) -> logging.Logger:
@@ -181,10 +64,52 @@ def setup_logging(
     if log_format is None:
         log_format = LogFormat.JSON
 
-    # Create formatters
-    json_formatter = JSONFormatter()
-    text_formatter = TextFormatter()
-    formatter = json_formatter if log_format == LogFormat.JSON else text_formatter
+    # The configuration was taken from structlog documentation
+    # https://www.structlog.org/en/stable/standard-library.html
+    # Specifically the section "Rendering Using structlog-based Formatters Within logging"
+
+    # Adds log level and timestamp to log entries
+    shared_processors = [
+        structlog.processors.add_log_level,
+        structlog.processors.TimeStamper(fmt="%Y-%m-%dT%H:%M:%S.%03dZ", utc=True),
+        structlog.processors.CallsiteParameterAdder(
+            [
+                structlog.processors.CallsiteParameter.MODULE,
+            ]
+        ),
+    ]
+    # Not sure why this is needed. I think it is a wrapper for the standard logging module.
+    # Should allow to log both with structlog and the standard logging module:
+    # import logging
+    # import structlog
+    # logging.getLogger("stdlog").info("woo")
+    # structlog.get_logger("structlog").info("amazing", events="oh yes")
+    structlog.configure(
+        processors=shared_processors
+        + [
+            # Prepare event dict for `ProcessorFormatter`.
+            structlog.stdlib.ProcessorFormatter.wrap_for_formatter,
+        ],
+        logger_factory=structlog.stdlib.LoggerFactory(),
+        cache_logger_on_first_use=True,
+    )
+
+    # The config aboves adds the following keys to all log entries: _record & _from_structlog.
+    # remove_processors_meta removes them.
+    processors = shared_processors + [structlog.stdlib.ProcessorFormatter.remove_processors_meta]
+    # Choose the processors based on the log format
+    if log_format == LogFormat.JSON:
+        processors = processors + [
+            structlog.processors.dict_tracebacks,
+            structlog.processors.JSONRenderer(),
+        ]
+    else:
+        processors = processors + [structlog.dev.ConsoleRenderer()]
+    formatter = structlog.stdlib.ProcessorFormatter(
+        # foreign_pre_chain run ONLY on `logging` entries that do NOT originate within structlog.
+        foreign_pre_chain=shared_processors,
+        processors=processors,
+    )
 
     # Create handlers for stdout and stderr
     stdout_handler = logging.StreamHandler(sys.stdout)
@@ -208,7 +133,7 @@ def setup_logging(
     root_logger.addHandler(stderr_handler)
 
     # Create a logger for our package
-    logger = logging.getLogger("codegate")
+    logger = structlog.get_logger("codegate")
     logger.debug(
         "Logging initialized",
         extra={
@@ -217,5 +142,3 @@ def setup_logging(
             "handlers": ["stdout", "stderr"],
         },
     )
-
-    return logger

src/codegate/codegate_logging.py

@@ -5,13 +5,14 @@
 from pathlib import Path
 from typing import Dict, Optional, Union
 
+import structlog
 import yaml
 
-from codegate.codegate_logging import LogFormat, LogLevel, setup_logging
+from codegate.codegate_logging import LogFormat, LogLevel
 from codegate.exceptions import ConfigurationError
 from codegate.prompts import PromptConfig
 
-logger = setup_logging()
+logger = structlog.get_logger("codegate")
 
 # Default provider URLs
 DEFAULT_PROVIDER_URLS = {

src/codegate/config.py

@@ -1,6 +1,7 @@
+import structlog
 from llama_cpp import Llama
 
-from codegate.codegate_logging import setup_logging
+logger = structlog.get_logger("codegate")
 
 
 class LlamaCppInferenceEngine:
@@ -21,7 +22,6 @@ def __new__(cls):
     def __init__(self):
         if not hasattr(self, "models"):
             self.__models = {}
-            self.__logger = setup_logging()
 
     def __del__(self):
         self.__close_models()
@@ -32,7 +32,7 @@ async def __get_model(self, model_path, embedding=False, n_ctx=512, n_gpu_layers
         is loaded and added to __models and returned.
         """
         if model_path not in self.__models:
-            self.__logger.info(
+            logger.info(
                 f"Loading model from {model_path} with parameters "
                 f"n_gpu_layers={n_gpu_layers} and n_ctx={n_ctx}"
             )

src/codegate/inference/inference_engine.py

@@ -1,9 +1,9 @@
+import structlog
 from litellm import ChatCompletionRequest
 
-from codegate.codegate_logging import setup_logging
 from codegate.pipeline.base import PipelineContext, PipelineResponse, PipelineResult, PipelineStep
 
-logger = setup_logging()
+logger = structlog.get_logger("codegate")
 
 
 class SecretAnalyzer(PipelineStep):