From 1c4e13270b930d4172b5e4ae5e3b64723c7f4406 Mon Sep 17 00:00:00 2001 From: wang-qisen Date: Thu, 27 Aug 2026 17:38:20 +0800 Subject: [PATCH] feat: add auto_image step for image resource caption notes --- reme/steps/evolve/__init__.py | 2 + reme/steps/evolve/auto_image.py | 450 ++++++++++++++++++ reme/steps/evolve/auto_image.yaml | 52 +++ reme/steps/file_io/__init__.py | 3 +- reme/steps/file_io/_path.py | 8 + tests/unit/test_auto_image_steps.py | 686 ++++++++++++++++++++++++++++ 6 files changed, 1200 insertions(+), 1 deletion(-) create mode 100644 reme/steps/evolve/auto_image.py create mode 100644 reme/steps/evolve/auto_image.yaml create mode 100644 tests/unit/test_auto_image_steps.py diff --git a/reme/steps/evolve/__init__.py b/reme/steps/evolve/__init__.py index 423b7598..fec242a4 100644 --- a/reme/steps/evolve/__init__.py +++ b/reme/steps/evolve/__init__.py @@ -1,6 +1,7 @@ """Evolve steps.""" from ._evolve import now +from .auto_image import AutoImageStep from .auto_memory import AutoMemoryStep from .auto_memory_cc import AutoMemoryCCStep from .auto_resource import AutoResourceStep @@ -9,6 +10,7 @@ from .dream import DreamExtractStep, DreamFinishStep, DreamIntegrateStep, DreamT __all__ = [ "now", + "AutoImageStep", "AutoMemoryStep", "AutoMemoryCCStep", "AutoResourceStep", diff --git a/reme/steps/evolve/auto_image.py b/reme/steps/evolve/auto_image.py new file mode 100644 index 00000000..ca643751 --- /dev/null +++ b/reme/steps/evolve/auto_image.py @@ -0,0 +1,450 @@ +"""auto_image — interpret image resource files into source-linked daily notes via a VLM.""" + +import base64 +import io +import json +import re +from pathlib import Path + +import aiofiles +from agentscope.message import Base64Source, DataBlock, TextBlock, UserMsg +from agentscope.model import ChatModelBase +from PIL import Image +from pydantic import BaseModel, Field + +from ..file_io import is_image_file, refresh_day_index +from ..file_io._path import IMAGE_MIME_BY_EXT +from .auto_resource import _SOURCE_RESOURCE_KEY, _sanitize_note_name, AutoResourceStep +from ...components import R +from ...enumeration import ComponentEnum + +try: + from pillow_heif import register_heif_opener + + register_heif_opener() +except ImportError: + # Without the plugin HEIC files stay passthrough and the provider rejects + # them per change; the pipeline itself keeps importing and running. + pass + +DEFAULT_MAX_IMAGE_INPUT_BYTES = 50 * 1024 * 1024 +MAX_IMAGE_REQUEST_DIMENSION = 2048 +_JPEG_QUALITY = 85 +# Suffixes re-encoded to PNG for VLM requests; the stored resource file is never modified. +_CONVERT_SUFFIXES = {".bmp", ".tiff", ".heic"} +_JSON_FENCE_RE = re.compile(r"^\s*```(?:json)?\s*(.*?)\s*```\s*$", re.DOTALL) + + +class _CaptionOutput(BaseModel): + """Structured caption contract enforced on the vision model.""" + + name: str = Field(description="short kebab-case topic stem for the note filename; never include dates") + description: str = Field(description="one-sentence summary that conveys the key information on its own") + caption: str = Field(description="complete description / verbatim transcription of meaningful visible text") + + +def _normalize_image_bytes(data: bytes, suffix: str) -> tuple[bytes, str] | None: + """Downscale or re-encode image bytes in memory for a VLM request. + + Returns ``None`` to pass the original bytes through untouched (already a + reasonable size and format, or content that PIL cannot decode — left for + the provider to reject). The stored resource file is never modified. + """ + try: + with Image.open(io.BytesIO(data)) as image: + needs_resize = image.width > MAX_IMAGE_REQUEST_DIMENSION or image.height > MAX_IMAGE_REQUEST_DIMENSION + needs_convert = suffix in _CONVERT_SUFFIXES + if not needs_resize and not needs_convert: + return None + has_alpha = image.mode in ("RGBA", "LA", "P") + frame = image.convert("RGBA" if has_alpha else "RGB") + if needs_resize: + frame.thumbnail((MAX_IMAGE_REQUEST_DIMENSION, MAX_IMAGE_REQUEST_DIMENSION), Image.LANCZOS) + buffer = io.BytesIO() + if frame.mode == "RGBA": + frame.save(buffer, format="PNG") + return buffer.getvalue(), "image/png" + frame.save(buffer, format="JPEG", quality=_JPEG_QUALITY) + return buffer.getvalue(), "image/jpeg" + except Exception: # pylint: disable=broad-except + return None + + +def _build_image_request_payload(data: bytes, suffix: str) -> dict: + """Return ``{"data_b64", "mime", "source_mime", "converted"}`` for a VLM request. + + ``mime`` is the format actually sent (after in-memory downscale/re-encode); + ``source_mime`` describes the stored resource file and is what notes record. + """ + source_mime = IMAGE_MIME_BY_EXT.get(suffix, "image/png") + normalized = _normalize_image_bytes(data, suffix) + if normalized is None: + return { + "data_b64": base64.b64encode(data).decode("ascii"), + "mime": source_mime, + "source_mime": source_mime, + "converted": False, + } + normalized_bytes, mime = normalized + return { + "data_b64": base64.b64encode(normalized_bytes).decode("ascii"), + "mime": mime, + "source_mime": source_mime, + "converted": normalized_bytes != data, + } + + +async def _response_text(result) -> str: + """Extract text blocks from a streaming or non-streaming ChatResponse.""" + if hasattr(type(result), "__aiter__"): + last = None + async for chunk in result: + last = chunk + result = last + if result is None: + return "" + parts: list[str] = [] + for block in result.content or []: + if isinstance(block, dict): + if block.get("type") == "text": + parts.append(str(block.get("text") or "")) + elif getattr(block, "type", None) == "text": + parts.append(str(getattr(block, "text", "") or "")) + return "".join(parts).strip() + + +def _normalize_caption_fields(parsed: dict) -> dict: + """Normalize parsed caption fields, cross-filling a missing ``caption`` + from a present ``description`` so raw JSON never reaches the note body.""" + caption = str(parsed.get("caption") or "").strip() + description = str(parsed.get("description") or "").strip() + if not caption and description: + caption = description + return { + "name": str(parsed.get("name") or "").strip(), + "description": description, + "caption": caption, + } + + +def _parse_caption_json(text: str) -> dict: + """Parse a plain-call caption response leniently. + + Used as the fallback when the schema-forced structured call fails: fenced + JSON and embedded ``{...}`` slices are tried before degrading the whole + response text to the caption. + """ + cleaned = text.strip() + fence = _JSON_FENCE_RE.match(cleaned) + if fence: + cleaned = fence.group(1) + for candidate in (cleaned, cleaned[cleaned.find("{") : cleaned.rfind("}") + 1]): + if not candidate: + continue + try: + parsed = json.loads(candidate) + except (json.JSONDecodeError, ValueError): + continue + if isinstance(parsed, dict): + normalized = _normalize_caption_fields(parsed) + if normalized["caption"] or normalized["description"]: + return normalized + return {"name": "", "description": "", "caption": text.strip()} + + +@R.register("auto_image_step") +class AutoImageStep(AutoResourceStep): + """Interpret image resource files into daily notes via a direct VLM call. + + Unlike text resources (agent + file tools), the image interpretation is a + single vision-model call. Images larger than the request budget or in + provider-unfriendly formats are downscaled/re-encoded in memory for the + request only; files under ``resource/`` are never modified. Note lookup, + renaming, deletion linkage, and day-index refresh reuse the inherited + AutoResourceStep helpers; only the interpretation differs. + """ + + def _skip_image_change(self, file_path: str) -> bool: + """This step interprets images, so the inherited image guard never applies.""" + return False + + def _max_image_bytes(self) -> int: + """Return the image read limit from Step or Job context.""" + value = self.kwargs.get("max_image_bytes") + if value is None and self.context is not None: + value = self.context.get("max_image_bytes") + return int(value) if value is not None else DEFAULT_MAX_IMAGE_INPUT_BYTES + + def _vision_model(self) -> ChatModelBase | None: + """Resolve the vision model: an explicit instance, then the ``vision`` + named as_llm component, then the ``default`` one.""" + for source in (self.kwargs, self.context or {}): + value = source.get("as_llm") + if isinstance(value, ChatModelBase): + return value + if self.app_context is None: + return None + models = self.app_context.components.get(ComponentEnum.AS_LLM, {}) + component = models.get("vision") or models.get("default") + if component is None: + return None + return getattr(component, "model", None) + + async def _caption_with_retry(self, model: ChatModelBase, user_message: UserMsg) -> dict: + """Return the caption fields from the vision model. + + Primary path is the schema-forced structured output (the SDK enforces + the ``name``/``description``/``caption`` contract and retries transport + errors). When that fails or yields no usable field, retry once with a + plain call parsed leniently. + """ + try: + structured = await model.generate_structured_output( + messages=[user_message], + structured_model=_CaptionOutput, + ) + content = structured.content if isinstance(structured.content, dict) else {} + normalized = _normalize_caption_fields(dict(content)) + if normalized["caption"] or normalized["description"]: + self.logger.info(f"[{self.name}] structured caption ok name={normalized['name']}") + return normalized + self.logger.warning(f"[{self.name}] structured caption empty; retrying with a plain call") + except Exception as exc: # pylint: disable=broad-except + self.logger.warning(f"[{self.name}] structured caption failed ({exc}); retrying with a plain call") + result = await model([user_message]) + return _parse_caption_json(await _response_text(result)) + + async def _read_image(self, file_path: str) -> dict | None: + """Read the image file and build the VLM request payload. + + Returns ``None`` when the change must be skipped (stat failure or + oversized file); the skip outcome is already recorded on the response. + """ + abs_path = self.workspace_path / file_path + max_image_bytes = self._max_image_bytes() + try: + size_bytes = abs_path.stat().st_size + except OSError as exc: + self.context.response.success = False + self.context.response.answer = f"Failed to inspect resource file: {file_path}: {exc}" + self.context.response.metadata.update( + { + "path": file_path, + "action": "failed", + "error": str(exc), + "modified": False, + }, + ) + self.logger.warning(f"[{self.name}] resource stat failed file_path={file_path} error={exc}") + return None + if size_bytes > max_image_bytes: + self.context.response.success = True + self.context.response.answer = ( + f"Skipped oversized image resource file: {file_path} ({size_bytes} > {max_image_bytes} bytes)" + ) + self.context.response.metadata.update( + { + "path": file_path, + "action": "skipped", + "reason": "file_too_large", + "oversized": True, + "size_bytes": size_bytes, + "max_image_bytes": max_image_bytes, + "modified": False, + }, + ) + self.logger.warning( + f"[{self.name}] skip oversized image resource file_path={file_path} " + f"size_bytes={size_bytes} max_image_bytes={max_image_bytes}", + ) + return None + + self.logger.info(f"[{self.name}] read image start file_path={file_path}") + async with aiofiles.open(abs_path, "rb") as f: + data = await f.read() + payload = _build_image_request_payload(data, Path(file_path).suffix.lower()) + self.logger.info( + f"[{self.name}] read image done file_path={file_path} size_bytes={size_bytes} " + f"mime={payload['mime']} converted={payload['converted']}", + ) + return payload + + async def _handle_change(self, file_path: str, raw_change) -> dict: + """Skip non-image changes; isolate per-change failures so the batch continues.""" + if file_path and not is_image_file(file_path): + file_path = self.to_workspace_relative(file_path) if Path(file_path).is_absolute() else file_path + self.context.response.metadata = {} + answer = f"Skipped non-image resource file: {file_path}" + self.context.response.success = True + self.context.response.answer = answer + self.context.response.metadata.update( + { + "path": file_path, + "action": "skipped", + "reason": "non_image_file", + "modified": False, + }, + ) + self.logger.info(f"[{self.name}] skip change file_path={file_path} reason=non_image_file") + return { + "success": True, + "path": file_path, + "change": str(raw_change), + "answer": answer, + "metadata": dict(self.context.response.metadata), + } + try: + return await super()._handle_change(file_path, raw_change) + except Exception as exc: # pylint: disable=broad-except + self.context.response.success = False + self.context.response.answer = f"Failed to caption image resource: {file_path}: {exc}" + self.context.response.metadata.update( + { + "path": file_path, + "action": "failed", + "error": str(exc), + "modified": False, + }, + ) + self.logger.warning(f"[{self.name}] caption failed file_path={file_path} error={exc}") + return { + "success": False, + "path": file_path, + "change": str(raw_change), + "answer": self.context.response.answer, + "metadata": dict(self.context.response.metadata), + } + + async def _handle_upsert(self, file_path: str, date_str: str, note_stem: str, added: bool) -> None: + """Caption the image and write/refresh its note (image counterpart of the text upsert).""" + daily_dir = self.config_value("daily_dir") + fallback_path = f"{daily_dir}/{date_str}/{note_stem}.md" + try: + note = await self._list_resource_note(date_str, file_path, fallback_path) + except RuntimeError as exc: + self.context.response.success = False + self.context.response.answer = str(exc) + self.logger.info(f"[{self.name}] list failed file_path={file_path} answer={str(exc)!r}") + return + + note_path = str(note["path"]) if note else fallback_path + note_created = note is None + before_note_path = note_path + before_note_bytes = self._note_bytes(note_path) + self.logger.info( + f"[{self.name}] upsert start file_path={file_path} date={date_str} " f"note_stem={note_stem} added={added}", + ) + + model = self._vision_model() + if model is None: + self.context.response.success = True + self.context.response.answer = f"Skipped image resource without a vision model: {file_path}" + self.context.response.metadata.update( + { + "path": file_path, + "action": "skipped", + "reason": "vision_model_not_configured", + "modified": False, + }, + ) + self.logger.warning(f"[{self.name}] no vision model configured file_path={file_path}") + return + + payload = await self._read_image(file_path) + if payload is None: + return + + user_message = UserMsg( + name="user", + content=[ + TextBlock(text=self.prompt_format("user_message", file_path=file_path, date=date_str)), + DataBlock( + source=Base64Source(data=payload["data_b64"], media_type=payload["mime"]), + name="image", + ), + ], + ) + parsed = await self._caption_with_retry(model, user_message) + name = _sanitize_note_name(str(parsed.get("name") or ""), note_stem) + caption = str(parsed.get("caption") or "").strip() + description = str(parsed.get("description") or "").strip() or caption[:120] + body = f"![[{file_path}]]\n\n## Caption\n\n{caption}\n" + + # The write job's ``name`` parameter is the note name; calling the job + # directly (instead of run_job) keeps it clear of run_job's + # positional-only job-selector argument. + write_job = self.get_job("write") + if write_job is None: + raise RuntimeError("Job write not found") + write_response = await write_job( + path=note_path, + name=name, + description=description, + content=body, + metadata={ + _SOURCE_RESOURCE_KEY: self._source_resource_link(file_path), + "kind": "image", + "media_type": payload["source_mime"], + }, + ) + if not write_response.success: + raise RuntimeError(f"write failed: {write_response.answer}") + + if note_created: + try: + note = await self._list_resource_note(date_str, file_path, fallback_path) + except RuntimeError as exc: + self.context.response.success = False + self.context.response.answer = str(exc) + self.context.response.metadata.update({"path": None, "created": note_created, "modified": False}) + self.logger.info(f"[{self.name}] post-create list failed file_path={file_path} answer={str(exc)!r}") + return + if note is None: + self.context.response.success = True + self.context.response.answer = f"Captioned image resource {file_path}" + self.context.response.metadata.update({"path": None, "created": False, "modified": False}) + self.logger.info(f"[{self.name}] done without note file_path={file_path} modified=False") + return + note_path = str(note["path"]) + + try: + await self._ensure_resource_frontmatter(note_path, file_path) + note_path = await self._rename_from_frontmatter_name( + note_path, + date_str, + file_path, + note_stem, + fallback_path, + allow_rename=note_created, + ) + except RuntimeError as exc: + self.context.response.success = False + self.context.response.answer = str(exc) + self.context.response.metadata.update( + { + "path": note_path, + "created": note_created, + "modified": self._note_modified(before_note_path, before_note_bytes, note_path), + }, + ) + self.logger.info(f"[{self.name}] post-write failed path={note_path} answer={str(exc)!r}") + return + + modified = self._note_modified(before_note_path, before_note_bytes, note_path) + index_payload = await refresh_day_index(self.file_store, date_str, daily_dir) + + self.context.response.success = True + self.context.response.answer = f"Captioned image resource {file_path} -> {note_path}" + self.context.response.metadata.update( + { + "path": note_path, + "created": note_created, + "modified": modified, + "session_id": note_stem, + "source_resource": self._source_resource_link(file_path), + "action": "added" if added else "modified", + "media_type": payload["source_mime"], + "index": index_payload, + }, + ) + self.logger.info(f"[{self.name}] done {note_path} modified={modified}") diff --git a/reme/steps/evolve/auto_image.yaml b/reme/steps/evolve/auto_image.yaml new file mode 100644 index 00000000..08e8e482 --- /dev/null +++ b/reme/steps/evolve/auto_image.yaml @@ -0,0 +1,52 @@ +user_message: | + Describe the attached image from the user's resource library for a memory knowledge base. + + Resource image file: {file_path} + Date: {date} + + ## What to Record + - **Visible facts**: people, objects, places, actions, and relationships that + may matter in future conversations. + - **Text in the image**: transcribe meaningful visible text verbatim (slides, + whiteboards, screenshots, signs, labels); keep the original language; do not + translate or correct it. + - **Numbers and dates**: quantities, measurements, versions, prices, timestamps. + - **Image type and purpose**: photo / screenshot / diagram / whiteboard / + document scan, and what it appears to be for. + + ## What to Ignore + - Pure visual style, composition, lighting, or aesthetic qualities with no + informational value. + - Speculation about anything not visible in the image. + + ## Completeness + The caption must let a person who cannot see the image understand its content. + For text-heavy images, verbatim transcription takes priority over summary; use + lists or line breaks to mirror the layout when helpful. + + ## Output + Return a JSON object with the fields `name` (short kebab-case topic stem for + the note filename; never include dates), `description` (one-sentence summary + that conveys the key information on its own), and `caption` (complete + description / transcription). Return only the JSON object. +user_message_zh: | + 为记忆知识库描述用户资源库中的这张图像。 + + 资源图像文件:{file_path} + 日期:{date} + + ## 记录什么 + - **可见事实**:人物、物体、地点、动作及相互关系——未来对话中可能重要的信息。 + - **图中的文字**:逐字转录有意义的可见文字(幻灯片、白板、截图、招牌、标签);保留原语言,不翻译、不纠错。 + - **数字与日期**:数量、度量、版本号、价格、时间戳。 + - **图像类型与用途**:照片/截图/图表/白板/文档扫描,以及它看起来是做什么用的。 + + ## 忽略什么 + - 纯视觉风格、构图、光照等无信息量的美学属性。 + - 对图中不存在内容的猜测。 + + ## 完整性 + caption 必须让看不到图的人理解其内容。文本密集的图,逐字转录优先于概括;可借用列表/换行还原版式。 + + ## 输出 + 返回只含以下字段的 JSON 对象:`name`(笔记文件名的简短 kebab-case 主题词;不要含任何日期)、`description`(一句话总结,单独读即可传达关键信息)、`caption`(完整描述/转录)。只返回 JSON 对象本身。 diff --git a/reme/steps/file_io/__init__.py b/reme/steps/file_io/__init__.py index 7417e9ed..2aca0e3e 100644 --- a/reme/steps/file_io/__init__.py +++ b/reme/steps/file_io/__init__.py @@ -2,7 +2,7 @@ from ._daily_index import extract_daily_date, parse_daily_date, refresh_day_index, validate_session_id from ._file_io import get_path_lock, write_file_safe -from ._path import validate_filename_component +from ._path import is_image_file, validate_filename_component from .daily_list import DailyListStep from .daily_reindex import DailyReindexStep from .daily_write import DailyWriteStep @@ -26,6 +26,7 @@ __all__ = [ "parse_daily_date", "validate_session_id", "validate_filename_component", + "is_image_file", "get_path_lock", "write_file_safe", "DailyListStep", diff --git a/reme/steps/file_io/_path.py b/reme/steps/file_io/_path.py index d5d217b7..ee5e042c 100644 --- a/reme/steps/file_io/_path.py +++ b/reme/steps/file_io/_path.py @@ -28,6 +28,14 @@ IMAGE_MIME_BY_EXT: dict[str, str] = { ".heic": "image/heic", } +IMAGE_SUFFIXES = frozenset(IMAGE_MIME_BY_EXT) + + +def is_image_file(path: str | Path) -> bool: + """Return True when ``path`` carries a known image file suffix.""" + return Path(path).suffix.lower() in IMAGE_SUFFIXES + + _INVALID_CHARS = re.compile(r'[<>:"/\\|?*\x00-\x1f]') _RESERVED_NAMES = { "CON", diff --git a/tests/unit/test_auto_image_steps.py b/tests/unit/test_auto_image_steps.py new file mode 100644 index 00000000..6a5c5fa5 --- /dev/null +++ b/tests/unit/test_auto_image_steps.py @@ -0,0 +1,686 @@ +"""Tests for AutoImageStep: image resource files become caption daily notes. + +The vision model boundary is faked (no network); test images are synthesized +with PIL inside a temporary workspace. +""" + +# pylint: disable=protected-access + +import asyncio +import base64 +import hashlib +import io +import json +import os +import tempfile +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock + +import frontmatter +import pytest +from agentscope.model import ChatModelBase +from PIL import Image + +from reme.components.agent_wrapper import BaseAgentWrapper +from reme.components.file_store import LocalFileStore +from reme.components.runtime_context import RuntimeContext +from reme.steps.evolve.auto_image import AutoImageStep, _parse_caption_json +from reme.steps.evolve.auto_resource import AutoResourceStep +from reme.steps.file_io import DailyListStep, FrontmatterUpdateStep, MoveStep, WriteStep + + +class temp_chdir: + """Context manager to temporarily chdir into a path and restore on exit.""" + + def __init__(self, path): + self.path = path + self.old = None + + def __enter__(self): + self.old = os.getcwd() + os.chdir(self.path) + return self + + def __exit__(self, *exc): + os.chdir(self.old) + + +class _FakeAgentWrapper(BaseAgentWrapper): + """Capture agent calls without invoking a real model.""" + + def __init__(self): + super().__init__() + self.inputs = "" + + async def reply(self, inputs, **kwargs) -> dict: + self.inputs = inputs + return {"result": "ok"} + + +class _FakeVisionModel(ChatModelBase): + """Capture VLM calls and return a canned text response (plain-call path).""" + + def __init__(self, text: str): + self.text = text + self.calls: list = [] + + async def generate_structured_output(self, messages, structured_model, **kwargs): # pylint: disable=unused-argument + """Force the plain-call path in tests.""" + raise NotImplementedError("structured path not faked") + + async def __call__(self, messages, **kwargs): + self.calls.append(messages) + return SimpleNamespace(content=[{"type": "text", "text": self.text}]) + + +class _FlakyVisionModel(ChatModelBase): + """Fail the first plain call, succeed afterwards.""" + + def __init__(self, text: str): + self.text = text + self.calls = 0 + + async def generate_structured_output(self, messages, structured_model, **kwargs): # pylint: disable=unused-argument + """Force the plain-call path in tests.""" + raise NotImplementedError("structured path not faked") + + async def __call__(self, messages, **kwargs): + self.calls += 1 + if self.calls == 1: + raise RuntimeError("vision backend unavailable") + return SimpleNamespace(content=[{"type": "text", "text": self.text}]) + + +class _StructuredVisionModel(ChatModelBase): + """Serve the schema-forced structured path; count fallback plain calls.""" + + def __init__(self, content: dict | None = None, error: Exception | None = None, plain_text: str = "plain"): + self.content = content + self.error = error + self.plain_text = plain_text + self.structured_calls: list = [] + self.plain_calls: list = [] + + async def generate_structured_output(self, messages, structured_model, **kwargs): # pylint: disable=unused-argument + """Serve the canned structured content (or raise the canned error).""" + self.structured_calls.append(messages) + if self.error is not None: + raise self.error + return SimpleNamespace(content=dict(self.content or {})) + + async def __call__(self, messages, **kwargs): # pylint: disable=unused-argument + """Serve the canned plain-call text response.""" + self.plain_calls.append(messages) + return SimpleNamespace(content=[{"type": "text", "text": self.plain_text}]) + + +class _StepJob: + """Tiny job adapter for unit tests that need BaseStep.run_job.""" + + def __init__(self, step_cls, app_context, file_store): + self.step_cls = step_cls + self.app_context = app_context + self.file_store = file_store + + async def __call__(self, **kwargs): + step = self.step_cls(app_context=self.app_context, file_store=self.file_store) + result = await step(**kwargs) + return result or step.context.response + + +def _make_app_context(workspace_path: Path): + """Create a mock app_context with app_config pointing to the given workspace.""" + ctx = MagicMock() + ctx.app_config.workspace_dir = str(workspace_path) + ctx.app_config.daily_dir = "daily" + ctx.app_config.digest_dir = "digest" + ctx.app_config.resource_dir = "resource" + ctx.app_config.session_dir = "session" + ctx.app_config.timezone = None + return ctx + + +def _install_file_jobs(app_context, file_store) -> None: + app_context.jobs = { + "daily_list": _StepJob(DailyListStep, app_context, file_store), + "frontmatter_update": _StepJob(FrontmatterUpdateStep, app_context, file_store), + "move": _StepJob(MoveStep, app_context, file_store), + "write": _StepJob(WriteStep, app_context, file_store), + } + + +def _png_bytes(width: int = 8, height: int = 8, color=(200, 30, 30)) -> bytes: + image = Image.new("RGB", (width, height), color) + buffer = io.BytesIO() + image.save(buffer, format="PNG") + return buffer.getvalue() + + +def _img_bytes(image_format: str, size=(8, 8), color=(200, 30, 30)) -> bytes: + """Synthesize an image in any PIL-supported format (incl. HEIF via pillow-heif).""" + if image_format == "HEIF": + from pillow_heif import register_heif_opener + + register_heif_opener() + image = Image.new("RGB", size, color) + buffer = io.BytesIO() + image.save(buffer, format=image_format) + return buffer.getvalue() + + +def _write_binary(path: Path, data: bytes) -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(data) + return path + + +def _write_note(path: Path, source_resource: str, body: str = "old caption") -> Path: + content = ( + f"---\nname: {path.stem}\ndescription: old\n" + f'source_resource: "{source_resource}"\nkind: image\n' + f"media_type: image/png\n---\n![[{source_resource[2:-2]}]]\n\n## Caption\n\n{body}\n" + ) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(content, encoding="utf-8") + return path + + +def _caption_json(name: str, description: str, caption: str) -> str: + return json.dumps({"name": name, "description": description, "caption": caption}) + + +def _run_step(step, changes, **context_kwargs): + return step(RuntimeContext(changes=changes, **context_kwargs)) + + +def test_auto_image_creates_caption_note(): + """An added image produces a renamed daily note with caption and embed link.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + cwd = Path.cwd() + app_ctx = _make_app_context(cwd) + fs = LocalFileStore(name="test_store", embedding_store="") + await fs.start() + _install_file_jobs(app_ctx, fs) + try: + source = _write_binary(cwd / "resource" / "2026-01-01" / "img.png", _png_bytes()) + model = _FakeVisionModel(_caption_json("red-square", "A red square", "An 8x8 solid red square.")) + step = AutoImageStep(app_context=app_ctx, file_store=fs, as_llm=model) + resp = await _run_step(step, [{"change": "added", "path": str(source)}]) + + assert resp.success is True + note_path = cwd / "daily" / "2026-01-01" / "red-square.md" + assert note_path.is_file() + post = frontmatter.loads(note_path.read_text(encoding="utf-8")) + assert post.metadata["source_resource"] == "[[resource/2026-01-01/img.png]]" + assert post.metadata["kind"] == "image" + assert post.metadata["media_type"] == "image/png" + assert post.metadata["name"] == "red-square" + assert "![[resource/2026-01-01/img.png]]" in post.content + assert "An 8x8 solid red square." in post.content + assert (cwd / "daily" / "2026-01-01.md").is_file() + assert len(model.calls) == 1 + finally: + await fs.close() + + asyncio.run(run()) + + +def test_auto_image_parses_fenced_json_and_falls_back_to_raw_text(): + """Fenced JSON is parsed; non-JSON output degrades to a raw-text caption.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + cwd = Path.cwd() + app_ctx = _make_app_context(cwd) + fs = LocalFileStore(name="test_store", embedding_store="") + await fs.start() + _install_file_jobs(app_ctx, fs) + try: + fenced = "```json\n" + _caption_json("fenced-note", "Fenced", "Fenced caption body.") + "\n```" + source = _write_binary(cwd / "resource" / "2026-01-01" / "fenced.png", _png_bytes()) + step = AutoImageStep(app_context=app_ctx, file_store=fs, as_llm=_FakeVisionModel(fenced)) + resp = await _run_step(step, [{"change": "added", "path": str(source)}]) + assert resp.success is True + fenced_post = frontmatter.loads((cwd / "daily" / "2026-01-01" / "fenced-note.md").read_text("utf-8")) + assert "Fenced caption body." in fenced_post.content + + raw = _write_binary(cwd / "resource" / "2026-01-01" / "photo.png", _png_bytes(color=(30, 30, 200))) + step = AutoImageStep( + app_context=app_ctx, + file_store=fs, + as_llm=_FakeVisionModel("A plain description."), + ) + resp = await _run_step(step, [{"change": "added", "path": str(raw)}]) + assert resp.success is True + raw_post = frontmatter.loads((cwd / "daily" / "2026-01-01" / "photo.md").read_text("utf-8")) + assert "A plain description." in raw_post.content + finally: + await fs.close() + + asyncio.run(run()) + + +def test_auto_image_updates_existing_note_in_place(): + """A modified image rewrites the same note found via source_resource.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + cwd = Path.cwd() + app_ctx = _make_app_context(cwd) + fs = LocalFileStore(name="test_store", embedding_store="") + await fs.start() + _install_file_jobs(app_ctx, fs) + try: + source = _write_binary(cwd / "resource" / "2026-01-01" / "img.png", _png_bytes()) + note_path = _write_note( + cwd / "daily" / "2026-01-01" / "red-square.md", + "[[resource/2026-01-01/img.png]]", + ) + model = _FakeVisionModel(_caption_json("red-square", "Updated", "The updated caption.")) + step = AutoImageStep(app_context=app_ctx, file_store=fs, as_llm=model) + resp = await _run_step(step, [{"change": "modified", "path": str(source)}]) + + assert resp.success is True + assert note_path.is_file() + post = frontmatter.loads(note_path.read_text(encoding="utf-8")) + assert "The updated caption." in post.content + assert not (cwd / "daily" / "2026-01-01" / "img.md").exists() + finally: + await fs.close() + + asyncio.run(run()) + + +def test_auto_image_deletes_linked_note(): + """Deleting the image resource removes its caption note.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + cwd = Path.cwd() + app_ctx = _make_app_context(cwd) + fs = LocalFileStore(name="test_store", embedding_store="") + await fs.start() + _install_file_jobs(app_ctx, fs) + try: + source = cwd / "resource" / "2026-01-01" / "img.png" + note_path = _write_note( + cwd / "daily" / "2026-01-01" / "red-square.md", + "[[resource/2026-01-01/img.png]]", + ) + step = AutoImageStep(app_context=app_ctx, file_store=fs, as_llm=_FakeVisionModel("{}")) + resp = await _run_step(step, [{"change": "deleted", "path": str(source)}]) + + assert resp.success is True + assert not note_path.exists() + finally: + await fs.close() + + asyncio.run(run()) + + +def test_auto_image_downscales_oversized_image_for_request_only(): + """Images beyond the request budget are downscaled in the request; storage is untouched.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + cwd = Path.cwd() + app_ctx = _make_app_context(cwd) + fs = LocalFileStore(name="test_store", embedding_store="") + await fs.start() + _install_file_jobs(app_ctx, fs) + try: + source = _write_binary( + cwd / "resource" / "2026-01-01" / "huge.png", + _png_bytes(width=3000, height=3000), + ) + stored_bytes = source.read_bytes() + model = _FakeVisionModel(_caption_json("huge-image", "Big", "A big image.")) + step = AutoImageStep(app_context=app_ctx, file_store=fs, as_llm=model) + resp = await _run_step(step, [{"change": "added", "path": str(source)}]) + + assert resp.success is True + assert len(model.calls) == 1 + data_block = model.calls[0][0].content[1] + assert data_block.source.media_type == "image/jpeg" + with Image.open(io.BytesIO(base64.b64decode(data_block.source.data))) as sent: + assert max(sent.size) <= 2048 + assert source.read_bytes() == stored_bytes + assert (cwd / "daily" / "2026-01-01" / "huge-image.md").is_file() + finally: + await fs.close() + + asyncio.run(run()) + + +def test_auto_image_skips_oversized_file(): + """Files beyond max_image_bytes are skipped without a VLM call.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + cwd = Path.cwd() + app_ctx = _make_app_context(cwd) + fs = LocalFileStore(name="test_store", embedding_store="") + await fs.start() + _install_file_jobs(app_ctx, fs) + try: + source = _write_binary(cwd / "resource" / "2026-01-01" / "img.png", _png_bytes()) + model = _FakeVisionModel(_caption_json("x", "y", "z")) + step = AutoImageStep(app_context=app_ctx, file_store=fs, as_llm=model) + resp = await _run_step(step, [{"change": "added", "path": str(source)}], max_image_bytes=8) + + result = resp.metadata["results"][0] + assert resp.success is True + assert result["metadata"]["reason"] == "file_too_large" + assert result["metadata"]["oversized"] is True + assert not model.calls + assert not (cwd / "daily" / "2026-01-01" / "img.md").exists() + finally: + await fs.close() + + asyncio.run(run()) + + +def test_auto_image_skips_without_vision_model(): + """Without any resolvable vision model the change is skipped with a reason.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + cwd = Path.cwd() + app_ctx = _make_app_context(cwd) + app_ctx.components = {} + fs = LocalFileStore(name="test_store", embedding_store="") + await fs.start() + _install_file_jobs(app_ctx, fs) + try: + source = _write_binary(cwd / "resource" / "2026-01-01" / "img.png", _png_bytes()) + step = AutoImageStep(app_context=app_ctx, file_store=fs) + resp = await _run_step(step, [{"change": "added", "path": str(source)}]) + + result = resp.metadata["results"][0] + assert resp.success is True + assert result["metadata"]["reason"] == "vision_model_not_configured" + assert not (cwd / "daily" / "2026-01-01" / "img.md").exists() + finally: + await fs.close() + + asyncio.run(run()) + + +def test_image_and_text_changes_are_routed_by_suffix(): + """auto_image skips text changes; auto_resource skips image changes.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + cwd = Path.cwd() + app_ctx = _make_app_context(cwd) + fs = LocalFileStore(name="test_store", embedding_store="") + await fs.start() + _install_file_jobs(app_ctx, fs) + try: + image_step = AutoImageStep(app_context=app_ctx, file_store=fs, as_llm=_FakeVisionModel("{}")) + resp = await _run_step( + image_step, + [{"change": "added", "path": str(cwd / "resource" / "2026-01-01" / "note.txt")}], + ) + assert resp.metadata["results"][0]["metadata"]["reason"] == "non_image_file" + + wrapper = _FakeAgentWrapper() + resource_step = AutoResourceStep(app_context=app_ctx, file_store=fs, agent_wrapper=wrapper) + resp = await _run_step( + resource_step, + [{"change": "added", "path": str(cwd / "resource" / "2026-01-01" / "img.png")}], + ) + assert resp.metadata["results"][0]["metadata"]["reason"] == "image_file" + assert wrapper.inputs == "" + finally: + await fs.close() + + asyncio.run(run()) + + +def test_auto_image_model_failure_is_isolated_per_change(): + """A failing VLM call marks one change failed while the rest of the batch continues.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + cwd = Path.cwd() + app_ctx = _make_app_context(cwd) + fs = LocalFileStore(name="test_store", embedding_store="") + await fs.start() + _install_file_jobs(app_ctx, fs) + try: + first = _write_binary(cwd / "resource" / "2026-01-01" / "img-a.png", _png_bytes()) + second = _write_binary(cwd / "resource" / "2026-01-01" / "img-b.png", _png_bytes(color=(20, 90, 200))) + model = _FlakyVisionModel(_caption_json("blue-square", "Blue", "A blue square.")) + step = AutoImageStep(app_context=app_ctx, file_store=fs, as_llm=model) + resp = await _run_step( + step, + [ + {"change": "added", "path": str(first)}, + {"change": "added", "path": str(second)}, + ], + ) + + assert resp.success is False + results = resp.metadata["results"] + assert results[0]["success"] is False + assert results[0]["metadata"]["action"] == "failed" + assert results[1]["success"] is True + assert (cwd / "daily" / "2026-01-01" / "blue-square.md").is_file() + assert first.exists() and second.exists() + finally: + await fs.close() + + asyncio.run(run()) + + +def test_auto_image_uniquifies_conflicting_note_name(): + """A name collision with an unrelated note falls back to the sha1-suffixed path.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + cwd = Path.cwd() + app_ctx = _make_app_context(cwd) + fs = LocalFileStore(name="test_store", embedding_store="") + await fs.start() + _install_file_jobs(app_ctx, fs) + try: + source = _write_binary(cwd / "resource" / "2026-01-01" / "img.png", _png_bytes()) + _write_note(cwd / "daily" / "2026-01-01" / "red-square.md", "[[resource/2026-01-01/other.png]]") + model = _FakeVisionModel(_caption_json("red-square", "A red square", "An 8x8 solid red square.")) + step = AutoImageStep(app_context=app_ctx, file_store=fs, as_llm=model) + resp = await _run_step(step, [{"change": "added", "path": str(source)}]) + + assert resp.success is True + suffix = hashlib.sha1(b"resource/2026-01-01/img.png").hexdigest()[:8] + assert (cwd / "daily" / "2026-01-01" / f"red-square--{suffix}.md").is_file() + finally: + await fs.close() + + asyncio.run(run()) + + +def test_parse_caption_json_cross_fills_missing_fields(): + """A description-only JSON payload cross-fills the caption instead of leaking raw JSON.""" + parsed = _parse_caption_json('{"description": "Waterfall in Iceland."}') + assert parsed["caption"] == "Waterfall in Iceland." + assert parsed["description"] == "Waterfall in Iceland." + assert parsed["name"] == "" + + parsed = _parse_caption_json('{"caption": "A red square."}') + assert parsed["caption"] == "A red square." + assert parsed["description"] == "" + + parsed = _parse_caption_json('```json\n{"name": "n", "description": "d", "caption": "c"}\n```') + assert parsed == {"name": "n", "description": "d", "caption": "c"} + + +def test_parse_caption_json_falls_back_to_raw_text(): + """Unusable payloads degrade to a raw-text caption.""" + parsed = _parse_caption_json("A plain description without json.") + assert parsed == {"name": "", "description": "", "caption": "A plain description without json."} + + parsed = _parse_caption_json('{"foo": 1}') + assert parsed["caption"] == '{"foo": 1}' + + +def test_auto_image_note_body_stays_clean_when_caption_field_missing(): + """Real-model regression: JSON with only a description must not enter the body verbatim.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + cwd = Path.cwd() + app_ctx = _make_app_context(cwd) + fs = LocalFileStore(name="test_store", embedding_store="") + await fs.start() + _install_file_jobs(app_ctx, fs) + try: + source = _write_binary(cwd / "resource" / "2026-01-01" / "img.png", _png_bytes()) + model = _FakeVisionModel('{"file": "resource/2026-01-01/img.png", "description": "A tall waterfall."}') + step = AutoImageStep(app_context=app_ctx, file_store=fs, as_llm=model) + resp = await _run_step(step, [{"change": "added", "path": str(source)}]) + + assert resp.success is True + note_path = cwd / "daily" / "2026-01-01" / "img.md" + assert note_path.is_file() + content = note_path.read_text(encoding="utf-8") + assert "A tall waterfall." in content + assert '{"file"' not in content + assert '"description"' not in content + finally: + await fs.close() + + asyncio.run(run()) + + +def test_auto_image_uses_structured_output_first(): + """The schema-forced structured call is the primary path; no plain fallback.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + cwd = Path.cwd() + app_ctx = _make_app_context(cwd) + fs = LocalFileStore(name="test_store", embedding_store="") + await fs.start() + _install_file_jobs(app_ctx, fs) + try: + source = _write_binary(cwd / "resource" / "2026-01-01" / "img.png", _png_bytes()) + model = _StructuredVisionModel( + content={"name": "red-square", "description": "A red square.", "caption": "An 8x8 red square."}, + ) + step = AutoImageStep(app_context=app_ctx, file_store=fs, as_llm=model) + resp = await _run_step(step, [{"change": "added", "path": str(source)}]) + + assert resp.success is True + assert len(model.structured_calls) == 1 + assert not model.plain_calls + note_path = cwd / "daily" / "2026-01-01" / "red-square.md" + assert note_path.is_file() + content = note_path.read_text(encoding="utf-8") + assert "An 8x8 red square." in content + finally: + await fs.close() + + asyncio.run(run()) + + +def test_auto_image_retries_with_plain_call_when_structured_fails(): + """A failing structured call retries once via the plain-call path.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + cwd = Path.cwd() + app_ctx = _make_app_context(cwd) + fs = LocalFileStore(name="test_store", embedding_store="") + await fs.start() + _install_file_jobs(app_ctx, fs) + try: + source = _write_binary(cwd / "resource" / "2026-01-01" / "img.png", _png_bytes()) + model = _StructuredVisionModel( + error=RuntimeError("provider rejects tool_choice"), + plain_text=_caption_json("plain-note", "Plain", "Plain-call caption."), + ) + step = AutoImageStep(app_context=app_ctx, file_store=fs, as_llm=model) + resp = await _run_step(step, [{"change": "added", "path": str(source)}]) + + assert resp.success is True + assert len(model.structured_calls) == 1 + assert len(model.plain_calls) == 1 + note_path = cwd / "daily" / "2026-01-01" / "plain-note.md" + content = note_path.read_text(encoding="utf-8") + assert "Plain-call caption." in content + finally: + await fs.close() + + asyncio.run(run()) + + +def test_auto_image_falls_back_when_structured_content_empty(): + """A structured response with no usable fields also triggers the plain retry.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + cwd = Path.cwd() + app_ctx = _make_app_context(cwd) + fs = LocalFileStore(name="test_store", embedding_store="") + await fs.start() + _install_file_jobs(app_ctx, fs) + try: + source = _write_binary(cwd / "resource" / "2026-01-01" / "img.png", _png_bytes()) + model = _StructuredVisionModel( + content={"name": "", "description": "", "caption": ""}, + plain_text=_caption_json("empty-note", "Empty", "Recovered by plain call."), + ) + step = AutoImageStep(app_context=app_ctx, file_store=fs, as_llm=model) + resp = await _run_step(step, [{"change": "added", "path": str(source)}]) + + assert resp.success is True + assert len(model.plain_calls) == 1 + content = (cwd / "daily" / "2026-01-01" / "empty-note.md").read_text(encoding="utf-8") + assert "Recovered by plain call." in content + finally: + await fs.close() + + asyncio.run(run()) + + +def test_auto_image_converts_bmp_tiff_heic_requests(): + """bmp/tiff/heic resources are re-encoded for the request; notes record the source media_type.""" + + async def run(): + pytest.importorskip("pillow_heif") + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + cwd = Path.cwd() + app_ctx = _make_app_context(cwd) + fs = LocalFileStore(name="test_store", embedding_store="") + await fs.start() + _install_file_jobs(app_ctx, fs) + try: + sources = [] + for stem, fmt, suffix in ( + ("photo", "BMP", ".bmp"), + ("scan", "TIFF", ".tiff"), + ("phone", "HEIF", ".heic"), + ): + sources.append(_write_binary(cwd / "resource" / "2026-01-01" / f"{stem}{suffix}", _img_bytes(fmt))) + model = _StructuredVisionModel(content={"name": "", "description": "d", "caption": "converted caption"}) + step = AutoImageStep(app_context=app_ctx, file_store=fs, as_llm=model) + resp = await _run_step(step, [{"change": "added", "path": str(p)} for p in sources]) + + assert resp.success is True + assert len(model.structured_calls) == 3 + sent_mimes = set() + for call in model.structured_calls: + sent_mimes.add(call[0].content[1].source.media_type) + assert sent_mimes <= {"image/png", "image/jpeg"} + for stem, expected_mime in (("photo", "image/bmp"), ("scan", "image/tiff"), ("phone", "image/heic")): + note = frontmatter.loads((cwd / "daily" / "2026-01-01" / f"{stem}.md").read_text(encoding="utf-8")) + assert note.metadata["media_type"] == expected_mime + assert "converted caption" in note.content + finally: + await fs.close() + + asyncio.run(run())