Teknium
460 lines
24 KiB
Python
460 lines
24 KiB
Python
"""MemoryStore — bounded, file-backed curated memory (MEMORY.md / USER.md).
|
|
Entries are joined by ``ENTRY_DELIMITER``; budgets are in chars (model-independent).
|
|
Module state that tests monkeypatch (``get_memory_dir``, ``fcntl``/``msvcrt``) stays
|
|
in ``tools.memory_tool`` and is read lazily."""
|
|
|
|
import logging
|
|
import time
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
from utils import atomic_write_text
|
|
from tools.threat_patterns import first_threat_message as _first_threat_message
|
|
|
|
logger = logging.getLogger("tools.memory_tool")
|
|
|
|
# Block header prefixes rendered by _render_block; agent/conversation_compression.py
|
|
# matches them to detect a leftover block for an emptied target — keep in lockstep.
|
|
MEMORY_BLOCK_HEADERS = {
|
|
"memory": "MEMORY (your personal notes)", "user": "USER PROFILE (who the user is)"}
|
|
|
|
ENTRY_DELIMITER = "\n§\n"
|
|
|
|
def _memory_dir() -> Path:
|
|
from tools import memory_tool
|
|
|
|
return memory_tool.get_memory_dir()
|
|
|
|
|
|
def _scan_memory_content(content: str) -> Optional[str]:
|
|
"""Error string if *content* matches injection/exfil patterns. Strict scope:
|
|
memory enters the system prompt, so a poisoned entry persists across sessions."""
|
|
return _first_threat_message(content, scope="strict")
|
|
|
|
|
|
def _drift_error(path: "Path", bak_path: str) -> Dict[str, Any]:
|
|
"""External drift: the file wouldn't round-trip, so flushing would discard content."""
|
|
return {
|
|
"success": False,
|
|
"error": (
|
|
f"Refusing to write {path.name}: file on disk has content that wouldn't round-trip "
|
|
f"through the memory tool (likely added by the patch tool, a shell append, a manual edit, "
|
|
f"or a concurrent session). A snapshot was saved to {bak_path}. Resolve the drift first — "
|
|
f"either rewrite the file as a clean §-delimited list of entries, or move the extra "
|
|
f"content out — then retry. This guard exists to prevent silent data loss (issue #26045)."
|
|
),
|
|
"drift_backup": bak_path,
|
|
"remediation": (
|
|
"Open the .bak file, integrate the missing entries into the memory tool one at a time via "
|
|
"memory(action=add, content=...), then remove or rewrite the original file to a clean state."
|
|
)}
|
|
|
|
|
|
def _read_failed_error(path: "Path") -> Dict[str, Any]:
|
|
"""Existing-but-unreadable file: saving from an assumed-empty view would wipe it."""
|
|
return {"success": False, "error": (
|
|
f"Refusing to write {path.name}: the file exists on disk but could not be read right now "
|
|
f"(temporarily locked by another program, a permission change, invalid/corrupt text encoding, "
|
|
f"or a filesystem error). Treating an unreadable file as empty and saving would wipe existing "
|
|
f"memory, so the write is refused. Nothing was changed — retry in a moment.")}
|
|
|
|
|
|
def _find_unique_match(entries: List[str], old_text: str) -> Tuple[Optional[int], bool]:
|
|
"""``(index, ambiguous)`` for entries containing *old_text*. Exact-duplicate
|
|
matches are safe (first wins); distinct matches → ``(None, True)``."""
|
|
matches = [i for i, e in enumerate(entries) if old_text in e]
|
|
if not matches:
|
|
return None, False
|
|
if len({entries[i] for i in matches}) > 1:
|
|
return None, True
|
|
return matches[0], False
|
|
|
|
|
|
class MemoryStore:
|
|
"""Bounded curated memory with file persistence; one instance per AIAgent.
|
|
``_system_prompt_snapshot`` is frozen at load time (prefix-cache stable);
|
|
``memory_entries`` / ``user_entries`` are live state persisted to disk."""
|
|
|
|
# Failed consolidation attempts (overflow / zero-match) allowed per turn before
|
|
# a TERMINAL "save skipped" result, so a fragile replace/add can't loop the turn
|
|
# to budget exhaustion and suppress the user's reply.
|
|
_MAX_CONSOLIDATION_FAILURES_PER_TURN = 3
|
|
|
|
def __init__(self, memory_char_limit: int = 2200, user_char_limit: int = 1375, *,
|
|
memory_enabled: bool = True, user_profile_enabled: bool = True):
|
|
self.memory_entries: List[str] = []
|
|
self.user_entries: List[str] = []
|
|
self.memory_char_limit = memory_char_limit
|
|
self.user_char_limit = user_char_limit
|
|
self.memory_enabled = memory_enabled
|
|
self.user_profile_enabled = user_profile_enabled
|
|
self._system_prompt_snapshot: Dict[str, str] = {"memory": "", "user": ""}
|
|
self._consolidation_failures = 0 # per turn; reset by reset_consolidation_failures()
|
|
|
|
def target_enabled(self, target: str) -> bool:
|
|
"""Return whether this session's selected built-in store is writable."""
|
|
return self.user_profile_enabled if target == "user" else self.memory_enabled
|
|
|
|
def reset_consolidation_failures(self) -> None:
|
|
"""Reset the per-turn consolidation-failure counter (call at turn start)."""
|
|
self._consolidation_failures = 0
|
|
|
|
def _consolidation_failure(self, response: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Count a consolidation failure: under the per-turn cap return ``response``
|
|
(it says how to retry); past it a TERMINAL result so the model stops looping."""
|
|
self._consolidation_failures += 1
|
|
if self._consolidation_failures <= self._MAX_CONSOLIDATION_FAILURES_PER_TURN:
|
|
return response
|
|
return {"success": False, "done": True, "error": (
|
|
f"Memory consolidation failed {self._consolidation_failures} times this turn. Stop retrying "
|
|
"memory calls — leave memory unchanged for now and continue with your reply to the user. "
|
|
"The fact can be saved in a later turn.")}
|
|
|
|
def load_from_disk(self):
|
|
"""Load MEMORY.md / USER.md and capture the frozen system-prompt snapshot.
|
|
Threat hits are replaced by a ``[BLOCKED: …]`` placeholder in the SNAPSHOT only;
|
|
live lists keep the raw text so the user can see and remove poisoned entries
|
|
(dropping them silently would hide the attack)."""
|
|
mem_dir = _memory_dir()
|
|
mem_dir.mkdir(parents=True, exist_ok=True)
|
|
# Deduplicate (order-preserving, first occurrence wins).
|
|
self.memory_entries = list(dict.fromkeys(self._read_file(mem_dir / "MEMORY.md")))
|
|
self.user_entries = list(dict.fromkeys(self._read_file(mem_dir / "USER.md")))
|
|
self._system_prompt_snapshot = {
|
|
target: self._render_block(target, self._sanitize_entries_for_snapshot(entries, filename))
|
|
for target, entries, filename in (
|
|
("memory", self.memory_entries, "MEMORY.md"), ("user", self.user_entries, "USER.md"))
|
|
}
|
|
|
|
@staticmethod
|
|
def _sanitize_entries_for_snapshot(entries: List[str], filename: str) -> List[str]:
|
|
"""*entries* with threat matches replaced by a ``[BLOCKED: …]`` placeholder
|
|
(strict scope, same as writes); empty / already-blocked entries pass through."""
|
|
from tools.threat_patterns import scan_for_threats
|
|
|
|
sanitized: List[str] = []
|
|
for entry in entries:
|
|
findings = scan_for_threats(entry, scope="strict") if entry and not entry.startswith("[BLOCKED:") else None
|
|
if not findings:
|
|
sanitized.append(entry)
|
|
continue
|
|
logger.warning("Memory entry from %s blocked at load time: %s", filename, ", ".join(findings))
|
|
sanitized.append(f"[BLOCKED: {filename} entry contained threat pattern(s): {', '.join(findings)}. "
|
|
f"Removed from system prompt; use memory(action=remove) to delete the original.]")
|
|
return sanitized
|
|
|
|
@staticmethod
|
|
@contextmanager
|
|
def _file_lock(path: Path):
|
|
"""Exclusive lock on a separate .lock file so the memory file itself can
|
|
still be atomically replaced."""
|
|
from tools import memory_tool as _mt # fcntl/msvcrt live (and are patched) there
|
|
|
|
fcntl, msvcrt = _mt.fcntl, _mt.msvcrt
|
|
lock_path = path.with_suffix(path.suffix + ".lock")
|
|
lock_path.parent.mkdir(parents=True, exist_ok=True)
|
|
if fcntl is None and msvcrt is None:
|
|
yield
|
|
return
|
|
fd = open(lock_path, "a+", encoding="utf-8")
|
|
|
|
def _flock(unlock: bool):
|
|
if fcntl:
|
|
fcntl.flock(fd, fcntl.LOCK_UN if unlock else fcntl.LOCK_EX)
|
|
else:
|
|
fd.seek(0)
|
|
msvcrt.locking(fd.fileno(), msvcrt.LK_UNLCK if unlock else msvcrt.LK_LOCK, 1)
|
|
|
|
try:
|
|
_flock(False)
|
|
yield
|
|
finally:
|
|
try:
|
|
_flock(True)
|
|
except OSError:
|
|
pass
|
|
fd.close()
|
|
|
|
@staticmethod
|
|
def _path_for(target: str) -> Path:
|
|
return _memory_dir() / ("USER.md" if target == "user" else "MEMORY.md")
|
|
|
|
def _reload_or_error(self, target: str, *, skip_drift: bool = False) -> Optional[Dict[str, Any]]:
|
|
"""Re-read entries from disk (under lock) before mutating; return the abort
|
|
error dict or None. Aborts on external drift (flushing would discard
|
|
un-roundtrippable content) and on an existing-but-unreadable file (even
|
|
append-only ``add`` rewrites the whole file). Drift check and parse use the
|
|
SAME raw snapshot — a failed second read used to count as "no drift"."""
|
|
path = self._path_for(target)
|
|
raw, read_ok = self._read_raw_checked(path)
|
|
if not read_ok:
|
|
return _read_failed_error(path)
|
|
bak = None if skip_drift else self._detect_external_drift(target, raw)
|
|
self._set_entries(target, list(dict.fromkeys(self._parse_entries(raw))))
|
|
return _drift_error(path, bak) if bak else None
|
|
|
|
def save_to_disk(self, target: str):
|
|
"""Persist entries to the appropriate file. Called after every mutation."""
|
|
_memory_dir().mkdir(parents=True, exist_ok=True)
|
|
self._write_file(self._path_for(target), self._entries_for(target))
|
|
|
|
def _entries_for(self, target: str) -> List[str]:
|
|
return self.user_entries if target == "user" else self.memory_entries
|
|
|
|
def _set_entries(self, target: str, entries: List[str]):
|
|
setattr(self, "user_entries" if target == "user" else "memory_entries", entries)
|
|
|
|
def _char_count(self, target: str) -> int:
|
|
return len(ENTRY_DELIMITER.join(self._entries_for(target)))
|
|
|
|
def _char_limit(self, target: str) -> int:
|
|
return self.user_char_limit if target == "user" else self.memory_char_limit
|
|
|
|
def _usage(self, target: str) -> str:
|
|
return f"{self._char_count(target):,}/{self._char_limit(target):,}"
|
|
|
|
def _usage_pct(self, target: str, current: int) -> str:
|
|
"""``"<pct>% — <current>/<limit> chars"`` for the given target."""
|
|
limit = self._char_limit(target)
|
|
pct = min(100, int((current / limit) * 100)) if limit > 0 else 0
|
|
return f"{pct}% — {current:,}/{limit:,} chars"
|
|
|
|
def _failure_with_entries(self, target: str, message: str) -> Dict[str, Any]:
|
|
"""Consolidation failure carrying the live entries so the model can consolidate."""
|
|
return self._consolidation_failure({"success": False, "error": message,
|
|
"current_entries": self._entries_for(target), "usage": self._usage(target)})
|
|
|
|
def _locate(self, target: str, old_text: str, verb: str):
|
|
"""Resolve *old_text* to a unique entry index, or an error dict."""
|
|
entries = self._entries_for(target)
|
|
idx, ambiguous = _find_unique_match(entries, old_text)
|
|
if ambiguous:
|
|
return None, {"success": False, "error": f"Multiple entries matched '{old_text}'. Be more specific.",
|
|
"matches": self._previews([e for e in entries if old_text in e])}
|
|
if idx is None:
|
|
return None, self._consolidation_failure({
|
|
"success": False,
|
|
"error": f"No entry matched '{old_text}'. Check current_entries below and retry with the exact text of the entry you want to {verb}.",
|
|
"current_entries": entries})
|
|
return idx, None
|
|
|
|
def _commit(self, target: str, entries: List[str], message: str) -> Dict[str, Any]:
|
|
self._set_entries(target, entries)
|
|
self.save_to_disk(target)
|
|
return self._success_response(target, message)
|
|
|
|
def add(self, target: str, content: str) -> Dict[str, Any]:
|
|
"""Append a new entry. Returns error if it would exceed the char limit."""
|
|
content = content.strip()
|
|
if not content:
|
|
return {"success": False, "error": "Content cannot be empty."}
|
|
scan_error = _scan_memory_content(content)
|
|
if scan_error:
|
|
return {"success": False, "error": scan_error}
|
|
with self._file_lock(self._path_for(target)):
|
|
# Append-only: skip the drift guard (appending never clobbers foreign
|
|
# content) but still refuse a failed read — add rewrites the WHOLE file.
|
|
err = self._reload_or_error(target, skip_drift=True)
|
|
if err:
|
|
return err
|
|
entries = self._entries_for(target)
|
|
limit = self._char_limit(target)
|
|
if content in entries:
|
|
return self._success_response(target, "Entry already exists (no duplicate added).")
|
|
if len(ENTRY_DELIMITER.join(entries + [content])) > limit:
|
|
return self._failure_with_entries(target, (
|
|
f"Memory at {self._char_count(target):,}/{limit:,} chars. Adding this entry "
|
|
f"({len(content)} chars) would exceed the limit. Consolidate now: use 'replace' to merge "
|
|
f"overlapping entries into shorter ones or 'remove' stale or less important entries (see "
|
|
f"current_entries below), then retry this add — all in this turn."))
|
|
entries.append(content)
|
|
return self._commit(target, entries, "Entry added.")
|
|
|
|
def replace(self, target: str, old_text: str, new_content: str) -> Dict[str, Any]:
|
|
"""Find entry containing old_text substring, replace it with new_content."""
|
|
new_content = new_content.strip()
|
|
if not old_text.strip():
|
|
return {"success": False, "error": "old_text cannot be empty."}
|
|
if not new_content:
|
|
return {"success": False, "error": "new_content cannot be empty. Use 'remove' to delete entries."}
|
|
scan_error = _scan_memory_content(new_content)
|
|
if scan_error:
|
|
return {"success": False, "error": scan_error}
|
|
return self._edit(target, old_text.strip(), new_content)
|
|
|
|
def remove(self, target: str, old_text: str) -> Dict[str, Any]:
|
|
"""Remove the entry containing old_text substring."""
|
|
if not old_text.strip():
|
|
return {"success": False, "error": "old_text cannot be empty."}
|
|
return self._edit(target, old_text.strip(), None)
|
|
|
|
def _edit(self, target: str, old_text: str, new_content: Optional[str]) -> Dict[str, Any]:
|
|
"""Locked replace (``new_content`` set) or remove (None) of the unique entry matching *old_text*."""
|
|
with self._file_lock(self._path_for(target)):
|
|
err = self._reload_or_error(target)
|
|
if err:
|
|
return err
|
|
idx, err = self._locate(target, old_text, "replace" if new_content else "remove")
|
|
if err:
|
|
return err
|
|
entries = self._entries_for(target)
|
|
if new_content is None:
|
|
entries.pop(idx)
|
|
return self._commit(target, entries, "Entry removed.")
|
|
limit = self._char_limit(target)
|
|
new_total = len(ENTRY_DELIMITER.join(entries[:idx] + [new_content] + entries[idx + 1:]))
|
|
if new_total > limit:
|
|
return self._failure_with_entries(target, (
|
|
f"Replacement would put memory at {new_total:,}/{limit:,} chars. Shorten the new content, "
|
|
f"or 'remove' other stale or less important entries to make room (see current_entries "
|
|
f"below), then retry — all in this turn."))
|
|
entries[idx] = new_content
|
|
return self._commit(target, entries, "Entry replaced.")
|
|
|
|
@staticmethod
|
|
def _apply_batch_op(working: List[str], act: str, content: str, old_text: str, pos: str) -> Optional[str]:
|
|
"""Apply one batch op to *working* in place; return an error message or None."""
|
|
if act == "add":
|
|
if not content:
|
|
return f"{pos}: content is required."
|
|
if content not in working: # idempotent -- skip duplicate, don't fail the batch
|
|
working.append(content)
|
|
return None
|
|
if act not in ("replace", "remove"):
|
|
return f"{pos}: unknown action. Use add, replace, or remove."
|
|
if not old_text:
|
|
return f"{pos}: old_text is required."
|
|
if act == "replace" and not content:
|
|
return f"{pos}: content is required (use action='remove' to delete)."
|
|
idx, ambiguous = _find_unique_match(working, old_text)
|
|
if ambiguous:
|
|
return f"{pos}: '{old_text}' matched multiple distinct entries -- be more specific."
|
|
if idx is None:
|
|
return f"{pos}: no entry matched '{old_text}'."
|
|
if act == "replace":
|
|
working[idx] = content
|
|
else:
|
|
working.pop(idx)
|
|
return None
|
|
|
|
def apply_batch(self, target: str, operations: List[Dict[str, Any]]) -> Dict[str, Any]:
|
|
"""Apply add/replace/remove ops to one target atomically against the FINAL
|
|
budget, so one call can free space and add entries. All-or-nothing: any
|
|
malformed / unmatched op or an over-limit result writes NOTHING and returns
|
|
the first failure plus live state."""
|
|
if not operations:
|
|
return {"success": False, "error": "operations list is empty."}
|
|
|
|
# Scan every add/replace content BEFORE touching disk -- one poisoned op rejects the batch.
|
|
for i, op in enumerate(operations):
|
|
op = op or {}
|
|
scan_error = op.get("action") in {"add", "replace"} and op.get("content") and _scan_memory_content(op["content"])
|
|
if scan_error:
|
|
return {"success": False, "error": f"Operation {i + 1}: {scan_error}"}
|
|
|
|
with self._file_lock(self._path_for(target)):
|
|
err = self._reload_or_error(target)
|
|
if err:
|
|
return err
|
|
# Work on a copy; only commit if the whole batch validates.
|
|
working: List[str] = list(self._entries_for(target))
|
|
limit = self._char_limit(target)
|
|
for i, op in enumerate(operations):
|
|
op = op or {}
|
|
act = op.get("action")
|
|
content = (op.get("content") or op.get("new_text") or "").strip()
|
|
old_text = (op.get("old_text") or "").strip()
|
|
pos = f"Operation {i + 1} ({act or 'unknown'})"
|
|
msg = self._apply_batch_op(working, act, content, old_text, pos)
|
|
if msg:
|
|
return self._failure_with_entries(
|
|
target, msg + " No operations were applied (batch is all-or-nothing).")
|
|
# Budget check against the FINAL state only.
|
|
new_total = len(ENTRY_DELIMITER.join(working))
|
|
if new_total > limit:
|
|
return self._failure_with_entries(target, (
|
|
f"After applying all {len(operations)} operations, memory would be at "
|
|
f"{new_total:,}/{limit:,} chars -- over the limit. Remove or shorten more "
|
|
f"entries in the same batch (see current_entries below), then retry."))
|
|
return self._commit(target, working, f"Applied {len(operations)} operation(s).")
|
|
|
|
def format_for_system_prompt(self, target: str) -> Optional[str]:
|
|
"""Frozen load-time snapshot for the system prompt (NOT live state — mid-session
|
|
writes don't touch it, preserving the prefix cache); None if empty."""
|
|
return self._system_prompt_snapshot.get(target, "") or None
|
|
|
|
@staticmethod
|
|
def _previews(entries: List[str], width: int = 80) -> List[str]:
|
|
"""Truncated one-line previews of entries for error feedback."""
|
|
return [e[:width] + ("..." if len(e) > width else "") for e in entries]
|
|
|
|
def _success_response(self, target: str, message: str = None) -> Dict[str, Any]:
|
|
# A successful write is progress: reset the per-turn (consecutive) failure budget.
|
|
self._consolidation_failures = 0
|
|
# TERMINAL and WITHOUT the entries list: echoing entries invites the model to
|
|
# "find more to fix" and re-issue the same ops. Entries only appear on errors.
|
|
return {"success": True, "done": True, "target": target,
|
|
"usage": self._usage_pct(target, self._char_count(target)),
|
|
"entry_count": len(self._entries_for(target)), **({"message": message} if message else {}),
|
|
"note": "Write saved. This update is complete — do not repeat it."}
|
|
|
|
def _render_block(self, target: str, entries: List[str]) -> str:
|
|
"""Render a system prompt block with header and usage indicator."""
|
|
if not entries:
|
|
return ""
|
|
content = ENTRY_DELIMITER.join(entries)
|
|
title = MEMORY_BLOCK_HEADERS["user" if target == "user" else "memory"]
|
|
separator = "═" * 46
|
|
return f"{separator}\n{title} [{self._usage_pct(target, len(content))}]\n{separator}\n{content}"
|
|
|
|
@staticmethod
|
|
def _read_raw_checked(path: Path) -> Tuple[str, bool]:
|
|
"""``(raw, read_ok)``; ``read_ok`` is False ONLY when the file EXISTS but can't
|
|
be read (absent → ``("", True)``). Decoding stays STRICT: ``errors="replace"``
|
|
would hand callers a lossy view that a save then persists. ``utf-8-sig`` strips
|
|
a Notepad BOM that otherwise glues U+FEFF onto the first entry forever."""
|
|
if not path.exists():
|
|
return "", True
|
|
try:
|
|
return path.read_text(encoding="utf-8-sig"), True
|
|
except (OSError, UnicodeDecodeError):
|
|
return "", False
|
|
|
|
@staticmethod
|
|
def _parse_entries(raw: str) -> List[str]:
|
|
"""Stripped, non-empty entries; splits on the FULL delimiter so a bare "§" survives."""
|
|
return [e for e in (x.strip() for x in raw.split(ENTRY_DELIMITER)) if e]
|
|
|
|
@staticmethod
|
|
def _read_file(path: Path) -> List[str]:
|
|
"""Entries of a memory file ([] on any error). Read-only callers only
|
|
(``load_from_disk``, learning_mutations); mutation paths must use
|
|
``_read_raw_checked`` so they can refuse to overwrite an unreadable file."""
|
|
return MemoryStore._parse_entries(MemoryStore._read_raw_checked(path)[0])
|
|
|
|
def _detect_external_drift(self, target: str, raw: str) -> Optional[str]:
|
|
"""Backup path if *raw* shows external drift, else None. Signals: round-trip
|
|
mismatch, or one entry over the whole-file limit (no tool-written entry can —
|
|
an external writer appended free-form text). Snapshots to ``.bak.<ts>``."""
|
|
if not raw.strip():
|
|
return None
|
|
parsed = self._parse_entries(raw)
|
|
if raw.strip() == ENTRY_DELIMITER.join(parsed) and max(map(len, parsed), default=0) <= self._char_limit(target):
|
|
return None
|
|
path = self._path_for(target)
|
|
bak_path = path.with_suffix(path.suffix + f".bak.{int(time.time())}")
|
|
try:
|
|
bak_path.write_text(raw, encoding="utf-8")
|
|
except OSError:
|
|
return str(bak_path) + " (BACKUP FAILED — file unchanged on disk)"
|
|
return str(bak_path)
|
|
|
|
@staticmethod
|
|
def _write_file(path: Path, entries: List[str]):
|
|
"""Atomic temp-file + rename: readers never see a truncated file."""
|
|
try:
|
|
atomic_write_text(path, ENTRY_DELIMITER.join(entries), tmp_prefix=".mem_")
|
|
except OSError as e:
|
|
raise RuntimeError(f"Failed to write memory file {path}: {e}")
|