"""auto_memory — record conversation facts into a daily note via an agent.""" import base64 import binascii 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, get_path_lock, parse_daily_date, refresh_day_index, write_file_safe from ..file_io import validate_filename_component, validate_session_id from ..file_io._path import IMAGE_MIME_BY_EXT, _check_path_permission, resolve_path from ..index import normalize_posix_path from ...components import R from ...utils.wikilink_handler import WikilinkHandler _SESSION_ID_KEY = "session_id" _SOURCE_CONVERSATION_KEY = "source_conversation" _MESSAGE_TIME_ALIASES = ("time_created", "timestamp", "createdAt", "timeCreated", "created_time") def _checked_write_path(workspace: Path, relative: str, allowed_paths) -> Path: """Use the file Jobs' path and permission checks for attachments and note metadata.""" target, error = resolve_path(workspace, relative) if error or target is None: raise ValueError(error or "Invalid write path") if not _check_path_permission(workspace, target, allowed_paths): raise PermissionError(f"No permission to write {relative}") return target def _check_existing_image(target: Path) -> bool: if target.is_file(): return True if target.exists(): raise ValueError(f"Image attachment path is not a file: {target.name}") return False async def _save_session_images( workspace: Path, session_dir: str, session_id: str, messages: list[Msg], images: dict[str, DataBlock], allowed_paths=None, ) -> list[str]: """Save Base64 bytes and annotate prepared image markers; URLs remain remote. ``messages`` must be the invocation-owned copies from image preparation. Image data and caller-owned messages remain unchanged. Image identity uses the message ID and block position, not the image block's optional identity. """ workspace = workspace.resolve() error = validate_filename_component(session_id, kind="session_id") if error: raise ValueError(error) if Path(session_dir).is_absolute(): raise ValueError("session_dir must be workspace-relative") pending: dict[Path, tuple[str, bytes]] = {} replacements = [] sources = [] for message in messages: for index, block in enumerate(message.content): if not isinstance(block, TextBlock) or block.text not in images: continue image = images[block.text] source = image.source if source.type == "url": reference = str(source.url) else: # A reversible encoding, not a hash: also distinguish IDs on # case-insensitive filesystems without trusting IDs as paths. encoded_id = message.id.encode("utf-8").hex() suffix = next( (ext for ext, mime in IMAGE_MIME_BY_EXT.items() if mime == source.media_type.lower()), ".bin", ) filename = f"msg-{encoded_id}-image-{index}{suffix}" if len(filename) > 255: raise ValueError("Message ID is too long for an image attachment filename") relative = (Path(session_dir) / "images" / session_id / filename).as_posix() error = WikilinkHandler.validate_src_dst(relative, relative) if error: raise ValueError(f"Image attachment path cannot be linked: {error}") target = _checked_write_path(workspace, relative, allowed_paths) # Wait for any in-process writer before reusing the attachment. # Stable message IDs and block positions identify an image. async with await get_path_lock(target): target = _checked_write_path(workspace, relative, allowed_paths) if not _check_existing_image(target): try: payload = base64.b64decode(source.data, validate=True) except (ValueError, binascii.Error) as exc: raise ValueError("Image source contains invalid Base64") from exc if not payload: raise ValueError("Image source contains empty Base64 data") if target in pending and pending[target][1] != payload: raise ValueError("Conflicting images share the same message ID and block position") pending[target] = (relative, payload) reference = f"[[{relative}]]" replacements.append((message, index, block, reference)) if reference not in sources: sources.append(reference) # Validate every input before creating files. Recheck each path under the # existing file-operation lock so concurrent saves cannot overwrite it. for relative, payload in pending.values(): target = _checked_write_path(workspace, relative, allowed_paths) async with await get_path_lock(target): target = _checked_write_path(workspace, relative, allowed_paths) if not _check_existing_image(target): await write_file_safe(target, payload) for message, index, block, reference in replacements: message.content[index] = block.model_copy(update={"text": f"Image source: {reference}\n{block.text}"}) return sources 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, image_sources: list[str] | None = None, ) -> None: metadata = { _SESSION_ID_KEY: session_id, _SOURCE_CONVERSATION_KEY: self._session_link(session_id), } if image_sources is not None and ( not isinstance(image_sources, list) or any(not isinstance(source, str) for source in image_sources) ): raise ValueError("source_images must be a list of strings") if image_sources: workspace = self.file_store.workspace_path.resolve() allowed_paths = self.context.get("_allowed_paths") target = _checked_write_path(workspace, path, allowed_paths) # Read and merge under the same lock as native file Jobs so a # concurrent write cannot lose its source links to a stale list. async with await get_path_lock(target): target = _checked_write_path(workspace, path, allowed_paths) if not target.is_file(): raise ValueError(f"Memory note not found: {path}") post = frontmatter.loads(target.read_text(encoding="utf-8")) current_sources = post.metadata.get("source_images", []) if not isinstance(current_sources, list) or any(not isinstance(s, str) for s in current_sources): raise ValueError("Existing source_images must be a list of strings") metadata["source_images"] = list(dict.fromkeys([*current_sources, *image_sources])) if not all(post.metadata.get(key) == value for key, value in metadata.items()): post.metadata.update(metadata) await write_file_safe(target, frontmatter.dumps(post)) return 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 def _image_user_message(self, prompt: str, images: dict[str, DataBlock]) -> UserMsg: """Add source instructions and restore images after rendering the memory prompt.""" prompt += "\n\n" + self.prompt_format("image_sources_instructions") 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 # Keep attachment provenance even when the Agent uses a full rewrite, # or this call updates an image-backed note with images disabled. image_sources = [] if before_note_bytes is not None: previous_sources = frontmatter.loads(before_note_bytes.decode("utf-8")).get("source_images", []) if isinstance(previous_sources, list) and all(isinstance(source, str) for source in previous_sources): image_sources = previous_sources elif images: raise ValueError("Existing source_images must be a list of strings") if images: image_sources += await _save_session_images( self.file_store.workspace_path, self._session_dir(), session_id, history_messages, images, self.context.get("_allowed_paths"), ) 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"]) if image_sources: await self._ensure_session_frontmatter(note_path, session_id, image_sources) else: try: await self._ensure_session_frontmatter(note_path, session_id, image_sources) 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}")