ReMe/reme/steps/evolve/auto_memory.py

"""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}")