observability.py

import json
import logging
import os
import sys
import threading
import time
from dataclasses import dataclass
from typing import Any

SENSITIVE_KEYS = {
    "api_key",
    "authorization",
    "cookie",
    "hf_token",
    "password",
    "secret",
    "token",
}


@dataclass
class GenerationStats:
    prompt_tokens: int = 0
    completion_tokens: int = 0
    queue_ms: float = 0
    inference_ms: float = 0
    first_token_ms: float | None = None
    chunks: int = 0
    output_chars: int = 0


class RuntimeMetrics:
    def __init__(self, started_at: float | None = None) -> None:
        self._lock = threading.Lock()
        self._started_at = time.time() if started_at is None else started_at
        self.http_requests = 0
        self.http_errors = 0
        self.completions = 0
        self.completion_errors = 0
        self.active_generations = 0
        self.prompt_tokens = 0
        self.completion_tokens = 0

    def record_http(self, status_code: int) -> None:
        with self._lock:
            self.http_requests += 1
            if status_code >= 500:
                self.http_errors += 1

    def generation_started(self) -> None:
        with self._lock:
            self.active_generations += 1

    def generation_finished(
        self,
        stats: GenerationStats,
        *,
        failed: bool,
    ) -> None:
        with self._lock:
            self.active_generations = max(0, self.active_generations - 1)
            self.completions += 1
            self.prompt_tokens += stats.prompt_tokens
            self.completion_tokens += stats.completion_tokens
            if failed:
                self.completion_errors += 1

    def snapshot(self) -> dict[str, Any]:
        with self._lock:
            return {
                "uptime_seconds": round(time.time() - self._started_at, 1),
                "http_requests": self.http_requests,
                "http_errors": self.http_errors,
                "completions": self.completions,
                "completion_errors": self.completion_errors,
                "active_generations": self.active_generations,
                "prompt_tokens": self.prompt_tokens,
                "completion_tokens": self.completion_tokens,
            }


def env_flag(name: str, default: bool = False) -> bool:
    value = os.getenv(name)
    if value is None:
        return default
    return value.strip().lower() in {"1", "true", "yes", "on"}


def is_sensitive_key(key: Any) -> bool:
    normalized = str(key).strip().lower().replace("-", "_")
    return (
        normalized in SENSITIVE_KEYS
        or normalized.endswith("_token")
        or "api_key" in normalized
        or "authorization" in normalized
        or "cookie" in normalized
        or "password" in normalized
        or "secret" in normalized
    )


def sanitize(value: Any, max_text_chars: int = 2_000) -> Any:
    if isinstance(value, dict):
        return {
            str(key): (
                "[REDACTED]"
                if is_sensitive_key(key)
                else sanitize(item, max_text_chars)
            )
            for key, item in value.items()
        }
    if isinstance(value, (list, tuple)):
        return [sanitize(item, max_text_chars) for item in value]
    if isinstance(value, str) and len(value) > max_text_chars:
        omitted = len(value) - max_text_chars
        return f"{value[:max_text_chars]}...[truncated {omitted} chars]"
    return value


class JsonFormatter(logging.Formatter):
    def format(self, record: logging.LogRecord) -> str:
        payload = {
            "timestamp": self.formatTime(record, "%Y-%m-%dT%H:%M:%S%z"),
            "level": record.levelname,
            "logger": record.name,
            "message": record.getMessage(),
        }
        event_fields = getattr(record, "event_fields", None)
        if isinstance(event_fields, dict):
            payload.update(event_fields)
        if record.exc_info:
            payload["exception"] = self.formatException(record.exc_info)
        return json.dumps(payload, ensure_ascii=True, default=str)


def configure_logger(name: str = "browso.server") -> logging.Logger:
    logger = logging.getLogger(name)
    level_name = os.getenv("LOG_LEVEL", "INFO").upper()
    logger.setLevel(getattr(logging, level_name, logging.INFO))
    logger.propagate = False

    if not logger.handlers:
        handler = logging.StreamHandler(sys.stdout)
        handler.setFormatter(JsonFormatter())
        logger.addHandler(handler)

    return logger


def log_event(
    logger: logging.Logger,
    event: str,
    *,
    level: int = logging.INFO,
    exc_info: bool = False,
    max_text_chars: int = 2_000,
    **fields: Any,
) -> None:
    logger.log(
        level,
        event,
        exc_info=exc_info,
        extra={
            "event_fields": {
                "event": event,
                **sanitize(fields, max_text_chars=max_text_chars),
            }
        },
    )


def summarize_messages(messages: list[dict[str, Any]]) -> dict[str, Any]:
    roles: dict[str, int] = {}
    text_chars = 0
    tool_calls = 0

    for message in messages:
        role = str(message.get("role", "unknown"))
        roles[role] = roles.get(role, 0) + 1
        content = message.get("content")
        if isinstance(content, str):
            text_chars += len(content)
        elif isinstance(content, list):
            for part in content:
                if isinstance(part, dict) and isinstance(part.get("text"), str):
                    text_chars += len(part["text"])
        if isinstance(message.get("tool_calls"), list):
            tool_calls += len(message["tool_calls"])

    return {
        "message_count": len(messages),
        "roles": roles,
        "text_chars": text_chars,
        "tool_call_count": tool_calls,
    }