storage/trace_logger.py

"""Agent trace untuk bonus quest Sharing is Caring — tanpa teks mimpi penuh."""
from __future__ import annotations

import datetime
import hashlib
import json
import os
import uuid
from pathlib import Path

from model.extractor import VERSI_PROMPT

TRACE_LOG_PATH = Path("./traces/oneiros_agent_traces.jsonl")
HF_DATASET_REPO = os.getenv(
    "ONEIROS_TRACE_DATASET",
    "adindamochamad/oneiros-agent-traces",
)
APP_VERSION = "0.1.0"


def log_trace(
    dream_text: str,
    raw_model_output: str,
    entities: dict,
    parse_ok: bool,
    parse_error: str | None,
    elapsed_extract: float,
    elapsed_render: float = 0.0,
    environment: str = "local",
    raw_percobaan_pertama: str | None = None,
) -> str:
    """Tulis satu baris JSONL; kembalikan trace_id."""
    TRACE_LOG_PATH.parent.mkdir(parents=True, exist_ok=True)
    id_jejak = str(uuid.uuid4())

    langkah_extract: dict = {
        "step": "extract",
        "prompt_template_id": VERSI_PROMPT,
        "raw_model_output": (raw_model_output or "")[:4000],
        "parse_ok": parse_ok,
        "parse_error": parse_error,
        "elapsed_seconds": round(elapsed_extract, 2),
    }
    if raw_percobaan_pertama:
        langkah_extract["raw_first_attempt"] = raw_percobaan_pertama[:2000]

    jejak = {
        "trace_id": id_jejak,
        "timestamp": datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "app_version": APP_VERSION,
        "environment": environment,
        "model_id": "Qwen2.5-7B-Instruct-Q4_K_M",
        "input": {
            "dream_char_count": len(dream_text),
            "dream_word_count": len(dream_text.split()),
            "dream_sha256_prefix": hashlib.sha256(dream_text.encode()).hexdigest()[:16],
        },
        "agent_steps": [
            langkah_extract,
            {
                "step": "render_map",
                "node_count": len(entities.get("characters", []))
                + len(entities.get("places", []))
                + len(entities.get("symbols", [])),
                "edge_count": len(entities.get("connections", [])),
                "elapsed_seconds": round(elapsed_render, 2),
            },
        ],
        "result_summary": {
            "mood": entities.get("mood"),
            "title": entities.get("title"),
            "entity_counts": {
                "characters": len(entities.get("characters", [])),
                "places": len(entities.get("places", [])),
                "symbols": len(entities.get("symbols", [])),
            },
        },
    }
    with open(TRACE_LOG_PATH, "a", encoding="utf-8") as f:
        f.write(json.dumps(jejak, ensure_ascii=False) + "\n")
    return id_jejak


def push_traces_to_hub() -> None:
    """Upload file trace ke Dataset publik di Hub (buat repo jika belum ada)."""
    from huggingface_hub import HfApi

    if not TRACE_LOG_PATH.is_file():
        raise FileNotFoundError(f"Tidak ada trace di {TRACE_LOG_PATH}")

    api = HfApi()
    try:
        api.repo_info(repo_id=HF_DATASET_REPO, repo_type="dataset")
    except Exception:
        api.create_repo(
            repo_id=HF_DATASET_REPO,
            repo_type="dataset",
            exist_ok=True,
            private=False,
        )

    api.upload_file(
        path_or_fileobj=str(TRACE_LOG_PATH),
        path_in_repo="oneiros_agent_traces.jsonl",
        repo_id=HF_DATASET_REPO,
        repo_type="dataset",
    )