src/utils/logger.py

# Structured logging for RAG and database operations
import hashlib
import json
import logging
import random
import re
import sys
import threading
import uuid
from datetime import datetime
from typing import Any, Dict, List, Optional

LOG_SAMPLING_RATE = 1.0
LOG_SLOW_QUERY_THRESHOLD_MS = 5000


class StructuredFormatter(logging.Formatter):
    """Formatter that outputs JSON structured logs."""
    
    def __init__(self, include_timestamp: bool = True):
        """Initialize formatter with configuration."""
        super().__init__()
        self._include_timestamp = include_timestamp
    
    def format(self, record: logging.LogRecord) -> str:
        """Format log record as JSON string."""
        log_data = {
            "level": record.levelname,
            "message": record.getMessage(),
            "logger": record.name,
            "module": record.module,
            "function": record.funcName,
            "line": record.lineno,
        }
        
        if self._include_timestamp:
            log_data["timestamp"] = datetime.utcnow().isoformat()
        
        if hasattr(record, "query"):
            log_data["query"] = record.query
        
        if hasattr(record, "execution_time"):
            log_data["execution_time"] = record.execution_time
        
        if hasattr(record, "extra_data"):
            log_data["data"] = record.extra_data
        
        if record.exc_info:
            log_data["exception"] = self.formatException(record.exc_info)
        
        return json.dumps(log_data, ensure_ascii=False, default=str)


class SimpleFormatter(logging.Formatter):
    """Simple formatter for console output."""
    
    DEFAULT_FORMAT = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
    
    def __init__(self, fmt: Optional[str] = None):
        """Initialize with optional format string."""
        super().__init__(fmt or self.DEFAULT_FORMAT)


def setup_logger(
    name: str,
    level: int = logging.INFO,
    structured: bool = False,
    log_file: Optional[str] = None,
) -> logging.Logger:
    """Set up and configure a logger instance."""
    logger = logging.getLogger(name)
    logger.setLevel(level)
    
    if logger.handlers:
        return logger
    
    if structured:
        formatter = StructuredFormatter()
    else:
        formatter = SimpleFormatter()
    
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setLevel(level)
    console_handler.setFormatter(formatter)
    logger.addHandler(console_handler)
    
    if log_file:
        file_handler = logging.FileHandler(log_file, encoding="utf-8")
        file_handler.setLevel(level)
        file_handler.setFormatter(StructuredFormatter())
        logger.addHandler(file_handler)
    
    return logger


def get_logger(name: str) -> logging.Logger:
    """Get or create a logger by name."""
    return logging.getLogger(name)


def log_query_execution(
    logger: logging.Logger,
    query: str,
    execution_time: float,
    mode: str,
    success: bool,
    error: Optional[str] = None,
) -> None:
    """Log query execution details."""
    extra = {
        "query": query,
        "execution_time": execution_time,
        "extra_data": {
            "mode": mode,
            "success": success,
            "error": error,
        },
    }
    
    if success:
        logger.info(
            f"Query executed: mode={mode}, time={execution_time:.3f}s",
            extra=extra,
        )
    else:
        logger.error(
            f"Query failed: mode={mode}, error={error}",
            extra=extra,
        )


def log_database_operation(
    logger: logging.Logger,
    operation: str,
    table: Optional[str] = None,
    execution_time: Optional[float] = None,
    row_count: Optional[int] = None,
    success: bool = True,
    error: Optional[str] = None,
) -> None:
    """Log database operation details."""
    extra = {
        "extra_data": {
            "operation": operation,
            "table": table,
            "execution_time": execution_time,
            "row_count": row_count,
            "success": success,
            "error": error,
        },
    }
    
    message_parts = [f"DB operation: {operation}"]
    if table:
        message_parts.append(f"table={table}")
    if execution_time is not None:
        message_parts.append(f"time={execution_time:.3f}s")
    if row_count is not None:
        message_parts.append(f"rows={row_count}")
    
    message = ", ".join(message_parts)
    
    if success:
        logger.info(message, extra=extra)
    else:
        logger.error(f"{message}, error={error}", extra=extra)


