diff --git a/agentic_security/lib.py b/agentic_security/lib.py index 67ddd32..21880e4 100644 --- a/agentic_security/lib.py +++ b/agentic_security/lib.py @@ -4,12 +4,12 @@ from datetime import datetime import colorama import tqdm.asyncio -from loguru import logger from rich.console import Console from rich.table import Table from tabulate import tabulate from agentic_security.config import SettingsMixin # Importing the configuration mixin +from agentic_security.logutils import logger from agentic_security.primitives import Scan from agentic_security.probe_data import REGISTRY from agentic_security.routes.scan import streaming_response_generator diff --git a/agentic_security/logutils.py b/agentic_security/logutils.py index 91d7ff1..f51d680 100644 --- a/agentic_security/logutils.py +++ b/agentic_security/logutils.py @@ -1,45 +1,22 @@ -import logging -from os import getenv +import sys -from rich.logging import RichHandler +from loguru import logger -LOGGER_NAME = "root" +# Define custom colors +BLUE = "#89CFF0" +BROWN = "#8B4513" # Brown for DEBUG +# Define custom log level colors +logger.level("DEBUG", color=f"") +logger.level("INFO", color=f"") -def get_logger(logger_name: str) -> logging.Logger: - # https://rich.readthedocs.io/en/latest/reference/logging.html#rich.logging.RichHandler - # https://rich.readthedocs.io/en/latest/logging.html#handle-exceptions - rich_handler = RichHandler( - show_time=False, - rich_tracebacks=False, - show_path=True if getenv("API_RUNTIME") == "dev" else False, - tracebacks_show_locals=False, - ) - rich_handler.setFormatter( - logging.Formatter( - fmt="%(message)s", - datefmt="[%X]", - ) - ) +# Define custom log format with aligned messages and colored levels +LOG_FORMAT = ( + "{level:<8} " # Properly formatted and colored log level + "{message:<100} " # Left-aligned message for readability + "{file.name}:{line}" # File name and line number in cyan +) - _logger = logging.getLogger(logger_name) - _logger.addHandler(rich_handler) - _logger.setLevel(logging.INFO) - _logger.propagate = True - return _logger - - -logger: logging.Logger = get_logger(LOGGER_NAME) - - -def set_log_level_to_debug(): - _logger = logging.getLogger(LOGGER_NAME) - _logger.setLevel(logging.DEBUG) - - -def set_log_level_to_info(): - _logger = logging.getLogger(LOGGER_NAME) - _logger.setLevel(logging.INFO) - - -set_log_level_to_debug() +# Remove default handlers and add a new one with custom formatting +logger.remove() +logger.add(sys.stdout, format=LOG_FORMAT, level="DEBUG", colorize=True) diff --git a/agentic_security/misc/banner.py b/agentic_security/misc/banner.py index f758c2f..dd41dac 100644 --- a/agentic_security/misc/banner.py +++ b/agentic_security/misc/banner.py @@ -80,6 +80,7 @@ def generate_banner( def init_banner(): + return ver = version("agentic_security") try: print(generate_banner(version=ver))