99 lines
3.4 KiB
Python
99 lines
3.4 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import shutil
|
|
from dataclasses import asdict, is_dataclass
|
|
from datetime import date, datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
class ArtifactStore:
|
|
def __init__(self, root: str | Path = "tasks/history") -> None:
|
|
self.root = Path(root)
|
|
|
|
def task_dir(self, task_id: str) -> Path:
|
|
return self.root / task_id
|
|
|
|
def write_step(
|
|
self,
|
|
*,
|
|
task_id: str,
|
|
index: int,
|
|
before_screenshot: bytes | None,
|
|
after_screenshot: bytes | None,
|
|
record: dict[str, Any],
|
|
) -> dict[str, str | None]:
|
|
task_dir = self.task_dir(task_id)
|
|
task_dir.mkdir(parents=True, exist_ok=True)
|
|
stem = f"{index:03d}"
|
|
before_screenshot_path: Path | None = None
|
|
after_screenshot_path: Path | None = None
|
|
if before_screenshot is not None:
|
|
before_screenshot_path = task_dir / f"{stem}-before.png"
|
|
before_screenshot_path.write_bytes(before_screenshot)
|
|
if after_screenshot is not None:
|
|
# Preserve the original filename as the compatibility alias for the
|
|
# screenshot captured after a step.
|
|
after_screenshot_path = task_dir / f"{stem}.png"
|
|
after_screenshot_path.write_bytes(after_screenshot)
|
|
|
|
json_path = task_dir / f"{stem}.json"
|
|
payload = {
|
|
**record,
|
|
"before_screenshot_path": (
|
|
str(before_screenshot_path) if before_screenshot_path else None
|
|
),
|
|
"after_screenshot_path": (
|
|
str(after_screenshot_path) if after_screenshot_path else None
|
|
),
|
|
"screenshot_path": (
|
|
str(after_screenshot_path) if after_screenshot_path else None
|
|
),
|
|
}
|
|
json_path.write_text(
|
|
json.dumps(_jsonable(payload), ensure_ascii=False, indent=2),
|
|
encoding="utf-8",
|
|
)
|
|
return {
|
|
"json_path": str(json_path),
|
|
"before_screenshot_path": (
|
|
str(before_screenshot_path) if before_screenshot_path else None
|
|
),
|
|
"after_screenshot_path": (
|
|
str(after_screenshot_path) if after_screenshot_path else None
|
|
),
|
|
"screenshot_path": (
|
|
str(after_screenshot_path) if after_screenshot_path else None
|
|
),
|
|
}
|
|
|
|
def read_steps(self, task_id: str) -> list[dict[str, Any]]:
|
|
task_dir = self.task_dir(task_id)
|
|
if not task_dir.exists():
|
|
return []
|
|
steps = []
|
|
for path in sorted(task_dir.glob("*.json")):
|
|
steps.append(json.loads(path.read_text(encoding="utf-8")))
|
|
return steps
|
|
|
|
def delete_task(self, task_id: str) -> None:
|
|
"""Remove all on-disk artifacts (JSON + screenshots) for a task."""
|
|
shutil.rmtree(self.task_dir(task_id), ignore_errors=True)
|
|
|
|
|
|
def _jsonable(value: Any) -> Any:
|
|
if hasattr(value, "to_dict"):
|
|
return value.to_dict()
|
|
if is_dataclass(value):
|
|
return asdict(value)
|
|
if isinstance(value, dict):
|
|
return {key: _jsonable(inner) for key, inner in value.items()}
|
|
if isinstance(value, list):
|
|
return [_jsonable(inner) for inner in value]
|
|
if isinstance(value, tuple):
|
|
return [_jsonable(inner) for inner in value]
|
|
if isinstance(value, (datetime, date)):
|
|
return value.isoformat()
|
|
return value
|