persistence/traces.py

"""Agent trace logging for Sharing is Caring badge."""

from __future__ import annotations

import json
import os
import time
from datetime import datetime, timezone
from pathlib import Path

TRACE_DIR = Path(os.environ.get("TRACE_DIR", "traces"))
TRACE_FILE = TRACE_DIR / "entity_generations.jsonl"
HF_TRACES_REPO = os.environ.get("HF_TRACES_REPO", "")


def log_entity_generation(
    *,
    user_input: str,
    system_prompt: str,
    user_prompt: str,
    raw_output: str,
    latency_ms: float,
    model_id: str,
    success: bool,
    entity_id: str | None = None,
    error: str | None = None,
) -> None:
    """Append one generation trace to local JSONL."""
    TRACE_DIR.mkdir(parents=True, exist_ok=True)
    record = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "task": "entity_generation",
        "model_id": model_id,
        "user_input": user_input,
        "system_prompt": system_prompt[:2000],
        "user_prompt": user_prompt[:4000],
        "raw_output": raw_output[:8000],
        "latency_ms": round(latency_ms, 1),
        "success": success,
        "entity_id": entity_id,
        "error": error,
    }
    with TRACE_FILE.open("a", encoding="utf-8") as f:
        f.write(json.dumps(record, ensure_ascii=False) + "\n")

    if HF_TRACES_REPO and os.environ.get("HF_TOKEN"):
        _maybe_push_traces()


def _maybe_push_traces() -> None:
    """Upload traces file to HF dataset if configured."""
    if not TRACE_FILE.exists():
        return
    try:
        from huggingface_hub import HfApi
        api = HfApi()
        api.upload_file(
            path_or_fileobj=str(TRACE_FILE),
            path_in_repo="entity_generations.jsonl",
            repo_id=HF_TRACES_REPO,
            repo_type="dataset",
            commit_message="Update agent traces",
        )
    except Exception:
        pass


def export_traces_summary() -> dict:
    if not TRACE_FILE.exists():
        return {"count": 0, "path": str(TRACE_FILE)}
    count = sum(1 for _ in TRACE_FILE.open())
    return {"count": count, "path": str(TRACE_FILE)}