def log_agent_step(
    logger: logging.Logger,
    step_name: str,
    step_type: str,
    tool_name: Optional[str] = None,
    execution_time: Optional[float] = None,
    success: bool = True,
    result_summary: Optional[str] = None,
) -> None:
    """Log agent execution step details."""
    extra = {
        "extra_data": {
            "step_name": step_name,
            "step_type": step_type,
            "tool_name": tool_name,
            "execution_time": execution_time,
            "success": success,
        },
    }
    
    message_parts = [f"Agent step: {step_name}", f"type={step_type}"]
    if tool_name:
        message_parts.append(f"tool={tool_name}")
    if execution_time is not None:
        message_parts.append(f"time={execution_time:.3f}s")
    if result_summary:
        message_parts.append(f"result={result_summary}")
    
    message = ", ".join(message_parts)
    
    if success:
        logger.info(message, extra=extra)
    else:
        logger.warning(f"{message} (failed)", extra=extra)


_request_id_local = threading.local()


def generate_request_id() -> str:
    request_id = str(uuid.uuid4())
    _request_id_local.request_id = request_id
    return request_id


def get_request_id() -> Optional[str]:
    return getattr(_request_id_local, "request_id", None)


def clear_request_id() -> None:
    _request_id_local.request_id = None


def _should_sample() -> bool:
    return random.random() < LOG_SAMPLING_RATE


def _build_log_data(operation: str, data: Dict[str, Any]) -> Dict[str, Any]:
    log_data = {"operation": operation, **data}
    request_id = get_request_id()
    if request_id:
        log_data["request_id"] = request_id
    return log_data


def log_query_rewrite(
    logger: logging.Logger,
    original: str,
    rewritten: str,
    method: str,
    time_ms: float,
    cache_hit: bool,
) -> None:
    data = _build_log_data("query_rewrite", {
        "original_query": original[:100],
        "rewritten_query": rewritten[:100],
        "method": method,
        "time_ms": round(time_ms, 2),
        "cache_hit": cache_hit,
        "changed": original != rewritten,
        "original_hash": hashlib.md5(original.encode()).hexdigest()[:8],
        "rewritten_hash": hashlib.md5(rewritten.encode()).hexdigest()[:8],
    })

    if not _should_sample():
        return

    logger.info(
        f"Query rewrite: method={method}, time={time_ms:.1f}ms, cache={cache_hit}, changed={original != rewritten}",
        extra={"extra_data": data},
    )


def log_conversation_summary(
    logger: logging.Logger,
    original_count: int,
    summarized_count: int,
    summary_text: str,
    time_ms: float,
) -> None:
    reduction = round(1 - (summarized_count / original_count), 4) if original_count > 0 else 0.0
    data = _build_log_data("conversation_summary", {
        "messages_before": original_count,
        "messages_after": summarized_count,
        "reduction_ratio": reduction,
        "summary_preview": summary_text[:200],
        "time_ms": round(time_ms, 2),
    })

    if not _should_sample():
        return

    logger.debug(
        f"Conversation summary: {original_count}->{summarized_count} msgs, reduction={reduction:.1%}, time={time_ms:.1f}ms",
        extra={"extra_data": data},
    )


def log_context_filtering(
    logger: logging.Logger,
    intent: str,
    original_length: int,
    filtered_length: int,
    filter_type: str,
) -> None:
    data = _build_log_data("context_filtering", {
        "intent": intent,
        "history_before": original_length,
        "history_after": filtered_length,
        "reduction": original_length - filtered_length,
        "filter_type": filter_type,
    })

    if not _should_sample():
        return

    logger.debug(
        f"Context filter: intent={intent}, {original_length}->{filtered_length} msgs, type={filter_type}",
        extra={"extra_data": data},
    )


