380 lines
13 KiB
Python
380 lines
13 KiB
Python
"""分层记忆系统 — SOUL/USER/MEMORY/DAILY 文件管理.
|
||
|
||
参考 Hermes/OpenClaw 架构，实现 Agent 人格、用户档案、工作笔记、
|
||
日志的持久化存储与 system prompt 注入。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
from dataclasses import dataclass, field
|
||
from datetime import datetime, timezone
|
||
from pathlib import Path
|
||
from typing import Any, Callable
|
||
|
||
|
||
class MemoryFile:
|
||
"""单个记忆文件的管理器，支持 section 级别 CRUD 和容量控制.
|
||
|
||
文件格式为 Markdown，使用 `## Section` 组织内容::
|
||
|
||
## 身份
|
||
我是小王，一个专业的 AI 助手。
|
||
|
||
## 性格
|
||
友好、耐心、注重细节
|
||
|
||
"""
|
||
|
||
def __init__(self, path: Path, char_budget: int | None = None,
|
||
protected_sections: set[str] | None = None):
|
||
self.path = Path(path)
|
||
self.char_budget = char_budget
|
||
self._protected_sections = protected_sections or set()
|
||
|
||
def read(self) -> str:
|
||
"""读取整个文件内容，文件不存在返回空字符串."""
|
||
if not self.path.exists():
|
||
return ""
|
||
return self.path.read_text(encoding="utf-8")
|
||
|
||
def write(self, content: str) -> None:
|
||
"""写入内容，自动创建父目录，超容量时自动裁剪.
|
||
|
||
在内存中完成裁剪后一次性写入，避免中间不一致状态。
|
||
"""
|
||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||
if self.char_budget and len(content) > self.char_budget:
|
||
content = self._trim_content(content, self._protected_sections or None)
|
||
self.path.write_text(content, encoding="utf-8")
|
||
|
||
def read_section(self, name: str) -> str:
|
||
"""读取指定 section 的内容（不含标题行）."""
|
||
content = self.read()
|
||
if not content:
|
||
return ""
|
||
# 匹配 ## name 后面的内容，直到下一个 ## 或文件末尾
|
||
pattern = rf"^## {re.escape(name)}\s*\n(.*?)(?=^## |\Z)"
|
||
match = re.search(pattern, content, re.MULTILINE | re.DOTALL)
|
||
if match:
|
||
return match.group(1).strip()
|
||
return ""
|
||
|
||
def add_section(self, name: str, content: str) -> None:
|
||
"""追加内容到指定 section，不存在则创建."""
|
||
existing = self.read()
|
||
section_content = self.read_section(name)
|
||
if section_content:
|
||
# 追加到已有 section
|
||
old_text = section_content
|
||
new_text = f"{old_text}\n{content}"
|
||
self.replace_section(name, old_text, new_text)
|
||
else:
|
||
# 创建新 section
|
||
new_section = f"\n## {name}\n{content}"
|
||
if existing and not existing.endswith("\n"):
|
||
new_section = "\n" + new_section
|
||
self.write(existing + new_section)
|
||
|
||
def replace_section(self, name: str, old_text: str, new_text: str) -> bool:
|
||
"""替换 section 内的文本，返回是否成功."""
|
||
section_content = self.read_section(name)
|
||
if old_text not in section_content:
|
||
return False
|
||
full_content = self.read()
|
||
# 替换 section 内的文本
|
||
pattern = rf"(^## {re.escape(name)}\s*\n)(.*?)(?=^## |\Z)"
|
||
match = re.search(pattern, full_content, re.MULTILINE | re.DOTALL)
|
||
if not match:
|
||
return False
|
||
original_section_body = match.group(2)
|
||
new_section_body = original_section_body.replace(old_text, new_text, 1)
|
||
updated = full_content[: match.start(2)] + new_section_body + full_content[match.end(2) :]
|
||
self.write(updated)
|
||
return True
|
||
|
||
def remove_section(self, name: str) -> None:
|
||
"""删除整个 section（含标题行）."""
|
||
content = self.read()
|
||
if not content:
|
||
return
|
||
pattern = rf"^## {re.escape(name)}\s*\n.*?(?=^## |\Z)"
|
||
new_content = re.sub(pattern, "", content, flags=re.MULTILINE | re.DOTALL).strip()
|
||
self.write(new_content)
|
||
|
||
def list_sections(self) -> list[str]:
|
||
"""列出所有 section 名称."""
|
||
content = self.read()
|
||
if not content:
|
||
return []
|
||
return re.findall(r"^## (.+)$", content, re.MULTILINE)
|
||
|
||
def trim_to_budget(self, protected_sections: set[str] | None = None) -> None:
|
||
"""裁剪内容到容量上限，按 section 边界截断.
|
||
|
||
保持原始 section 顺序，仅从后向前移除非保护 section。
|
||
protected_sections 中的 section 始终保留，不参与裁剪。
|
||
"""
|
||
if not self.char_budget:
|
||
return
|
||
content = self.read()
|
||
if len(content) <= self.char_budget:
|
||
return
|
||
trimmed = self._trim_content(content, protected_sections)
|
||
self.path.write_text(trimmed, encoding="utf-8")
|
||
|
||
def _trim_content(self, content: str, protected_sections: set[str] | None = None) -> str:
|
||
"""在内存中裁剪内容到容量上限，返回裁剪后的字符串（不写文件）.
|
||
|
||
保持原始 section 顺序，仅从后向前移除非保护 section。
|
||
"""
|
||
if not self.char_budget or len(content) <= self.char_budget:
|
||
return content
|
||
|
||
protected = protected_sections or set()
|
||
|
||
# 解析所有 section 及其位置
|
||
sections: list[tuple[str, int, int]] = [] # (name, start, end)
|
||
for match in re.finditer(r"^## (.+)$", content, re.MULTILINE):
|
||
name = match.group(1).strip()
|
||
start = match.start()
|
||
next_match = re.search(r"^## ", content[match.end():], re.MULTILINE)
|
||
if next_match:
|
||
end = match.end() + next_match.start()
|
||
else:
|
||
end = len(content)
|
||
sections.append((name, start, end))
|
||
|
||
if not sections:
|
||
return content[:self.char_budget]
|
||
|
||
# 保持原始顺序，标记每个 section 是否受保护
|
||
ordered: list[tuple[str, str, bool]] = [] # (name, text, is_protected)
|
||
for name, start, end in sections:
|
||
ordered.append((name, content[start:end], name in protected))
|
||
|
||
# 从后向前移除非保护 section，直到总长度在预算内
|
||
while ordered:
|
||
total = len("\n\n".join(s[1] for s in ordered))
|
||
if total <= self.char_budget:
|
||
break
|
||
# 从后向前找第一个非保护 section 移除
|
||
for i in range(len(ordered) - 1, -1, -1):
|
||
if not ordered[i][2]:
|
||
ordered.pop(i)
|
||
break
|
||
else:
|
||
break # 所有剩余 section 都是受保护的
|
||
|
||
return "\n\n".join(s[1] for s in ordered).strip()
|
||
|
||
|
||
@dataclass
|
||
class MemorySnapshot:
|
||
"""一次加载的所有记忆文件快照."""
|
||
|
||
soul: str = ""
|
||
user: str = ""
|
||
memory: str = ""
|
||
daily: str = ""
|
||
total_chars: int = 0
|
||
|
||
def is_empty(self) -> bool:
|
||
return not any([self.soul, self.user, self.memory, self.daily])
|
||
|
||
|
||
# 容量上限常量（字符数）
|
||
SOUL_BUDGET = 2000
|
||
USER_BUDGET = 1400
|
||
MEMORY_BUDGET = 2200
|
||
DAILY_BUDGET = 1000 # 每天日志上限
|
||
|
||
# 默认 SOUL.md 内容
|
||
DEFAULT_SOUL = """## 身份
|
||
我是 AgentKit，一个专业的 AI 助手。
|
||
|
||
## 性格
|
||
专业、友好、注重细节
|
||
|
||
## 说话方式
|
||
简洁清晰，偶尔使用比喻帮助理解
|
||
|
||
## 做事准则
|
||
- 准确回答用户问题
|
||
- 主动记住用户提到的偏好和信息
|
||
- 不确定时坦诚说明
|
||
"""
|
||
|
||
|
||
class MemoryStore:
|
||
"""管理 SOUL/USER/MEMORY/DAILY 四类记忆文件.
|
||
|
||
存储路径::
|
||
|
||
base_dir/
|
||
├── SOUL.md
|
||
├── memories/
|
||
│ ├── USER.md
|
||
│ ├── MEMORY.md
|
||
│ └── daily/
|
||
│ ├── 2026-06-07.md
|
||
│ └── 2026-06-08.md
|
||
|
||
"""
|
||
|
||
def __init__(self, base_dir: Path | str | None = None,
|
||
on_change: Callable[[str], None] | None = None):
|
||
if base_dir is None:
|
||
base_dir = Path.home() / ".agentkit"
|
||
self.base_dir = Path(base_dir)
|
||
self.base_dir.mkdir(parents=True, exist_ok=True)
|
||
self._on_change = on_change
|
||
self._base_prompt: str = ""
|
||
|
||
# 初始化四个 MemoryFile
|
||
self._soul = MemoryFile(
|
||
self.base_dir / "SOUL.md",
|
||
char_budget=SOUL_BUDGET,
|
||
protected_sections={"版本", "更新历史"},
|
||
)
|
||
self._user = MemoryFile(self.base_dir / "memories" / "USER.md", char_budget=USER_BUDGET)
|
||
self._memory = MemoryFile(self.base_dir / "memories" / "MEMORY.md", char_budget=MEMORY_BUDGET)
|
||
self._daily_dir = self.base_dir / "memories" / "daily"
|
||
self._daily_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
def get_file(self, file_key: str) -> MemoryFile:
|
||
"""获取指定类型的 MemoryFile.
|
||
|
||
Args:
|
||
file_key: "soul" | "user" | "memory" | "daily"
|
||
"""
|
||
mapping = {
|
||
"soul": self._soul,
|
||
"user": self._user,
|
||
"memory": self._memory,
|
||
}
|
||
if file_key in mapping:
|
||
return mapping[file_key]
|
||
if file_key == "daily":
|
||
# daily 返回今天的日志文件
|
||
today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
||
return MemoryFile(self._daily_dir / f"{today}.md", char_budget=DAILY_BUDGET)
|
||
raise ValueError(f"Invalid file_key: {file_key}. Must be soul/user/memory/daily")
|
||
|
||
def ensure_defaults(self) -> None:
|
||
"""首次运行时创建默认 SOUL.md."""
|
||
if not self._soul.read():
|
||
self._soul.write(DEFAULT_SOUL.strip())
|
||
|
||
def load_all(self) -> MemorySnapshot:
|
||
"""加载所有记忆文件."""
|
||
soul = self._soul.read()
|
||
user = self._user.read()
|
||
memory = self._memory.read()
|
||
daily = self.load_daily_logs()
|
||
total = len(soul) + len(user) + len(memory) + len(daily)
|
||
return MemorySnapshot(
|
||
soul=soul,
|
||
user=user,
|
||
memory=memory,
|
||
daily=daily,
|
||
total_chars=total,
|
||
)
|
||
|
||
def load_daily_logs(self, days: int = 2) -> str:
|
||
"""加载最近 N 天的日志."""
|
||
parts: list[str] = []
|
||
for i in range(days):
|
||
from datetime import timedelta
|
||
|
||
date = datetime.now(timezone.utc) - timedelta(days=i)
|
||
filename = f"{date.strftime('%Y-%m-%d')}.md"
|
||
daily_file = MemoryFile(self._daily_dir / filename)
|
||
content = daily_file.read()
|
||
if content:
|
||
parts.append(f"### {date.strftime('%Y-%m-%d')}\n{content}")
|
||
return "\n\n".join(parts)
|
||
|
||
def archive_old_dailies(self, keep_days: int = 2) -> int:
|
||
"""归档超过 N 天的日志（删除旧文件）."""
|
||
from datetime import timedelta
|
||
|
||
count = 0
|
||
cutoff = datetime.now(timezone.utc) - timedelta(days=keep_days)
|
||
if not self._daily_dir.exists():
|
||
return 0
|
||
for f in self._daily_dir.glob("*.md"):
|
||
# 从文件名解析日期
|
||
try:
|
||
date_str = f.stem # e.g. "2026-06-07"
|
||
file_date = datetime.strptime(date_str, "%Y-%m-%d").replace(tzinfo=timezone.utc)
|
||
if file_date < cutoff:
|
||
f.unlink()
|
||
count += 1
|
||
except ValueError:
|
||
continue
|
||
return count
|
||
|
||
def build_system_prompt(self, snapshot: MemorySnapshot, base_prompt: str = "") -> str:
|
||
"""将记忆注入 system prompt.
|
||
|
||
格式::
|
||
|
||
<agent-identity>
|
||
[SOUL.md content]
|
||
</agent-identity>
|
||
|
||
<user-profile>
|
||
[USER.md content]
|
||
</user-profile>
|
||
|
||
<agent-notes>
|
||
[MEMORY.md content]
|
||
</agent-notes>
|
||
|
||
<recent-activity>
|
||
[DAILY.md content]
|
||
</recent-activity>
|
||
|
||
[base_prompt]
|
||
"""
|
||
# 保存 base_prompt 供后续刷新使用
|
||
if base_prompt:
|
||
self._base_prompt = base_prompt
|
||
|
||
parts: list[str] = []
|
||
|
||
if snapshot.soul:
|
||
parts.append(f"<agent-identity>\n{snapshot.soul}\n</agent-identity>")
|
||
if snapshot.user:
|
||
parts.append(f"<user-profile>\n{snapshot.user}\n</user-profile>")
|
||
if snapshot.memory:
|
||
parts.append(f"<agent-notes>\n{snapshot.memory}\n</agent-notes>")
|
||
if snapshot.daily:
|
||
parts.append(f"<recent-activity>\n{snapshot.daily}\n</recent-activity>")
|
||
|
||
if base_prompt:
|
||
parts.append(base_prompt)
|
||
|
||
return "\n\n".join(parts) if parts else base_prompt
|
||
|
||
def refresh_system_prompt(self) -> str:
|
||
"""重新加载所有记忆文件并构建 system prompt.
|
||
|
||
在 MemoryTool 写入记忆后调用，确保 agent 的 _system_prompt
|
||
反映最新的记忆内容。
|
||
"""
|
||
snapshot = self.load_all()
|
||
return self.build_system_prompt(snapshot, self._base_prompt)
|
||
|
||
def notify_change(self) -> None:
|
||
"""记忆文件变更后通知回调，刷新所有订阅者的 system prompt."""
|
||
if self._on_change is None:
|
||
return
|
||
try:
|
||
new_prompt = self.refresh_system_prompt()
|
||
self._on_change(new_prompt)
|
||
except Exception:
|
||
import logging
|
||
logging.getLogger(__name__).warning("Memory notify_change failed", exc_info=True)
|