opensource/AstrBot
3
0
Fork 0
mirror of https://github.com/AstrBotDevs/AstrBot.git synced 2026-10-16 06:58:51 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: centralize model image preparation and attachment notices

Browse Source
This commit is contained in:
Soulter
2026-09-17 23:52:14 +08:00
parent 35b404c495
commit a793318f50
12 changed files with 1030 additions and 270 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrbot/core/astr_main_agent.py
+175 -29
View File
@@ -107,8 +107,13 @@ from astrbot.core.utils.astrbot_path import (
get_astrbot_workspaces_path,
)
from astrbot.core.utils.file_extract import extract_file_moonshotai
from astrbot.core.utils.image_input import prepare_request_images
from astrbot.core.utils.llm_metadata import LLM_METADATAS
from astrbot.core.utils.media_utils import is_file_uri, is_recoverable_image_error
from astrbot.core.utils.media_utils import (
is_file_uri,
is_recoverable_image_error,
normalize_model_image_max_size,
)
from astrbot.core.utils.quoted_message.settings import (
SETTINGS as DEFAULT_QUOTED_MESSAGE_SETTINGS,
)
@@ -126,6 +131,14 @@ from astrbot.core.workspace import (
)
LLM_ERROR_MESSAGE_EXTRA_KEY = "_llm_error_message"
ANIMATION_CAPTION_NOTICE = (
"\n<system_notice>\n"
"Input images at positions {indices} (1-based) are animations (e.g. GIFs), "
"each converted to a single image of frames in reading order. "
"Describe them as animations, including motion or changes; "
"do not mention the conversion or frame layout.\n"
"</system_notice>"
)
WEEKDAY_NAMES = (
"Monday",
"Tuesday",
@@ -692,7 +705,23 @@ async def _request_img_caption(
cfg: dict,
image_urls: list[str],
plugin_context: Context,
montage_refs: set[str] | None = None,
) -> str:
"""Describe prepared images, preserving animation semantics.
Args:
provider_id: Caption provider ID.
cfg: Provider settings containing the caption prompt.
image_urls: Prepared images in their input order.
plugin_context: Services used to resolve the caption provider.
montage_refs: Prepared paths containing animation frame montages.
Returns:
The generated image description.
Raises:
ValueError: If the caption provider is missing or invalid.
"""
prov = plugin_context.get_provider_by_id(provider_id)
if prov is None:
raise ValueError(
@@ -707,6 +736,15 @@ async def _request_img_caption(
"image_caption_prompt",
"Please describe the image.",
)
montage_indices = [
str(index)
for index, ref in enumerate(image_urls, start=1)
if ref in (montage_refs or ())
]
if montage_indices:
img_cap_prompt += ANIMATION_CAPTION_NOTICE.format(
indices=", ".join(montage_indices)
)
logger.debug("Processing image caption with provider: %s", provider_id)
llm_resp = await prov.text_chat(
prompt=img_cap_prompt,
@@ -721,30 +759,41 @@ async def _ensure_img_caption(
cfg: dict,
plugin_context: Context,
image_caption_provider: str,
) -> None:
montage_refs: set[str] | None = None,
) -> set[str]:
"""Append a caption and return the references actually described.
Args:
event: Current event.
req: Request whose image_urls will be replaced with a text description.
cfg: Provider settings for caption generation.
plugin_context: Services used to resolve the caption provider.
image_caption_provider: Configured caption provider ID.
montage_refs: Prepared paths containing animation frame montages.
Returns:
Successfully described image references, or an empty set on failure.
"""
image_refs = set(req.image_urls)
try:
caption = await _request_img_caption(
image_caption_provider,
cfg,
req.image_urls,
plugin_context,
montage_refs=montage_refs,
)
if caption:
req.extra_user_content_parts.append(
TextPart(text=f"<image_caption>{caption}</image_caption>")
)
req.image_urls = []
return image_refs
except Exception as exc: # noqa: BLE001
logger.error("处理图片描述失败: %s", exc)
req.extra_user_content_parts.append(TextPart(text="[Image Captioning Failed]"))
finally:
req.image_urls = []
def _append_quoted_image_attachment(req: ProviderRequest, image_path: str) -> None:
req.extra_user_content_parts.append(
TextPart(text=f"[Image Attachment in quoted message: path {image_path}]")
)
return set()
def _append_audio_attachment(req: ProviderRequest, audio_path: str) -> None:
@@ -819,7 +868,8 @@ async def _process_quote_message(
main_provider_supports_image: bool = False,
skip_quote_image_caption: bool = False,
image_ref: str | None = None,
) -> None:
image_is_montage: bool = False,
) -> str | None:
"""Append quoted text and optionally describe an explicitly supplied image.
Args:
@@ -830,8 +880,11 @@ async def _process_quote_message(
quoted_message_settings: Existing quote extraction limits and options.
main_provider_supports_image: Whether the main model can see images.
skip_quote_image_caption: Whether the main caption branch handles it.
image_ref: Already collected image reference; the local stage supplies
a prepared path. No image is re-read from the event here.
image_ref: Prepared image path supplied by the builder.
image_is_montage: Whether the prepared image represents an animation.
Returns:
The successfully described image reference, or None if no caption was produced.
"""
quote = None
for comp in event.message_obj.message:
@@ -841,6 +894,7 @@ async def _process_quote_message(
if not quote:
return
captioned_ref = None
content_parts = []
sender_info = f"({quote.sender_nickname}): " if quote.sender_nickname else ""
message_str = (
@@ -878,11 +932,15 @@ async def _process_quote_message(
)
if prov and isinstance(prov, Provider):
caption_prompt = "Please describe the image content."
if image_is_montage:
caption_prompt += ANIMATION_CAPTION_NOTICE.format(indices="1")
llm_resp = await prov.text_chat(
prompt="Please describe the image content.",
prompt=caption_prompt,
image_urls=[image_ref],
)
if llm_resp.completion_text:
captioned_ref = image_ref
content_parts.append(
f"[Image Caption in quoted message]: {llm_resp.completion_text}"
)
@@ -894,6 +952,7 @@ async def _process_quote_message(
quoted_content = "\n".join(content_parts)
quoted_text = f"<Quoted Message>\n{quoted_content}\n</Quoted Message>"
req.extra_user_content_parts.append(TextPart(text=quoted_text))
return captioned_ref
def _append_system_reminders(
@@ -945,7 +1004,21 @@ async def _decorate_llm_request(
plugin_context: Context,
config: MainAgentBuildConfig,
provider: Provider | None = None,
) -> None:
montage_refs: set[str] | None = None,
) -> set[str]:
"""Decorate the request and return image references successfully captioned.
Args:
event: Current event.
req: Working request receiving prompts and any image description.
plugin_context: Services used to resolve configuration and providers.
config: Main-agent build configuration.
provider: Main chat provider, used to determine vision support.
montage_refs: Prepared paths containing animation frame montages.
Returns:
References successfully described by the image-captioning provider.
"""
cfg = config.provider_settings or plugin_context.get_config(
umo=event.unified_msg_origin
).get("provider_settings", {})
@@ -959,14 +1032,16 @@ async def _decorate_llm_request(
await _ensure_persona_and_skills(req, cfg, plugin_context, event)
captioned_refs = set()
if req.conversation:
if img_cap_prov_id and req.image_urls and not main_provider_supports_image:
await _ensure_img_caption(
captioned_refs = await _ensure_img_caption(
event,
req,
cfg,
plugin_context,
img_cap_prov_id,
montage_refs=montage_refs,
)
tz = config.timezone
@@ -974,6 +1049,7 @@ async def _decorate_llm_request(
tz = plugin_context.get_config().get("timezone")
_append_system_reminders(event, req, cfg, tz)
await _apply_workspace_extra_prompt(event, req, plugin_context)
return captioned_refs
def _plugin_tool_fix(event: AstrMessageEvent, req: ProviderRequest) -> None:
@@ -1327,6 +1403,8 @@ async def collect_initial_request(
plugin_context: Context,
config: MainAgentBuildConfig,
req: ProviderRequest | None = None,
*,
quoted_image_refs: set[str] | None = None,
) -> tuple[ProviderRequest | None, str | None]:
"""Collect raw attachments without applying model image policy.
@@ -1336,11 +1414,14 @@ async def collect_initial_request(
config: Existing request collection settings.
req: Explicit request for direct callers; its attachments are already
collected, so only its quote reference needs to be resolved.
quoted_image_refs: Optional output set populated from quoted image components.
Returns:
The initial request and the first embedded quote image reference used by
the dedicated quote caption branch. A rejected wake prefix returns None.
"""
if quoted_image_refs is None:
quoted_image_refs = set()
attachment_paths: list[str] = []
if req is None:
if event.get_extra("provider_request"):
@@ -1399,9 +1480,6 @@ async def collect_initial_request(
source_is_local = False
if not source_is_local and Path(image_path).is_file():
event.track_temporary_local_file(image_path)
req.extra_user_content_parts.append(
TextPart(text=f"[Image Attachment: path {image_path}]")
)
elif isinstance(comp, Record):
audio_path = await comp.convert_to_file_path()
req.audio_urls.append(audio_path)
@@ -1455,7 +1533,7 @@ async def collect_initial_request(
source_is_local = False
if not source_is_local and Path(image_path).is_file():
event.track_temporary_local_file(image_path)
_append_quoted_image_attachment(req, image_path)
quoted_image_refs.add(image_path)
elif isinstance(reply_comp, Record):
audio_path = await reply_comp.convert_to_file_path()
req.audio_urls.append(audio_path)
@@ -1511,7 +1589,7 @@ async def collect_initial_request(
continue
req.image_urls.append(image_ref)
fallback_quoted_image_count += 1
_append_quoted_image_attachment(req, image_ref)
quoted_image_refs.add(image_ref)
except Exception as exc: # noqa: BLE001
logger.warning(
"Failed to resolve fallback quoted images for umo=%s, reply_id=%s: %s",
@@ -1550,10 +1628,22 @@ async def build_main_agent(
provider: Provider | None = None,
req: ProviderRequest | None = None,
apply_reset: bool = True,
prepared_images: dict[str, dict] | None = None,
) -> MainAgentBuildResult | None:
"""构建主对话代理（Main Agent），并且自动 reset。
"""Collect and prepare input, configure the main agent, and optionally reset it.
If apply_reset is False, will not call reset on the agent runner.
Args:
event: Incoming platform event.
plugin_context: Services used to build the request and agent.
config: Main-agent configuration for this session.
provider: Selected provider, or None to select one for the event.
req: Existing request, or None to collect one from the event.
apply_reset: Whether to initialize the runner before returning.
prepared_images: Optional request-local cache reused to process only new
image references after the request hook.
Returns:
The configured agent and request, or None when the request is rejected.
"""
provider = provider or await _select_provider(event, plugin_context)
if provider is None:
@@ -1573,13 +1663,41 @@ async def build_main_agent(
)
)
quote_image_ref = None
quoted_image_refs: set[str] = set()
if collected_request:
req, quote_image_ref = await collect_initial_request(
event, plugin_context, config, req=req
event, plugin_context, config, req=req, quoted_image_refs=quoted_image_refs
)
if req is None:
return None
cfg = config.provider_settings or plugin_context.get_config(
umo=event.unified_msg_origin
).get("provider_settings", {})
options = cfg.get("image_compress_options", {})
max_size = normalize_model_image_max_size(
options.get("max_size") if isinstance(options, dict) else None
)
if prepared_images is None:
prepared_images = {}
supports_image = _provider_supports_modality(provider, "image")
caption_provider_id = cfg.get("default_image_caption_provider_id") or ""
if (
supports_image
or not caption_provider_id
or (req.conversation and req.image_urls)
):
quote_image_ref = None
await prepare_request_images(
req,
event,
max_size=max_size,
prepared=prepared_images,
quote_image_ref=quote_image_ref,
quoted_refs=quoted_image_refs,
finalize=False,
)
if isinstance(req.contexts, str):
req.contexts = json.loads(req.contexts)
thread_selected_text = event.get_extra("thread_selected_text")
@@ -1610,11 +1728,14 @@ async def build_main_agent(
else:
return None
captioned_refs: set[str] = set()
montage_refs = {
result["path"]
for result in prepared_images.values()
if result["path"] and result["montage"]
}
if collected_request:
cfg = config.provider_settings or plugin_context.get_config(
umo=event.unified_msg_origin
).get("provider_settings", {})
await _process_quote_message(
captioned_ref = await _process_quote_message(
event,
req,
cfg.get("default_image_caption_provider_id") or "",
@@ -1622,9 +1743,25 @@ async def build_main_agent(
_get_quoted_message_parser_settings(cfg),
main_provider_supports_image=_provider_supports_modality(provider, "image"),
skip_quote_image_caption=bool(req.conversation and req.image_urls),
image_ref=quote_image_ref,
image_ref=prepared_images[quote_image_ref]["path"]
if quote_image_ref
else None,
image_is_montage=bool(
quote_image_ref and prepared_images[quote_image_ref]["montage"]
),
)
await _decorate_llm_request(event, req, plugin_context, config, provider=provider)
if captioned_ref:
captioned_refs.add(captioned_ref)
captioned_refs.update(
await _decorate_llm_request(
event,
req,
plugin_context,
config,
provider=provider,
montage_refs=montage_refs,
)
)
await _apply_kb(event, req, plugin_context, config)
@@ -1730,6 +1867,15 @@ async def build_main_agent(
_apply_web_search_citation_prompt(event, req)
await prepare_request_images(
req,
event,
max_size=max_size,
prepared=prepared_images,
captioned_refs=captioned_refs,
supports_image=_provider_supports_modality(provider, "image"),
)
reset_coro = agent_runner.reset(
provider=provider,
request=req,
astrbot/core/pipeline/process_stage/method/agent_sub_stages/image_input.py
-104
View File
@@ -1,104 +0,0 @@
"""Prepare current local-agent input; history and tool results are not inputs here."""
from pathlib import Path
from astrbot.core.agent.message import ImageURLPart, TextPart
from astrbot.core.platform.astr_message_event import AstrMessageEvent
from astrbot.core.provider.entities import ProviderRequest
from astrbot.core.utils.media_utils import prepare_model_image
from astrbot.core.utils.string_utils import normalize_and_dedupe_strings
async def prepare_request_images(
req: ProviderRequest,
event: AstrMessageEvent,
*,
max_size: int,
output_dir: Path,
prepared: dict[str, str | None],
quote_image_ref: str | None = None,
) -> None:
"""Replace current images on a working request and track their owned files.
Args:
req: Working request; shared lists and image blocks are copied on write.
event: Owner of downloaded source files and prepared working files.
max_size: Normalized longest-edge limit for this request.
output_dir: Event working file directory.
prepared: Per-request mapping reused after the request hook.
quote_image_ref: Optional input for the dedicated quote caption branch.
"""
req.image_urls = normalize_and_dedupe_strings(req.image_urls)
refs = list(req.image_urls)
for part in req.extra_user_content_parts:
if isinstance(part, ImageURLPart):
refs.append(part.image_url.url)
elif isinstance(part, dict) and part.get("type") == "image_url":
refs.append(part["image_url"]["url"])
if quote_image_ref:
refs.append(quote_image_ref)
failed = False
has_montage = False
for ref in dict.fromkeys(refs):
if ref not in prepared:
path = None
image = await prepare_model_image(
ref,
max_size=max_size,
output_dir=output_dir,
)
if image:
path, is_montage, needs_cleanup = image
if needs_cleanup:
event.track_temporary_local_file(path)
has_montage |= is_montage
prepared[ref] = path
if path:
prepared[path] = path
failed |= prepared[ref] is None
req.image_urls = normalize_and_dedupe_strings(
[prepared[ref] for ref in req.image_urls if prepared[ref] is not None]
)
parts = []
for part in req.extra_user_content_parts:
if isinstance(part, ImageURLPart):
path = prepared[part.image_url.url]
if path is None:
continue
part = part.model_copy(
update={"image_url": part.image_url.model_copy(update={"url": path})}
)
elif isinstance(part, dict) and part.get("type") == "image_url":
path = prepared[part["image_url"]["url"]]
if path is None:
continue
part = {**part, "image_url": {**part["image_url"], "url": path}}
parts.append(part)
req.extra_user_content_parts = parts
if (
failed
and not (req.prompt or "").strip()
and not req.image_urls
and not req.audio_urls
):
if not any(
(part.get("type") != "text" or part.get("text", "").strip())
if isinstance(part, dict)
else (part.type != "text" or part.text.strip())
for part in parts
):
req.prompt = "[Image unavailable]"
if has_montage:
# Transient per-request hint; the montage file itself is not persisted.
req.extra_user_content_parts = [
*req.extra_user_content_parts,
TextPart(
text="<system_notice>\n"
"The input includes a GIF converted into a single image with frames "
"in reading order. Treat it as an animation; "
"do not mention the conversion or frame layout.\n"
"</system_notice>"
).mark_as_temp(),
]
astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py
+12 -65
View File
@@ -4,7 +4,6 @@ import asyncio
import base64
from collections.abc import AsyncGenerator
from dataclasses import replace
from pathlib import Path
from astrbot.core import db_helper, logger
from astrbot.core.agent.message import (
@@ -18,12 +17,8 @@ from astrbot.core.astr_main_agent import (
LLM_ERROR_MESSAGE_EXTRA_KEY,
MainAgentBuildConfig,
MainAgentBuildResult,
_get_quoted_message_parser_settings,
_process_quote_message,
_provider_supports_modality,
_select_provider,
build_main_agent,
collect_initial_request,
)
from astrbot.core.config.agent_runner import resolve_context_compression_config
from astrbot.core.message.components import File, Image, Record, Reply, Video
@@ -42,7 +37,7 @@ from astrbot.core.provider.entities import (
ProviderRequest,
)
from astrbot.core.star.star_handler import EventType
from astrbot.core.utils.astrbot_path import get_astrbot_temp_path
from astrbot.core.utils.image_input import prepare_request_images
from astrbot.core.utils.media_utils import normalize_model_image_max_size
from astrbot.core.utils.metrics import Metric
from astrbot.core.utils.session_lock import session_lock_manager
@@ -57,7 +52,6 @@ from ...follow_up import (
try_capture_follow_up,
unregister_active_runner,
)
from .image_input import prepare_request_images
class InternalAgentSubStage(Stage):
@@ -224,67 +218,13 @@ class InternalAgentSubStage(Stage):
)
plugin_context = self.ctx.plugin_manager.context
provider = await _select_provider(event, plugin_context)
if provider is None:
await self._send_llm_error_message(
event,
event.get_extra(LLM_ERROR_MESSAGE_EXTRA_KEY)
or "LLM 请求失败：未找到任何可用的对话模型（提供商）。请先在 WebUI 中配置并启用可用模型。",
)
return
req, quote_image_ref = await collect_initial_request(
event, plugin_context, build_cfg
)
if req is None:
return
settings = self.ctx.astrbot_config["provider_settings"]
options = settings.get("image_compress_options", {})
max_size = normalize_model_image_max_size(
options.get("max_size") if isinstance(options, dict) else None
)
output_dir = Path(get_astrbot_temp_path())
prepared: dict[str, str | None] = {}
supports_image = _provider_supports_modality(provider, "image")
caption_provider_id = (
settings.get("default_image_caption_provider_id") or ""
)
# Plugin requests may replace the event's image inputs. Only
# materialize a separate quote when its caption will be used.
if (
supports_image
or not caption_provider_id
or (req.conversation and req.image_urls)
):
quote_image_ref = None
await prepare_request_images(
req,
event,
max_size=max_size,
output_dir=output_dir,
prepared=prepared,
quote_image_ref=quote_image_ref,
)
await _process_quote_message(
event,
req,
caption_provider_id,
plugin_context,
_get_quoted_message_parser_settings(settings),
main_provider_supports_image=supports_image,
skip_quote_image_caption=bool(
req.conversation and req.image_urls
),
image_ref=prepared.get(quote_image_ref)
if quote_image_ref
else None,
)
prepared: dict[str, dict] = {}
build_result: MainAgentBuildResult | None = await build_main_agent(
event=event,
plugin_context=plugin_context,
config=build_cfg,
req=req,
provider=provider,
apply_reset=False,
prepared_images=prepared,
)
if build_result is None:
@@ -321,12 +261,19 @@ class InternalAgentSubStage(Stage):
if await call_event_hook(event, EventType.OnLLMRequestEvent, req):
return
options = build_cfg.provider_settings.get(
"image_compress_options", {}
)
await prepare_request_images(
req,
event,
max_size=max_size,
output_dir=output_dir,
max_size=normalize_model_image_max_size(
options.get("max_size")
if isinstance(options, dict)
else None
),
prepared=prepared,
supports_image=_provider_supports_modality(provider, "image"),
)
# apply reset
if reset_coro:
astrbot/core/utils/image_input.py
+206
View File
@@ -0,0 +1,206 @@
"""Prepare current local-agent images and generate their original-path notices."""
from pathlib import Path
from astrbot.core.agent.message import ImageURLPart, TextPart
from astrbot.core.platform.astr_message_event import AstrMessageEvent
from astrbot.core.provider.entities import ProviderRequest
from astrbot.core.utils.astrbot_path import get_astrbot_temp_path
from astrbot.core.utils.media_utils import (
MODEL_IMAGE_MAX_INPUT_BYTES,
ImageInputTooLargeError,
describe_media_ref,
prepare_model_image,
)
from astrbot.core.utils.string_utils import normalize_and_dedupe_strings
async def prepare_request_images(
req: ProviderRequest,
event: AstrMessageEvent,
*,
max_size: int,
output_dir: Path | None = None,
prepared: dict[str, dict],
quote_image_ref: str | None = None,
quoted_refs: set[str] | None = None,
captioned_refs: set[str] | None = None,
finalize: bool = True,
supports_image: bool = True,
) -> None:
"""Prepare new references and generate notices from actual processing results.
Args:
req: Working request; shared lists and image blocks are copied on write.
event: Owner of prepared working files.
max_size: Normalized longest-edge limit for this request.
output_dir: Event working file directory, defaulting to AstrBot temp.
prepared: Request-local cache of paths, processing outcomes, and owned
notice objects. Original and prepared references share each result.
quote_image_ref: Optional input for the dedicated quote caption branch.
quoted_refs: References collected directly from quoted message components.
captioned_refs: References successfully described by a caption provider.
finalize: Generate notices after captioning or the request hook finishes.
supports_image: Whether the final provider accepts visual input.
"""
output_dir = output_dir if output_dir is not None else Path(get_astrbot_temp_path())
req.image_urls = normalize_and_dedupe_strings(req.image_urls)
refs = list(req.image_urls)
persistent_refs = set(refs)
for part in req.extra_user_content_parts:
if isinstance(part, ImageURLPart):
refs.append(part.image_url.url)
if not part._no_save:
persistent_refs.add(part.image_url.url)
elif isinstance(part, dict) and part.get("type") == "image_url":
refs.append(part["image_url"]["url"])
if not part.get("_no_save"):
persistent_refs.add(part["image_url"]["url"])
if quote_image_ref:
refs.append(quote_image_ref)
persistent_refs.add(quote_image_ref)
for ref in dict.fromkeys(refs):
if ref in prepared:
if ref in persistent_refs:
prepared[ref]["temporary"] = False
continue
path = None
original_path = describe_media_ref(ref)
status = None
is_montage = False
try:
image = await prepare_model_image(
ref, max_size=max_size, output_dir=output_dir
)
except ImageInputTooLargeError as exc:
image = None
original_path = str(exc)
status = "oversized"
if image:
path, is_montage, needs_cleanup, original_path = image
if needs_cleanup:
event.track_temporary_local_file(path)
elif status is None:
status = "unavailable"
result = {
"path": path,
"original_path": original_path,
"quoted": ref == quote_image_ref or ref in (quoted_refs or ()),
"montage": is_montage,
"status": status,
"temporary": ref not in persistent_refs,
"notices": [],
}
prepared[ref] = result
if path:
prepared[path] = result
for ref in captioned_refs or ():
if ref in prepared:
prepared[ref]["status"] = "captioned"
req.image_urls = normalize_and_dedupe_strings(
[prepared[ref]["path"] for ref in req.image_urls if prepared[ref]["path"]]
)
parts = []
visuals = []
for part in req.extra_user_content_parts:
if isinstance(part, ImageURLPart):
result = prepared[part.image_url.url]
path = result["path"]
if path is None:
continue
part = part.model_copy(
update={"image_url": part.image_url.model_copy(update={"url": path})}
)
visuals.append((result, part._no_save))
elif isinstance(part, dict) and part.get("type") == "image_url":
result = prepared[part["image_url"]["url"]]
path = result["path"]
if path is None:
continue
part = {**part, "image_url": {**part["image_url"], "url": path}}
visuals.append((result, bool(part.get("_no_save"))))
parts.append(part)
visuals.extend((prepared[path], False) for path in req.image_urls)
req.extra_user_content_parts = parts
if refs and not visuals and not (req.prompt or "").strip() and not req.audio_urls:
if not any(
(part.get("type") != "text" or part.get("text", "").strip())
if isinstance(part, dict)
else (part.type != "text" or part.text.strip())
for part in parts
):
req.prompt = "[Image unavailable]"
if not finalize:
return
results = list({id(result): result for result in prepared.values()}.values())
# Only replace objects produced by this request. Plugin text is never parsed
# or removed, even when it is identical to one of our notices.
owned = {id(part) for result in results for part in result["notices"]}
parts = [part for part in parts if id(part) not in owned]
for result in results:
result["notices"] = []
if not supports_image:
visuals = []
active = {id(result) for result, _ in visuals}
rows = [
(result, None, result["temporary"])
for result in results
if id(result) not in active
]
# Extra image parts precede image_urls in the provider's visual order.
rows.extend(
(result, index, temporary)
for index, (result, temporary) in enumerate(visuals, 1)
)
oversized = None
montage = None
for result, index, temporary in rows:
name = f"Image {index}" if index is not None else "Image Attachment"
origin = " in quoted message" if result["quoted"] else ""
status = ""
if index is not None:
if result["montage"]:
status = "; animation converted to a 3x3 frame montage"
montage = result
elif result["status"] == "oversized":
status = (
f"; skipped: exceeds {MODEL_IMAGE_MAX_INPUT_BYTES // (1024 * 1024)} MiB"
)
oversized = result
elif result["status"] == "unavailable":
status = "; skipped: image unavailable"
elif result["status"] == "captioned":
status = "; sent to image-captioning model; description included as text"
else:
status = "; not included in this request"
label = TextPart(
text=f"[{name}{origin}: original path {result['original_path']}{status}]"
)
if temporary:
label.mark_as_temp()
result["notices"].append(label)
parts.append(label)
if oversized is not None:
advice = TextPart(
text="<system_notice>\n"
"For skipped images, use astrbot_file_read_tool if available, "
"or ask the user to resend a smaller image.\n"
"</system_notice>"
).mark_as_temp()
oversized["notices"].append(advice)
parts.append(advice)
if montage is not None:
advice = TextPart(
text="<system_notice>\n"
"Images labeled as animation montages contain frames in reading order. "
"Treat them as animations; do not mention the conversion or frame layout.\n"
"</system_notice>"
).mark_as_temp()
montage["notices"].append(advice)
parts.append(advice)
req.extra_user_content_parts = parts
astrbot/core/utils/media_utils.py
+21 -12
View File
@@ -1198,12 +1198,16 @@ def _extract_animation_montage_sync(image: PILImage.Image, max_size: int) -> byt
return encoded
class ImageInputTooLargeError(ValueError):
"""Raised with the retained source path when an image exceeds the input cap."""
async def prepare_model_image(
image_ref: str,
*,
max_size: int,
output_dir: Path,
) -> tuple[str, bool, bool] | None:
) -> tuple[str, bool, bool, str] | None:
"""Prepare an image, reusing compliant local files without copying them.
Args:
@@ -1212,10 +1216,13 @@ async def prepare_model_image(
output_dir: Directory for event-owned working files.
Returns:
The image path, whether it is an animation montage, and whether the caller
must delete the file after use. Existing compliant local files are borrowed;
newly materialized sources and encoded previews are owned by the caller.
Returns None for a recoverable input or write failure.
The image path, whether it is an animation montage, whether the caller must
delete the preview after use, and the retained original path. Compliant
originals are reused directly. Returns None for a recoverable failure.
Raises:
ImageInputTooLargeError: The input exceeds the size cap. Its original path
is retained for file-tool access and returned as the error message.
"""
try:
async with MediaResolver(image_ref, media_type="image").as_path() as source:
@@ -1226,17 +1233,17 @@ async def prepare_model_image(
input_size,
source.path,
)
return None
source.detach()
raise ImageInputTooLargeError(str(source.path))
image_bytes = await asyncio.to_thread(source.read_bytes)
converted_bytes, is_montage = await asyncio.to_thread(
_prepare_model_image_sync, image_bytes, max_size
)
original_path = str(source.path)
# Final image labels expose this original for later file-tool access.
source.detach()
if converted_bytes is image_bytes:
# The encoder returns the original bytes object for passthroughs.
# Transfer ownership only if the resolver materialized this file.
needs_cleanup = bool(source.cleanup_paths)
source.detach()
return str(source.path), is_montage, needs_cleanup
return original_path, is_montage, False, original_path
# Publish the working file synchronously after encoding, so cancellation
# cannot leave an untracked background write alive after this call.
output_dir.mkdir(parents=True, exist_ok=True)
@@ -1251,7 +1258,9 @@ async def prepare_model_image(
except BaseException:
output_path.unlink(missing_ok=True)
raise
return str(output_path), is_montage, True
return str(output_path), is_montage, True, original_path
except ImageInputTooLargeError:
raise
except Exception as exc:
if not is_recoverable_image_error(exc):
raise
changelogs/unreleased.md
+2 -1
View File
@@ -1,5 +1,6 @@
# Unreleased
- Raise the model image input cap from 32 MiB to 64 MiB. Larger originals are skipped before reading image bytes; accepted inputs still produce images strictly below 512 KiB.
- Generate numbered original-path notices in the image preparation stage during main-agent construction, including quoted images. After request hooks, apply the same size and format policy only to newly introduced image references and refresh their notices. Skipped and captioned attachments retain explicit status without taking a visual index; animation labels identify frame montages.
- Raise the model image input cap from 32 MiB to 64 MiB. Larger originals are skipped before reading image bytes, with a model notice retaining their paths and suggesting the file-reading tool or a smaller upload; accepted inputs still produce images strictly below 512 KiB.
- Local Agent input images are always prepared as JPEG/PNG files strictly below 512 KiB. Compliant local images are reused without copying. Transparent previews retain PNG alpha; animations become 3×3 montages. Original attachment paths remain available to tools, and event-owned previews are deleted after use without a shared conversion cache.
- Configuration mapping: **Enable image compression** (`provider_settings.image_compress_enabled`) is replaced by always-on preparation; **JPEG quality** (`provider_settings.image_compress_options.quality`) is replaced by automatic size control; **Maximum edge length** (`provider_settings.image_compress_options.max_size`) remains available. User attachments in CUA sessions follow the same limits.
docs/en/providers/image-formats.md
+6 -4
View File
@@ -1,12 +1,12 @@
# Model input images
The local Agent always prepares current input images before sending them to the model. This covers ordinary attachments, quoted images, and plugin `ProviderRequest` images, including images added or replaced by `OnLLMRequestEvent`.
The main-agent builder (`build_main_agent`) prepares initial input images before captioning or agent initialization. This covers ordinary attachments, quoted images, plugin `ProviderRequest` images, and direct calls to the builder. After `OnLLMRequestEvent`, the pipeline processes only newly introduced image references; existing references reuse their prepared results without another file check.
## Limits and formats
- Original image files larger than **64 MiB** are skipped before reading their contents or decoding pixels. Files up to and including 64 MiB are processed normally.
- Original image files larger than **64 MiB** are skipped before reading their contents or decoding pixels. The omission reason and retained original path are saved in history. A temporary notice suggests the file-reading tool if available or asking the user to resend a smaller image; this advice is not saved. Files up to and including 64 MiB are processed normally.
- Each prepared image file is **strictly smaller than 512 KiB** (524,288 bytes). This limit applies to image bytes before Base64 encoding, not the entire request.
- **Maximum edge length**, under **Configuration**, remains configurable through `provider_settings.image_compress_options.max_size` (default 1280 pixels). Images keep their aspect ratio and are never enlarged. They may be reduced further to meet the byte limit.
- **Maximum edge length**, under **Configuration**, remains configurable through `provider_settings.image_compress_options.max_size` (default 1280 pixels). Initial images keep their aspect ratio and are never enlarged. They may be reduced further to meet the byte limit. Images newly introduced by the request hook follow the same configured edge cap, format and animation handling, 64 MiB input limit, and strict 512 KiB output limit.
- Correctly oriented JPEG/PNG images already within both limits retain their original bytes. Other still images are orientation-corrected and encoded as JPEG, or PNG when they have transparency. Transparent images retain their alpha channel even when further reduction is necessary.
- GIF, animated WebP, and APNG inputs become a white 3×3 frame montage. Up to nine evenly spaced frames include the first and last animation frames; unused cells stay blank. An independent APNG cover is excluded.
- Unneeded metadata is removed when re-encoding. A color profile is retained when compatible with the output and the byte limit.
@@ -25,6 +25,8 @@ Values of the removed settings in existing configuration files no longer affect
## Original files and preview lifetime
The image preparation stage adds ordinary `Image N` text notices with original paths and quoted-message origins. Numbers follow the image-block order after hooks and filtering, including image content parts before `image_urls`. No attachment metadata is added to request or message models. Quoted origins and successful caption results are passed directly from their processing steps; notice text is never parsed to infer state. Only notice objects generated by this request are refreshed after hooks, preserving plugin-supplied text. Skipped images keep an unnumbered reason and original path; animation labels identify the 3×3 montage. Images handled by a captioning model are marked as described in text, without taking a visual index in the main request.
`image_urls` and image content parts reuse compliant local JPEG/PNG files directly; no working copy is created. Images requiring conversion use prepared preview files. The attachment paths shown to the model for file access continue to point to the **original images**, so tools can read the originals. Original files and event image components are not overwritten.
Preview files are owned by the current event and deleted when the event finishes, including failure or cancellation. There is no shared image-conversion cache. Repeated references within one request reuse the same preview, including the preparation pass after the request hook.
@@ -33,4 +35,4 @@ Successfully localized original attachments, including quoted images, survive ev
Conversation history captures the prepared image content as Base64 data URIs before preview files expire. Old history is not reprocessed. An unreadable image is skipped; valid text and other images remain. Cancellation, resource exhaustion, and programming errors are not treated as bad images.
Third-party Agent backends, direct calls that bypass the local pipeline, and tool-result images are outside this input preparation flow.
Third-party Agent backends, direct provider calls that bypass the main-agent builder, and tool-result images are outside this input preparation flow.
docs/zh/providers/image-formats.md
+6 -4
View File
@@ -1,12 +1,12 @@
# 图片处理
使用内置 Agent 处理新消息时，AstrBot 始终在发送给模型前处理本次输入的图片，包括普通附件、引用图片，以及插件 `ProviderRequest` 中的图片。请求 Hook 新增或替换的图片也会处理。
`build_main_agent` 在图片描述和 Agent 初始化前处理初始图片，包括普通附件、引用图片及插件 `ProviderRequest` 中的图片，直接调用构建函数也适用。请求 Hook 执行后，流水线只检查新增的图片引用；已有引用复用处理结果，不再次检查文件。
## 大小限制与格式
- 原始图片文件超过 **64 MiB** 时，在读取内容和解码像素前跳过；小于等于 64 MiB 的图片正常处理。
- 原始图片文件超过 **64 MiB** 时，在读取内容和解码像素前跳过；原因与保留的原图路径会保存到历史。另向模型提供临时建议：使用可用的文件读取工具，或请用户重发较小的图片，这段建议不保存到历史。小于等于 64 MiB 的图片正常处理。
- 模型使用的单张图片文件**严格小于 512 KiB**（524,288 字节）。限制针对 Base64 编码前的图片文件，不是整个请求体。
- 在管理面板的**配置文件**页面搜索**最大边长**，对应 `provider_settings.image_compress_options.max_size`，默认 1280 像素。图片保持宽高比，小图不会放大；为满足文件体积限制，必要时会进一步缩小。
- 在管理面板的**配置文件**页面搜索**最大边长**，对应 `provider_settings.image_compress_options.max_size`，默认 1280 像素。初始图片保持宽高比，小图不会放大；为满足文件体积限制，必要时会进一步缩小。Hook 新增图片同样按配置限制最长边、处理格式和动图，并应用 64 MiB 输入上限和严格小于 512 KiB 的输出限制。
- 方向正常、像素尺寸及文件体积均合规的 JPEG/PNG 保留原始字节。其余静图修正方向后编码为 JPEG；含透明度的图片输出 PNG，即使需要进一步缩小也保留透明通道。
- GIF、动画 WebP、APNG 等动图均匀选取最多 9 帧，包含首尾动画帧，生成白底 3×3 拼图。不足 9 帧的格子留白，APNG 的独立封面不参与采样。
- 重新编码时移除不必要的元数据；与输出色彩空间兼容、且不影响体积达标的色彩配置保留。
@@ -25,6 +25,8 @@ CUA 会话中的用户输入图片也遵守这些限制；截图及其他工具
## 原图路径与文件清理
图片处理阶段统一生成普通的 `Image N` 文本说明，包含原图路径，并标明是否来自引用消息。序号按 Hook 和过滤之后的图片顺序生成：额外图片内容块在前，`image_urls` 在后。请求实体和消息模型不新增附件元数据。引用来源和成功生成描述的图片由对应处理步骤直接传递，不解析提示文本推断状态；Hook 后只更新本次流程自己生成的提示对象，保留插件传入的文本。跳过的图片只保留原因和原图路径，不占图片序号；动图会标明已转换为九宫格。交给图片描述模型处理的图片标为已转文字描述，不占主请求的图片序号。
`image_urls` 及图片内容块直接复用合规的本地 JPEG/PNG 文件，不创建工作副本；需要转换的图片才使用处理后的预览文件。提供给模型、供工具读取文件的附件说明仍然指向**原始图片路径**，不会替换为预览路径。原始文件和事件中的图片组件均不被覆盖。
模型预览文件归当前事件所有，事件完成、失败或取消后清理，不再保留跨请求图片转换缓存。同一请求中重复出现的图片引用复用同一份预览，请求 Hook 前后的准备过程也会去重。
@@ -33,4 +35,4 @@ CUA 会话中的用户输入图片也遵守这些限制；截图及其他工具
对话历史会在预览文件删除前保存模型使用的图片内容，采用 Base64 data URI；旧历史不会重新处理。无法读取的图片会跳过，不影响其他文字和正常图片；取消、资源耗尽及程序错误不会被当作坏图忽略。
第三方 Agent 后端、绕过本地流水线的直接调用，以及工具返回的图片，不属于本次输入图片处理范围。
第三方 Agent 后端、绕过主 Agent 构建函数直接调用 provider，以及工具返回的图片，不属于本次输入图片处理范围。
tests/test_media_utils.py
+7 -5
View File
@@ -925,11 +925,12 @@ async def test_prepare_model_image_skips_oversized_input(tmp_path, monkeypatch):
pytest.fail("Oversized inputs must be rejected before reading image bytes")
monkeypatch.setattr(media_utils.ResolvedMediaFile, "read_bytes", fail_read)
result = await media_utils.prepare_model_image(
str(image_path), max_size=1280, output_dir=tmp_path
)
with pytest.raises(media_utils.ImageInputTooLargeError) as error:
await media_utils.prepare_model_image(
str(image_path), max_size=1280, output_dir=tmp_path
)
assert result is None
assert str(error.value) == str(image_path)
assert image_path.is_file()
@@ -949,7 +950,8 @@ async def test_prepare_model_image_accepts_inputs_up_to_64_mib(tmp_path, input_s
)
assert result is not None
output_path, is_montage, needs_cleanup = result
output_path, is_montage, needs_cleanup, original_path = result
assert original_path == str(image_path)
assert not is_montage and needs_cleanup
assert Path(output_path).stat().st_size < 512 * 1024
assert image_path.stat().st_size == input_size
tests/test_model_image_preparation.py
+1 -1
View File
@@ -53,7 +53,7 @@ async def test_compliant_stills_reuse_source_without_copying(tmp_path, fmt, file
prepared = await media.prepare_model_image(
source.as_uri() if file_uri else str(source), max_size=1280, output_dir=output
)
assert prepared == (str(source), False, False)
assert prepared == (str(source), False, False, str(source))
assert not output.exists()
assert source.read_bytes() == original
tests/test_process_stage_images.py
+553 -20
View File
@@ -24,7 +24,6 @@ from astrbot.core.config.default import DEFAULT_CONFIG
from astrbot.core.message.components import Image, Plain, Reply
from astrbot.core.pipeline.preprocess_stage import stage as preprocess
from astrbot.core.pipeline.process_stage.method.agent_sub_stages import (
image_input,
internal,
)
from astrbot.core.pipeline.process_stage.stage import ProcessStage
@@ -35,6 +34,7 @@ from astrbot.core.platform.platform_metadata import PlatformMetadata
from astrbot.core.provider.entities import LLMResponse, ProviderRequest
from astrbot.core.provider.provider import Provider
from astrbot.core.star.star_handler import EventType
from astrbot.core.utils import image_input
from astrbot.core.utils import media_utils as media
@@ -85,7 +85,7 @@ def source_image(tmp_path, fmt="GIF"):
@pytest.fixture
def harness(tmp_path, monkeypatch):
work = tmp_path / "work"
for module in (media, internal, preprocess):
for module in (media, image_input, preprocess):
monkeypatch.setattr(module, "get_astrbot_temp_path", lambda: str(work))
config = copy.deepcopy(DEFAULT_CONFIG)
config["provider_settings"].update(
@@ -339,15 +339,17 @@ async def test_animation_montage_notice_reaches_model(harness, tmp_path):
part
for part in harness.captured[-1].req.extra_user_content_parts
if isinstance(part, TextPart)
and part.text.startswith("<system_notice>\nThe input includes a GIF")
and part.text.startswith(
"<system_notice>\nImages labeled as animation montages"
)
]
assert len(notices) == int(expected)
if expected:
assert notices[0]._no_save
assert notices[0].text.startswith("<system_notice>\n")
assert notices[0].text.endswith("\n</system_notice>")
assert "converted into a single image with frames" in notices[0].text
assert "Treat it as an animation" in notices[0].text
assert "contain frames in reading order" in notices[0].text
assert "Treat them as animations" in notices[0].text
assert "do not mention the conversion or frame layout" in notices[0].text
@@ -579,15 +581,27 @@ async def test_provider_and_history_keep_bounded_previews_and_original_paths(
@pytest.mark.asyncio
@pytest.mark.parametrize("plugin", [False, True])
async def test_caption_first_call_uses_prepared_quote(harness, tmp_path, plugin):
source = source_image(tmp_path)
@pytest.mark.parametrize("fmt", ["GIF", "JPEG"])
async def test_caption_first_call_uses_prepared_quote(harness, tmp_path, plugin, fmt):
source = source_image(tmp_path, fmt)
harness.provider.provider_config["modalities"] = ["text"]
harness.config["provider_settings"]["default_image_caption_provider_id"] = "caption"
harness.config["provider_settings"]["image_caption_prompt"] = "Describe in Chinese."
caption = MagicMock(spec=Provider)
captured = []
async def describe(**kwargs):
assert "image_settings" not in kwargs
prompt = kwargs["prompt"]
if not plugin:
assert prompt.startswith("Describe in Chinese.")
if fmt == "GIF":
assert "positions 1 (1-based)" in prompt
assert "single image of frames in reading order" in prompt
assert "Describe them as animations" in prompt
assert "do not mention the conversion or frame layout" in prompt
else:
assert "<system_notice>" not in prompt
for ref in kwargs["image_urls"]:
with PILImage.open(ref) as image:
assert image.format == "JPEG" and getattr(image, "n_frames", 1) == 1
@@ -612,6 +626,30 @@ async def test_caption_first_call_uses_prepared_quote(harness, tmp_path, plugin)
)
@pytest.mark.asyncio
async def test_caption_animation_notice_uses_actual_image_order(harness, tmp_path):
still = source_image(tmp_path, "JPEG")
animated = source_image(tmp_path, "GIF")
harness.provider.provider_config["modalities"] = ["text"]
harness.config["provider_settings"]["default_image_caption_provider_id"] = "caption"
caption = MagicMock(spec=Provider)
caption.text_chat = AsyncMock(
return_value=LLMResponse(role="assistant", completion_text="caption")
)
harness.context.get_provider_by_id.return_value = caption
event = make_event([Image(file=str(still)), Image(file=str(animated))])
await process_event(harness, event, preprocess_first=True)
caption.text_chat.assert_awaited_once()
args = caption.text_chat.await_args.kwargs
assert len(args["image_urls"]) == 2
assert args["image_urls"][0] == str(still)
assert "positions 2 (1-based)" in args["prompt"]
assert "Describe them as animations" in args["prompt"]
assert "do not mention the conversion or frame layout" in args["prompt"]
@pytest.mark.asyncio
async def test_reply_fallback_limit_and_deduplication(harness, tmp_path, monkeypatch):
paths = []
@@ -645,6 +683,117 @@ async def test_failed_images_keep_valid_input(harness, tmp_path, text, good):
assert harness.captured[0].run_context.messages[-1].content
@pytest.mark.asyncio
@pytest.mark.parametrize("origin", ["ordinary", "quote", "plugin", "hook"])
@pytest.mark.parametrize("with_valid_image", [False, True])
async def test_oversized_images_explain_omission_and_keep_original_path(
harness, tmp_path, monkeypatch, origin, with_valid_image
):
source = tmp_path / "oversized.png"
PILImage.new("RGB", (4, 4)).save(source)
with source.open("ab") as file:
file.truncate(64 * 1024 * 1024 + 1)
refs = [str(source)]
if with_valid_image:
refs.append(str(source_image(tmp_path, "JPEG")))
attachments = [Image(file=ref) for ref in refs]
event = make_event(text="describe")
if origin == "ordinary":
event.message_obj.message = attachments
elif origin == "quote":
event.message_obj.message = [Reply(id="large", chain=attachments)]
elif origin == "plugin":
event.set_extra(
"provider_request",
ProviderRequest(
prompt="describe",
image_urls=refs,
conversation=harness.context.conversation_manager.get_conversation.return_value,
),
)
else:
async def hook(event, kind, *args):
if kind == EventType.OnLLMRequestEvent:
args[0].image_urls.extend(refs)
return False
monkeypatch.setattr(internal, "call_event_hook", hook)
await process_event(harness, event, preprocess_first=True)
req = harness.captured[0].req
assert len(req.image_urls) == int(with_valid_image)
payload = harness.provider.text_chat.await_args.kwargs["contexts"][-1]["content"]
notices = [
part["text"]
for part in payload
if part["type"] == "text" and "skipped: exceeds" in part["text"]
]
assert len(notices) == 1
assert "64 MiB" in notices[0] and str(source) in notices[0]
advice = [
part["text"]
for part in payload
if part["type"] == "text"
and "astrbot_file_read_tool if available" in part["text"]
]
assert len(advice) == 1
assert "resend a smaller image" in advice[0]
saved = harness.context.conversation_manager.update_conversation.await_args.kwargs[
"history"
]
assert any(
part.get("text") == notices[0]
for message in saved
if message["role"] == "user"
for part in message["content"]
)
assert all(
"resend a smaller image" not in part.get("text", "")
for message in saved
if message["role"] == "user"
for part in message["content"]
)
event.cleanup_temporary_local_files()
assert source.stat().st_size == 64 * 1024 * 1024 + 1
assert not list(harness.work.rglob("model_image_*"))
@pytest.mark.asyncio
async def test_oversized_plugin_download_is_retained_for_file_tool(
harness, tmp_path, monkeypatch
):
source = tmp_path / "source.png"
PILImage.new("RGB", (4, 4)).save(source)
original = source.read_bytes()
async def download(url, target):
path = Path(target)
path.write_bytes(original)
with path.open("ab") as file:
file.truncate(64 * 1024 * 1024 + 1)
monkeypatch.setattr(media, "download_file", download)
event = make_event()
req = ProviderRequest(image_urls=["https://example.com/large.png"])
prepared = {}
for _ in range(2):
req.image_urls = ["https://example.com/large.png"]
await image_input.prepare_request_images(
req, event, max_size=1280, output_dir=harness.work, prepared=prepared
)
assert not req.image_urls
assert len(req.extra_user_content_parts) == 2
assert not req.extra_user_content_parts[0]._no_save
assert req.extra_user_content_parts[1]._no_save
context = await req.assemble_context()
notices = [part["text"] for part in context["content"] if part["type"] == "text"]
downloaded = list(harness.work.glob("media_image_*"))
assert len(downloaded) == 1
assert any(str(downloaded[0]) in notice for notice in notices)
event.cleanup_temporary_local_files()
assert downloaded[0].stat().st_size == 64 * 1024 * 1024 + 1
@pytest.mark.asyncio
@pytest.mark.parametrize(
"streaming,fallback", [(False, False), (True, False), (False, True), (True, True)]
@@ -781,10 +930,10 @@ async def test_compliant_plugin_images_reuse_localized_file_with_correct_ownersh
assert path.read_bytes() == original
borrowed = reference in {"file", "path"}
assert (path == source) == borrowed
assert event._temporary_local_files == ([] if borrowed else [str(path)])
assert event._temporary_local_files == []
assert set(harness.work.rglob("*")) == (set() if borrowed else {path})
event.cleanup_temporary_local_files()
assert path.exists() == borrowed
assert path.exists()
assert source.read_bytes() == original
@@ -924,14 +1073,18 @@ async def test_attachment_sources_survive_event_cleanup(
monkeypatch.setattr(internal, "call_event_hook", hook)
await process_event(harness, event, preprocess_first=True)
req = harness.captured[0].req
label = "Image Attachment in quoted message" if quoted else "Image Attachment"
prefix = f"[{label}: path "
label = "Image 1 in quoted message" if quoted else "Image 1"
prefix = f"[{label}: original path "
text = next(
part.text
for part in req.extra_user_content_parts
if isinstance(part, TextPart) and part.text.startswith(prefix)
)
attachment_path = Path(text[len(prefix) : -1])
attachment_path = Path(
text[len(prefix) : -1].removesuffix(
"; animation converted to a 3x3 frame montage"
)
)
visual_path = Path(req.image_urls[0])
assert attachment_path.read_bytes() == original
assert str(attachment_path) not in event._temporary_local_files
@@ -989,7 +1142,7 @@ async def test_stage_failure_preserves_sources_and_cleans_working_files(
@pytest.mark.asyncio
@pytest.mark.parametrize("explicit", [False, True])
async def test_direct_build_keeps_raw_images_and_quote_collection(
async def test_direct_build_prepares_images_and_quote_collection(
harness, tmp_path, monkeypatch, explicit
):
source = source_image(tmp_path)
@@ -1003,11 +1156,8 @@ async def test_direct_build_keeps_raw_images_and_quote_collection(
),
]
)
monkeypatch.setattr(
image_input,
"prepare_model_image",
AsyncMock(side_effect=AssertionError("direct build must not convert")),
)
convert = AsyncMock(wraps=image_input.prepare_model_image)
monkeypatch.setattr(image_input, "prepare_model_image", convert)
req = (
ProviderRequest(prompt="direct", image_urls=[str(source)]) if explicit else None
)
@@ -1020,7 +1170,9 @@ async def test_direct_build_keeps_raw_images_and_quote_collection(
provider=harness.provider,
req=req,
)
assert result and result.provider_request.image_urls == [str(source)]
assert result and len(result.provider_request.image_urls) == 1
assert convert.await_count == 1
assert result.provider_request.image_urls[0] != str(source)
quote_parts = [
part
for part in result.provider_request.extra_user_content_parts
@@ -1028,7 +1180,11 @@ async def test_direct_build_keeps_raw_images_and_quote_collection(
]
assert len(quote_parts) == 1 and "quoted text" in quote_parts[0].text
with PILImage.open(result.provider_request.image_urls[0]) as image:
assert image.format == "GIF" and image.n_frames == 2
assert image.format == "JPEG" and image.size == (90, 45)
assert Path(result.provider_request.image_urls[0]).stat().st_size < 512 * 1024
event.cleanup_temporary_local_files()
assert source.is_file()
assert not Path(result.provider_request.image_urls[0]).exists()
@pytest.mark.asyncio
@@ -1067,3 +1223,380 @@ async def test_materialized_sources_remain_owned_when_collection_fails(
assert owned.read_bytes() == source.read_bytes()
event.cleanup_temporary_local_files()
assert not owned.exists() and source.exists()
@pytest.mark.asyncio
@pytest.mark.parametrize("remove_image", [False, True])
@pytest.mark.parametrize("rewrite_text_parts", [False, True])
async def test_final_image_labels_follow_mixed_visuals_after_hook(
harness, tmp_path, monkeypatch, remove_image, rewrite_text_parts
):
oversized = tmp_path / "oversized.png"
PILImage.new("RGB", (4, 4)).save(oversized)
with oversized.open("ab") as file:
file.truncate(64 * 1024 * 1024 + 1)
still = source_image(tmp_path, "JPEG")
animation = source_image(tmp_path, "GIF")
inserted = source_image(tmp_path, "BMP")
event = make_event(
[
Image(file=str(oversized)),
Image(file=str(still)),
Reply(id="animation", chain=[Image(file=str(animation))]),
]
)
convert = AsyncMock(wraps=image_input.prepare_model_image)
monkeypatch.setattr(image_input, "prepare_model_image", convert)
async def hook(event, kind, *args):
if kind == EventType.OnLLMRequestEvent:
req = args[0]
if rewrite_text_parts:
req.extra_user_content_parts = []
# Extra image parts precede image_urls in the actual provider payload.
req.extra_user_content_parts.insert(
0, ImageURLPart(image_url=ImageURLPart.ImageURL(url=str(inserted)))
)
req.extra_user_content_parts.append(
TextPart(text="[Image captions supplied by a plugin]")
)
req.image_urls = ([] if remove_image else [str(animation)]) + [
str(still),
str(oversized),
str(still),
]
return False
monkeypatch.setattr(internal, "call_event_hook", hook)
await process_event(harness, event, preprocess_first=True)
assert convert.await_count == 4
req = harness.captured[0].req
payload = harness.provider.text_chat.await_args.kwargs["contexts"][-1]["content"]
expected = [inserted, still] if remove_image else [inserted, animation, still]
images = [
index for index, part in enumerate(payload) if part["type"] == "image_url"
]
assert len(images) == len(expected)
visual_paths = [
part.image_url.url
for part in req.extra_user_content_parts
if isinstance(part, ImageURLPart)
] + req.image_urls
for number, (position, original) in enumerate(zip(images, expected), 1):
label = next(
part
for part in payload
if part["type"] == "text"
and part["text"].startswith(f"[Image {number}:")
or part["type"] == "text"
and part["text"].startswith(f"[Image {number} in quoted message:")
)
assert f"original path {original}" in label["text"]
assert ("in quoted message" in label["text"]) == (original == animation)
assert ("3x3 frame montage" in label["text"]) == (original == animation)
data = base64.b64decode(payload[position]["image_url"]["url"].split(",", 1)[1])
assert data == Path(visual_paths[number - 1]).read_bytes()
texts = [part["text"] for part in payload if part["type"] == "text"]
assert "[Image captions supplied by a plugin]" in texts
for original in (oversized, still, animation, inserted):
assert sum(str(original) in text for text in texts) == 1
assert any(
str(oversized) in text and "skipped: exceeds 64 MiB" in text for text in texts
)
if remove_image:
assert any(
str(animation) in text and "not included in this request" in text
for text in texts
)
history = (
harness.context.conversation_manager.update_conversation.await_args.kwargs[
"history"
]
)
saved = next(message["content"] for message in history if message["role"] == "user")
assert [
p["text"]
for p in saved
if p["type"] == "text" and p["text"].startswith("[Image")
] == [text for text in texts if text.startswith("[Image")]
event.cleanup_temporary_local_files()
assert all(path.is_file() for path in (oversized, still, animation, inserted))
@pytest.mark.asyncio
@pytest.mark.parametrize("add_uncaptioned_image", [False, True])
async def test_captioned_quote_has_status_without_visual_index(
harness, tmp_path, monkeypatch, add_uncaptioned_image
):
source = source_image(tmp_path)
harness.provider.provider_config["modalities"] = ["text"]
harness.config["provider_settings"]["default_image_caption_provider_id"] = "caption"
caption = MagicMock(spec=Provider)
caption.text_chat = AsyncMock(
return_value=LLMResponse(
role="assistant", completion_text="An animated red square."
)
)
harness.context.get_provider_by_id.return_value = caption
event = make_event([Reply(id="quoted", chain=[Image(file=str(source))])])
inserted = source_image(tmp_path, "JPEG")
async def hook(event, kind, *args):
if kind == EventType.OnLLMRequestEvent and add_uncaptioned_image:
args[0].image_urls.append(str(inserted))
return False
monkeypatch.setattr(internal, "call_event_hook", hook)
await process_event(harness, event, preprocess_first=True)
assert caption.text_chat.await_count == 1
payload = harness.provider.text_chat.await_args.kwargs["contexts"][-1]["content"]
assert not any(part["type"] == "image_url" for part in payload)
labels = [
part["text"]
for part in payload
if part["type"] == "text" and str(source) in part["text"]
]
assert len(labels) == 1
assert "in quoted message" in labels[0]
assert "description included as text" in labels[0]
assert "[Image 1" not in labels[0]
if add_uncaptioned_image:
label = next(
part["text"] for part in payload if str(inserted) in part.get("text", "")
)
assert "not included in this request" in label
assert "description included as text" not in label
@pytest.mark.asyncio
async def test_temporary_skipped_image_does_not_persist_its_path(harness, tmp_path):
from astrbot.core.agent.message import Message
source = tmp_path / "private.png"
PILImage.new("RGB", (4, 4)).save(source)
with source.open("ab") as file:
file.truncate(64 * 1024 * 1024 + 1)
req = ProviderRequest(
prompt="describe",
extra_user_content_parts=[
ImageURLPart(
image_url=ImageURLPart.ImageURL(url=str(source))
).mark_as_temp()
],
)
await image_input.prepare_request_images(
req, make_event(), max_size=1280, output_dir=harness.work, prepared={}
)
context = await req.assemble_context()
assert any(str(source) in part.get("text", "") for part in context["content"])
saved = dump_messages_with_checkpoints([Message.model_validate(context)])
assert saved[0]["content"] == [{"type": "text", "text": "describe"}]
@pytest.mark.asyncio
async def test_duplicate_current_and_quoted_image_share_one_visual_label(
harness, tmp_path
):
source = source_image(tmp_path, "JPEG")
event = make_event(
[
Image(file=str(source)),
Reply(id="same", chain=[Image(file=str(source))]),
]
)
await process_event(harness, event, preprocess_first=True)
content = harness.provider.text_chat.await_args.kwargs["contexts"][-1]["content"]
images = [i for i, part in enumerate(content) if part["type"] == "image_url"]
assert len(images) == 1
label = next(
part["text"]
for part in content
if part.get("text", "").startswith("[Image 1 in quoted message:")
)
assert label == f"[Image 1 in quoted message: original path {source}]"
assert sum(str(source) in part.get("text", "") for part in content) == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("origin", ["image_urls", "image_part", "dict_part"])
@pytest.mark.parametrize(
"size_case", ["compliant", "boundary", "compress", "oversized"]
)
async def test_hook_only_new_refs_receive_full_image_policy(
harness, tmp_path, monkeypatch, origin, size_case
):
initial = tmp_path / "initial.jpg"
PILImage.new("RGB", (300, 150), "blue").save(initial)
added = tmp_path / "added.jpg"
PILImage.new("RGB", (600, 300), "red").save(added)
if size_case == "compress":
added = tmp_path / "added.png"
pixels = random.Random(42).randbytes(1000 * 800 * 3)
PILImage.frombytes("RGB", (1000, 800), pixels).save(added)
assert 512 * 1024 < added.stat().st_size < 64 * 1024 * 1024
elif size_case in {"boundary", "oversized"}:
with added.open("ab") as file:
file.truncate(
512 * 1024 if size_case == "boundary" else 64 * 1024 * 1024 + 1
)
original_bytes = added.read_bytes() if size_case != "oversized" else None
convert = AsyncMock(wraps=image_input.prepare_model_image)
monkeypatch.setattr(image_input, "prepare_model_image", convert)
async def hook(event, kind, *args):
if kind == EventType.OnLLMRequestEvent:
req = args[0]
assert convert.await_count == 1
with PILImage.open(req.image_urls[0]) as image:
assert image.size == (90, 45)
# Reintroducing an already processed original must also reuse its preview.
req.image_urls.append(str(initial))
if origin == "image_urls":
req.image_urls.append(str(added))
elif origin == "image_part":
req.extra_user_content_parts.append(
ImageURLPart(image_url=ImageURLPart.ImageURL(url=str(added)))
)
else:
req.extra_user_content_parts.append(
{"type": "image_url", "image_url": {"url": str(added)}}
)
return False
monkeypatch.setattr(internal, "call_event_hook", hook)
event = make_event([Image(file=str(initial))])
await process_event(harness, event, preprocess_first=True)
assert [call.args[0] for call in convert.await_args_list] == [
str(initial),
str(added),
]
assert [call.kwargs["max_size"] for call in convert.await_args_list] == [90, 90]
payload = harness.provider.text_chat.await_args.kwargs["contexts"][-1]["content"]
images = [part for part in payload if part["type"] == "image_url"]
assert len(images) == (1 if size_case == "oversized" else 2)
for image in images:
assert (
len(base64.b64decode(image["image_url"]["url"].split(",", 1)[1]))
< 512 * 1024
)
if size_case == "oversized":
assert any(
str(added) in part.get("text", "")
and "skipped: exceeds 64 MiB" in part["text"]
for part in payload
)
else:
req = harness.captured[0].req
if origin == "image_urls":
result_path = req.image_urls[-1]
else:
part = next(
part
for part in req.extra_user_content_parts
if isinstance(part, ImageURLPart)
or isinstance(part, dict)
and part.get("type") == "image_url"
)
result_path = (
part["image_url"]["url"]
if isinstance(part, dict)
else part.image_url.url
)
with PILImage.open(result_path) as image:
assert max(image.size) <= 90
assert result_path != str(added)
assert Path(result_path).stat().st_size < 512 * 1024
event.cleanup_temporary_local_files()
assert initial.exists() and added.exists()
if original_bytes is not None:
assert added.read_bytes() == original_bytes
@pytest.mark.asyncio
@pytest.mark.parametrize("with_image", [False, True])
@pytest.mark.parametrize("as_dict", [False, True])
async def test_plugin_image_notices_are_not_parsed(
harness, tmp_path, monkeypatch, with_image, as_dict
):
source = source_image(tmp_path, "JPEG")
texts = [
"[Image Attachment: path /plugin/reference-only.png]",
f"[Image 1: original path {source}]",
"<system_notice>\nFor skipped images, keep this plugin instruction.\n</system_notice>",
"<system_notice>\nImages labeled as animation montages contain frames in reading order. "
"Treat them as animations; do not mention the conversion or frame layout.\n</system_notice>",
"<image_caption>This text is supplied by a plugin.</image_caption>",
]
plugin_parts = [
{"type": "text", "text": text} if as_dict else TextPart(text=text)
for text in texts
]
req = ProviderRequest(
prompt="Use these references if needed",
image_urls=[str(source)] if with_image else [],
extra_user_content_parts=list(plugin_parts),
)
prepared = {}
convert = AsyncMock(wraps=image_input.prepare_model_image)
monkeypatch.setattr(image_input, "prepare_model_image", convert)
for _ in range(2):
await image_input.prepare_request_images(
req, make_event(), max_size=90, prepared=prepared, output_dir=harness.work
)
for part in plugin_parts:
assert any(current is part for current in req.extra_user_content_parts)
content = (await req.assemble_context())["content"]
for text in texts:
assert any(part.get("text") == text for part in content)
assert len(req.extra_user_content_parts) == len(plugin_parts) + int(with_image)
assert convert.await_count == int(with_image)
@pytest.mark.asyncio
@pytest.mark.parametrize("drop_extra_in_hook", [False, True])
async def test_only_actually_captioned_images_receive_caption_status(
harness, tmp_path, monkeypatch, drop_extra_in_hook
):
main_image = source_image(tmp_path, "JPEG")
extra_image = source_image(tmp_path, "PNG")
harness.provider.provider_config["modalities"] = ["text"]
harness.config["provider_settings"]["default_image_caption_provider_id"] = "caption"
caption = MagicMock(spec=Provider)
caption.text_chat = AsyncMock(
return_value=LLMResponse(role="assistant", completion_text="A red square.")
)
harness.context.get_provider_by_id.return_value = caption
event = make_event()
event.set_extra(
"provider_request",
ProviderRequest(
prompt="describe both",
image_urls=[str(main_image)],
extra_user_content_parts=[
ImageURLPart(image_url=ImageURLPart.ImageURL(url=str(extra_image)))
],
conversation=harness.context.conversation_manager.get_conversation.return_value,
),
)
async def hook(event, kind, *args):
if kind == EventType.OnLLMRequestEvent and drop_extra_in_hook:
req = args[0]
req.extra_user_content_parts = [
part
for part in req.extra_user_content_parts
if not isinstance(part, ImageURLPart)
]
return False
monkeypatch.setattr(internal, "call_event_hook", hook)
await process_event(harness, event)
assert caption.text_chat.await_count == 1
assert caption.text_chat.await_args.kwargs["image_urls"] == [str(main_image)]
content = harness.provider.text_chat.await_args.kwargs["contexts"][-1]["content"]
assert not any(part["type"] == "image_url" for part in content)
for path, captioned in ((main_image, True), (extra_image, False)):
labels = [part["text"] for part in content if str(path) in part.get("text", "")]
assert len(labels) == 1
assert ("description included as text" in labels[0]) == captioned
assert ("not included in this request" in labels[0]) != captioned
tests/unit/test_astr_main_agent.py
+41 -25
View File
@@ -32,6 +32,16 @@ from astrbot.core.star.context import Context
from astrbot.core.star.star import StarMetadata
@pytest.fixture
def valid_image_path(tmp_path):
"""Create a real image for builder tests that exercise provider selection."""
from PIL import Image as PILImage
path = tmp_path / "input.jpg"
PILImage.new("RGB", (8, 8), "red").save(path)
return str(path)
@pytest.fixture
def mock_provider():
"""Create a mock provider."""
@@ -1885,14 +1895,14 @@ class TestBuildMainAgent:
quoted,
compression_enabled,
):
"""Direct builders keep raw attachments regardless of pipeline settings."""
"""Direct builders prepare both ordinary and quoted images before reset."""
from PIL import Image as PILImage
from astrbot.core.utils import media_utils
from astrbot.core.utils import image_input
module = ama
monkeypatch.setattr(media_utils, "get_astrbot_temp_path", lambda: str(tmp_path))
monkeypatch.setattr(image_input, "get_astrbot_temp_path", lambda: str(tmp_path))
source_path = tmp_path / "image.jpg"
PILImage.new("RGB", (8, 8), (255, 0, 0)).save(source_path)
original = source_path.read_bytes()
@@ -1923,23 +1933,23 @@ class TestBuildMainAgent:
tool_call_timeout=60,
provider_settings={
"image_compress_enabled": compression_enabled,
"image_compress_options": {"max_size": 2},
"image_compress_options": {"max_size": 4},
},
),
)
assert result is not None
request = result.provider_request
label = "Image Attachment in quoted message" if quoted else "Image Attachment"
assert f"[{label}: path {source_path}]" in [
label = "Image 1 in quoted message" if quoted else "Image 1"
assert f"[{label}: original path {source_path}]" in [
part.text for part in request.extra_user_content_parts
]
assert len(request.image_urls) == 1
visual_path = Path(request.image_urls[0])
assert visual_path == source_path
assert visual_path != source_path
with PILImage.open(visual_path) as visual_image:
assert visual_image.size == (8, 8)
mock_event.track_temporary_local_file.assert_not_called()
assert visual_image.size == (4, 4)
mock_event.track_temporary_local_file.assert_called_once_with(str(visual_path))
mock_event.untrack_temporary_local_file.assert_called_once_with(
str(source_path)
)
@@ -1955,11 +1965,11 @@ class TestBuildMainAgent:
@pytest.mark.asyncio
async def test_build_main_agent_skips_caption_when_main_provider_supports_images(
self, mock_event, mock_context, mock_provider
self, mock_event, mock_context, mock_provider, valid_image_path
):
"""Test image-capable chat providers receive quoted images directly."""
module = ama
mock_image = Image(file="file:///tmp/quoted.jpg")
mock_image = Image(file=Path(valid_image_path).as_uri())
mock_reply = Reply(
id="reply-1",
chain=[Plain(text="quoted text"), mock_image],
@@ -1981,7 +1991,7 @@ class TestBuildMainAgent:
patch.object(
Image,
"convert_to_file_path",
AsyncMock(return_value="/tmp/quoted.jpg"),
AsyncMock(return_value=valid_image_path),
),
):
mock_runner = MagicMock()
@@ -2001,7 +2011,7 @@ class TestBuildMainAgent:
)
assert result is not None
assert result.provider_request.image_urls == ["/tmp/quoted.jpg"]
assert result.provider_request.image_urls == [valid_image_path]
assert not any(
"Image Caption" in part.text or "<image_caption>" in part.text
for part in result.provider_request.extra_user_content_parts
@@ -2010,7 +2020,7 @@ class TestBuildMainAgent:
@pytest.mark.asyncio
async def test_build_main_agent_does_not_caption_quoted_image_twice(
self, mock_event, mock_context
self, mock_event, mock_context, valid_image_path
):
"""Quoted images should not be captioned again after request image captioning."""
module = ama
@@ -2028,7 +2038,10 @@ class TestBuildMainAgent:
mock_reply = Reply(
id="reply-1",
chain=[Plain(text="quoted text"), Image(file="file:///tmp/quoted.jpg")],
chain=[
Plain(text="quoted text"),
Image(file=Path(valid_image_path).as_uri()),
],
sender_nickname="Alice",
message_str="quoted text",
)
@@ -2047,7 +2060,7 @@ class TestBuildMainAgent:
patch.object(
Image,
"convert_to_file_path",
AsyncMock(return_value="/tmp/quoted.jpg"),
AsyncMock(return_value=valid_image_path),
),
):
mock_runner = MagicMock()
@@ -2077,7 +2090,7 @@ class TestBuildMainAgent:
@pytest.mark.asyncio
async def test_build_main_agent_does_not_retry_quoted_image_caption_when_empty(
self, mock_event, mock_context
self, mock_event, mock_context, valid_image_path
):
"""Quoted images already sent to image captioning should not be retried."""
module = ama
@@ -2095,7 +2108,10 @@ class TestBuildMainAgent:
mock_reply = Reply(
id="reply-1",
chain=[Plain(text="quoted text"), Image(file="file:///tmp/quoted.jpg")],
chain=[
Plain(text="quoted text"),
Image(file=Path(valid_image_path).as_uri()),
],
sender_nickname="Alice",
message_str="quoted text",
)
@@ -2114,7 +2130,7 @@ class TestBuildMainAgent:
patch.object(
Image,
"convert_to_file_path",
AsyncMock(return_value="/tmp/quoted.jpg"),
AsyncMock(return_value=valid_image_path),
),
):
mock_runner = MagicMock()
@@ -2144,7 +2160,7 @@ class TestBuildMainAgent:
@pytest.mark.asyncio
async def test_build_main_agent_uses_image_fallback_provider(
self, mock_event, mock_context
self, mock_event, mock_context, valid_image_path
):
"""Test image requests use a fallback provider that supports images."""
module = ama
@@ -2166,7 +2182,7 @@ class TestBuildMainAgent:
req = ProviderRequest(
prompt="describe this",
image_urls=["/tmp/image.jpg"],
image_urls=[valid_image_path],
model="text-model",
)
mock_context.get_provider_by_id.side_effect = lambda provider_id: (
@@ -2198,14 +2214,14 @@ class TestBuildMainAgent:
assert result is not None
assert result.provider is image_provider
assert result.provider_request.image_urls == ["/tmp/image.jpg"]
assert result.provider_request.image_urls == [valid_image_path]
assert result.provider_request.model is None
assert mock_runner.reset.call_args.kwargs["provider"] is image_provider
assert mock_runner.reset.call_args.kwargs["fallback_providers"] == []
@pytest.mark.asyncio
async def test_build_main_agent_keeps_text_provider_without_image_fallback(
self, mock_event, mock_context
self, mock_event, mock_context, valid_image_path
):
"""Test image requests fall back to existing sanitizing when no image provider exists."""
module = ama
@@ -2219,7 +2235,7 @@ class TestBuildMainAgent:
req = ProviderRequest(
prompt="describe this",
image_urls=["/tmp/image.jpg"],
image_urls=[valid_image_path],
)
mock_context.get_provider_by_id.return_value = None
mock_context.get_config.return_value = {}
@@ -2250,7 +2266,7 @@ class TestBuildMainAgent:
assert result is not None
assert result.provider is text_provider
assert result.provider_request.image_urls == ["/tmp/image.jpg"]
assert result.provider_request.image_urls == [valid_image_path]
assert mock_runner.reset.call_args.kwargs["provider"] is text_provider
@pytest.mark.asyncio