scriptoriumadmin/ReMe
1
0
Fork 0
mirror of https://github.com/agentscope-ai/ReMe.git synced 2026-11-15 09:11:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

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

Browse source
This commit is contained in:
wang-qisen 2026-09-11 00:27:10 +08:00
parent c92543a560
commit a681dc2368
11 changed files with 379 additions and 744 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
docs/en/auto_memory.md
View file

@ -89,7 +89,8 @@ reme auto_memory session_id=session-a include_images=true supports_vision=true m
For a one-shot CLI invocation, use `reme start job=auto_memory` with the same arguments.
No YAML edit or rebuild is needed. `include_images` controls whether Auto Memory considers images;
`supports_vision` is the caller's declaration about the selected model, not automatic capability detection.
`supports_vision` declares whether the AgentScope wrapper's bound `as_llm` supports images; it does not detect capabilities
or select a different model.
Only `include_images=true supports_vision=true` sends images. Choose a model and AgentScope formatter that support them.
`include_images=false` keeps the existing text-only behavior on every wrapper, regardless of `supports_vision`.
@ -119,9 +120,9 @@ When no image-count limit is explicitly configured, direct mode reserves room fo
context. An explicit `context_config.max_image_num` below the incoming count causes a visible whole-input text fallback,
not silent removal of older images. Model/provider context limits still apply.
When sending images, Auto Memory uses `components.as_llm.vision` if configured, otherwise the wrapper's currently
bound `as_llm`. `supports_vision=true` must describe that selected model. This selection applies only to the current
memory Agent invocation; calls without image inputs keep the wrapper's original model.
Text-only and image-bearing Auto Memory calls use the same model bound through the AgentScope wrapper's `as_llm`.
`supports_vision=true` must describe that model. The optional `components.as_llm.vision` used for image resources
does not change Auto Memory's model.
Auto Memory does not create resource image files or separate caption cards. Relevant image facts may become part of
the normal daily memory note. The existing AgentScope wrapper also saves its internal Agent state under

11
docs/zh/auto_memory.md
View file

