修改了nanobot,往Hermes agent的风格走,进度1/3
This commit is contained in:
463
app-instance/backend/beaver/memory/curated/store.py
Normal file
463
app-instance/backend/beaver/memory/curated/store.py
Normal file
@ -0,0 +1,463 @@
|
||||
"""Beaver 的精炼长期记忆存储层。
|
||||
|
||||
这个文件实现的是以 Hermes-agent 为基线的 curated memory 模型,目标不是
|
||||
“把所有历史都存下来”,而是只保存跨会话仍然值得保留的稳定事实。
|
||||
|
||||
核心设计:
|
||||
1. 只保留两个持久化记忆桶:
|
||||
- ``memory``: agent 自己对环境、项目、工具 quirks 的长期备注
|
||||
- ``user``: 对用户偏好、习惯、身份信息的长期理解
|
||||
2. ``replace`` / ``remove`` 不使用 UUID,而是使用短语义片段做子串匹配。
|
||||
这是为了适配 LLM 更擅长“记住一句话片段”而不是“追踪一个随机 ID”的现实。
|
||||
3. 写入前先做安全扫描,避免把 prompt injection / secrets exfiltration
|
||||
一类危险内容写入长期记忆,再在未来会话中反向污染 system prompt。
|
||||
4. 写入协议严格遵守:
|
||||
- scan
|
||||
- lock
|
||||
- reload
|
||||
- validate
|
||||
- atomic write
|
||||
5. 本文件维护两份状态:
|
||||
- live state: 当前内存中的真实条目,tool 写入后立刻变化
|
||||
- frozen snapshot: 会话开始时冻结的一份 prompt 注入快照
|
||||
|
||||
其中最重要的一点是:本会话中新增的记忆会立刻写盘,但不会反向修改本会话
|
||||
已经冻结的 system prompt。这样可以保住 prefix cache,也避免“会话中途 prompt
|
||||
变了导致行为抖动”的问题。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
import tempfile
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
import fcntl
|
||||
except ImportError: # pragma: no cover - Windows fallback
|
||||
fcntl = None
|
||||
|
||||
try:
|
||||
import msvcrt
|
||||
except ImportError: # pragma: no cover - Unix platforms
|
||||
msvcrt = None
|
||||
|
||||
ENTRY_DELIMITER = "\n§\n"
|
||||
DEFAULT_MEMORY_FILENAME = "MEMORY.md"
|
||||
DEFAULT_USER_FILENAME = "USER.md"
|
||||
|
||||
_MEMORY_THREAT_PATTERNS: list[tuple[str, str]] = [
|
||||
(r"ignore\s+(previous|all|above|prior)\s+instructions", "prompt_injection"),
|
||||
(r"you\s+are\s+now\s+", "role_hijack"),
|
||||
(r"do\s+not\s+tell\s+the\s+user", "deception_hide"),
|
||||
(r"system\s+prompt\s+override", "sys_prompt_override"),
|
||||
(r"disregard\s+(your|all|any)\s+(instructions|rules|guidelines)", "disregard_rules"),
|
||||
(r"act\s+as\s+(if|though)\s+you\s+(have\s+no|don't\s+have)\s+(restrictions|limits|rules)", "bypass_restrictions"),
|
||||
(r"curl\s+[^\n]*\$\{?\w*(KEY|TOKEN|SECRET|PASSWORD|CREDENTIAL|API)", "exfil_curl"),
|
||||
(r"wget\s+[^\n]*\$\{?\w*(KEY|TOKEN|SECRET|PASSWORD|CREDENTIAL|API)", "exfil_wget"),
|
||||
(r"cat\s+[^\n]*(\.env|credentials|\.netrc|\.pgpass|\.npmrc|\.pypirc)", "read_secrets"),
|
||||
(r"authorized_keys", "ssh_backdoor"),
|
||||
(r"\$HOME/\.ssh|\~/\.ssh", "ssh_access"),
|
||||
(r"\$HOME/\.beaver/\.env|\~/\.beaver/\.env", "beaver_env"),
|
||||
]
|
||||
|
||||
_INVISIBLE_CHARS = {
|
||||
"\u200b",
|
||||
"\u200c",
|
||||
"\u200d",
|
||||
"\u2060",
|
||||
"\ufeff",
|
||||
"\u202a",
|
||||
"\u202b",
|
||||
"\u202c",
|
||||
"\u202d",
|
||||
"\u202e",
|
||||
}
|
||||
|
||||
|
||||
def scan_memory_content(content: str) -> str | None:
|
||||
"""扫描待写入内容,拦截明显危险的记忆条目。
|
||||
|
||||
这里不是在做完备的安全审计,而是在做“进入长期记忆之前的最低限度闸门”。
|
||||
因为长期记忆会在未来会话中重新注入 system prompt,所以一旦把恶意文本写进去,
|
||||
风险远高于普通临时上下文。
|
||||
"""
|
||||
|
||||
for char in _INVISIBLE_CHARS:
|
||||
if char in content:
|
||||
return (
|
||||
f"Blocked: content contains invisible unicode character "
|
||||
f"U+{ord(char):04X}."
|
||||
)
|
||||
|
||||
for pattern, pattern_id in _MEMORY_THREAT_PATTERNS:
|
||||
if re.search(pattern, content, re.IGNORECASE):
|
||||
return (
|
||||
f"Blocked: content matches threat pattern '{pattern_id}'. "
|
||||
"Memory entries are injected into future system prompts."
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
class MemoryStore:
|
||||
"""带容量上限的长期记忆存储。
|
||||
|
||||
这个类负责:
|
||||
1. 从磁盘加载 `MEMORY.md` / `USER.md`
|
||||
2. 在 session 启动时冻结 prompt snapshot
|
||||
3. 为 `add / replace / remove` 提供安全写接口
|
||||
4. 维护 live state 与 frozen snapshot 的边界
|
||||
|
||||
它不负责:
|
||||
1. 自动从对话里抽取要记住的内容
|
||||
2. session transcript 检索
|
||||
3. skills 的学习和发布
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
root: str | Path,
|
||||
*,
|
||||
memory_char_limit: int = 2200,
|
||||
user_char_limit: int = 1375,
|
||||
) -> None:
|
||||
self.root = Path(root)
|
||||
self.memory_char_limit = memory_char_limit
|
||||
self.user_char_limit = user_char_limit
|
||||
self.memory_entries: list[str] = []
|
||||
self.user_entries: list[str] = []
|
||||
self._system_prompt_snapshot: dict[str, str] = {"memory": "", "user": ""}
|
||||
|
||||
def load_from_disk(self) -> None:
|
||||
"""从磁盘加载 live state,并冻结当前 session 的 prompt snapshot。
|
||||
|
||||
调用时机应该是“会话启动时”,而不是每次工具写入后。
|
||||
如果在每次写入后都重新 load 并更新 system prompt,就会破坏 frozen snapshot
|
||||
这个设计,导致本轮会话 prompt 前缀发生变化。
|
||||
"""
|
||||
|
||||
self.root.mkdir(parents=True, exist_ok=True)
|
||||
self.memory_entries = list(dict.fromkeys(self._read_file(self._path_for("memory"))))
|
||||
self.user_entries = list(dict.fromkeys(self._read_file(self._path_for("user"))))
|
||||
self._system_prompt_snapshot = {
|
||||
"memory": self._render_block("memory", self.memory_entries),
|
||||
"user": self._render_block("user", self.user_entries),
|
||||
}
|
||||
|
||||
@contextmanager
|
||||
def _file_lock(self, path: Path):
|
||||
"""对目标记忆文件加排他锁。
|
||||
|
||||
锁文件使用 sibling `.lock` 文件,而不是直接锁业务文件本身。
|
||||
原因是业务文件使用的是“临时文件写入 + os.replace 原子替换”,如果直接锁目标
|
||||
文件,替换时会让锁语义和文件句柄关系变得更脆弱。
|
||||
"""
|
||||
|
||||
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
|
||||
|
||||
if msvcrt and (not lock_path.exists() or lock_path.stat().st_size == 0):
|
||||
lock_path.write_text(" ", encoding="utf-8")
|
||||
|
||||
fd = open(lock_path, "r+" if msvcrt else "a+", encoding="utf-8")
|
||||
try:
|
||||
if fcntl is not None:
|
||||
fcntl.flock(fd, fcntl.LOCK_EX)
|
||||
elif msvcrt is not None: # pragma: no cover - Windows fallback
|
||||
fd.seek(0)
|
||||
msvcrt.locking(fd.fileno(), msvcrt.LK_LOCK, 1)
|
||||
yield
|
||||
finally:
|
||||
if fcntl is not None:
|
||||
fcntl.flock(fd, fcntl.LOCK_UN)
|
||||
elif msvcrt is not None: # pragma: no cover - Windows fallback
|
||||
try:
|
||||
fd.seek(0)
|
||||
msvcrt.locking(fd.fileno(), msvcrt.LK_UNLCK, 1)
|
||||
except OSError:
|
||||
pass
|
||||
fd.close()
|
||||
|
||||
def _path_for(self, target: str) -> Path:
|
||||
"""根据目标桶返回实际文件路径。"""
|
||||
if target == "user":
|
||||
return self.root / DEFAULT_USER_FILENAME
|
||||
return self.root / DEFAULT_MEMORY_FILENAME
|
||||
|
||||
def _entries_for(self, target: str) -> list[str]:
|
||||
"""读取某个目标桶当前的 live entries。"""
|
||||
if target == "user":
|
||||
return self.user_entries
|
||||
return self.memory_entries
|
||||
|
||||
def _set_entries(self, target: str, entries: list[str]) -> None:
|
||||
"""更新某个目标桶在内存中的 live entries。"""
|
||||
if target == "user":
|
||||
self.user_entries = entries
|
||||
else:
|
||||
self.memory_entries = entries
|
||||
|
||||
def _char_limit(self, target: str) -> int:
|
||||
"""返回目标桶的字符预算。
|
||||
|
||||
这里使用字符数而不是 token 数,是因为字符预算更稳定,也不依赖具体模型。
|
||||
"""
|
||||
return self.user_char_limit if target == "user" else self.memory_char_limit
|
||||
|
||||
def _char_count(self, target: str) -> int:
|
||||
"""返回目标桶当前 live state 的字符占用。"""
|
||||
entries = self._entries_for(target)
|
||||
return len(ENTRY_DELIMITER.join(entries)) if entries else 0
|
||||
|
||||
def _reload_target(self, target: str) -> None:
|
||||
"""在持锁状态下重新从磁盘读取目标桶。
|
||||
|
||||
这是并发安全协议里最关键的一步之一。
|
||||
必须在拿到锁之后 reload,才能确保当前进程不会覆盖掉其他并发会话刚刚写入
|
||||
的最新内容。
|
||||
"""
|
||||
fresh = list(dict.fromkeys(self._read_file(self._path_for(target))))
|
||||
self._set_entries(target, fresh)
|
||||
|
||||
def save_to_disk(self, target: str) -> None:
|
||||
"""把当前 live entries 持久化到磁盘。"""
|
||||
self.root.mkdir(parents=True, exist_ok=True)
|
||||
self._write_file(self._path_for(target), self._entries_for(target))
|
||||
|
||||
def add(self, target: str, content: str) -> dict[str, Any]:
|
||||
"""追加一条新的长期记忆。
|
||||
|
||||
规则:
|
||||
1. 空内容拒绝
|
||||
2. 安全扫描不通过拒绝
|
||||
3. 精确重复拒绝
|
||||
4. 超出字符预算拒绝
|
||||
5. 否则追加并立即写盘
|
||||
"""
|
||||
|
||||
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)):
|
||||
self._reload_target(target)
|
||||
entries = self._entries_for(target)
|
||||
if content in entries:
|
||||
return self._success_response(target, "Entry already exists (skipped duplicate).")
|
||||
|
||||
new_entries = entries + [content]
|
||||
new_total = len(ENTRY_DELIMITER.join(new_entries))
|
||||
limit = self._char_limit(target)
|
||||
if new_total > limit:
|
||||
current = self._char_count(target)
|
||||
return {
|
||||
"success": False,
|
||||
"error": (
|
||||
f"Memory at {current:,}/{limit:,} chars. "
|
||||
f"Adding this entry ({len(content)} chars) would exceed the limit."
|
||||
),
|
||||
"current_entries": list(entries),
|
||||
"usage": f"{current:,}/{limit:,}",
|
||||
}
|
||||
|
||||
entries.append(content)
|
||||
self._set_entries(target, entries)
|
||||
self.save_to_disk(target)
|
||||
|
||||
return self._success_response(target, "Entry added.")
|
||||
|
||||
def replace(self, target: str, old_text: str, new_content: str) -> dict[str, Any]:
|
||||
"""用新的内容替换一条已有记忆。
|
||||
|
||||
这里按 `old_text in entry` 做子串匹配,而不是要求调用方提供完整条目或 UUID。
|
||||
如果命中多条且它们内容不同,会要求调用方给出更精确的片段,避免误替换。
|
||||
"""
|
||||
|
||||
old_text = old_text.strip()
|
||||
new_content = new_content.strip()
|
||||
if not old_text:
|
||||
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}
|
||||
|
||||
with self._file_lock(self._path_for(target)):
|
||||
self._reload_target(target)
|
||||
entries = self._entries_for(target)
|
||||
matches = [(index, entry) for index, entry in enumerate(entries) if old_text in entry]
|
||||
if not matches:
|
||||
return {"success": False, "error": f"No entry matched '{old_text}'."}
|
||||
|
||||
if len(matches) > 1:
|
||||
unique_texts = {entry for _, entry in matches}
|
||||
if len(unique_texts) > 1:
|
||||
return {
|
||||
"success": False,
|
||||
"error": f"Multiple entries matched '{old_text}'. Be more specific.",
|
||||
"matches": [
|
||||
entry[:80] + ("..." if len(entry) > 80 else "")
|
||||
for _, entry in matches
|
||||
],
|
||||
}
|
||||
|
||||
index = matches[0][0]
|
||||
candidate_entries = list(entries)
|
||||
candidate_entries[index] = new_content
|
||||
new_total = len(ENTRY_DELIMITER.join(candidate_entries))
|
||||
limit = self._char_limit(target)
|
||||
if new_total > limit:
|
||||
return {
|
||||
"success": False,
|
||||
"error": (
|
||||
f"Replacement would put memory at {new_total:,}/{limit:,} chars. "
|
||||
"Shorten the new content or remove other entries first."
|
||||
),
|
||||
}
|
||||
|
||||
entries[index] = new_content
|
||||
self._set_entries(target, entries)
|
||||
self.save_to_disk(target)
|
||||
|
||||
return self._success_response(target, "Entry replaced.")
|
||||
|
||||
def remove(self, target: str, old_text: str) -> dict[str, Any]:
|
||||
"""删除一条已有记忆。
|
||||
|
||||
删除和替换共享同样的匹配策略:优先服务于 LLM 可操作性,而不是数据库式的强 ID。
|
||||
"""
|
||||
|
||||
old_text = old_text.strip()
|
||||
if not old_text:
|
||||
return {"success": False, "error": "old_text cannot be empty."}
|
||||
|
||||
with self._file_lock(self._path_for(target)):
|
||||
self._reload_target(target)
|
||||
entries = self._entries_for(target)
|
||||
matches = [(index, entry) for index, entry in enumerate(entries) if old_text in entry]
|
||||
if not matches:
|
||||
return {"success": False, "error": f"No entry matched '{old_text}'."}
|
||||
|
||||
if len(matches) > 1:
|
||||
unique_texts = {entry for _, entry in matches}
|
||||
if len(unique_texts) > 1:
|
||||
return {
|
||||
"success": False,
|
||||
"error": f"Multiple entries matched '{old_text}'. Be more specific.",
|
||||
"matches": [
|
||||
entry[:80] + ("..." if len(entry) > 80 else "")
|
||||
for _, entry in matches
|
||||
],
|
||||
}
|
||||
|
||||
entries.pop(matches[0][0])
|
||||
self._set_entries(target, entries)
|
||||
self.save_to_disk(target)
|
||||
|
||||
return self._success_response(target, "Entry removed.")
|
||||
|
||||
def format_for_system_prompt(self, target: str) -> str | None:
|
||||
"""返回 session 启动时冻结下来的 prompt block。
|
||||
|
||||
这里明确返回的是 frozen snapshot,而不是 live state。
|
||||
所以如果 session 中途调用 `add()` 写入了新记忆,这里不会立刻变化。
|
||||
"""
|
||||
|
||||
block = self._system_prompt_snapshot.get(target, "")
|
||||
return block or None
|
||||
|
||||
def _success_response(self, target: str, message: str | None = None) -> dict[str, Any]:
|
||||
"""统一生成 memory tool 的成功响应。
|
||||
|
||||
响应里返回 live entries 和占用信息,目的是让模型能“看到自己刚写进去什么”,
|
||||
即使 system prompt 仍然保持冻结不变。
|
||||
"""
|
||||
current = self._char_count(target)
|
||||
limit = self._char_limit(target)
|
||||
percent = min(100, int((current / limit) * 100)) if limit > 0 else 0
|
||||
payload: dict[str, Any] = {
|
||||
"success": True,
|
||||
"target": target,
|
||||
"entries": list(self._entries_for(target)),
|
||||
"entry_count": len(self._entries_for(target)),
|
||||
"usage": f"{percent}% — {current:,}/{limit:,} chars",
|
||||
}
|
||||
if message:
|
||||
payload["message"] = message
|
||||
return payload
|
||||
|
||||
def _render_block(self, target: str, entries: list[str]) -> str:
|
||||
"""把条目渲染成适合注入 system prompt 的块。"""
|
||||
if not entries:
|
||||
return ""
|
||||
|
||||
current = len(ENTRY_DELIMITER.join(entries))
|
||||
limit = self._char_limit(target)
|
||||
percent = min(100, int((current / limit) * 100)) if limit > 0 else 0
|
||||
if target == "user":
|
||||
header = f"USER PROFILE (who the user is) [{percent}% — {current:,}/{limit:,} chars]"
|
||||
else:
|
||||
header = f"MEMORY (your personal notes) [{percent}% — {current:,}/{limit:,} chars]"
|
||||
separator = "═" * 46
|
||||
return f"{separator}\n{header}\n{separator}\n{ENTRY_DELIMITER.join(entries)}"
|
||||
|
||||
@staticmethod
|
||||
def _read_file(path: Path) -> list[str]:
|
||||
"""读取记忆文件并按 entry delimiter 拆分。
|
||||
|
||||
这里不额外加读锁,因为写入采用的是原子替换:读者只会看到旧完整文件或新完整文件,
|
||||
不会看到半写入状态。
|
||||
"""
|
||||
if not path.exists():
|
||||
return []
|
||||
try:
|
||||
raw = path.read_text(encoding="utf-8")
|
||||
except OSError:
|
||||
return []
|
||||
if not raw.strip():
|
||||
return []
|
||||
return [entry for entry in (item.strip() for item in raw.split(ENTRY_DELIMITER)) if entry]
|
||||
|
||||
@staticmethod
|
||||
def _write_file(path: Path, entries: list[str]) -> None:
|
||||
"""以原子方式写入记忆文件。
|
||||
|
||||
这里不能直接 `open(path, "w")`,因为那会先截断原文件,再写新内容。
|
||||
如果恰好此时别的进程正在读,就可能读到空文件或半成品。
|
||||
|
||||
正确方式是:
|
||||
1. 在同目录创建临时文件
|
||||
2. 写入并 fsync
|
||||
3. 使用 `os.replace()` 原子替换
|
||||
"""
|
||||
content = ENTRY_DELIMITER.join(entries) if entries else ""
|
||||
fd, tmp_path = tempfile.mkstemp(dir=str(path.parent), suffix=".tmp", prefix=".mem_")
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as handle:
|
||||
handle.write(content)
|
||||
handle.flush()
|
||||
os.fsync(handle.fileno())
|
||||
os.replace(tmp_path, path)
|
||||
except BaseException:
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
Reference in New Issue
Block a user