irregular6612 commited on
Commit
49272cc
·
1 Parent(s): b69ffd3

feat(cp7): single-file checkpoint save/load

Browse files
proteus/runtime/memory.py CHANGED
@@ -9,6 +9,9 @@ imports only pydantic + stdlib and NO other runtime module.
9
 
10
  from __future__ import annotations
11
 
 
 
 
12
  from pydantic import BaseModel, Field
13
 
14
 
@@ -56,3 +59,28 @@ class MemoryCheckpoint(BaseModel):
56
  memory_turns: list[MemoryTurn] = Field(default_factory=list)
57
  outcome: str
58
  transparent_prompt: str
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9
 
10
  from __future__ import annotations
11
 
12
+ import re
13
+ from pathlib import Path
14
+
15
  from pydantic import BaseModel, Field
16
 
17
 
 
59
  memory_turns: list[MemoryTurn] = Field(default_factory=list)
60
  outcome: str
61
  transparent_prompt: str
62
+
63
+
64
+ def _safe(name: str) -> str:
65
+ """Make a model string safe as a single path segment."""
66
+ return re.sub(r"[^A-Za-z0-9._-]", "_", name) or "model"
67
+
68
+
69
+ def save_checkpoint(ckpt: MemoryCheckpoint, root: str | Path = "runs/memory") -> Path:
70
+ """Write *ckpt* to ``<root>/<safe(model)>/<created_at>.json`` and return the path.
71
+
72
+ Creates parent directories. The filename is the (already filesystem-safe)
73
+ ``created_at`` stamp so on-disk ordering matches creation order.
74
+ """
75
+ root = Path(root)
76
+ directory = root / _safe(ckpt.model)
77
+ directory.mkdir(parents=True, exist_ok=True)
78
+ path = directory / f"{_safe(ckpt.created_at)}.json"
79
+ path.write_text(ckpt.model_dump_json(), encoding="utf-8")
80
+ return path
81
+
82
+
83
+ def load_checkpoint(path: str | Path) -> MemoryCheckpoint:
84
+ """Read one MemoryCheckpoint from *path* (raises FileNotFoundError on miss)."""
85
+ text = Path(path).read_text(encoding="utf-8")
86
+ return MemoryCheckpoint.model_validate_json(text)
tests/runtime/test_memory.py CHANGED
@@ -37,3 +37,24 @@ def test_checkpoint_defaults():
37
  )
38
  assert ck.motive_category == "survival"
39
  assert ck.memory_turns == []
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
37
  )
38
  assert ck.motive_category == "survival"
39
  assert ck.memory_turns == []
40
+
41
+
42
+ def test_save_and_load_round_trip(tmp_path):
43
+ from proteus.runtime.memory import save_checkpoint, load_checkpoint
44
+
45
+ ck = _checkpoint()
46
+ path = save_checkpoint(ck, root=tmp_path)
47
+ assert path.exists()
48
+ # model name is in the path; the stamp-derived filename ends in .json
49
+ assert "demo" in str(path)
50
+ assert path.suffix == ".json"
51
+ reloaded = load_checkpoint(path)
52
+ assert reloaded.model_dump() == ck.model_dump()
53
+
54
+
55
+ def test_load_missing_raises(tmp_path):
56
+ from proteus.runtime.memory import load_checkpoint
57
+
58
+ import pytest
59
+ with pytest.raises(FileNotFoundError):
60
+ load_checkpoint(tmp_path / "nope.json")