solar_eval/stores/jsonl.py

"""JSONL file-based result store for local CLI runs."""

import json
import logging
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
from uuid import uuid4

logger = logging.getLogger(__name__)


class JsonlStore:
    """Persist run results to JSONL files in a local directory.

    Directory structure:
        output_dir/
        ├── run.json          (run metadata)
        ├── results.jsonl     (inference results, one per line)
        ├── evaluation.json   (aggregated evaluation)
        └── eval_details.jsonl (per-sample eval details)
    """

    def __init__(self, output_dir: str | Path) -> None:
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

    async def update_run(self, run_id: str, data: dict[str, Any]) -> None:
        run_file = self.output_dir / "run.json"
        existing: dict[str, Any] = {}
        if run_file.exists():
            existing = json.loads(run_file.read_text())
        # Serialize datetimes
        for k, v in data.items():
            if isinstance(v, datetime):
                data[k] = v.isoformat()
        existing.update(data)
        existing["run_id"] = run_id
        run_file.write_text(json.dumps(existing, ensure_ascii=False, indent=2))

    async def insert_run_result(self, data: dict[str, Any]) -> str:
        result_id = str(uuid4())[:8]
        results_file = self.output_dir / "results.jsonl"
        data["id"] = result_id
        with open(results_file, "a") as f:
            f.write(json.dumps(data, ensure_ascii=False) + "\n")
        return result_id

    async def create_evaluation(self, data: dict[str, Any]) -> str:
        eval_id = str(uuid4())[:8]
        eval_file = self.output_dir / "evaluation.json"
        data["id"] = eval_id
        # Serialize datetimes
        for k, v in data.items():
            if isinstance(v, datetime):
                data[k] = v.isoformat()
        payload = json.dumps(data, ensure_ascii=False, indent=2)
        # tmp 파일에 먼저 쓰고 원자적으로 교체한다(F3) -- 재계산 도중 죽어도
        # evaluation.json 은 이전 내용 그대로거나 새 내용 그대로다. 절반짜리가 안 남는다.
        tmp_file = eval_file.with_name(eval_file.name + ".tmp")
        tmp_file.write_text(payload)
        tmp_file.replace(eval_file)
        return eval_id

    def load_existing_results(self) -> list[dict[str, Any]]:
        """Load existing inference results from results.jsonl."""
        results_file = self.output_dir / "results.jsonl"
        if not results_file.exists():
            return []
        results = []
        for line in results_file.read_text().strip().split("\n"):
            if line:
                results.append(json.loads(line))
        return results

    def load_run_meta(self) -> dict[str, Any]:
        """Load run metadata from run.json."""
        run_file = self.output_dir / "run.json"
        if not run_file.exists():
            return {}
        return json.loads(run_file.read_text())

    async def insert_eval_details(self, details: list[dict[str, Any]]) -> int:
        if not details:
            return 0
        details_file = self.output_dir / "eval_details.jsonl"
        with open(details_file, "a") as f:
            for d in details:
                f.write(json.dumps(d, ensure_ascii=False) + "\n")
        return len(details)