src/compliment_forest/trace.py

from __future__ import annotations

import json
import re
from pathlib import Path
from typing import Any

_REDACT_KEYS = {
    "api_key",
    "authorization",
    "cohere_api_key",
    "hf_token",
    "name",
    "password",
    "secret",
    "situation",
    "token",
    "user_name",
}


def _sensitive_values(value: Any) -> tuple[str, ...]:
    found: set[str] = set()
    if isinstance(value, dict):
        for item_key, item_value in value.items():
            if (
                item_key.casefold() in {"name", "situation", "user_name"}
                and isinstance(item_value, str)
                and item_value.strip()
            ):
                found.add(item_value.strip())
            found.update(_sensitive_values(item_value))
    elif isinstance(value, (list, tuple)):
        for item in value:
            found.update(_sensitive_values(item))
    return tuple(sorted(found, key=len, reverse=True))


def sanitize_trace(
    value: Any,
    key: str = "",
    sensitive_values: tuple[str, ...] | None = None,
) -> Any:
    if sensitive_values is None:
        sensitive_values = _sensitive_values(value)
    normalized_key = key.casefold()
    if normalized_key in _REDACT_KEYS or any(
        marker in normalized_key for marker in ("token", "api_key", "password", "secret")
    ):
        return "[redacted]"
    if isinstance(value, dict):
        return {
            item_key: sanitize_trace(item_value, item_key, sensitive_values)
            for item_key, item_value in value.items()
        }
    if isinstance(value, list):
        return [sanitize_trace(item, sensitive_values=sensitive_values) for item in value]
    if isinstance(value, tuple):
        return [sanitize_trace(item, sensitive_values=sensitive_values) for item in value]
    if isinstance(value, str):
        sanitized = value
        for sensitive in sensitive_values:
            sanitized = re.sub(
                re.escape(sensitive),
                "[redacted]",
                sanitized,
                flags=re.IGNORECASE,
            )
        return sanitized
    return value


class TraceRecorder:
    def __init__(self, path: str | Path) -> None:
        self.path = Path(path)

    def append(self, record: dict[str, Any]) -> None:
        self.path.parent.mkdir(parents=True, exist_ok=True)
        safe_record = sanitize_trace(record)
        with self.path.open("a", encoding="utf-8") as trace_file:
            trace_file.write(json.dumps(safe_record, ensure_ascii=True) + "\n")