feat: add auto_image step for image resource caption notes

This commit is contained in:
wang-qisen 2026-08-27 17:38:20 +08:00
parent ef3f99f019
commit 1c4e13270b
6 changed files with 1200 additions and 1 deletions

View file

@ -1,6 +1,7 @@
"""Evolve steps.""" """Evolve steps."""
from ._evolve import now from ._evolve import now
from .auto_image import AutoImageStep
from .auto_memory import AutoMemoryStep from .auto_memory import AutoMemoryStep
from .auto_memory_cc import AutoMemoryCCStep from .auto_memory_cc import AutoMemoryCCStep
from .auto_resource import AutoResourceStep from .auto_resource import AutoResourceStep
@ -9,6 +10,7 @@ from .dream import DreamExtractStep, DreamFinishStep, DreamIntegrateStep, DreamT
__all__ = [ __all__ = [
"now", "now",
"AutoImageStep",
"AutoMemoryStep", "AutoMemoryStep",
"AutoMemoryCCStep", "AutoMemoryCCStep",
"AutoResourceStep", "AutoResourceStep",

View file

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

View file

@ -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 对象本身。

View file

@ -2,7 +2,7 @@
from ._daily_index import extract_daily_date, parse_daily_date, refresh_day_index, validate_session_id 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 ._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_list import DailyListStep
from .daily_reindex import DailyReindexStep from .daily_reindex import DailyReindexStep
from .daily_write import DailyWriteStep from .daily_write import DailyWriteStep
@ -26,6 +26,7 @@ __all__ = [
"parse_daily_date", "parse_daily_date",
"validate_session_id", "validate_session_id",
"validate_filename_component", "validate_filename_component",
"is_image_file",
"get_path_lock", "get_path_lock",
"write_file_safe", "write_file_safe",
"DailyListStep", "DailyListStep",

View file

@ -28,6 +28,14 @@ IMAGE_MIME_BY_EXT: dict[str, str] = {
".heic": "image/heic", ".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]') _INVALID_CHARS = re.compile(r'[<>:"/\\|?*\x00-\x1f]')
_RESERVED_NAMES = { _RESERVED_NAMES = {
"CON", "CON",

View file

@ -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())