fix: Logging module always producing JSON

Right now even if we specify that we want TEXT logging output Codegate
will use JSON. This fixes it. It also introduces `structlog` as
a convenient already built formatter instead of making our own
formatters hand-made

Author: aponcedeleonch

pyproject.toml

@@ -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"
 

src/codegate/cli.py

@@ -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/codegate_logging.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,47 @@ 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),
+    ]
+    # 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 +128,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 +137,3 @@ def setup_logging(
             "handlers": ["stdout", "stderr"],
         },
     )
-
-    return logger

src/codegate/config.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/inference/inference_engine.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/pipeline/codegate_system_prompt/codegate.py

@@ -17,8 +17,7 @@ class CodegateSystemPrompt(PipelineStep):
 
     def __init__(self, system_prompt_message: Optional[str] = None):
         self._system_message = ChatCompletionSystemMessage(
-            content=system_prompt_message,
-            role="system"
+            content=system_prompt_message, role="system"
         )
 
     @property
@@ -29,7 +28,7 @@ def name(self) -> str:
         return "codegate-system-prompt"
 
     async def process(
-            self, request: ChatCompletionRequest, context: PipelineContext
+        self, request: ChatCompletionRequest, context: PipelineContext
     ) -> PipelineResult:
         """
         Process the completion request and add a system prompt if the user message contains

src/codegate/pipeline/fim/secret_analyzer.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):