ai-exchange/scripts/run_training_case_package.py

from pathlib import Path
import argparse
import json
import sys
from typing import List, Optional

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from src.services.training import TrainingCaseService


DEFAULT_FIXTURE = ROOT / "tests" / "fixtures" / "training_case_package_m6.json"


class _PublicCliError(Exception):
    def __init__(self, code: str) -> None:
        super().__init__(code)
        self.code = code


class _OfflineTrainingRepository:
    def __init__(self) -> None:
        self._cases = []
        self._tasks = []
        self._next_case_id = 6000
        self._next_task_id = 9000

    def create_case(self, **kwargs) -> int:
        self._next_case_id += 1
        row = dict(kwargs)
        row["id"] = self._next_case_id
        self._cases.append(row)
        return self._next_case_id

    def add_task(self, **kwargs) -> int:
        self._next_task_id += 1
        row = dict(kwargs)
        row["id"] = self._next_task_id
        self._tasks.append(row)
        return self._next_task_id

    def list_cases(
        self,
        status: Optional[str] = None,
        skill_tag: Optional[str] = None,
        instrument_id: Optional[int] = None,
        limit: int = 100,
    ) -> List[dict]:
        rows = []
        for row in self._cases:
            if status is not None and row.get("status") != status:
                continue
            if skill_tag is not None and row.get("skill_tag") != skill_tag:
                continue
            if instrument_id is not None and row.get("instrument_id") != instrument_id:
                continue
            rows.append(dict(row))
        return rows[:limit]


def run_fixture_smoke(input_path: Path, output_path: Optional[Path] = None) -> dict:
    payload = _load_package(input_path)
    service = TrainingCaseService(repository=_OfflineTrainingRepository())
    import_result = service.import_cases(payload)
    export_result = service.export_cases(limit=100)

    if output_path is not None:
        _write_public_export(output_path, export_result)

    return {
        "passed": True,
        "input": _display_path(input_path),
        "output": _display_path(output_path) if output_path is not None else None,
        "imported": import_result["imported"],
        "exported": export_result["metadata"]["item_count"],
    }


def main(argv=None) -> int:
    parser = _build_parser()
    args = parser.parse_args(argv)
    input_path = Path(args.input)
    output_path = Path(args.output) if args.output else None
    try:
        summary = run_fixture_smoke(input_path, output_path=output_path)
    except (_PublicCliError, OSError, ValueError, json.JSONDecodeError) as exc:
        print(json.dumps(_error_summary(input_path, output_path, exc), ensure_ascii=False))
        return 2

    print(json.dumps(summary, ensure_ascii=False))
    return 0


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="scripts/run_training_case_package.py",
        description="Offline package smoke: Training Case import/export smoke with local files.",
        epilog=(
            "Imports a public package from --input and writes public export/public output "
            "to --output when set. On public errors, stdout is one-line JSON with "
            "passed, input, output, and error. Stable public error codes include "
            "input_file_not_found, input_path_is_directory, output_path_is_directory, "
            "and output_path_unwritable."
        ),
    )
    input_help = f"Public package JSON path to import (default: {DEFAULT_FIXTURE.relative_to(ROOT)})."
    parser.add_argument("--input", default=str(DEFAULT_FIXTURE), help=input_help)
    parser.add_argument("--output", help="Optional public export JSON path; writes public output only.")
    return parser


def _load_package(input_path: Path) -> dict:
    with input_path.open("r", encoding="utf-8") as handle:
        payload = json.load(handle)
    if not isinstance(payload, dict):
        raise ValueError("invalid_json_body")
    return payload


def _write_public_export(output_path: Path, export_result: dict) -> None:
    try:
        if output_path.exists() and output_path.is_dir():
            raise _PublicCliError("output_path_is_directory")
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(
            json.dumps(export_result, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
            encoding="utf-8",
        )
    except _PublicCliError:
        raise
    except OSError as exc:
        raise _PublicCliError("output_path_unwritable") from exc


def _display_path(path: Optional[Path]) -> Optional[str]:
    if path is None:
        return None
    try:
        return str(path.resolve().relative_to(ROOT))
    except ValueError:
        return str(path)


def _public_error(exc: Exception) -> str:
    if isinstance(exc, _PublicCliError):
        return exc.code
    if isinstance(exc, IsADirectoryError):
        return "input_path_is_directory"
    if isinstance(exc, FileNotFoundError):
        return "input_file_not_found"
    if isinstance(exc, json.JSONDecodeError):
        return "invalid_json_body"
    message = str(exc)
    if "answer" + "_key" in message:
        return "invalid_json_field"
    return message


def _is_output_path_error(exc: Exception) -> bool:
    return isinstance(exc, _PublicCliError) and exc.code.startswith("output_path_")


def _display_public_output_path(path: Optional[Path]) -> Optional[str]:
    if path is None:
        return None
    try:
        return str(path.resolve().relative_to(ROOT))
    except (OSError, ValueError):
        if path.is_absolute():
            return path.name or "output"
        return str(path)


def _error_summary(input_path: Path, output_path: Optional[Path], exc: Exception) -> dict:
    displayed_output = (
        _display_public_output_path(output_path)
        if _is_output_path_error(exc)
        else _display_path(output_path)
    )
    return {
        "passed": False,
        "input": _display_path(input_path),
        "output": displayed_output,
        "error": _public_error(exc),
    }


if __name__ == "__main__":
    raise SystemExit(main())