def log_sql_generation(
    logger: logging.Logger,
    query: str,
    sql: str,
    schema_size: int,
    chat_history_size: int,
    generation_time_ms: float,
    success: bool = True,
    error: Optional[str] = None,
) -> None:
    data = _build_log_data("sql_generation", {
        "user_query": query[:100],
        "generated_sql": sql[:200],
        "schema_context_bytes": schema_size,
        "chat_history_messages": chat_history_size,
        "generation_time_ms": round(generation_time_ms, 2),
        "sql_hash": hashlib.md5(sql.encode()).hexdigest()[:12] if sql else None,
        "success": success,
        "error": error,
    })

    if success:
        if _should_sample():
            logger.info(
                f"SQL generated: schema={schema_size}B, history={chat_history_size}msgs, time={generation_time_ms:.1f}ms",
                extra={"extra_data": data},
            )
    else:
        logger.error(
            f"SQL generation failed: query={query[:80]}, error={error}",
            extra={"extra_data": data},
        )


def log_database_query_execution(
    logger: logging.Logger,
    sql: str,
    row_count: int,
    execution_time_ms: float,
    success: bool,
    error: Optional[str] = None,
) -> None:
    data = _build_log_data("database_execution", {
        "sql_hash": hashlib.md5(sql.encode()).hexdigest()[:12] if sql else None,
        "sql_preview": sql[:100] if sql else "",
        "row_count": row_count,
        "execution_time_ms": round(execution_time_ms, 2),
        "success": success,
        "error": error if not success else None,
    })

    if not success:
        logger.error(
            f"DB execution failed: sql={sql[:60]}, error={error}",
            extra={"extra_data": data},
        )
        return

    if execution_time_ms > LOG_SLOW_QUERY_THRESHOLD_MS:
        logger.warning(
            f"Slow DB query: time={execution_time_ms:.1f}ms, rows={row_count}, sql={sql[:60]}",
            extra={"extra_data": data},
        )
        return

    if _should_sample():
        logger.info(
            f"DB executed: time={execution_time_ms:.1f}ms, rows={row_count}",
            extra={"extra_data": data},
        )


def filter_sensitive_data(log_data: Dict[str, Any]) -> Dict[str, Any]:
    filtered = {}
    email_pattern = re.compile(r'[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}')
    phone_pattern = re.compile(r'\+?\d[\d\s\-()]{7,}\d')
    token_pattern = re.compile(r'(Bearer\s+\S+|api[_-]?key[=:]\s*\S+|token[=:]\s*\S+)', re.IGNORECASE)

    for key, value in log_data.items():
        if isinstance(value, str):
            sanitized = email_pattern.sub("[EMAIL_REDACTED]", value)
            sanitized = phone_pattern.sub("[PHONE_REDACTED]", sanitized)
            sanitized = token_pattern.sub("[TOKEN_REDACTED]", sanitized)
            filtered[key] = sanitized
        elif isinstance(value, dict):
            filtered[key] = filter_sensitive_data(value)
        else:
            filtered[key] = value

    return filtered


def aggregate_logs(log_entries: List[Dict[str, Any]], time_range_hours: int = 24) -> Dict[str, Any]:
    total = len(log_entries)
    by_operation: Dict[str, Dict[str, Any]] = {}
    errors: List[Dict[str, Any]] = []
    slow_ops: List[Dict[str, Any]] = []

    for entry in log_entries:
        op = entry.get("operation", "unknown")
        if op not in by_operation:
            by_operation[op] = {"count": 0, "total_time_ms": 0.0, "errors": 0}
        by_operation[op]["count"] += 1

        time_ms = entry.get("time_ms") or entry.get("execution_time_ms") or entry.get("generation_time_ms", 0)
        by_operation[op]["total_time_ms"] += time_ms

        if entry.get("error") or entry.get("success") is False:
            by_operation[op]["errors"] += 1
            errors.append(entry)

        if time_ms > LOG_SLOW_QUERY_THRESHOLD_MS:
            slow_ops.append(entry)

    summary = {}
    for op, data in by_operation.items():
        avg_time = data["total_time_ms"] / data["count"] if data["count"] > 0 else 0.0
        success_count = data["count"] - data["errors"]
        success_rate = success_count / data["count"] if data["count"] > 0 else 1.0
        summary[op] = {
            "count": data["count"],
            "avg_time_ms": round(avg_time, 2),
            "success_rate": round(success_rate, 4),
        }

    return {
        "total_operations": total,
        "time_range_hours": time_range_hours,
        "by_operation_type": summary,
        "errors": errors[:50],
        "slow_operations": slow_ops[:50],
    }