src/app_kit/tracing.py

from __future__ import annotations

import json
import uuid
from dataclasses import asdict, is_dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

_CANONICAL_INPUT_KEYS = (
    'inputs',
    'input',
    'pack',
    'pack_id',
    'pack_name',
    'pack_path',
    'project',
    'scenario_id',
    'scenario_count',
    'query',
    'prompt',
    'transcript',
    'text',
)
_MODEL_HINT_KEYS = (
    'model_name',
    'base_model_id',
    'base_model',
    'model_id',
    'adapter_name',
    'loaded_from',
    'adapter_path',
)


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat(timespec='seconds')


def _json_safe(value: Any) -> Any:
    if is_dataclass(value):
        return _json_safe(asdict(value))
    if isinstance(value, Path):
        return str(value)
    if isinstance(value, dict):
        return {str(key): _json_safe(item) for key, item in value.items()}
    if isinstance(value, (list, tuple)):
        return [_json_safe(item) for item in value]
    return value


def _nonempty(value: Any) -> bool:
    return value not in (None, '', [], {}, ())


def _infer_input_payload(payload: dict[str, Any]) -> Any:
    existing_inputs = payload.get('inputs')
    if _nonempty(existing_inputs):
        return existing_inputs
    inputs: dict[str, Any] = {}
    for key in _CANONICAL_INPUT_KEYS:
        if key in payload and key != 'inputs' and _nonempty(payload[key]):
            inputs[key] = payload[key]
    if inputs:
        return inputs
    for key in ('project', 'pack_id', 'pack_name', 'pack_path', 'pack'):
        if key in payload and _nonempty(payload[key]):
            inputs[key] = payload[key]
    return inputs


def _infer_model_name(value: Any) -> str | None:
    if isinstance(value, dict):
        for key in _MODEL_HINT_KEYS:
            candidate = value.get(key)
            if isinstance(candidate, str) and candidate.strip():
                return candidate.strip()
        for item in value.values():
            candidate = _infer_model_name(item)
            if candidate:
                return candidate
    elif isinstance(value, list):
        for item in value:
            candidate = _infer_model_name(item)
            if candidate:
                return candidate
    return None


def canonicalize_trace_payload(payload: dict[str, Any]) -> dict[str, Any]:
    normalized = _json_safe(payload)
    timestamp = normalized.get('timestamp') or normalized.get('finished_at') or normalized.get('started_at') or utc_now()
    parsed_outputs = normalized.get('parsed_outputs')
    if not _nonempty(parsed_outputs):
        parsed_outputs = normalized.get('result')
    if not _nonempty(parsed_outputs):
        parsed_outputs = normalized.get('results')
    if not _nonempty(parsed_outputs):
        parsed_outputs = normalized.get('output')
    if not _nonempty(parsed_outputs):
        parsed_outputs = {}
    model_name = normalized.get('model_name')
    if not _nonempty(model_name):
        model_name = _infer_model_name(parsed_outputs) or _infer_model_name(normalized)
    if not _nonempty(model_name):
        model_name = 'unknown'
    normalized['timestamp'] = str(timestamp)
    normalized['inputs'] = _infer_input_payload(normalized)
    normalized['parsed_outputs'] = parsed_outputs
    normalized['model_name'] = str(model_name)
    return normalized


def write_trace_artifact(artifact_dir: str | Path, payload: dict[str, Any]) -> Path:
    artifact_dir = Path(artifact_dir)
    trace_dir = artifact_dir / 'traces'
    trace_dir.mkdir(parents=True, exist_ok=True)
    run_id = str(payload.get('run_id') or uuid.uuid4().hex)
    kind = str(payload.get('kind', 'trace')).strip().replace('/', '_').replace(' ', '_') or 'trace'
    trace_path = trace_dir / f'{kind}-{run_id}.json'
    normalized = canonicalize_trace_payload(payload)
    normalized['run_id'] = run_id
    normalized['trace_path'] = str(trace_path)
    trace_path.write_text(json.dumps(normalized, indent=2, ensure_ascii=False, sort_keys=True, default=str), encoding='utf-8')
    return trace_path