Skip to content

azad.logging_formatter Module

azad.logging_formatter

Logging formatters for structured logging output.

Classes

CustomJsonFormatter

CustomJsonFormatter(*args, **kwargs)

Bases: Formatter

Custom JSON formatter with pretty console output for specific events.

Initialize the formatter with custom options.

Source code in azad/logging_formatter.py
def __init__(self, *args, **kwargs):
    """Initialize the formatter with custom options."""
    super().__init__(*args, **kwargs)
    self.width = 80  # Terminal width for formatting
    self.separator_top = "═" * self.width
    self.separator_bottom = "─" * self.width
    self.default_time_format = '%Y-%m-%d %H:%M:%S'
    self.default_msec_format = '%s.%03d'
Attributes
separator_top instance-attribute
separator_top = '═' * width
separator_bottom instance-attribute
separator_bottom = '─' * width
default_time_format instance-attribute
default_time_format = '%Y-%m-%d %H:%M:%S'
default_msec_format instance-attribute
default_msec_format = '%s.%03d'
Functions
format
format(record: LogRecord) -> str

Format the log record with JSON structure.

Source code in azad/logging_formatter.py
def format(self, record: logging.LogRecord) -> str:
    """Format the log record with JSON structure."""
    # Handle error logs specially
    if record.levelno == logging.ERROR:
        return self._format_error(record)

    # For non-error logs, proceed with normal formatting
    message = record.getMessage()
    # Only truncate messages in PROD mode
    if record.levelno == LogLevel.PROD and len(message) > 150:
        message = self._truncate_text(message)

    # Base log data
    log_data = {
        "level": record.levelname,
        "message": message
    }

    # Add filename and line number for DEBUG level
    if record.levelno == logging.DEBUG:
        log_data["filename"] = record.filename
        log_data["lineno"] = str(record.lineno)

    # Add all extra attributes from the record
    extra_attrs = {
        key: value 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"] and key not in log_data
    }
    if extra_attrs:
        log_data.update(extra_attrs)

    # For PROD level, apply special formatting
    if record.levelno == LogLevel.PROD:
        if hasattr(record, 'event'):
            event = getattr(record, 'event')

            if event == "task":
                formatted = self._format_task(message, record)
                if formatted:  # Only use special formatting if we got a result
                    return formatted
            elif event == "initialization":
                formatted = self._format_initialization(record)
                if formatted:  # Only use special formatting if we got a result
                    return formatted
            elif event == "tool_request":
                params = getattr(record, 'parameters', None)
                tool_name = getattr(record, 'tool_name', None)
                return self._format_tool_request(message, params, tool_name)
            elif event == "tool_response":
                response_data = getattr(record, 'response_data', {})  # Default to empty dict instead of empty string
                return self._format_tool_response(response_data)
            elif event == "separator":
                return record.getMessage()

        # If no special formatting applied, return JSON for PROD
        return json.dumps(log_data, default=str)

    # Remove taskName from extra_attrs for all log levels
    if 'taskName' in extra_attrs:
        del extra_attrs['taskName']

    # Set color based on log level
    if record.levelno == logging.DEBUG:
        level_color = Colors.GREEN
    elif record.levelno == logging.INFO:
        level_color = Colors.CYAN
    elif record.levelno == logging.WARNING:
        level_color = Colors.YELLOW
    elif record.levelno == logging.ERROR:
        level_color = Colors.RED
    else:
        level_color = Colors.RESET

    level_name = f"{level_color}{record.levelname}{Colors.RESET}".ljust(8)

    # Extract standard keys
    standard_keys = ['event', 'component', 'type', 'status', 'model']
    standard_attrs = []

    # Add colored values to extra_attrs for all levels
    if hasattr(record, 'event'):
        extra_attrs['event'] = f"{Colors.CYAN}{getattr(record, 'event')}{Colors.RESET}"
    if hasattr(record, 'component'):
        extra_attrs['component'] = f"{Colors.YELLOW}{getattr(record, 'component')}{Colors.RESET}"
    if hasattr(record, 'type'):
        extra_attrs['type'] = f"{Colors.GREEN}{getattr(record, 'type')}{Colors.RESET}"
    if hasattr(record, 'status'):
        extra_attrs['status'] = f"{Colors.BLUE}{getattr(record, 'status')}{Colors.RESET}"
    if hasattr(record, 'model'):
        extra_attrs['model'] = f"{Colors.CYAN}{getattr(record, 'model')}{Colors.RESET}"

    # Special handling for AI Network setting model message
    if record.name == "azad.ai_network" and message.startswith("SETTING MODEL NOW:"):
        model_name = message.split(": ")[1]
        output = [f"{Colors.BLUE}*{Colors.RESET} {level_name} Setting model={model_name}"]
    else:
        # Start with basic log line without standard attributes in brackets
        output = [f"{Colors.BLUE}*{Colors.RESET} {level_name} {message}"]

    # Remove taskName from any remaining output parts for all log levels
    output = [line for line in output if 'taskName=' not in line]

    # Format common keys for the [] block
    common_keys = ['event', 'status', 'component']
    common_parts = []
    remaining_attrs = {}

    for key in common_keys:
        if key in extra_attrs:
            value = extra_attrs[key]
            if isinstance(value, str) and any(color in str(value) for color in [Colors.CYAN, Colors.YELLOW, Colors.GREEN, Colors.BLUE]):
                common_parts.append(f"{key}={value}")
            else:
                common_parts.append(f"{key}={value}")
            del extra_attrs[key]

    # Add remaining attributes to remaining_attrs
    remaining_attrs.update(extra_attrs)

    # Add the common keys block to the output line
    if common_parts:
        output[0] = f"{output[0]} [{' '.join(common_parts)}]"

    # Format remaining extras
    if remaining_attrs:
        # Try to format as one line first
        try:
            inline_parts = []
            for key, value in remaining_attrs.items():
                if isinstance(value, str) and any(color in value for color in [Colors.CYAN, Colors.YELLOW, Colors.GREEN, Colors.BLUE]):
                    inline_parts.append(f"{key}={value}")
                else:
                    inline_parts.append(f"{key}={json.dumps(value, default=str)}")

            inline_str = " ".join(inline_parts)
            # Calculate visible length without color codes
            visible_length = len(inline_str)
            for color in [Colors.CYAN, Colors.YELLOW, Colors.GREEN, Colors.BLUE, Colors.RESET]:
                visible_length -= inline_str.count(color) * len(color)

            if visible_length <= 100 and '\n' not in str(remaining_attrs):
                # If it fits on one line, add it indented
                output.append(f"        {inline_str}")
            else:
                # Format as pretty-printed
                extra_lines = []
                for key, value in remaining_attrs.items():
                    if isinstance(value, (dict, list)):
                        formatted_value = self._format_json_value(value, depth=2)
                        if len(formatted_value) > 1000:
                            # Summarize long values
                            formatted_value = f"{formatted_value[:1000]}... (truncated, total length: {len(formatted_value)})"
                        extra_lines.append(f"        {key}={formatted_value}")
                    elif isinstance(value, str) and '\n' in value:
                        formatted_value = value.replace('\n', '\n            ')
                        extra_lines.append(f"        {key}={formatted_value}")
                    elif isinstance(value, str) and any(color in value for color in [Colors.CYAN, Colors.YELLOW, Colors.GREEN, Colors.BLUE]):
                        extra_lines.append(f"        {key}={value}")
                    else:
                        extra_lines.append(f"        {key}={json.dumps(value, default=str)}")
                output.extend(extra_lines)
        except Exception:
            # If JSON serialization fails, fall back to multi-line format
            extra_lines = []
            for key, value in extra_attrs.items():
                if isinstance(value, (dict, list)):
                    formatted_value = self._format_json_value(value, depth=2)
                    extra_lines.append(f"    {key}: {formatted_value}")
                elif isinstance(value, str) and '\n' in value:
                    formatted_value = value.replace('\n', '\n    ')
                    extra_lines.append(f"    {key}: {formatted_value}")
                else:
                    extra_lines.append(f"    {key}: {json.dumps(value, default=str)}")
            if extra_lines:
                output.append("  extra:")
                output.extend(extra_lines)

    return "\n".join(output)

Functions

setup_logging

setup_logging(level: str | int = PROD)

Configure logging with custom formatter and levels.

Source code in azad/logging_formatter.py
def setup_logging(level: str | int = LogLevel.PROD):
    """Configure logging with custom formatter and levels."""
    # Set up custom log levels
    setup_log_levels()

    # Convert string level to numeric level if needed
    if isinstance(level, str):
        level = getattr(logging, level.upper())

    # Create console handler with custom formatter
    console_handler = logging.StreamHandler(sys.stdout)
    formatter = CustomJsonFormatter(
        fmt="%(levelname)s %(message)s"
    )
    console_handler.setFormatter(formatter)

    # Configure root logger
    root_logger = logging.getLogger()
    for handler in root_logger.handlers[:]:
        root_logger.removeHandler(handler)
    root_logger.addHandler(console_handler)
    root_logger.setLevel(level)

    # Configure websockets loggers to use appropriate level
    websockets_loggers = [
        'websockets.server',
        'websockets.protocol',
        'websockets.client'
    ]

    for logger_name in websockets_loggers:
        logger = logging.getLogger(logger_name)
        # Just set the level, let websockets handle its own logging
        logger.setLevel(level)