ReMe/reme/steps/evolve/auto_memory.py
WQS 6125fc197d
feat(auto-memory): add session image support (#532)
* feat(auto-memory): add opt-in image input

* feat(auto-memory): caption session images into source-linked notes

* fix(auto-memory): version Pillow 10 compatible image preparation

* fix(auto-memory): harden image evidence and retry boundaries

* fix(auto-memory): preserve image evidence across replay and concurrent writes

Keep persisted image positions through disabled history backfills and transcript filtering. Merge note links atomically, map 16-bit grayscale without clipping, recheck restored caption owners, and reuse unchanged identity metadata during batch publication. Add regressions and document conservative custom-rename behavior.

* refactor(auto-memory): restore main baseline for image modes v2

* refactor(images): share resource caption preprocessing and model calls

* feat(watch): support scoped exclusions for managed session images

* feat(auto-memory): add opt-in resource and caption-only image input

* refactor(auto-memory): keep caption-only mode with text fallback

* refactor(auto-memory): make caption-only mode dispatch explicit

* docs(auto-memory): focus image guide on caption-only mode

* feat(auto-memory): add direct multimodal image extraction

* refactor(auto-memory): route direct images through the vision model

* test(auto-memory): consolidate overlapping image regressions

* feat(auto-memory): interleave direct images with conversation text

* docs(auto-memory): clarify history rendering scope for direct inputs

* feat(auto-memory): require vision declaration for direct-only images

* refactor(auto-memory): keep bound model and native image helpers

* refactor(auto-memory): keep image inputs native and opt-in

* refactor(auto-memory): simplify image docs and tests

* refactor(auto-memory): confine image adaptation to image requests

* refactor(auto-memory): remove image switch type validation
2026-09-20 15:35:00 +08:00

486 lines
21 KiB
Python

"""auto_memory — record conversation facts into a daily note via an agent."""
import datetime
from pathlib import Path
import re
from urllib.parse import urlsplit
from uuid import uuid4
import zoneinfo
import aiofiles
import frontmatter
from agentscope.agent import ContextConfig
from agentscope.message import DataBlock, Msg, TextBlock, UserMsg
from ._evolve import agent_reply_result_text, format_history, now
from ..base_step import BaseStep
from ..file_io import extract_daily_date, parse_daily_date, refresh_day_index
from ..file_io import validate_filename_component, validate_session_id
from ..index import normalize_posix_path
from ...components import R
_SESSION_ID_KEY = "session_id"
_SOURCE_CONVERSATION_KEY = "source_conversation"
_MESSAGE_TIME_ALIASES = ("time_created", "timestamp", "createdAt", "timeCreated", "created_time")
def _sanitize_msg_for_save(msg: Msg) -> Msg:
new_content = []
changed = False
for block in msg.content:
# Tool results often contain recalled memory/search/read output. Keeping
# them in saved conversation history lets retrieved facts masquerade as
# user-provided context in future auto-memory runs.
if block.type == "tool_result":
changed = True
continue
if block.type == "data" and hasattr(block, "source") and getattr(block.source, "type", None) == "base64":
changed = True
continue
new_content.append(block)
if not changed:
return msg
return msg.model_copy(update={"content": new_content})
def _normalize_msg_timestamp(item: dict) -> dict:
"""Map common message timestamp aliases to AgentScope's ``created_at`` field."""
if item.get("created_at"):
return item
for key in _MESSAGE_TIME_ALIASES:
value = item.get(key)
if value:
return {**item, "created_at": value}
metadata = item.get("metadata")
if isinstance(metadata, dict):
for key in _MESSAGE_TIME_ALIASES:
value = metadata.get(key)
if value:
return {**item, "created_at": value}
return item
@R.register("auto_memory_step")
class AutoMemoryStep(BaseStep):
"""Record conversation facts into a daily note via an Agent."""
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.create_tools: list[str] = ["daily_write"]
self.update_tools: list[str] = ["read", "edit", "frontmatter_update", "write"]
def _session_dir(self) -> str:
return normalize_posix_path(str(self.config_value("session_dir")))
def _session_path(self, session_id: str) -> Path:
return self.file_store.workspace_path / self._session_dir() / "dialog" / f"{session_id}.jsonl"
def _session_source_path(self, session_id: str) -> str:
return normalize_posix_path(f"{self._session_dir()}/dialog/{session_id}.jsonl")
def _session_link(self, session_id: str) -> str:
return f"[[{self._session_source_path(session_id)}]]"
def _daily_note_path(self, day: str, name: str) -> str:
return f"{self.config_value('daily_dir')}/{day}/{name}.md"
def _frontmatter(self, path: str) -> dict:
post = frontmatter.loads((self.file_store.workspace_path / path).read_text(encoding="utf-8"))
return dict(post.metadata or {})
def _note_bytes(self, path: str) -> bytes | None:
note_path = self.file_store.workspace_path / path
if not note_path.is_file():
return None
return note_path.read_bytes()
def _note_modified(self, before_path: str, before_bytes: bytes | None, after_path: str) -> bool:
if not after_path:
return False
after_bytes = self._note_bytes(after_path)
if after_bytes is None:
return before_bytes is not None
return after_path != before_path or before_bytes != after_bytes
def _find_session_note(self, notes: list[dict], session_id: str) -> dict | None:
source = self._session_link(session_id)
for note in notes:
if str(note.get(_SESSION_ID_KEY, "")).strip() == session_id:
return note
for note in notes:
if str(note.get(_SOURCE_CONVERSATION_KEY, "")).strip() == source:
return note
return None
async def _list_session_note(self, day: str, session_id: str) -> dict | None:
list_response = await self.run_job("daily_list", date=day)
if not list_response.success:
raise RuntimeError(f"daily_list failed: {list_response.answer}")
notes = list_response.metadata.get("notes") or []
return self._find_session_note(notes, session_id)
async def _ensure_session_frontmatter(self, path: str, session_id: str) -> None:
metadata = {
_SESSION_ID_KEY: session_id,
_SOURCE_CONVERSATION_KEY: self._session_link(session_id),
}
current = self._frontmatter(path)
if all(current.get(key) == value for key, value in metadata.items()):
return
response = await self.run_job(
"frontmatter_update",
path=path,
metadata=metadata,
)
if not response.success:
raise RuntimeError(f"frontmatter_update failed: {response.answer}")
async def _rename_from_frontmatter_name(self, path: str, day: str) -> str:
meta = self._frontmatter(path)
name = str(meta.get("name", "")).strip()
if not name:
return path
if err := validate_filename_component(name, kind="name"):
raise RuntimeError(err)
target_path = self._daily_note_path(day, name)
if target_path == path:
return path
move_response = await self.run_job(
"move",
src_path=path,
dst_path=target_path,
overwrite=False,
retarget=True,
)
if not move_response.success:
raise RuntimeError(f"move failed: {move_response.answer}")
return target_path
async def _save_session_messages(self, session_id: str, messages: list[Msg]) -> None:
if not session_id or not messages:
return
path = self._session_path(session_id)
self.logger.info(
f"[{self.name}] save session start session_id={session_id!r} messages={len(messages)} path={path}",
)
existing: list[Msg] = []
if path.exists():
async with aiofiles.open(path, encoding="utf-8") as f:
content = await f.read()
for line in content.splitlines():
line = line.strip()
if line:
try:
existing.append(Msg.model_validate_json(line))
except Exception:
pass
by_id: dict[str, Msg] = {}
for msg in existing:
by_id[msg.id] = msg
for msg in messages:
by_id[msg.id] = msg
merged = sorted(by_id.values(), key=lambda m: m.created_at)
can_append = 0 < len(existing) <= len(merged) and all(
merged[i].id == existing[i].id for i in range(len(existing))
)
path.parent.mkdir(parents=True, exist_ok=True)
if can_append:
new_msgs = merged[len(existing) :]
if new_msgs:
async with aiofiles.open(path, "a", encoding="utf-8") as f:
for msg in new_msgs:
await f.write(_sanitize_msg_for_save(msg).model_dump_json() + "\n")
self.logger.info(
f"[{self.name}] save session appended session_id={session_id!r} "
f"existing={len(existing)} appended={len(new_msgs)} total={len(merged)}",
)
else:
self.logger.info(
f"[{self.name}] save session unchanged session_id={session_id!r} "
f"existing={len(existing)} total={len(merged)}",
)
else:
async with aiofiles.open(path, "w", encoding="utf-8") as f:
for msg in merged:
await f.write(_sanitize_msg_for_save(msg).model_dump_json() + "\n")
self.logger.info(
f"[{self.name}] save session rewrote session_id={session_id!r} "
f"existing={len(existing)} total={len(merged)}",
)
@staticmethod
def _to_msg(item) -> Msg:
if isinstance(item, Msg):
return item
if isinstance(item, dict):
item = _normalize_msg_timestamp(item)
if isinstance(item, dict) and isinstance(item.get("content"), str):
item = {**item, "content": [{"type": "text", "text": item["content"]}]}
return Msg.model_validate(item)
@staticmethod
def _messages_day(messages: list[Msg], timezone: str | None = None) -> str | None:
days = [day for msg in messages if (day := AutoMemoryStep._message_day(msg.created_at, timezone))]
return max(days) if days else None
@staticmethod
def _message_day(value, timezone: str | None) -> str | None:
"""Resolve an absolute timestamp to the workspace's calendar date."""
fallback = extract_daily_date(value)
text = str(value or "").strip()
if not fallback or len(text) <= 10:
return fallback
try:
timestamp = datetime.datetime.fromisoformat(text.replace("Z", "+00:00"))
except ValueError:
return fallback
if timestamp.tzinfo is None:
return fallback
if timezone:
timestamp = timestamp.astimezone(zoneinfo.ZoneInfo(timezone))
else:
timestamp = timestamp.astimezone()
return timestamp.date().isoformat()
def _build_messages(self, raw_messages: list) -> list[Msg]:
"""Convert raw message payloads into ``Msg`` objects.
Overridable hook: subclasses can preprocess ``raw_messages`` (e.g. fill
in missing timestamps) before conversion.
"""
return [self._to_msg(item) for item in raw_messages]
def _reply_extra_kwargs(self, day: str) -> dict: # pylint: disable=unused-argument
"""Extra keyword arguments for ``agent_wrapper.reply``.
Overridable hook: subclasses can inject additional reply options such
as per-tool defaults keyed on ``day``.
"""
return {}
def _format_history(self, messages: list[Msg]) -> str:
"""Render the conversation history injected into the prompt.
Overridable hook: subclasses can annotate messages (e.g. with their
source line numbers in the session file) before rendering.
"""
return format_history(messages)
def _prepare_image_history(
self,
messages: list[Msg],
day: str,
) -> tuple[list[Msg], dict[str, DataBlock], dict | None]:
"""Validate image inputs before saving, without reading or changing their sources."""
include_images = self.context.get("include_images", False)
if include_images is False:
return messages, {}, None
images = [
(message_index, block_index, block)
for message_index, message in enumerate(messages)
for block_index, block in enumerate(message.content)
if isinstance(block, DataBlock) and block.source.media_type.startswith("image/")
]
if not images:
return messages, {}, None
wrapper = self.agent_wrapper
if wrapper is None or wrapper.backend != "agentscope":
raise NotImplementedError("Auto Memory image inputs require the AgentScope wrapper")
for _, _, block in images:
if block.source.type == "url" and urlsplit(str(block.source.url)).scheme not in {"http", "https"}:
raise ValueError("Image URLs must use HTTP(S); convert local files to Base64Source before calling")
reply_kwargs = dict(self._reply_extra_kwargs(day))
context_config = reply_kwargs.get("context_config", wrapper.kwargs.get("context_config")) or {}
limit = ContextConfig(**context_config).max_image_num
if len(images) > limit:
raise ValueError(
f"Session has {len(images)} images, exceeding context_config.max_image_num={limit}; "
"configure the AgentScope wrapper's image limit explicitly",
)
prepared = [message.model_copy(deep=True) for message in messages]
image_blocks = {}
prefix = f"__reme_image_{uuid4().hex}_"
for number, (message_index, block_index, _) in enumerate(images):
marker = f"{prefix}{number}__"
image_blocks[marker] = prepared[message_index].content[block_index]
prepared[message_index].content[block_index] = TextBlock(text=marker)
return prepared, image_blocks, reply_kwargs
@staticmethod
def _image_user_message(prompt: str, images: dict[str, DataBlock]) -> UserMsg:
"""Restore images after the existing templates and history hooks have rendered."""
parts = re.split("(" + "|".join(map(re.escape, images)) + ")", prompt)
if [part for part in parts if part in images] != list(images):
raise ValueError("Memory prompt must preserve every image once in conversation order")
return UserMsg(
name="user",
content=[images[part] if part in images else TextBlock(text=part) for part in parts if part],
)
# pylint: disable=too-many-return-statements
async def execute(self):
assert self.context is not None
self.context["changes"] = []
raw_messages = self.context.get("messages") or []
session_id: str = self.context.get("session_id", "")
memory_hint: str = self.context.get("memory_hint", "")
raw_date = self.context.get("date", "")
tz = self.app_context.app_config.timezone if self.app_context is not None else None
current = now(tz)
messages: list[Msg] = self._build_messages(raw_messages)
self.logger.info(
f"[{self.name}] start session_id={session_id!r} raw_messages={len(raw_messages)} "
f"messages={len(messages)} hint={bool(memory_hint)}",
)
if session_id and (err := validate_session_id(session_id)):
self.context.response.success = False
self.context.response.answer = f"Error: {err}"
self.logger.warning(f"[{self.name}] invalid session_id={session_id!r} err={err}")
return
if not session_id:
self.context.response.success = False
self.context.response.answer = "Error: session_id is required"
self.logger.warning(f"[{self.name}] missing session_id")
return
day = (
parse_daily_date(raw_date) if raw_date else self._messages_day(messages, tz) or current.strftime("%Y-%m-%d")
)
if raw_date and day is None:
self.context.response.success = False
self.context.response.answer = "Error: date must be YYYY-MM-DD"
self.context.response.metadata.update({"date": raw_date, "modified": False, "n_messages": len(messages)})
self.logger.warning(f"[{self.name}] invalid date={raw_date!r}")
return
history_messages, images, reply_kwargs = self._prepare_image_history(messages, day)
await self._save_session_messages(session_id, messages)
if not messages:
self.context.response.success = True
self.context.response.answer = "Skipped: no messages"
self.context.response.metadata.update({"date": day, "modified": False, "n_messages": 0})
self.logger.info(f"[{self.name}] Skipped: no messages session_id={session_id!r} modified=False")
return
try:
note = await self._list_session_note(day, session_id)
except RuntimeError as exc:
self.context.response.success = False
self.context.response.answer = str(exc)
self.context.response.metadata.update({"date": day, "modified": False, "n_messages": len(messages)})
self.logger.info(f"[{self.name}] list failed session_id={session_id!r} answer={str(exc)!r}")
return
note_path = str(note["path"]) if note else ""
created = note is None
before_note_path = note_path
before_note_bytes = self._note_bytes(note_path) if note_path else None
self.logger.info(
f"[{self.name}] note lookup session_id={session_id!r} path={note_path!r} "
f"created={created} msgs={len(messages)} hint={bool(memory_hint)}",
)
template_key = "user_message_create" if created else "user_message_update"
user_message = self.prompt_format(
template_key,
today=day,
note=memory_hint or "(none)",
note_path=note_path,
session_id=session_id,
session_file=self._session_source_path(session_id),
history=self._format_history(history_messages),
)
if images:
user_message = self._image_user_message(user_message, images)
self.logger.info(f"[{self.name}] agent start path={note_path} template={template_key}")
# Existing-note updates are restricted to the resolved note path. New
# notes retain the upstream ``daily_write`` date behavior, where the
# model supplies the date from the prompt.
if reply_kwargs is None:
reply_kwargs = self._reply_extra_kwargs(day)
if not created:
reply_kwargs["injected_job_kwargs"] = {"_allowed_paths": [note_path]}
result = await self.agent_wrapper.reply(
user_message,
system_prompt=self.prompt_format("system_prompt"),
job_tools=self.create_tools if created else self.update_tools,
**reply_kwargs,
)
self.logger.info(f"[{self.name}] agent done path={note_path} has_result={bool(result.get('result'))}")
if created:
try:
note = await self._list_session_note(day, session_id)
except RuntimeError as exc:
self.context.response.success = False
self.context.response.answer = str(exc)
self.context.response.metadata.update(
{"date": day, "path": None, "created": created, "modified": False, "n_messages": len(messages)},
)
self.logger.info(f"[{self.name}] post-create list failed session_id={session_id!r} answer={str(exc)!r}")
return
if note is None:
self.context.response.success = True
self.context.response.answer = agent_reply_result_text(result)
self.context.response.metadata.update(
{"date": day, "path": None, "created": False, "modified": False, "n_messages": len(messages)},
)
self.logger.info(f"[{self.name}] done without note session_id={session_id!r} modified=False")
return
note_path = str(note["path"])
else:
try:
await self._ensure_session_frontmatter(note_path, session_id)
note_path = await self._rename_from_frontmatter_name(note_path, day)
except RuntimeError as exc:
self.context.response.success = False
self.context.response.answer = str(exc)
self.context.response.metadata.update(
{
"date": day,
"path": note_path,
"created": created,
"modified": self._note_modified(before_note_path, before_note_bytes, note_path),
"n_messages": len(messages),
},
)
self.logger.info(f"[{self.name}] post-update failed path={note_path} answer={str(exc)!r}")
return
modified = self._note_modified(before_note_path, before_note_bytes, note_path)
if modified:
self.context["changes"] = [{"change": "added" if created else "modified", "path": note_path}]
daily_dir = self.config_value("daily_dir")
self.logger.info(f"[{self.name}] refresh index start date={day} daily_dir={daily_dir}")
index_payload = await refresh_day_index(self.file_store, day, daily_dir)
self.logger.info(f"[{self.name}] refresh index done path={note_path}")
source_conversation = self._session_link(session_id)
self.context.response.success = True
self.context.response.answer = agent_reply_result_text(result)
self.context.response.metadata.update(
{
"date": day,
"path": note_path,
"created": created,
"modified": modified,
"n_messages": len(messages),
"source_conversation": source_conversation,
"index": index_payload,
},
)
self.logger.info(f"[{self.name}] done {note_path} modified={modified}")