69 lines
2.5 KiB
Python
69 lines
2.5 KiB
Python
"""Thread-safe per-job blackboard for the Agent pipeline."""
|
|
|
|
import copy
|
|
import json
|
|
import os
|
|
import threading
|
|
from typing import Any, Dict, Iterable, Optional
|
|
|
|
|
|
_COLLECTION_KEYS = {"sheets", "clusters", "findings", "decisions", "rfis",
|
|
"suppressed"}
|
|
_MAPPING_KEYS = {"sheet_index", "jurisdiction", "object_graph"}
|
|
_MEMORY_KEYS = _COLLECTION_KEYS | _MAPPING_KEYS
|
|
|
|
|
|
class ProjectMemory:
|
|
"""Owns intermediate Agent-mode state and optional debug snapshots."""
|
|
|
|
def __init__(self, artifact_dir: Optional[str] = None) -> None:
|
|
self.artifact_dir = artifact_dir
|
|
self._lock = threading.RLock()
|
|
self._data: Dict[str, Any] = {
|
|
**{key: [] for key in _COLLECTION_KEYS},
|
|
**{key: {} for key in _MAPPING_KEYS},
|
|
}
|
|
if artifact_dir:
|
|
os.makedirs(artifact_dir, exist_ok=True)
|
|
|
|
def replace(self, key: str, value: Any) -> None:
|
|
"""Replace one named memory section."""
|
|
self._validate_key(key)
|
|
with self._lock:
|
|
self._data[key] = copy.deepcopy(value)
|
|
|
|
def append(self, key: str, value: Dict[str, Any]) -> None:
|
|
"""Append one artifact to a list-backed memory section."""
|
|
if key not in _COLLECTION_KEYS:
|
|
raise KeyError(f"{key!r} is not an appendable memory section")
|
|
with self._lock:
|
|
self._data[key].append(copy.deepcopy(value))
|
|
|
|
def extend(self, key: str, values: Iterable[Dict[str, Any]]) -> None:
|
|
"""Append several artifacts under one lock."""
|
|
if key not in _COLLECTION_KEYS:
|
|
raise KeyError(f"{key!r} is not an appendable memory section")
|
|
with self._lock:
|
|
self._data[key].extend(copy.deepcopy(list(values)))
|
|
|
|
def snapshot(self) -> Dict[str, Any]:
|
|
"""Return a detached, JSON-serializable view of current state."""
|
|
with self._lock:
|
|
return copy.deepcopy(self._data)
|
|
|
|
def dump(self, filename: str = "memory.json") -> Optional[str]:
|
|
"""Persist a snapshot when this job has an artifact directory."""
|
|
if not self.artifact_dir:
|
|
return None
|
|
path = os.path.join(self.artifact_dir, filename)
|
|
temp_path = f"{path}.tmp"
|
|
with open(temp_path, "w", encoding="utf-8") as f:
|
|
json.dump(self.snapshot(), f, indent=2)
|
|
os.replace(temp_path, path)
|
|
return path
|
|
|
|
@staticmethod
|
|
def _validate_key(key: str) -> None:
|
|
if key not in _MEMORY_KEYS:
|
|
raise KeyError(f"Unknown project memory section: {key!r}")
|