feat(agent): two-stage memory system with Dream consolidation

Replace single-stage MemoryConsolidator with a two-stage architecture:

- Consolidator: lightweight token-budget triggered summarization,
  appends to HISTORY.md with cursor-based tracking
- Dream: cron-scheduled two-phase processor that analyzes HISTORY.md
  and updates SOUL.md, USER.md, MEMORY.md via AgentRunner with
  edit_file tools for surgical, fault-tolerant updates

New files: MemoryStore (pure file I/O), Dream class, DreamConfig,
/dream and /dream-log commands. 89 tests covering all components.
This commit is contained in:
chengyongru
2026-04-02 22:42:25 +08:00
committed by chengyongru
parent 7113ad34f4
commit b9616674f0
18 changed files with 856 additions and 728 deletions
+395 -184
View File
@@ -1,4 +1,4 @@
"""Memory system for persistent agent memory."""
"""Memory system: pure file I/O store, lightweight Consolidator, and Dream processor."""
from __future__ import annotations
@@ -11,94 +11,181 @@ from typing import TYPE_CHECKING, Any, Callable
from loguru import logger
from nanobot.utils.helpers import ensure_dir, estimate_message_tokens, estimate_prompt_tokens_chain
from nanobot.utils.helpers import ensure_dir, estimate_message_tokens, estimate_prompt_tokens_chain, strip_think
from nanobot.agent.runner import AgentRunSpec, AgentRunner
from nanobot.agent.tools.registry import ToolRegistry
if TYPE_CHECKING:
from nanobot.providers.base import LLMProvider
from nanobot.session.manager import Session, SessionManager
_SAVE_MEMORY_TOOL = [
{
"type": "function",
"function": {
"name": "save_memory",
"description": "Save the memory consolidation result to persistent storage.",
"parameters": {
"type": "object",
"properties": {
"history_entry": {
"type": "string",
"description": "A paragraph summarizing key events/decisions/topics. "
"Start with [YYYY-MM-DD HH:MM]. Include detail useful for grep search.",
},
"memory_update": {
"type": "string",
"description": "Full updated long-term memory as markdown. Include all existing "
"facts plus new ones. Return unchanged if nothing new.",
},
},
"required": ["history_entry", "memory_update"],
},
},
}
]
def _ensure_text(value: Any) -> str:
"""Normalize tool-call payload values to text for file storage."""
return value if isinstance(value, str) else json.dumps(value, ensure_ascii=False)
def _normalize_save_memory_args(args: Any) -> dict[str, Any] | None:
"""Normalize provider tool-call arguments to the expected dict shape."""
if isinstance(args, str):
args = json.loads(args)
if isinstance(args, list):
return args[0] if args and isinstance(args[0], dict) else None
return args if isinstance(args, dict) else None
_TOOL_CHOICE_ERROR_MARKERS = (
"tool_choice",
"toolchoice",
"does not support",
'should be ["none", "auto"]',
)
def _is_tool_choice_unsupported(content: str | None) -> bool:
"""Detect provider errors caused by forced tool_choice being unsupported."""
text = (content or "").lower()
return any(m in text for m in _TOOL_CHOICE_ERROR_MARKERS)
# ---------------------------------------------------------------------------
# MemoryStore — pure file I/O layer
# ---------------------------------------------------------------------------
class MemoryStore:
"""Two-layer memory: MEMORY.md (long-term facts) + HISTORY.md (grep-searchable log)."""
"""Pure file I/O for memory files: MEMORY.md, history.jsonl, SOUL.md, USER.md."""
_MAX_FAILURES_BEFORE_RAW_ARCHIVE = 3
_DEFAULT_MAX_HISTORY = 1000
def __init__(self, workspace: Path):
def __init__(self, workspace: Path, max_history_entries: int = _DEFAULT_MAX_HISTORY):
self.workspace = workspace
self.max_history_entries = max_history_entries
self.memory_dir = ensure_dir(workspace / "memory")
self.memory_file = self.memory_dir / "MEMORY.md"
self.history_file = self.memory_dir / "HISTORY.md"
self._consecutive_failures = 0
self.history_file = self.memory_dir / "history.jsonl"
self.soul_file = workspace / "SOUL.md"
self.user_file = workspace / "USER.md"
self._dream_log_file = self.memory_dir / ".dream-log.md"
self._cursor_file = self.memory_dir / ".cursor"
self._dream_cursor_file = self.memory_dir / ".dream_cursor"
def read_long_term(self) -> str:
if self.memory_file.exists():
return self.memory_file.read_text(encoding="utf-8")
return ""
# -- generic helpers -----------------------------------------------------
def write_long_term(self, content: str) -> None:
@staticmethod
def read_file(path: Path) -> str:
try:
return path.read_text(encoding="utf-8")
except FileNotFoundError:
return ""
# -- MEMORY.md (long-term facts) -----------------------------------------
def read_memory(self) -> str:
return self.read_file(self.memory_file)
def write_memory(self, content: str) -> None:
self.memory_file.write_text(content, encoding="utf-8")
def append_history(self, entry: str) -> None:
with open(self.history_file, "a", encoding="utf-8") as f:
f.write(entry.rstrip() + "\n\n")
# -- SOUL.md -------------------------------------------------------------
def read_soul(self) -> str:
return self.read_file(self.soul_file)
def write_soul(self, content: str) -> None:
self.soul_file.write_text(content, encoding="utf-8")
# -- USER.md -------------------------------------------------------------
def read_user(self) -> str:
return self.read_file(self.user_file)
def write_user(self, content: str) -> None:
self.user_file.write_text(content, encoding="utf-8")
# -- context injection (used by context.py) ------------------------------
def get_memory_context(self) -> str:
long_term = self.read_long_term()
long_term = self.read_memory()
return f"## Long-term Memory\n{long_term}" if long_term else ""
# -- history.jsonl — append-only, JSONL format ---------------------------
def append_history(self, entry: str) -> int:
"""Append *entry* to history.jsonl and return its auto-incrementing cursor."""
cursor = self._next_cursor()
ts = datetime.now().strftime("%Y-%m-%d %H:%M")
record = {"cursor": cursor, "timestamp": ts, "content": strip_think(entry.rstrip()) or entry.rstrip()}
with open(self.history_file, "a", encoding="utf-8") as f:
f.write(json.dumps(record, ensure_ascii=False) + "\n")
self._cursor_file.write_text(str(cursor), encoding="utf-8")
return cursor
def _next_cursor(self) -> int:
"""Read the current cursor counter and return next value."""
if self._cursor_file.exists():
try:
return int(self._cursor_file.read_text(encoding="utf-8").strip()) + 1
except (ValueError, OSError):
pass
# Fallback: read last line's cursor from the JSONL file.
last = self._read_last_entry()
if last:
return last["cursor"] + 1
return 1
def read_unprocessed_history(self, since_cursor: int) -> list[dict[str, Any]]:
"""Return history entries with cursor > *since_cursor*."""
return [e for e in self._read_entries() if e["cursor"] > since_cursor]
def compact_history(self) -> None:
"""Drop oldest entries if the file exceeds *max_history_entries*."""
if self.max_history_entries <= 0:
return
entries = self._read_entries()
if len(entries) <= self.max_history_entries:
return
kept = entries[-self.max_history_entries:]
self._write_entries(kept)
# -- JSONL helpers -------------------------------------------------------
def _read_entries(self) -> list[dict[str, Any]]:
"""Read all entries from history.jsonl."""
entries: list[dict[str, Any]] = []
try:
with open(self.history_file, "r", encoding="utf-8") as f:
for line in f:
line = line.strip()
if line:
try:
entries.append(json.loads(line))
except json.JSONDecodeError:
continue
except FileNotFoundError:
pass
return entries
def _read_last_entry(self) -> dict[str, Any] | None:
"""Read the last entry from the JSONL file efficiently."""
try:
with open(self.history_file, "rb") as f:
f.seek(0, 2)
size = f.tell()
if size == 0:
return None
read_size = min(size, 4096)
f.seek(size - read_size)
data = f.read().decode("utf-8")
lines = [l for l in data.split("\n") if l.strip()]
if not lines:
return None
return json.loads(lines[-1])
except (FileNotFoundError, json.JSONDecodeError):
return None
def _write_entries(self, entries: list[dict[str, Any]]) -> None:
"""Overwrite history.jsonl with the given entries."""
with open(self.history_file, "w", encoding="utf-8") as f:
for entry in entries:
f.write(json.dumps(entry, ensure_ascii=False) + "\n")
# -- dream cursor --------------------------------------------------------
def get_last_dream_cursor(self) -> int:
if self._dream_cursor_file.exists():
try:
return int(self._dream_cursor_file.read_text(encoding="utf-8").strip())
except (ValueError, OSError):
pass
return 0
def set_last_dream_cursor(self, cursor: int) -> None:
self._dream_cursor_file.write_text(str(cursor), encoding="utf-8")
# -- dream log -----------------------------------------------------------
def read_dream_log(self) -> str:
return self.read_file(self._dream_log_file)
def append_dream_log(self, entry: str) -> None:
with open(self._dream_log_file, "a", encoding="utf-8") as f:
f.write(f"{entry.rstrip()}\n\n")
# -- message formatting utility ------------------------------------------
@staticmethod
def _format_messages(messages: list[dict]) -> str:
lines = []
@@ -111,107 +198,10 @@ class MemoryStore:
)
return "\n".join(lines)
async def consolidate(
self,
messages: list[dict],
provider: LLMProvider,
model: str,
) -> bool:
"""Consolidate the provided message chunk into MEMORY.md + HISTORY.md."""
if not messages:
return True
current_memory = self.read_long_term()
prompt = f"""Process this conversation and call the save_memory tool with your consolidation.
## Current Long-term Memory
{current_memory or "(empty)"}
## Conversation to Process
{self._format_messages(messages)}"""
chat_messages = [
{"role": "system", "content": "You are a memory consolidation agent. Call the save_memory tool with your consolidation of the conversation."},
{"role": "user", "content": prompt},
]
try:
forced = {"type": "function", "function": {"name": "save_memory"}}
response = await provider.chat_with_retry(
messages=chat_messages,
tools=_SAVE_MEMORY_TOOL,
model=model,
tool_choice=forced,
)
if response.finish_reason == "error" and _is_tool_choice_unsupported(
response.content
):
logger.warning("Forced tool_choice unsupported, retrying with auto")
response = await provider.chat_with_retry(
messages=chat_messages,
tools=_SAVE_MEMORY_TOOL,
model=model,
tool_choice="auto",
)
if not response.has_tool_calls:
logger.warning(
"Memory consolidation: LLM did not call save_memory "
"(finish_reason={}, content_len={}, content_preview={})",
response.finish_reason,
len(response.content or ""),
(response.content or "")[:200],
)
return self._fail_or_raw_archive(messages)
args = _normalize_save_memory_args(response.tool_calls[0].arguments)
if args is None:
logger.warning("Memory consolidation: unexpected save_memory arguments")
return self._fail_or_raw_archive(messages)
if "history_entry" not in args or "memory_update" not in args:
logger.warning("Memory consolidation: save_memory payload missing required fields")
return self._fail_or_raw_archive(messages)
entry = args["history_entry"]
update = args["memory_update"]
if entry is None or update is None:
logger.warning("Memory consolidation: save_memory payload contains null required fields")
return self._fail_or_raw_archive(messages)
entry = _ensure_text(entry).strip()
if not entry:
logger.warning("Memory consolidation: history_entry is empty after normalization")
return self._fail_or_raw_archive(messages)
self.append_history(entry)
update = _ensure_text(update)
if update != current_memory:
self.write_long_term(update)
self._consecutive_failures = 0
logger.info("Memory consolidation done for {} messages", len(messages))
return True
except Exception:
logger.exception("Memory consolidation failed")
return self._fail_or_raw_archive(messages)
def _fail_or_raw_archive(self, messages: list[dict]) -> bool:
"""Increment failure count; after threshold, raw-archive messages and return True."""
self._consecutive_failures += 1
if self._consecutive_failures < self._MAX_FAILURES_BEFORE_RAW_ARCHIVE:
return False
self._raw_archive(messages)
self._consecutive_failures = 0
return True
def _raw_archive(self, messages: list[dict]) -> None:
def raw_archive(self, messages: list[dict]) -> None:
"""Fallback: dump raw messages to HISTORY.md without LLM summarization."""
ts = datetime.now().strftime("%Y-%m-%d %H:%M")
self.append_history(
f"[{ts}] [RAW] {len(messages)} messages\n"
f"[RAW] {len(messages)} messages\n"
f"{self._format_messages(messages)}"
)
logger.warning(
@@ -219,8 +209,14 @@ class MemoryStore:
)
class MemoryConsolidator:
"""Owns consolidation policy, locking, and session offset updates."""
# ---------------------------------------------------------------------------
# Consolidator — lightweight token-budget triggered consolidation
# ---------------------------------------------------------------------------
class Consolidator:
"""Lightweight consolidation: summarizes evicted messages, appends to HISTORY.md."""
_MAX_CONSOLIDATION_ROUNDS = 5
@@ -228,7 +224,7 @@ class MemoryConsolidator:
def __init__(
self,
workspace: Path,
store: MemoryStore,
provider: LLMProvider,
model: str,
sessions: SessionManager,
@@ -237,7 +233,7 @@ class MemoryConsolidator:
get_tool_definitions: Callable[[], list[dict[str, Any]]],
max_completion_tokens: int = 4096,
):
self.store = MemoryStore(workspace)
self.store = store
self.provider = provider
self.model = model
self.sessions = sessions
@@ -245,16 +241,14 @@ class MemoryConsolidator:
self.max_completion_tokens = max_completion_tokens
self._build_messages = build_messages
self._get_tool_definitions = get_tool_definitions
self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary()
self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = (
weakref.WeakValueDictionary()
)
def get_lock(self, session_key: str) -> asyncio.Lock:
"""Return the shared consolidation lock for one session."""
return self._locks.setdefault(session_key, asyncio.Lock())
async def consolidate_messages(self, messages: list[dict[str, object]]) -> bool:
"""Archive a selected message chunk into persistent memory."""
return await self.store.consolidate(messages, self.provider, self.model)
def pick_consolidation_boundary(
self,
session: Session,
@@ -294,14 +288,48 @@ class MemoryConsolidator:
self._get_tool_definitions(),
)
async def archive_messages(self, messages: list[dict[str, object]]) -> bool:
"""Archive messages with guaranteed persistence (retries until raw-dump fallback)."""
async def archive(self, messages: list[dict]) -> bool:
"""Summarize messages via LLM and append to HISTORY.md.
Returns True on success (or degraded success), False if nothing to do.
"""
if not messages:
return False
try:
formatted = MemoryStore._format_messages(messages)
response = await self.provider.chat_with_retry(
model=self.model,
messages=[
{
"role": "system",
"content": (
"Extract key facts from this conversation. "
"Only output items matching these categories, skip everything else:\n"
"- User facts: personal info, preferences, stated opinions, habits\n"
"- Decisions: choices made, conclusions reached\n"
"- Events: plans, deadlines, notable occurrences\n"
"- Preferences: communication style, tool preferences\n\n"
"Priority: user corrections and preferences > decisions > events > environment facts. "
"The most valuable memory prevents the user from having to repeat themselves.\n\n"
"Skip: code patterns derivable from source, git history, debug steps already in code, "
"or anything already captured in existing memory.\n\n"
"Output as concise bullet points, one fact per line. "
"No preamble, no commentary.\n"
"If nothing noteworthy happened, output: (nothing)"
),
},
{"role": "user", "content": formatted},
],
tools=None,
tool_choice=None,
)
summary = response.content or "[no summary]"
self.store.append_history(summary)
return True
except Exception:
logger.warning("Consolidation LLM call failed, raw-dumping to history")
self.store.raw_archive(messages)
return True
for _ in range(self.store._MAX_FAILURES_BEFORE_RAW_ARCHIVE):
if await self.consolidate_messages(messages):
return True
return True
async def maybe_consolidate_by_tokens(self, session: Session) -> None:
"""Loop: archive old messages until prompt fits within safe budget.
@@ -356,7 +384,7 @@ class MemoryConsolidator:
source,
len(chunk),
)
if not await self.consolidate_messages(chunk):
if not await self.archive(chunk):
return
session.last_consolidated = end_idx
self.sessions.save(session)
@@ -364,3 +392,186 @@ class MemoryConsolidator:
estimated, source = self.estimate_session_prompt_tokens(session)
if estimated <= 0:
return
# ---------------------------------------------------------------------------
# Dream — heavyweight cron-scheduled memory consolidation
# ---------------------------------------------------------------------------
class Dream:
"""Two-phase memory processor: analyze HISTORY.md, then edit files via AgentRunner.
Phase 1 produces an analysis summary (plain LLM call).
Phase 2 delegates to AgentRunner with read_file / edit_file tools so the
LLM can make targeted, incremental edits instead of replacing entire files.
"""
_PHASE1_SYSTEM = (
"Compare conversation history against current memory files. "
"Output one line per finding:\n"
"[FILE] atomic fact or change description\n\n"
"Files: USER (identity, preferences, habits), "
"SOUL (bot behavior, tone), "
"MEMORY (knowledge, project context, tool patterns)\n\n"
"Rules:\n"
"- Only new or conflicting information — skip duplicates and ephemera\n"
"- Prefer atomic facts: \"has a cat named Luna\" not \"discussed pet care\"\n"
"- Corrections: [USER] location is Tokyo, not Osaka\n"
"- Also capture confirmed approaches: if the user validated a non-obvious choice, note it\n\n"
"If nothing needs updating: [SKIP] no new information"
)
_PHASE2_SYSTEM = (
"Update memory files based on the analysis below.\n\n"
"## Quality standards\n"
"- Every line must carry standalone value — no filler\n"
"- Concise bullet points under clear headers\n"
"- Remove outdated or contradicted information\n\n"
"## Editing\n"
"- File contents provided below — edit directly, no read_file needed\n"
"- Batch changes to the same file into one edit_file call\n"
"- Surgical edits only — never rewrite entire files\n"
"- Do NOT overwrite correct entries — only add, update, or remove\n"
"- If nothing to update, stop without calling tools"
)
def __init__(
self,
store: MemoryStore,
provider: LLMProvider,
model: str,
max_batch_size: int = 20,
max_iterations: int = 10,
):
self.store = store
self.provider = provider
self.model = model
self.max_batch_size = max_batch_size
self.max_iterations = max_iterations
self._runner = AgentRunner(provider)
self._tools = self._build_tools()
# -- tool registry -------------------------------------------------------
def _build_tools(self) -> ToolRegistry:
"""Build a minimal tool registry for the Dream agent."""
from nanobot.agent.tools.filesystem import EditFileTool, ReadFileTool
tools = ToolRegistry()
workspace = self.store.workspace
tools.register(ReadFileTool(workspace=workspace, allowed_dir=workspace))
tools.register(EditFileTool(workspace=workspace, allowed_dir=workspace))
return tools
# -- main entry ----------------------------------------------------------
async def run(self) -> bool:
"""Process unprocessed history entries. Returns True if work was done."""
last_cursor = self.store.get_last_dream_cursor()
entries = self.store.read_unprocessed_history(since_cursor=last_cursor)
if not entries:
return False
batch = entries[: self.max_batch_size]
logger.info(
"Dream: processing {} entries (cursor {}{}), batch={}",
len(entries), last_cursor, batch[-1]["cursor"], len(batch),
)
# Build history text for LLM
history_text = "\n".join(
f"[{e['timestamp']}] {e['content']}" for e in batch
)
# Current file contents
current_memory = self.store.read_memory() or "(empty)"
current_soul = self.store.read_soul() or "(empty)"
current_user = self.store.read_user() or "(empty)"
file_context = (
f"## Current MEMORY.md\n{current_memory}\n\n"
f"## Current SOUL.md\n{current_soul}\n\n"
f"## Current USER.md\n{current_user}"
)
# Phase 1: Analyze
phase1_prompt = (
f"## Conversation History\n{history_text}\n\n{file_context}"
)
try:
phase1_response = await self.provider.chat_with_retry(
model=self.model,
messages=[
{"role": "system", "content": self._PHASE1_SYSTEM},
{"role": "user", "content": phase1_prompt},
],
tools=None,
tool_choice=None,
)
analysis = phase1_response.content or ""
logger.debug("Dream Phase 1 complete ({} chars)", len(analysis))
except Exception:
logger.exception("Dream Phase 1 failed")
return False
# Phase 2: Delegate to AgentRunner with read_file / edit_file
phase2_prompt = f"## Analysis Result\n{analysis}\n\n{file_context}"
tools = self._tools
messages: list[dict[str, Any]] = [
{"role": "system", "content": self._PHASE2_SYSTEM},
{"role": "user", "content": phase2_prompt},
]
try:
result = await self._runner.run(AgentRunSpec(
initial_messages=messages,
tools=tools,
model=self.model,
max_iterations=self.max_iterations,
fail_on_tool_error=True,
))
logger.debug(
"Dream Phase 2 complete: stop_reason={}, tool_events={}",
result.stop_reason, len(result.tool_events),
)
except Exception:
logger.exception("Dream Phase 2 failed")
result = None
# Build changelog from tool events
changelog: list[str] = []
if result and result.tool_events:
for event in result.tool_events:
if event["status"] == "ok":
changelog.append(f"{event['name']}: {event['detail']}")
# Advance cursor — always, to avoid re-processing Phase 1
new_cursor = batch[-1]["cursor"]
self.store.set_last_dream_cursor(new_cursor)
self.store.compact_history()
if result and result.stop_reason == "completed":
logger.info(
"Dream done: {} change(s), cursor advanced to {}",
len(changelog), new_cursor,
)
else:
reason = result.stop_reason if result else "exception"
logger.warning(
"Dream incomplete ({}): cursor advanced to {}",
reason, new_cursor,
)
# Write dream log
ts = datetime.now().strftime("%Y-%m-%d %H:%M")
if changelog:
log_entry = f"## {ts}\n"
for change in changelog:
log_entry += f"- {change}\n"
self.store.append_dream_log(log_entry)
else:
self.store.append_dream_log(f"## {ts}\nNo changes.\n")
return True