File size: 1,838 Bytes
331ec0d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
import json
import time
import shutil
from pathlib import Path
import tarfile
class PersistenceEngine:
def __init__(self, hive_home):
self.hive_home = Path(hive_home)
self.conversations_dir = self.hive_home / "conversations"
self.conversations_dir.mkdir(parents=True, exist_ok=True)
self.snapshots_dir = self.hive_home / "snapshots"
self.snapshots_dir.mkdir(parents=True, exist_ok=True)
def save_turn(self, turn):
conv_file = self.conversations_dir / f"{turn.session_id}.jsonl"
turn_data = {
"turn_id": turn.turn_id,
"user_id": turn.user_id,
"timestamp": turn.timestamp,
"input": turn.raw_input,
"output": turn.postprocessed_output
}
with open(conv_file, "a") as f:
f.write(json.dumps(turn_data) + "\n")
def save_snapshot(self):
timestamp = int(time.time())
snapshot_path = self.snapshots_dir / f"snapshot_{timestamp}.tar.gz"
# Create snapshot
with tarfile.open(snapshot_path, "w:gz") as tar:
tar.add(self.hive_home, arcname="hive_data")
# Keep only latest snapshot
for snap in self.snapshots_dir.glob("snapshot_*.tar.gz"):
if snap != snapshot_path:
snap.unlink()
return f"Snapshot saved: {snapshot_path.name}"
def rollback(self):
snapshots = list(self.snapshots_dir.glob("snapshot_*.tar.gz"))
if not snapshots:
return "No snapshots found"
latest = max(snapshots, key=lambda x: x.stat().st_mtime)
# Extract snapshot
with tarfile.open(latest, "r:gz") as tar:
tar.extractall(path=self.hive_home.parent)
return f"Rolled back to: {latest.name}" |