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/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):

src/codegate/pipeline/secrets/secrets.py

@@ -1,8 +1,8 @@
 import re
 
+import structlog
 from litellm import ChatCompletionRequest
 
-from codegate.codegate_logging import setup_logging
 from codegate.pipeline.base import (
     PipelineContext,
     PipelineResult,
@@ -11,6 +11,8 @@
 from codegate.pipeline.secrets.gatecrypto import CodeGateCrypto
 from codegate.pipeline.secrets.signatures import CodegateSignatures
 
+logger = structlog.get_logger("codegate")
+
 
 class CodegateSecrets(PipelineStep):
     """Pipeline step that handles secret information requests."""
@@ -21,7 +23,6 @@ def __init__(self):
         self.crypto = CodeGateCrypto()
         self._session_store = {}
         self._encrypted_to_session = {}  # Reverse lookup index
-        self.__logger = setup_logging()
 
     @property
     def name(self) -> str:
@@ -87,7 +88,7 @@ def _redeact_text(self, text: str) -> str:
         if not matches:
             return text
 
-        self.__logger.debug(f"Found {len(matches)} secrets in the user message")
+        logger.debug(f"Found {len(matches)} secrets in the user message")
 
         # Convert line positions to absolute positions and extend boundaries
         absolute_matches = []
@@ -133,7 +134,7 @@ def _redeact_text(self, text: str) -> str:
             self._encrypted_to_session[encrypted_value] = session_id
 
             # Print the session store
-            self.__logger.info(f"Session store: {self._session_store}")
+            logger.info(f"Session store: {self._session_store}")
 
             # Create the replacement string
             replacement = f"REDACTED<${encrypted_value}>"
@@ -155,12 +156,12 @@ def _redeact_text(self, text: str) -> str:
         protected_string = "".join(protected_text)
 
         # Log the findings
-        self.__logger.info("\nFound secrets:")
+        logger.info("\nFound secrets:")
         for secret in found_secrets:
-            self.__logger.info(f"\nService: {secret['service']}")
-            self.__logger.info(f"Type: {secret['type']}")
-            self.__logger.info(f"Original: {secret['original']}")
-            self.__logger.info(f"Encrypted: REDACTED<${secret['encrypted']}>")
+            logger.info(f"\nService: {secret['service']}")
+            logger.info(f"Type: {secret['type']}")
+            logger.info(f"Original: {secret['original']}")
+            logger.info(f"Encrypted: REDACTED<${secret['encrypted']}>")
 
         (f"\nProtected text:\n{protected_string}")
         return protected_string
@@ -181,7 +182,7 @@ def _get_original_value(self, encrypted_value: str) -> str:
             if session_id:
                 return self._session_store[session_id]["original"]
         except Exception as e:
-            self.__logger.error(f"Error looking up original value: {e}")
+            logger.error(f"Error looking up original value: {e}")
         return encrypted_value
 
     def get_by_session_id(self, session_id: str) -> dict | None:
@@ -197,7 +198,7 @@ def get_by_session_id(self, session_id: str) -> dict | None:
         try:
             return self._session_store.get(session_id)
         except Exception as e:
-            self.__logger.error(f"Error looking up by session ID: {e}")
+            logger.error(f"Error looking up by session ID: {e}")
             return None
 
     def _cleanup_session_store(self):
@@ -215,9 +216,9 @@ def _cleanup_session_store(self):
             self._session_store.clear()
             self._encrypted_to_session.clear()
 
-            self.__logger.info("Session stores securely wiped")
+            logger.info("Session stores securely wiped")
         except Exception as e:
-            self.__logger.error(f"Error during secure cleanup: {e}")
+            logger.error(f"Error during secure cleanup: {e}")
 
     def _unredact_text(self, protected_text: str) -> str:
         """
@@ -253,7 +254,7 @@ def _unredact_text(self, protected_text: str) -> str:
 
         # Join all parts together
         unprotected_text = "".join(result)
-        self.__logger.info(f"\nUnprotected text:\n{unprotected_text}")
+        logger.info(f"\nUnprotected text:\n{unprotected_text}")
         return unprotected_text
 
     async def process(
@@ -293,7 +294,7 @@ async def process(
 
             return PipelineResult(request=request)
         except Exception as e:
-            self.__logger.error(f"CodegateSecrets operation failed: {e}")
+            logger.error(f"CodegateSecrets operation failed: {e}")
 
         finally:
             # Clean up sensitive data