This commit is contained in:
thanhnv
@@ -0,0 +1,211 @@
|
||||
"""Persistence for model assessments — the routing feature's "loader/writer".
|
||||
|
||||
Assessments can be large and are rewritten wholesale on every reassess, so they
|
||||
live in their OWN file (``~/.cowork_local/assessments.json`` by default) rather
|
||||
than bloating the main ``config.json``. Two robustness guarantees the task
|
||||
requires:
|
||||
|
||||
* **Atomic write** — results are written to a temp file in the same directory
|
||||
and then ``os.replace``'d over the target, so an interrupted reassess can
|
||||
never leave a half-written / corrupt store behind.
|
||||
* **Versioned history** — before each overwrite, the previous store is copied
|
||||
to ``assessments_history/<timestamp>.json`` so a model that *degrades*
|
||||
between runs can be detected after the fact.
|
||||
|
||||
On-disk shape (JSON, mirrors the task's YAML ``assessments`` block)::
|
||||
|
||||
{
|
||||
"last_updated": "2026-07-22T09:30:00+00:00",
|
||||
"policy": "balanced",
|
||||
"results": {
|
||||
"anthropic/claude-opus-4-8": { <ModelAssessment> },
|
||||
...
|
||||
}
|
||||
}
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional
|
||||
|
||||
from .models import ModelAssessment, Policy
|
||||
|
||||
# Default location under the app's config dir. Imported lazily so tests can
|
||||
# point the store anywhere without touching the real home directory.
|
||||
_DEFAULT_STORE_NAME = "assessments.json"
|
||||
_HISTORY_DIR_NAME = "assessments_history"
|
||||
|
||||
|
||||
def utc_now_iso() -> str:
|
||||
"""Current UTC time as an ISO-8601 string (used for ``last_updated``)."""
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def _fs_safe_stamp() -> str:
|
||||
"""A filesystem-safe timestamp for history filenames (no ``:``)."""
|
||||
return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S%fZ")
|
||||
|
||||
|
||||
class AssessmentStore:
|
||||
"""Reads/writes the assessment JSON file with atomic writes + history.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
store_path:
|
||||
Path to the assessments JSON file. If ``None``, defaults to
|
||||
``~/.cowork_local/assessments.json``.
|
||||
history_dir:
|
||||
Directory for pre-overwrite backups. Defaults to a sibling
|
||||
``assessments_history/`` next to ``store_path``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
store_path: Optional[Path] = None,
|
||||
history_dir: Optional[Path] = None,
|
||||
) -> None:
|
||||
if store_path is None:
|
||||
from ...config import CONFIG_DIR # lazy: avoids import cost in tests
|
||||
store_path = CONFIG_DIR / _DEFAULT_STORE_NAME
|
||||
self.store_path = Path(store_path)
|
||||
self.history_dir = Path(
|
||||
history_dir or self.store_path.parent / _HISTORY_DIR_NAME
|
||||
)
|
||||
|
||||
# -- read ----------------------------------------------------------- #
|
||||
def load_raw(self) -> Dict:
|
||||
"""Return the raw JSON dict, or an empty skeleton if absent/corrupt.
|
||||
|
||||
A corrupt store must never crash the app (same philosophy as
|
||||
``AppConfig.load``) — we fall back to an empty result set so the next
|
||||
reassess simply rebuilds it.
|
||||
"""
|
||||
if not self.store_path.exists():
|
||||
return {"last_updated": None, "policy": Policy.BALANCED.value, "results": {}}
|
||||
try:
|
||||
return json.loads(self.store_path.read_text(encoding="utf-8"))
|
||||
except (json.JSONDecodeError, OSError):
|
||||
return {"last_updated": None, "policy": Policy.BALANCED.value, "results": {}}
|
||||
|
||||
def load(self) -> Dict[str, ModelAssessment]:
|
||||
"""Return ``{candidate_key: ModelAssessment}`` parsed from disk.
|
||||
|
||||
Individual malformed entries are skipped rather than failing the whole
|
||||
load — one bad row shouldn't hide every good assessment.
|
||||
"""
|
||||
raw = self.load_raw()
|
||||
out: Dict[str, ModelAssessment] = {}
|
||||
for key, payload in (raw.get("results") or {}).items():
|
||||
try:
|
||||
out[key] = ModelAssessment.model_validate(payload)
|
||||
except Exception: # noqa: BLE001 — skip a single corrupt entry
|
||||
continue
|
||||
return out
|
||||
|
||||
def last_updated(self) -> Optional[str]:
|
||||
return self.load_raw().get("last_updated")
|
||||
|
||||
def policy(self) -> str:
|
||||
return self.load_raw().get("policy") or Policy.BALANCED.value
|
||||
|
||||
# -- write ---------------------------------------------------------- #
|
||||
def save(
|
||||
self,
|
||||
results: Dict[str, ModelAssessment],
|
||||
policy: Policy | str = Policy.BALANCED,
|
||||
*,
|
||||
last_updated: Optional[str] = None,
|
||||
backup: bool = True,
|
||||
) -> Path:
|
||||
"""Atomically write ``results`` to the store, backing up the previous
|
||||
version to history first.
|
||||
|
||||
Returns the store path. Never leaves a partially-written file: the
|
||||
payload is fully serialized to a temp file and only then swapped into
|
||||
place with ``os.replace`` (atomic on the same filesystem, incl. NTFS).
|
||||
"""
|
||||
if backup:
|
||||
self._backup_existing()
|
||||
|
||||
policy_val = policy.value if isinstance(policy, Policy) else str(policy)
|
||||
payload = {
|
||||
"last_updated": last_updated or utc_now_iso(),
|
||||
"policy": policy_val,
|
||||
"results": {
|
||||
key: assessment.model_dump(mode="json")
|
||||
for key, assessment in results.items()
|
||||
},
|
||||
}
|
||||
|
||||
self.store_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
text = json.dumps(payload, indent=2, ensure_ascii=False)
|
||||
|
||||
# Temp file MUST be on the same volume as the target for os.replace to
|
||||
# be atomic — so create it in the target's own directory.
|
||||
fd, tmp_name = tempfile.mkstemp(
|
||||
dir=str(self.store_path.parent),
|
||||
prefix=".assessments-",
|
||||
suffix=".tmp",
|
||||
)
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as fh:
|
||||
fh.write(text)
|
||||
fh.flush()
|
||||
os.fsync(fh.fileno()) # durability: survive a crash right after
|
||||
os.replace(tmp_name, self.store_path) # atomic swap
|
||||
except BaseException:
|
||||
# Clean up the temp file on any failure so we never leave litter.
|
||||
try:
|
||||
os.unlink(tmp_name)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
return self.store_path
|
||||
|
||||
def _backup_existing(self) -> Optional[Path]:
|
||||
"""Copy the current store into history as ``<timestamp>.json``.
|
||||
|
||||
No-op when there is nothing to back up. Best-effort: a failed backup
|
||||
must not block the (more important) new write.
|
||||
"""
|
||||
if not self.store_path.exists():
|
||||
return None
|
||||
try:
|
||||
self.history_dir.mkdir(parents=True, exist_ok=True)
|
||||
# The wall clock can be coarse (Windows ~15ms), so two rapid
|
||||
# backups may share a timestamp — disambiguate with a counter so a
|
||||
# snapshot is never silently overwritten.
|
||||
stamp = _fs_safe_stamp()
|
||||
dst = self.history_dir / f"{stamp}.json"
|
||||
n = 1
|
||||
while dst.exists():
|
||||
dst = self.history_dir / f"{stamp}-{n}.json"
|
||||
n += 1
|
||||
dst.write_text(
|
||||
self.store_path.read_text(encoding="utf-8"), encoding="utf-8"
|
||||
)
|
||||
return dst
|
||||
except OSError:
|
||||
return None
|
||||
|
||||
def history_files(self) -> list[Path]:
|
||||
"""All history snapshots, oldest first (for degrade detection / UI)."""
|
||||
if not self.history_dir.exists():
|
||||
return []
|
||||
return sorted(self.history_dir.glob("*.json"))
|
||||
|
||||
def prune_history(self, keep: int = 30) -> None:
|
||||
"""Keep only the newest ``keep`` history snapshots; delete the rest."""
|
||||
files = self.history_files()
|
||||
for old in files[:-keep] if keep > 0 else files:
|
||||
try:
|
||||
old.unlink()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
__all__ = ["AssessmentStore", "utc_now_iso"]
|
||||
Reference in New Issue
Block a user