diff --git a/src/mlops_nlp/logging_config.py b/src/mlops_nlp/logging_config.py index af7829d..19be250 100644 --- a/src/mlops_nlp/logging_config.py +++ b/src/mlops_nlp/logging_config.py @@ -1,12 +1,51 @@ +import json import logging -from typing import Optional +from datetime import datetime, timezone +from typing import Any, Optional + + +class JSONFormatter(logging.Formatter): + """ + Custom JSON formatter for structured logging. + """ + def format(self, record: logging.LogRecord) -> str: + log_record: dict[str, Any] = { + "timestamp": datetime.fromtimestamp(record.created, tz=timezone.utc).isoformat(), + "level": record.levelname, + "name": record.name, + "message": record.getMessage(), + } + + # Include extra fields passed via 'extra' parameter + # Standard LogRecord attributes to ignore + standard_attrs = { + "name", "msg", "args", "levelname", "levelno", "pathname", "filename", + "module", "exc_info", "exc_text", "stack_info", "lineno", "funcName", + "created", "msecs", "relativeCreated", "thread", "threadName", + "processName", "process", "message" + } + for key, value in record.__dict__.items(): + if key not in standard_attrs: + log_record[key] = value + + if record.exc_info: + log_record["exception"] = self.formatException(record.exc_info) + + return json.dumps(log_record) def configure_logging(level: str = "INFO") -> None: - logging.basicConfig( - level=getattr(logging, level.upper(), logging.INFO), - format="%(asctime)s | %(levelname)s | %(name)s | %(message)s", - ) + handler = logging.StreamHandler() + handler.setFormatter(JSONFormatter()) + + root_logger = logging.getLogger() + root_logger.setLevel(getattr(logging, level.upper(), logging.INFO)) + + # Remove existing handlers to avoid duplicate logs + for h in root_logger.handlers[:]: + root_logger.removeHandler(h) + + root_logger.addHandler(handler) def get_logger(name: Optional[str] = None) -> logging.Logger: diff --git a/tests/verify_logging.py b/tests/verify_logging.py new file mode 100644 index 0000000..fb83bd4 --- /dev/null +++ b/tests/verify_logging.py @@ -0,0 +1,19 @@ +from mlops_nlp.logging_config import configure_logging, get_logger +import json + +def test_logging(): + configure_logging("INFO") + logger = get_logger("test_logger") + + # We'll use a string buffer or just check the output manually + # For automated check, we can use capsys in pytest, but here we'll just run it. + logger.info("This is a test message") + logger.info("Message with extra context", extra={"request_id": "12345", "user_id": "abc"}) + + try: + 1/0 + except Exception: + logger.exception("An error occurred") + +if __name__ == "__main__": + test_logging()