@ -81,8 +81,9 @@ reme auto_memory session_id=session-a include_images=true supports_vision=true m
```
单次 CLI 调用可使用 `reme start job=auto_memory`，其余参数相同。
不需要修改 YAML 或重新编译。`include_images` 控制是否考虑图像；`supports_vision` 是调用方对所选模型的视觉能力声明，
不是自动探测。仅当 `include_images=true supports_vision=true` 时才发送图像；请选用支持这些图像的模型与 AgentScope formatter。
不需要修改 YAML 或重新编译。`include_images` 控制是否考虑图像；`supports_vision` 声明 AgentScope wrapper 绑定的 `as_llm`
是否支持图像，不会自动探测能力或切换模型。仅当 `include_images=true supports_vision=true` 时才发送图像；
请选用支持这些图像的模型与 AgentScope formatter。
`include_images=false` 时，无论 `supports_vision` 如何设置，各种 wrapper 都保持原有纯文本行为。
纯文本仍通过已有的 `agent_wrapper.reply()` 字符串输入路径执行，无需为其他 wrapper 改造多模态接口。
@ -109,9 +110,9 @@ reme auto_memory session_id=session-a include_images=true supports_vision=true m
若显式设置的 `context_config.max_image_num` 小于输入图像数，会整次显式回退为纯文本，不静默移除较早图像。
模型和 provider 自身的上下文限制仍然适用。
发送图像时优先使用已配置的 `components.as_llm.vision`，未配置时使用 wrapper 当前绑定的 `as_llm`。
`supports_vision=true` 应当对应这个实际选中的模型。模型选择仅用于当前这次记忆 Agent 调用；
没有图像输入的调用仍使用 wrapper 原有模型。
纯文本与图文 Auto Memory 均使用 AgentScope wrapper 通过 `as_llm` 绑定的同一个模型。
`supports_vision=true` 应当描述这个模型的能力。图像资源处理使用的可选 `components.as_llm.vision`
不会改变 Auto Memory 的模型。
Auto Memory 不创建 resource 图像文件或独立 caption 卡片；
相关图像事实仍可被提取进普通 daily 记忆笔记。AgentScope wrapper 仍按原有行为在 `mem_session/agentscope` 保存

18
reme/components/agent_wrapper/as_agent_wrapper.py
View file

@ -37,7 +37,6 @@ from agentscope.event import (
ToolResultTextDeltaEvent,
)
from agentscope.message import TextBlock, ToolResultState, UserMsg
from agentscope.model import ChatModelBase
from agentscope.permission import PermissionBehavior, PermissionContext, PermissionDecision, PermissionMode
from agentscope.state import AgentState
from agentscope.tool import (
@ -300,20 +299,11 @@ class AsAgentWrapper(BaseAgentWrapper):
"""Resolve configured skill names to AgentScope local skill directories."""
return [str(path) for path in self._resolve_project_skills(skills).values()]
def _resolve_model(self, kwargs: dict[str, Any]) -> ChatModelBase:
"""Use an internal, caller-owned model override for this invocation only."""
if (model := kwargs.get("_model")) is not None:
if not isinstance(model, ChatModelBase):
raise TypeError("_model must be an initialized AgentScope ChatModelBase instance.")
return model
async def _build_agent(self, inputs: Any, **kwargs) -> tuple[Agent, Any]:
"""Build an Agent instance from kwargs. Returns (agent, processed_inputs)."""
model = self.as_llm.model if self.as_llm else None
if model is None:
raise ValueError("AsAgentWrapper requires a bound as_llm component with a valid model.")
return model
async def _build_agent(self, inputs: Any, **kwargs) -> tuple[Agent, Any]:
"""Build an Agent instance from kwargs. Returns (agent, processed_inputs)."""
model = self._resolve_model(kwargs)
self._cleanup_expired_sessions()
@ -373,7 +363,9 @@ class AsAgentWrapper(BaseAgentWrapper):
output_schema: dict | None = kwargs.get("output_schema")
if output_schema is not None:
model = self._resolve_model(kwargs)
assert self.as_llm is not None, "AsAgentWrapper requires a bound as_llm component with a valid model."
model = self.as_llm.model
assert model is not None, "AsAgentWrapper requires a bound as_llm component with a valid model."
res = await model.generate_structured_output(
messages=agent.state.context,
structured_model=output_schema,

9
reme/config/default.yaml
View file

@ -186,7 +186,7 @@ jobs:
default: false
supports_vision:
type: boolean
description: "Caller declares that the selected memory model supports images; does not enable images by itself"
description: "Caller declares that the AgentScope wrapper's bound as_llm supports images; does not enable images by itself"
default: false
image_mode:
type: string
@ -827,10 +827,9 @@ components:
max_tokens: 65536
thinking_enable: false
# Optional vision model for image resources and Auto Memory image inputs.
# Direct Auto Memory uses it only for the current image-bearing Agent call;
# when absent, it keeps the model already bound to that Agent wrapper.
# Image resource/caption processing otherwise falls back to "default".
# Optional dedicated vision model for image resources (auto_image_resource_step).
# Falls back to the "default" instance above when absent; uncomment to
# decouple the vision model from the main LLM.
# vision:
# backend: ${VLM_BACKEND:-openai}
# model: ${VLM_MODEL_NAME:-}

14
reme/steps/evolve/_auto_memory_image.py
View file

@ -12,10 +12,9 @@ import httpx
from agentscope.agent import ContextConfig
from agentscope.message import Base64Source, DataBlock, Msg, TextBlock, UserMsg
from ._image_caption import DEFAULT_MAX_IMAGE_INPUT_BYTES, _build_image_request_payload
from .auto_image_resource import DEFAULT_MAX_IMAGE_INPUT_BYTES, _build_image_request_payload
from ..file_io._path import _check_path_permission, resolve_path
from ...components.agent_wrapper.as_agent_wrapper import AsAgentWrapper
from ...enumeration import ComponentEnum
async def prepare_direct_message(
@ -49,15 +48,8 @@ async def prepare_direct_message(
f"[{step.name}] supports_vision is false; continuing with text-only memory without reading images.",
)
return None
stage = "model"
stage = "context-image-limit"
try:
component = wrapper.as_llm
if step.app_context is not None:
component = step.app_context.components.get(ComponentEnum.AS_LLM, {}).get("vision", component)
model = component.model if component is not None else None
if model is None:
raise ValueError("Direct images require a started memory model")
stage = "context-image-limit"
context_config = dict(
(reply_kwargs or {}).get("context_config", wrapper.kwargs.get("context_config")) or {},
)
@ -99,8 +91,6 @@ async def prepare_direct_message(
return None
if reply_kwargs is not None:
reply_kwargs["context_config"] = context_config
# Override only this Agent invocation; never rebind the shared wrapper.
reply_kwargs["_model"] = model
metadata["status"] = "prepared"
return result

335
reme/steps/evolve/_image_caption.py
View file

@ -1,335 +0,0 @@
"""Shared image preprocessing and caption generation, without workspace writes."""
import base64
import io
import json
import re
import warnings
from agentscope.message import Base64Source, DataBlock, TextBlock, UserMsg
from agentscope.model import ChatModelBase
from pydantic import BaseModel, Field
from ...enumeration import ComponentEnum
DEFAULT_MAX_IMAGE_INPUT_BYTES = 50 * 1024 * 1024
DEFAULT_MAX_IMAGE_PIXELS = 40_000_000
MAX_IMAGE_REQUEST_DIMENSION = 2048
_JPEG_QUALITY = 85
# Decoded formats outside this set are re-encoded to provider-friendly
# PNG/JPEG for VLM requests; the stored resource file is never modified.
_PASSTHROUGH_IMAGE_MIMES = frozenset({"image/png", "image/jpeg", "image/webp", "image/gif"})
_HEIF_BRANDS = frozenset(
{b"heic", b"heif", b"heix", b"heim", b"heis", b"hevc", b"hevx", b"hevm", b"hevs", b"mif1", b"msf1"},
)
_MAX_FTYP_SCAN_BYTES = 4096
_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 based on visible content; filename is only a weak naming hint",
)
description: str = Field(description="one-sentence summary of visible image content that stands on its own")
caption: str = Field(
description="complete description / verbatim transcription of meaningful content visible in the image",
)
def _load_pillow():
"""Load the core image dependency only when image processing runs."""
try:
from PIL import Image, ImageOps # pylint: disable=import-outside-toplevel
except ImportError as exc:
raise RuntimeError("Image captioning requires Pillow; install reme-ai[core]") from exc
return Image, ImageOps
def _looks_like_heif(data: bytes) -> bool:
"""Return whether an ISO-BMFF header declares a HEIC/HEIF brand."""
if len(data) < 12 or data[4:8] != b"ftyp":
return False
box_size = int.from_bytes(data[:4], "big")
if box_size < 12:
return False
end = min(box_size, len(data), _MAX_FTYP_SCAN_BYTES)
if data[8:12] in _HEIF_BRANDS:
return True
return any(data[index : index + 4] in _HEIF_BRANDS for index in range(16, end - 3, 4))
def _register_heif_opener() -> None:
"""Load and register HEIC support only for bytes that declare HEIF."""
try:
from pillow_heif import register_heif_opener # pylint: disable=import-outside-toplevel
except ImportError as exc:
raise RuntimeError("HEIC image captioning requires pillow-heif; install reme-ai[image-heif]") from exc
try:
register_heif_opener()
except Exception as exc: # pylint: disable=broad-except
raise RuntimeError(f"Failed to initialize HEIC image support: {exc}") from exc
def _normalize_image_bytes(
data: bytes,
suffix: str,
*,
max_image_pixels: int = DEFAULT_MAX_IMAGE_PIXELS,
) -> tuple[bytes | None, str, str]:
"""Validate and optionally normalize image bytes for a VLM request.
The returned tuple is ``(normalized_bytes, request_mime, source_mime)``.
``normalized_bytes`` is ``None`` only when the original bytes can be sent
unchanged. MIME values come from the decoded image rather than its suffix.
Missing dependencies, unsafe pixel counts, and decode/convert failures are
explicit. The stored resource file is never modified.
"""
try:
pixel_limit = int(max_image_pixels)
except (TypeError, ValueError) as exc:
raise ValueError(f"max_image_pixels must be a positive integer: {max_image_pixels!r}") from exc
if pixel_limit <= 0:
raise ValueError(f"max_image_pixels must be a positive integer: {max_image_pixels!r}")
image_module, image_ops = _load_pillow()
if _looks_like_heif(data):
_register_heif_opener()
try:
with warnings.catch_warnings():
warnings.simplefilter("error", image_module.DecompressionBombWarning)
image = image_module.open(io.BytesIO(data))
except (image_module.DecompressionBombWarning, image_module.DecompressionBombError) as exc:
raise RuntimeError(
f"Image rejected by Pillow decompression-bomb protection ({suffix or 'unknown suffix'})",
) from exc
except Exception as exc: # pylint: disable=broad-except
raise RuntimeError(f"Failed to decode image ({suffix or 'unknown suffix'}): {exc}") from exc
with image:
width, height = image.size
pixel_count = width * height
if pixel_count > pixel_limit:
raise RuntimeError(
f"Image exceeds max_image_pixels before decode: " f"{width}x{height}={pixel_count} > {pixel_limit}",
)
source_mime = str(image.get_format_mimetype() or "").strip().lower()
if not source_mime.startswith("image/"):
raise RuntimeError(
f"Cannot determine decoded image MIME type ({suffix or 'unknown suffix'}, format={image.format!r})",
)
needs_resize = width > MAX_IMAGE_REQUEST_DIMENSION or height > MAX_IMAGE_REQUEST_DIMENSION
if source_mime == "image/jpeg" and needs_resize:
max_dimension = max(width, height)
decoder_size = (
max(1, (width * MAX_IMAGE_REQUEST_DIMENSION + max_dimension - 1) // max_dimension),
max(1, (height * MAX_IMAGE_REQUEST_DIMENSION + max_dimension - 1) // max_dimension),
)
try:
# JPEG supports power-of-two decoder scaling. ``draft`` picks
# the smallest decoded frame that still covers decoder_size,
# reducing peak memory before the final LANCZOS thumbnail.
image.draft(None, decoder_size)
except Exception as exc: # pylint: disable=broad-except
raise RuntimeError(f"Failed to prepare JPEG decoder downsampling ({width}x{height}): {exc}") from exc
try:
with warnings.catch_warnings():
warnings.simplefilter("error", image_module.DecompressionBombWarning)
image.load()
orientation = int(image.getexif().get(274, 1) or 1)
image_ops.exif_transpose(image, in_place=True)
except (image_module.DecompressionBombWarning, image_module.DecompressionBombError) as exc:
raise RuntimeError(
f"Image rejected by Pillow decompression-bomb protection ({suffix or 'unknown suffix'})",
) from exc
except Exception as exc: # pylint: disable=broad-except
raise RuntimeError(f"Failed to decode image ({suffix or 'unknown suffix'}): {exc}") from exc
needs_convert = source_mime not in _PASSTHROUGH_IMAGE_MIMES
needs_orientation = orientation in range(2, 9)
if not needs_resize and not needs_convert and not needs_orientation:
return None, source_mime, source_mime
resize_frame = None
try:
frame = image
if needs_resize:
# Pillow forces NEAREST for palette and bilevel images, even
# when LANCZOS is requested. Expand these modes within the
# checked pixel budget so resizing retains fine strokes and
# palette transparency. Other modes resize before conversion.
if image.mode in ("P", "1"):
resize_frame = image.convert("RGBA" if image.mode == "P" else "L")
frame = resize_frame
# Pillow 10 cannot apply LANCZOS directly to 16-bit integer
# modes; NEAREST keeps that path bounded without a full-size
# RGB conversion first.
resize_filter = image_module.Resampling.LANCZOS
if frame.mode.startswith("I;16"):
resize_filter = image_module.Resampling.NEAREST
frame.thumbnail((MAX_IMAGE_REQUEST_DIMENSION, MAX_IMAGE_REQUEST_DIMENSION), resize_filter)
has_alpha = frame.mode in ("RGBA", "LA", "P")
frame = frame.convert("RGBA" if has_alpha else "RGB")
try:
buffer = io.BytesIO()
if frame.mode == "RGBA":
frame.save(buffer, format="PNG")
return buffer.getvalue(), "image/png", source_mime
frame.save(buffer, format="JPEG", quality=_JPEG_QUALITY)
return buffer.getvalue(), "image/jpeg", source_mime
finally:
frame.close()
except Exception as exc: # pylint: disable=broad-except
raise RuntimeError(f"Failed to convert/resize image ({suffix or 'unknown suffix'}): {exc}") from exc
finally:
if resize_frame is not None:
resize_frame.close()
def _build_image_request_payload(
data: bytes,
suffix: str,
*,
max_image_pixels: int = DEFAULT_MAX_IMAGE_PIXELS,
) -> 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`` is the decoded format of the stored resource file and is
what notes record. Both are based on actual bytes, not the filename suffix.
"""
normalized_bytes, mime, source_mime = _normalize_image_bytes(
data,
suffix,
max_image_pixels=max_image_pixels,
)
request_bytes = data if normalized_bytes is None else normalized_bytes
return {
"data_b64": base64.b64encode(request_bytes).decode("ascii"),
"mime": mime,
"source_mime": source_mime,
"converted": normalized_bytes is not None,
}
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)
parsed_json = False
for candidate in (cleaned, cleaned[cleaned.find("{") : cleaned.rfind("}") + 1]):
if not candidate:
continue
try:
parsed = json.loads(candidate)
except (json.JSONDecodeError, ValueError):
continue
parsed_json = True
if isinstance(parsed, dict):
normalized = _normalize_caption_fields(parsed)
if normalized["caption"] or normalized["description"]:
return normalized
if parsed_json:
return {"name": "", "description": "", "caption": ""}
return {"name": "", "description": "", "caption": cleaned.strip()}
def resolve_vision_model(step) -> ChatModelBase | None:
"""Resolve explicit ``as_llm`` through Ref, otherwise prefer vision/default."""
context_model = step.context.get("as_llm") if step.context is not None else None
if "as_llm" in step.kwargs or isinstance(context_model, ChatModelBase):
return step.as_llm
if step.app_context is None:
return None
models = step.app_context.components.get(ComponentEnum.AS_LLM, {})
for name in ("vision", "default"):
if name in models:
step.kwargs["as_llm"] = name
return step.as_llm
return None
async def generate_image_caption(
model: ChatModelBase,
payload: dict,
prompt: str,
*,
logger,
name: str,
) -> dict:
"""Caption a prepared image, preserving the resource model-call contract.
Try structured output first, then one lenient plain-call fallback. The
caller owns model resolution, prompt rendering, and any persistence.
"""
user_message = UserMsg(
name="user",
content=[
TextBlock(text=prompt),
DataBlock(
source=Base64Source(data=payload["data_b64"], media_type=payload["mime"]),
name="image",
),
],
)
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"]:
logger.info(f"[{name}] structured caption ok name={normalized['name']}")
return normalized
logger.warning(f"[{name}] structured caption empty; retrying with a plain call")
except Exception as exc: # pylint: disable=broad-except
logger.warning(f"[{name}] structured caption failed ({exc}); retrying with a plain call")
result = await model([user_message])
parsed = _parse_caption_json(await _response_text(result))
if not parsed["caption"] and not parsed["description"]:
raise RuntimeError("Vision model returned no usable caption")
return parsed

348
reme/steps/evolve/auto_image_resource.py
View file

@ -1,20 +1,285 @@
"""Image resource processor for the unified auto-resource router."""
import base64
import io
import json
import re
import warnings
from pathlib import Path, PurePosixPath
import aiofiles
from agentscope.message import Base64Source, DataBlock, TextBlock, UserMsg
from agentscope.model import ChatModelBase
from pydantic import BaseModel, Field
from ..file_io._path import IMAGE_SUFFIXES
from ._image_caption import (
DEFAULT_MAX_IMAGE_INPUT_BYTES,
DEFAULT_MAX_IMAGE_PIXELS,
_build_image_request_payload,
generate_image_caption,
resolve_vision_model,
)
from .base_auto_resource import _SOURCE_RESOURCE_KEY, _sanitize_note_name, BaseAutoResourceStep
from ...components import R
from ...enumeration import ComponentEnum
DEFAULT_MAX_IMAGE_INPUT_BYTES = 50 * 1024 * 1024
DEFAULT_MAX_IMAGE_PIXELS = 40_000_000
MAX_IMAGE_REQUEST_DIMENSION = 2048
_JPEG_QUALITY = 85
# Decoded formats outside this set are re-encoded to provider-friendly
# PNG/JPEG for VLM requests; the stored resource file is never modified.
_PASSTHROUGH_IMAGE_MIMES = frozenset({"image/png", "image/jpeg", "image/webp", "image/gif"})
_HEIF_BRANDS = frozenset(
{b"heic", b"heif", b"heix", b"heim", b"heis", b"hevc", b"hevx", b"hevm", b"hevs", b"mif1", b"msf1"},
)
_MAX_FTYP_SCAN_BYTES = 4096
_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 based on visible content; filename is only a weak naming hint",
)
description: str = Field(description="one-sentence summary of visible image content that stands on its own")
caption: str = Field(
description="complete description / verbatim transcription of meaningful content visible in the image",
)
def _load_pillow():
"""Load the core image dependency only when image processing runs."""
try:
from PIL import Image, ImageOps # pylint: disable=import-outside-toplevel
except ImportError as exc:
raise RuntimeError("Image captioning requires Pillow; install reme-ai[core]") from exc
return Image, ImageOps
def _looks_like_heif(data: bytes) -> bool:
"""Return whether an ISO-BMFF header declares a HEIC/HEIF brand."""
if len(data) < 12 or data[4:8] != b"ftyp":
return False
box_size = int.from_bytes(data[:4], "big")
if box_size < 12:
return False
end = min(box_size, len(data), _MAX_FTYP_SCAN_BYTES)
if data[8:12] in _HEIF_BRANDS:
return True
return any(data[index : index + 4] in _HEIF_BRANDS for index in range(16, end - 3, 4))
def _register_heif_opener() -> None:
"""Load and register HEIC support only for bytes that declare HEIF."""
try:
from pillow_heif import register_heif_opener # pylint: disable=import-outside-toplevel
except ImportError as exc:
raise RuntimeError("HEIC image captioning requires pillow-heif; install reme-ai[image-heif]") from exc
try:
register_heif_opener()
except Exception as exc: # pylint: disable=broad-except
raise RuntimeError(f"Failed to initialize HEIC image support: {exc}") from exc
def _normalize_image_bytes(
data: bytes,
suffix: str,
*,
max_image_pixels: int = DEFAULT_MAX_IMAGE_PIXELS,
) -> tuple[bytes | None, str, str]:
"""Validate and optionally normalize image bytes for a VLM request.
The returned tuple is ``(normalized_bytes, request_mime, source_mime)``.
``normalized_bytes`` is ``None`` only when the original bytes can be sent
unchanged. MIME values come from the decoded image rather than its suffix.
Missing dependencies, unsafe pixel counts, and decode/convert failures are
explicit. The stored resource file is never modified.
"""
try:
pixel_limit = int(max_image_pixels)
except (TypeError, ValueError) as exc:
raise ValueError(f"max_image_pixels must be a positive integer: {max_image_pixels!r}") from exc
if pixel_limit <= 0:
raise ValueError(f"max_image_pixels must be a positive integer: {max_image_pixels!r}")
image_module, image_ops = _load_pillow()
if _looks_like_heif(data):
_register_heif_opener()
try:
with warnings.catch_warnings():
warnings.simplefilter("error", image_module.DecompressionBombWarning)
image = image_module.open(io.BytesIO(data))
except (image_module.DecompressionBombWarning, image_module.DecompressionBombError) as exc:
raise RuntimeError(
f"Image rejected by Pillow decompression-bomb protection ({suffix or 'unknown suffix'})",
) from exc
except Exception as exc: # pylint: disable=broad-except
raise RuntimeError(f"Failed to decode image ({suffix or 'unknown suffix'}): {exc}") from exc
with image:
width, height = image.size
pixel_count = width * height
if pixel_count > pixel_limit:
raise RuntimeError(
f"Image exceeds max_image_pixels before decode: " f"{width}x{height}={pixel_count} > {pixel_limit}",
)
source_mime = str(image.get_format_mimetype() or "").strip().lower()
if not source_mime.startswith("image/"):
raise RuntimeError(
f"Cannot determine decoded image MIME type ({suffix or 'unknown suffix'}, format={image.format!r})",
)
needs_resize = width > MAX_IMAGE_REQUEST_DIMENSION or height > MAX_IMAGE_REQUEST_DIMENSION
if source_mime == "image/jpeg" and needs_resize:
max_dimension = max(width, height)
decoder_size = (
max(1, (width * MAX_IMAGE_REQUEST_DIMENSION + max_dimension - 1) // max_dimension),
max(1, (height * MAX_IMAGE_REQUEST_DIMENSION + max_dimension - 1) // max_dimension),
)
try:
# JPEG supports power-of-two decoder scaling. ``draft`` picks
# the smallest decoded frame that still covers decoder_size,
# reducing peak memory before the final LANCZOS thumbnail.
image.draft(None, decoder_size)
except Exception as exc: # pylint: disable=broad-except
raise RuntimeError(f"Failed to prepare JPEG decoder downsampling ({width}x{height}): {exc}") from exc
try:
with warnings.catch_warnings():
warnings.simplefilter("error", image_module.DecompressionBombWarning)
image.load()
orientation = int(image.getexif().get(274, 1) or 1)
image_ops.exif_transpose(image, in_place=True)
except (image_module.DecompressionBombWarning, image_module.DecompressionBombError) as exc:
raise RuntimeError(
f"Image rejected by Pillow decompression-bomb protection ({suffix or 'unknown suffix'})",
) from exc
except Exception as exc: # pylint: disable=broad-except
raise RuntimeError(f"Failed to decode image ({suffix or 'unknown suffix'}): {exc}") from exc
needs_convert = source_mime not in _PASSTHROUGH_IMAGE_MIMES
needs_orientation = orientation in range(2, 9)
if not needs_resize and not needs_convert and not needs_orientation:
return None, source_mime, source_mime
resize_frame = None
try:
frame = image
if needs_resize:
# Pillow forces NEAREST for palette and bilevel images, even
# when LANCZOS is requested. Expand these modes within the
# checked pixel budget so resizing retains fine strokes and
# palette transparency. Other modes resize before conversion.
if image.mode in ("P", "1"):
resize_frame = image.convert("RGBA" if image.mode == "P" else "L")
frame = resize_frame
# Pillow 10 cannot apply LANCZOS directly to 16-bit integer
# modes; NEAREST keeps that path bounded without a full-size
# RGB conversion first.
resize_filter = image_module.Resampling.LANCZOS
if frame.mode.startswith("I;16"):
resize_filter = image_module.Resampling.NEAREST
frame.thumbnail((MAX_IMAGE_REQUEST_DIMENSION, MAX_IMAGE_REQUEST_DIMENSION), resize_filter)
has_alpha = frame.mode in ("RGBA", "LA", "P")
frame = frame.convert("RGBA" if has_alpha else "RGB")
try:
buffer = io.BytesIO()
if frame.mode == "RGBA":
frame.save(buffer, format="PNG")
return buffer.getvalue(), "image/png", source_mime
frame.save(buffer, format="JPEG", quality=_JPEG_QUALITY)
return buffer.getvalue(), "image/jpeg", source_mime
finally:
frame.close()
except Exception as exc: # pylint: disable=broad-except
raise RuntimeError(f"Failed to convert/resize image ({suffix or 'unknown suffix'}): {exc}") from exc
finally:
if resize_frame is not None:
resize_frame.close()
def _build_image_request_payload(
data: bytes,
suffix: str,
*,
max_image_pixels: int = DEFAULT_MAX_IMAGE_PIXELS,
) -> 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`` is the decoded format of the stored resource file and is
what notes record. Both are based on actual bytes, not the filename suffix.
"""
normalized_bytes, mime, source_mime = _normalize_image_bytes(
data,
suffix,
max_image_pixels=max_image_pixels,
)
request_bytes = data if normalized_bytes is None else normalized_bytes
return {
"data_b64": base64.b64encode(request_bytes).decode("ascii"),
"mime": mime,
"source_mime": source_mime,
"converted": normalized_bytes is not None,
}
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)
parsed_json = False
for candidate in (cleaned, cleaned[cleaned.find("{") : cleaned.rfind("}") + 1]):
if not candidate:
continue
try:
parsed = json.loads(candidate)
except (json.JSONDecodeError, ValueError):
continue
parsed_json = True
if isinstance(parsed, dict):
normalized = _normalize_caption_fields(parsed)
if normalized["caption"] or normalized["description"]:
return normalized
if parsed_json:
return {"name": "", "description": "", "caption": ""}
return {"name": "", "description": "", "caption": cleaned.strip()}
@R.register("auto_image_resource_step")
@ -54,7 +319,44 @@ class AutoImageResourceStep(BaseAutoResourceStep):
def _vision_model(self) -> ChatModelBase | None:
"""Resolve explicit ``as_llm`` through Ref, otherwise prefer vision/default."""
return resolve_vision_model(self)
context_model = self.context.get("as_llm") if self.context is not None else None
if "as_llm" in self.kwargs or isinstance(context_model, ChatModelBase):
return self.as_llm
if self.app_context is None:
return None
models = self.app_context.components.get(ComponentEnum.AS_LLM, {})
for name in ("vision", "default"):
if name in models:
self.kwargs["as_llm"] = name
return self.as_llm
return 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])
parsed = _parse_caption_json(await _response_text(result))
if not parsed["caption"] and not parsed["description"]:
raise RuntimeError("Vision model returned no usable caption")
return parsed
async def _read_image(self, file_path: str, source_path: Path) -> dict | None:
"""Read the image file and build the VLM request payload.
@ -156,19 +458,25 @@ class AutoImageResourceStep(BaseAutoResourceStep):
if payload is None:
return
parsed = await generate_image_caption(
model,
payload,
self.prompt_format(
"user_message",
file_path=file_path,
filename=PurePosixPath(file_path).name,
stem=note_stem,
date=date_str,
),
logger=self.logger,
name=self.name,
user_message = UserMsg(
name="user",
content=[
TextBlock(
text=self.prompt_format(
"user_message",
file_path=file_path,
filename=PurePosixPath(file_path).name,
stem=note_stem,
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]

164
tests/unit/test_as_agent_wrapper_model_override.py
View file

@ -1,164 +0,0 @@
"""Internal model overrides stay local to one AgentScope invocation."""
# pylint: disable=protected-access,missing-function-docstring
import asyncio
from copy import deepcopy
from types import SimpleNamespace
from unittest.mock import AsyncMock
from agentscope.agent import Agent
from agentscope.message import Msg
from agentscope.model import ChatModelBase
import pytest
import pytest_asyncio
from reme.components.agent_wrapper.as_agent_wrapper import AsAgentWrapper
from reme.components.application_context import ApplicationContext
from reme.components.as_llm import BaseAsLLM
from reme.enumeration import ComponentEnum
class _Model(ChatModelBase):
"""Local model double with no provider client or network calls."""
def __init__(self, name):
self.model = name
self.structured_calls = []
async def generate_structured_output(self, **kwargs):
self.structured_calls.append(kwargs)
return SimpleNamespace(
content={"model": self.model},
usage=SimpleNamespace(input_tokens=2, output_tokens=1),
)
@pytest_asyncio.fixture(name="models")
async def model_fixture(tmp_path):
context = ApplicationContext(workspace_dir=str(tmp_path))
default = BaseAsLLM(name="default")
default.model = _Model("default")
vision = BaseAsLLM(name="vision")
vision.model = _Model("vision")
context.components = {ComponentEnum.AS_LLM: {"default": default, "vision": vision}}
wrapper = AsAgentWrapper(
app_context=context,
system_prompt="Configured prompt",
model_config={"max_retries": 1},
)
await default.start()
await vision.start()
await wrapper.start()
yield SimpleNamespace(wrapper=wrapper, default=default, vision=vision, context=context)
await wrapper.close()
await vision.close()
await default.close()
@pytest.mark.asyncio
@pytest.mark.parametrize("override", [False, True, None])
async def test_build_agent_binds_call_model_without_replacing_default(models, override):
wrapper = models.wrapper
before_kwargs = deepcopy(wrapper.kwargs)
before_bindings = wrapper.dependency_bindings
kwargs = {} if override is False else {"_model": models.vision.model if override else None}
agent, inputs = await wrapper._build_agent("hello", **kwargs)
assert isinstance(agent, Agent)
assert isinstance(inputs, Msg)
assert inputs.role == "user"
assert agent.model is (models.vision.model if override else models.default.model)
ordinary_agent, _ = await wrapper._build_agent("ordinary")
assert ordinary_agent.model is models.default.model
assert wrapper.as_llm is models.default
assert wrapper.kwargs == before_kwargs
assert wrapper.dependency_bindings == before_bindings
assert wrapper._owned == []
assert models.context.components == {ComponentEnum.AS_LLM: {"default": models.default, "vision": models.vision}}
assert kwargs == ({} if override is False else {"_model": models.vision.model if override else None})
@pytest.mark.asyncio
@pytest.mark.parametrize("override", ["vision", {}, object()])
async def test_invalid_internal_model_override_fails_before_agent_construction(models, override):
with pytest.raises(TypeError, match="ChatModelBase"):
await models.wrapper._build_agent("hello", _model=override)
assert models.wrapper.as_llm is models.default
@pytest.mark.asyncio
async def test_concurrent_replies_keep_agent_and_structured_output_on_each_call_model(models, monkeypatch):
wrapper = models.wrapper
barrier = asyncio.Barrier(3)
observed = {}
before_kwargs = deepcopy(wrapper.kwargs)
before_bindings = wrapper.dependency_bindings
async def observe(agent, inputs):
agent.state.context = [inputs]
async def reply(agent):
label = agent.state.context[0].get_text_content()
observed[label] = agent.model
await asyncio.wait_for(barrier.wait(), timeout=5)
assert wrapper.as_llm is models.default
return SimpleNamespace(
usage=SimpleNamespace(input_tokens=3, output_tokens=1),
model_dump=lambda: {"text": label},
get_text_content=lambda: label,
)
monkeypatch.setattr(Agent, "observe", observe)
monkeypatch.setattr(Agent, "reply", reply)
monkeypatch.setattr(wrapper, "_dump_state", AsyncMock())
schema = {"type": "object"}
kwargs = {"_model": models.vision.model, "output_schema": schema}
vision_result, ordinary_result, second_vision_result = await asyncio.gather(
wrapper.reply("vision-first", **kwargs),
wrapper.reply("ordinary", output_schema=schema),
wrapper.reply("vision-second", **kwargs),
)
assert observed == {
"vision-first": models.vision.model,
"ordinary": models.default.model,
"vision-second": models.vision.model,
}
for result, expected in ((vision_result, "vision"), (ordinary_result, "default"), (second_vision_result, "vision")):
assert result["structured_output"] == {"model": expected}
assert result["usage"] == {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}
assert [call["messages"][0].get_text_content() for call in models.default.model.structured_calls] == ["ordinary"]
assert {call["messages"][0].get_text_content() for call in models.vision.model.structured_calls} == {
"vision-first",
"vision-second",
}
assert wrapper.as_llm is models.default
assert wrapper.kwargs == before_kwargs
assert wrapper.dependency_bindings == before_bindings
assert models.context.components == {ComponentEnum.AS_LLM: {"default": models.default, "vision": models.vision}}
assert kwargs == {"_model": models.vision.model, "output_schema": {"type": "object"}}
assert models.default.is_started and models.vision.is_started
await wrapper.close()
assert models.default.is_started and models.vision.is_started
@pytest.mark.asyncio
@pytest.mark.parametrize("override", [False, True])
async def test_stream_builds_with_the_call_model_without_changing_default(models, monkeypatch, override):
observed = []
async def reply_stream(agent, inputs):
observed.append((agent.model, inputs.get_text_content()))
if inputs is None:
yield None
monkeypatch.setattr(Agent, "reply_stream", reply_stream)
monkeypatch.setattr(models.wrapper, "_dump_state", AsyncMock())
kwargs = {"_model": models.vision.model} if override else {}
assert [chunk async for chunk in models.wrapper.reply_stream("hello", **kwargs)] == []
assert observed == [(models.vision.model if override else models.default.model, "hello")]
assert models.wrapper.as_llm is models.default

45
tests/unit/test_auto_image_steps.py
View file

@ -9,7 +9,6 @@ with PIL inside a temporary workspace.
import base64
import hashlib
import io
import logging
import struct
import subprocess
import sys
@ -29,14 +28,13 @@ from reme.components.component_registry import ComponentRegistry
from reme.components.job import BaseJob
from reme.components.runtime_context import RuntimeContext
from reme.enumeration import ComponentEnum
from reme.steps.evolve._image_caption import (
from reme.steps.evolve.auto_image_resource import (
AutoImageResourceStep,
DEFAULT_MAX_IMAGE_PIXELS,
_build_image_request_payload,
generate_image_caption,
_normalize_image_bytes,
_parse_caption_json,
)
from reme.steps.evolve.auto_image_resource import AutoImageResourceStep
from reme.steps.evolve.auto_resource import AutoResourceStep
from reme.steps.evolve.auto_text_resource import AutoTextResourceStep
from .auto_resource_test_support import (
@ -1031,45 +1029,6 @@ def test_parse_caption_json(text, expected):
assert _parse_caption_json(text) == expected
@pytest.mark.parametrize("plain_fallback", [False, True])
@pytest.mark.asyncio
async def test_shared_caption_helper_preserves_request_without_workspace(plain_fallback, caplog):
"""The shared VLM path needs only bytes, a model and an already rendered prompt."""
expected = {"name": "image", "description": "Visible text", "caption": "Invoice 482."}
model = _StructuredVisionModel(
content=expected,
error=RuntimeError("structured output unavailable") if plain_fallback else None,
plain_text="```json\n" + _caption_json(**expected) + "\n```",
)
original = _png_bytes()
payload = _build_image_request_payload(original, ".png")
prompt = "Describe visible facts. Keep invoice numbers verbatim."
result = await generate_image_caption(
model,
payload,
prompt,
logger=logging.getLogger(__name__),
name="shared-caption-test",
)
assert result == expected
assert len(model.structured_calls) == 1
assert len(model.plain_calls) == int(plain_fallback)
message = model.structured_calls[0][0]
assert message.content[0].text == prompt
assert message.content[1].source.media_type == "image/png"
assert base64.b64decode(message.content[1].source.data) == original
if plain_fallback:
assert model.plain_calls[0][0] is message
assert caplog.messages == [
"[shared-caption-test] structured caption failed (structured output unavailable); "
"retrying with a plain call",
]
else:
assert not caplog.messages
@pytest.mark.parametrize(
("structured_mode", "note_name", "caption", "expected_plain_calls"),
[

164
tests/unit/test_auto_memory_direct_images.py
View file

@ -1,4 +1,4 @@
"""Direct image inputs select a per-call AgentScope model without caption calls."""
"""Direct image inputs preserve native model binding and session contracts."""
# pylint: disable=protected-access,missing-function-docstring
@ -7,11 +7,9 @@ import base64
import copy
import io
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
from agentscope.agent import Agent, ContextConfig
from agentscope.formatter import OpenAIChatFormatter
from agentscope.message import Base64Source, Msg, TextBlock, URLSource
from PIL import Image
import pytest
@ -20,7 +18,7 @@ import yaml
from reme.components.as_llm import BaseAsLLM
from reme.components.runtime_context import RuntimeContext
from reme.enumeration import ComponentEnum
from reme.steps.evolve import _auto_memory_image, _image_caption
from reme.steps.evolve import _auto_memory_image, auto_image_resource
from .test_auto_memory_image_modes import (
_DAY,
@ -63,7 +61,7 @@ def test_default_config_keeps_images_disabled_and_selects_direct_mode():
@pytest.mark.asyncio
async def test_enabled_default_is_direct_with_a_single_memory_reply(harness, monkeypatch):
caption = AsyncMock(side_effect=AssertionError("Direct mode must not caption images"))
monkeypatch.setattr(_image_caption, "generate_image_caption", caption)
monkeypatch.setattr(auto_image_resource.AutoImageResourceStep, "_caption_with_retry", caption)
message = _message()
message.content[0].text += " Keep literal [Image 1]; image ID: gallery-before."
message.content.insert(2, TextBlock(text="Between observations; image ID: gallery-after."))
@ -102,7 +100,7 @@ async def test_enabled_default_is_direct_with_a_single_memory_reply(harness, mon
assert harness.session_path.read_bytes() == _main_saved_line(message)
assert not (harness.workspace / "resource").exists()
caption.assert_not_called()
assert kwargs["_model"] is harness.wrapper.as_llm.model
assert "_model" not in kwargs
assert "supports_vision" not in kwargs
formatted = await harness.wrapper.as_llm.model.formatter.format([inputs])
@ -128,9 +126,6 @@ async def test_enabled_default_is_direct_with_a_single_memory_reply(harness, mon
async def test_disabled_keeps_original_string_and_does_not_prepare(harness, monkeypatch, include_images):
prepare = AsyncMock(side_effect=AssertionError("Disabled images must not be inspected"))
monkeypatch.setattr("reme.steps.evolve.auto_memory.prepare_direct_message", prepare)
original_component = harness.wrapper.as_llm
original_model = original_component.model
harness.app_context.components[ComponentEnum.AS_LLM]["vision"] = SimpleNamespace(model=_DirectModel("vision"))
step = harness.step()
kwargs = {} if include_images is None else {"include_images": include_images}
message = _message()
@ -143,16 +138,12 @@ async def test_disabled_keeps_original_string_and_does_not_prepare(harness, monk
assert "[Image" not in harness.wrapper.calls[0][0]
assert "auto_memory_images" not in step.context.response.metadata
assert "_model" not in harness.wrapper.calls[0][1]
assert harness.wrapper.as_llm is original_component
assert original_component.model is original_model
assert harness.session_path.read_bytes() == _main_saved_line(message)
prepare.assert_not_called()
@pytest.mark.asyncio
async def test_no_images_skips_model_selection_and_source_checks(harness, monkeypatch):
original_model = harness.wrapper.as_llm.model
harness.app_context.components[ComponentEnum.AS_LLM]["vision"] = None
async def test_no_images_skips_source_checks(harness, monkeypatch):
loader = AsyncMock(side_effect=AssertionError("No image should be loaded"))
monkeypatch.setattr(_auto_memory_image, "_image_bytes", loader)
@ -163,7 +154,6 @@ async def test_no_images_skips_model_selection_and_source_checks(harness, monkey
assert response.metadata["auto_memory_images"]["status"] == "skipped"
assert response.metadata["auto_memory_images"]["image_count"] == 0
assert "_model" not in harness.wrapper.calls[0][1]
assert harness.wrapper.as_llm.model is original_model
harness.logger.warning.assert_not_called()
loader.assert_not_called()
@ -203,7 +193,7 @@ async def test_non_agentscope_wrapper_still_accepts_original_text_only_calls(har
@pytest.mark.asyncio
@pytest.mark.parametrize("declared,images", [(None, True), (False, True), (None, False)])
async def test_missing_vision_declaration_does_not_read_images_or_choose_a_model(
async def test_missing_vision_declaration_does_not_read_images(
harness,
monkeypatch,
declared,
@ -213,8 +203,6 @@ async def test_missing_vision_declaration_does_not_read_images_or_choose_a_model
monkeypatch.setattr(_auto_memory_image, "_image_bytes", loader)
step = harness.step()
step.kwargs.pop("supports_vision")
harness.wrapper.as_llm.model = None
harness.app_context.components[ComponentEnum.AS_LLM]["vision"] = None
message = _message(images=images)
if images:
message.content[1].source.data = "invalid-base64"
@ -245,117 +233,28 @@ async def test_missing_vision_declaration_does_not_read_images_or_choose_a_model
@pytest.mark.asyncio
async def test_model_catalog_and_formatter_declarations_do_not_gate_direct_inputs(harness, monkeypatch):
model = _DirectModel("custom-deployment")
model.formatter = OpenAIChatFormatter(input_types=["text/plain"])
catalog = Mock(side_effect=AssertionError("Model capabilities belong to the caller"))
monkeypatch.setattr(model, "list_models", catalog)
harness.wrapper.as_llm.model = model
response = await _run(harness, [_message()])
assert response.metadata["auto_memory_images"]["status"] == "completed"
assert isinstance(harness.wrapper.calls[0][0], Msg)
assert harness.wrapper.calls[0][1]["_model"] is model
catalog.assert_not_called()
harness.logger.warning.assert_not_called()
@pytest.mark.asyncio
@pytest.mark.parametrize("wrapper_binding", ["default", "custom"])
async def test_dedicated_vision_model_is_preferred_without_rebinding_shared_components(harness, wrapper_binding):
@pytest.mark.parametrize("include_images", [False, True])
async def test_memory_uses_bound_as_llm_even_when_another_vision_component_exists(harness, include_images):
models = harness.app_context.components[ComponentEnum.AS_LLM]
default_component = models["default"]
custom_component = BaseAsLLM()
custom_component.model = _DirectModel("custom-memory")
vision_component = BaseAsLLM()
vision_component.model = _DirectModel("dedicated-vision")
models.update(custom=custom_component, vision=vision_component)
harness.wrapper.as_llm = models[wrapper_binding]
original_components = dict(models)
original_models = {name: component.model for name, component in models.items()}
wrapper_options = {"context_config": {"trigger_ratio": 0.7}, "model_config": {"max_retries": 2}}
harness.wrapper.kwargs.update(wrapper_options)
response = await _run(harness, [_message()])
assert response.metadata["auto_memory_images"]["status"] == "completed"
assert harness.wrapper.calls[0][1]["_model"] is vision_component.model
assert harness.wrapper.as_llm is original_components[wrapper_binding]
assert all(models[name] is component for name, component in original_components.items())
assert all(models[name].model is model for name, model in original_models.items())
assert harness.wrapper.kwargs["context_config"] is wrapper_options["context_config"]
assert harness.wrapper.kwargs["model_config"] is wrapper_options["model_config"]
assert wrapper_options == {"context_config": {"trigger_ratio": 0.7}, "model_config": {"max_retries": 2}}
assert default_component.model is original_models["default"]
@pytest.mark.asyncio
@pytest.mark.parametrize("application_context", [True, False])
async def test_absent_vision_uses_actual_custom_wrapper_binding(harness, application_context):
default_component = harness.app_context.components[ComponentEnum.AS_LLM]["default"]
custom_component = BaseAsLLM()
custom_component.model = _DirectModel("custom-memory")
harness.wrapper.as_llm = custom_component
bound = BaseAsLLM()
bound.model = _DirectModel("bound-memory")
vision = BaseAsLLM()
vision.model = _DirectModel("unused-resource-vision")
models.update(custom=bound, vision=vision)
harness.wrapper.as_llm = bound
step = harness.step()
if not application_context:
step.app_context = None
await step(session_id=_SESSION, date=_DAY, messages=[_message()], include_images=True)
await step(session_id=_SESSION, date=_DAY, messages=[_message()], include_images=include_images)
assert step.context.response.metadata["auto_memory_images"]["status"] == "completed"
assert harness.wrapper.calls[0][1]["_model"] is custom_component.model
assert harness.wrapper.calls[0][1]["_model"] is not default_component.model
assert harness.wrapper.as_llm is custom_component
@pytest.mark.asyncio
@pytest.mark.parametrize("unavailable", ["bound-model", "vision-model", "vision-component"])
async def test_selected_model_must_be_started_without_silent_model_substitution(harness, monkeypatch, unavailable):
if unavailable == "bound-model":
harness.wrapper.as_llm.model = None
elif unavailable == "vision-model":
harness.app_context.components[ComponentEnum.AS_LLM]["vision"] = BaseAsLLM()
else:
harness.app_context.components[ComponentEnum.AS_LLM]["vision"] = None
original_component = harness.wrapper.as_llm
original_model = original_component.model
loader = AsyncMock(side_effect=AssertionError("Missing model must be rejected before image IO"))
monkeypatch.setattr(_auto_memory_image, "_image_bytes", loader)
response = await _run(harness, [_message()])
assert response.metadata["auto_memory_images"]["status"] == "fallback"
assert response.metadata["auto_memory_images"]["reason"] == "model: ValueError"
assert isinstance(harness.wrapper.calls[0][0], str)
assert "_model" not in harness.wrapper.calls[0][1]
assert "context_config" not in harness.wrapper.calls[0][1]
assert harness.wrapper.as_llm is original_component
assert original_component.model is original_model
loader.assert_not_called()
@pytest.mark.asyncio
async def test_reply_hook_model_options_are_read_once_and_passed_unchanged(harness, monkeypatch):
fallback = _DirectModel()
fallback.formatter = OpenAIChatFormatter(input_types=["text/plain"])
harness.wrapper.kwargs["model_config"] = {"max_retries": 2}
model_config = {"fallback_model": fallback}
options = {"model_config": model_config}
step = harness.step()
hook = Mock(return_value=options)
monkeypatch.setattr(step, "_reply_extra_kwargs", hook)
await step(session_id=_SESSION, date=_DAY, messages=[_message()], include_images=True)
assert step.context.response.metadata["auto_memory_images"]["status"] == "completed"
assert isinstance(harness.wrapper.calls[0][0], Msg)
assert harness.wrapper.calls[0][1]["model_config"] is model_config
assert harness.wrapper.calls[0][1]["_model"] is harness.wrapper.as_llm.model
assert harness.wrapper.kwargs["model_config"] == {"max_retries": 2}
assert model_config["fallback_model"] is fallback
assert options == {"model_config": model_config}
hook.assert_called_once_with(_DAY)
inputs, kwargs = harness.wrapper.calls[0]
assert isinstance(inputs, Msg if include_images else str)
assert "_model" not in kwargs and "supports_vision" not in kwargs
# Exercise native Agent construction without resolving the test's file tools.
agent, _ = await harness.wrapper._build_agent(inputs, **{**kwargs, "job_tools": []})
assert agent.model is bound.model
assert harness.wrapper.as_llm is bound
assert models["default"] is not bound
assert models["vision"] is vision
def _six_image_message():
@ -518,10 +417,8 @@ async def test_interleaving_keeps_localized_create_update_prompt_boundaries(harn
@pytest.mark.parametrize("template", ["No history requested.", "{history}\n{history}"])
async def test_missing_or_repeated_history_falls_back_without_direct_overrides(harness, monkeypatch, template):
step = harness.step(prompt_dict={"user_message_create": template})
original_model = harness.wrapper.as_llm.model
original_options = {"context_config": {"trigger_ratio": 0.7}, "model_config": {"max_retries": 2}}
monkeypatch.setattr(step, "_reply_extra_kwargs", Mock(return_value=original_options))
harness.app_context.components[ComponentEnum.AS_LLM]["vision"] = SimpleNamespace(model=_DirectModel("vision"))
message = _message()
await step(session_id=_SESSION, date=_DAY, messages=[message], include_images=True)
@ -536,7 +433,6 @@ async def test_missing_or_repeated_history_falls_back_without_direct_overrides(h
assert "_model" not in options
assert options["context_config"] is original_options["context_config"]
assert original_options == {"context_config": {"trigger_ratio": 0.7}, "model_config": {"max_retries": 2}}
assert harness.wrapper.as_llm.model is original_model
harness.logger.warning.assert_called_once()
assert harness.session_path.read_bytes() == _main_saved_line(message)
@ -570,7 +466,7 @@ async def test_direct_uses_shared_normalized_provider_bytes(harness):
before = message.model_dump()
step = harness.step()
step.context = RuntimeContext()
expected = _image_caption._build_image_request_payload(original_data, "")
expected = auto_image_resource._build_image_request_payload(original_data, "")
direct_message = await _auto_memory_image.prepare_direct_message(step, [message], step._format_history)
@ -578,7 +474,7 @@ async def test_direct_uses_shared_normalized_provider_bytes(harness):
assert image_input.media_type == expected["mime"] == "image/jpeg"
assert image_input.data == expected["data_b64"]
with Image.open(io.BytesIO(base64.b64decode(image_input.data))) as image:
assert max(image.size) == _image_caption.MAX_IMAGE_REQUEST_DIMENSION
assert max(image.size) == auto_image_resource.MAX_IMAGE_REQUEST_DIMENSION
assert message.model_dump() == before
@ -607,10 +503,6 @@ async def test_second_image_failure_discards_all_prepared_inputs(harness, failur
message = _six_image_message()
source_config = {"trigger_ratio": 0.7}
harness.wrapper.kwargs["context_config"] = source_config
original_component = harness.wrapper.as_llm
original_model = original_component.model
vision_model = _DirectModel("dedicated-vision")
harness.app_context.components[ComponentEnum.AS_LLM]["vision"] = SimpleNamespace(model=vision_model)
kwargs = {}
if failure == "invalid-base64":
message.content[2].source.data = "INVALID_BASE64!!!"
@ -637,11 +529,7 @@ async def test_second_image_failure_discards_all_prepared_inputs(harness, failur
assert "Remember this observation." in harness.wrapper.calls[0][0]
assert response.metadata["auto_memory_images"]["captioned_images"] == 0
assert "context_config" not in harness.wrapper.calls[0][1]
assert "_model" not in harness.wrapper.calls[0][1]
assert source_config == {"trigger_ratio": 0.7}
assert harness.wrapper.as_llm is original_component
assert original_component.model is original_model
assert harness.app_context.components[ComponentEnum.AS_LLM]["vision"].model is vision_model
harness.logger.warning.assert_called_once()
assert message.model_dump() == before
assert harness.session_path.read_bytes() == _main_saved_line(message)

6
tests/unit/test_auto_memory_image_modes.py
View file

@ -112,17 +112,13 @@ class _MemoryWrapper(BaseAgentWrapper):
class _DirectModel(ChatModelBase):
"""Local model with no API client or model-catalog requirement."""
"""Local model with no API client."""
def __init__(self, name="custom-memory-model"):
self.model = name
self.context_size = 200000
self.formatter = OpenAIChatFormatter()
def list_models(self, custom_yaml_dir=None):
del custom_yaml_dir
raise AssertionError("Direct image routing must not inspect model catalogs")
class _DirectWrapper(AsAgentWrapper):
"""Keep a real wrapper type and binding while observing the reply boundary."""