Create-knowledge-find-decode/pipeline/tracing.py

"""Best-effort tracing helpers for the formal creation pipeline."""
from __future__ import annotations

import hashlib
import logging
import os
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Protocol
from uuid import UUID

from core.config import CreationDbConfig, env_value, load_env_file
from core.text_limits import clip_text

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class TraceConfig:
    enabled: bool = True
    capture_request: bool = True
    capture_response: bool = True
    max_payload_chars: int = 50_000

    @classmethod
    def from_env(cls, env_file: str | Path = ".env") -> "TraceConfig":
        file_env = load_env_file(env_file)
        return cls(
            enabled=_env_bool(env_value("CK_TRACE_ENABLED", file_env, "true")),
            capture_request=_env_bool(env_value("CK_TRACE_CAPTURE_REQUEST", file_env, "true")),
            capture_response=_env_bool(env_value("CK_TRACE_CAPTURE_RESPONSE", file_env, "true")),
            max_payload_chars=_env_int(env_value("CK_TRACE_MAX_PAYLOAD_CHARS", file_env, "50000"), 50_000),
        )


@dataclass(frozen=True)
class TraceContext:
    pipeline_run_id: UUID | None = None
    pipeline_job_id: UUID | None = None
    acquisition_run_id: UUID | None = None
    acquisition_job_id: UUID | None = None
    query_id: UUID | None = None
    item_id: UUID | None = None
    decode_job_id: UUID | None = None
    decode_result_id: UUID | None = None
    payload_draft_id: UUID | None = None
    ingest_record_id: UUID | None = None
    platform: str | None = None
    stage: str | None = None
    substage: str | None = None
    metadata: dict[str, Any] = field(default_factory=dict)

    def child(self, **updates: Any) -> "TraceContext":
        values = {
            "pipeline_run_id": self.pipeline_run_id,
            "pipeline_job_id": self.pipeline_job_id,
            "acquisition_run_id": self.acquisition_run_id,
            "acquisition_job_id": self.acquisition_job_id,
            "query_id": self.query_id,
            "item_id": self.item_id,
            "decode_job_id": self.decode_job_id,
            "decode_result_id": self.decode_result_id,
            "payload_draft_id": self.payload_draft_id,
            "ingest_record_id": self.ingest_record_id,
            "platform": self.platform,
            "stage": self.stage,
            "substage": self.substage,
            "metadata": dict(self.metadata),
        }
        metadata = updates.pop("metadata", None)
        values.update(updates)
        if metadata:
            values["metadata"].update(metadata)
        return TraceContext(**values)

    def as_payload(self) -> dict[str, Any]:
        out: dict[str, Any] = {
            "pipeline_run_id": self.pipeline_run_id,
            "pipeline_job_id": self.pipeline_job_id,
            "acquisition_run_id": self.acquisition_run_id,
            "acquisition_job_id": self.acquisition_job_id,
            "query_id": self.query_id,
            "item_id": self.item_id,
            "decode_job_id": self.decode_job_id,
            "decode_result_id": self.decode_result_id,
            "payload_draft_id": self.payload_draft_id,
            "ingest_record_id": self.ingest_record_id,
            "platform": self.platform,
            "stage": self.stage,
            "substage": self.substage,
            **self.metadata,
        }
        return {key: str(value) if isinstance(value, UUID) else value for key, value in out.items() if value is not None}


class TraceWriter(Protocol):
    def event(self, *, context: TraceContext, stage: str, event_type: str, **kwargs: Any) -> None:
        ...

    def candidate_hit(self, *, context: TraceContext, **kwargs: Any) -> None:
        ...

    def llm_call(self, *, context: TraceContext, stage: str, substage: str | None = None, **kwargs: Any) -> None:
        ...


class NoopTraceWriter:
    def event(self, *, context: TraceContext, stage: str, event_type: str, **kwargs: Any) -> None:
        return None

    def candidate_hit(self, *, context: TraceContext, **kwargs: Any) -> None:
        return None

    def llm_call(self, *, context: TraceContext, stage: str, substage: str | None = None, **kwargs: Any) -> None:
        return None


def new_trace_writer(db_config: CreationDbConfig, *, env_file: str | Path = ".env") -> TraceWriter:
    config = TraceConfig.from_env(env_file)
    if not config.enabled:
        return NoopTraceWriter()
    from pipeline.postgres import PostgresTraceWriter

    return PostgresTraceWriter(db_config=db_config, config=config)


def hash_prompt(value: str | None) -> str | None:
    if not value:
        return None
    return hashlib.sha256(value.encode("utf-8")).hexdigest()


def redact_headers(headers: dict[str, Any] | None) -> dict[str, Any]:
    if not headers:
        return {}
    out = dict(headers)
    for key in list(out):
        if key.lower() in {"authorization", "api-key", "x-api-key"}:
            out[key] = "<redacted>"
    return out


def sanitize_payload(value: Any, *, max_chars: int) -> Any:
    if value is None or isinstance(value, (bool, int, float)):
        return value
    if isinstance(value, UUID):
        return str(value)
    if isinstance(value, str):
        if value.startswith("data:") and ";base64," in value[:80]:
            head, _, data = value.partition(",")
            digest = hashlib.sha256(data.encode("utf-8")).hexdigest() if data else ""
            return {
                "omitted": "base64_data_url",
                "media_type": head[:80],
                "chars": len(value),
                "sha256": digest,
            }
        return clip_text(value, max_chars)
    if isinstance(value, dict):
        return {
            str(key): sanitize_payload(item, max_chars=max_chars)
            for key, item in value.items()
        }
    if isinstance(value, (list, tuple)):
        return [sanitize_payload(item, max_chars=max_chars) for item in value]
    if hasattr(value, "model_dump"):
        return sanitize_payload(value.model_dump(mode="json"), max_chars=max_chars)
    return clip_text(str(value), max_chars)


def timed_ms(start: float) -> int:
    return max(0, int((time.perf_counter() - start) * 1000))


def _env_bool(value: str) -> bool:
    return value.strip().lower() not in {"0", "false", "no", "off"}


def _env_int(value: str, default: int) -> int:
    try:
        parsed = int(value)
    except (TypeError, ValueError):
        logger.warning("invalid trace integer config %r; using %s", value, default)
        return default
    return parsed if parsed > 0 else default


def current_env_file(default: str = ".env") -> str:
    return os.getenv("CK_ENV_FILE", default)