diff --git a/astrbot/core/agent/context/compressor.py b/astrbot/core/agent/context/compressor.py index 7aaa7adf5..6330b13bf 100644 --- a/astrbot/core/agent/context/compressor.py +++ b/astrbot/core/agent/context/compressor.py @@ -7,14 +7,16 @@ from astrbot.core.utils.image_media_store import ( ImageMediaStore, materialize_image_media_refs, ) -from astrbot.core.utils.media_utils import ImagePayloadTooLargeError +from astrbot.core.utils.media_utils import ( + ImagePayloadTooLargeError, + get_image_preparation_options, +) from ...provider.modalities import ( log_context_sanitize_stats, sanitize_contexts_by_modalities, ) from ..message import Message -from .image_budget import get_image_encoded_byte_limit, validate_context_image_bytes from .token_counter import EstimateTokenCounter, TokenCounter if TYPE_CHECKING: @@ -286,14 +288,13 @@ class LLMSummaryCompressor: # Generate summary try: - limit = get_image_encoded_byte_limit( - getattr(self.provider, "provider_settings", {}) - ) - validate_context_image_bytes(sanitized_summary_contexts, limit) sanitized_summary_contexts = await materialize_image_media_refs( sanitized_summary_contexts, self.image_media_store or ImageMediaStore(Path(get_astrbot_data_path()) / "media"), + options=get_image_preparation_options( + getattr(self.provider, "provider_settings", {}) + ), ) response = await self.provider.text_chat( contexts=sanitized_summary_contexts, diff --git a/astrbot/core/agent/context/image_budget.py b/astrbot/core/agent/context/image_budget.py deleted file mode 100644 index d02eb807f..000000000 --- a/astrbot/core/agent/context/image_budget.py +++ /dev/null @@ -1,87 +0,0 @@ -"""Validate image bytes independently from text-token estimates.""" - -from collections.abc import Sequence - -from astrbot.core.agent.message import ImageMediaRefPart, ImageURLPart, Message -from astrbot.core.utils.media_utils import ( - IMAGE_COMPRESS_DEFAULT_MAX_ENCODED_BYTES, - ImagePayloadTooLargeError, -) - - -def get_image_encoded_byte_limit(provider_settings: dict | None) -> int: - """Read the configured per-image limit with the same rules for every request. - - Args: - provider_settings: Settings belonging to the actual request provider. - - Returns: - A positive byte limit, or the default when the setting is invalid. - """ - options = ( - provider_settings.get("image_compress_options", {}) - if isinstance(provider_settings, dict) - else {} - ) - limit = options.get("max_encoded_bytes") if isinstance(options, dict) else None - if isinstance(limit, int) and not isinstance(limit, bool) and limit > 0: - return limit - return IMAGE_COMPRESS_DEFAULT_MAX_ENCODED_BYTES - - -def validate_context_image_bytes( - messages: Sequence[Message | dict], - max_encoded_bytes: int = IMAGE_COMPRESS_DEFAULT_MAX_ENCODED_BYTES, -) -> int: - """Validate selected images without loading or rewriting historical bytes. - - Args: - messages: Selected main or summary request messages. - max_encoded_bytes: Maximum Base64 bytes for a single image. - - Returns: - Total known encoded-image bytes, excluding JSON and data URI headers. - Remote URLs have unknown size until resolved and are not counted. - - Raises: - ImagePayloadTooLargeError: A selected image exceeds the single-image cap. - ValueError: The configured cap is invalid. - """ - if isinstance(max_encoded_bytes, bool) or max_encoded_bytes < 1: - raise ValueError("Image byte budget must be a positive integer") - total = 0 - for message in messages: - parts = ( - message.content if isinstance(message, Message) else message.get("content") - ) - if not isinstance(parts, list): - continue - for part in parts: - size = 0 - url = None - if isinstance(part, ImageMediaRefPart): - size = 4 * ((part.byte_size + 2) // 3) - elif isinstance(part, ImageURLPart): - url = part.image_url.url - elif isinstance(part, dict): - if part.get("type") == "image_media_ref": - size = 4 * ((int(part["byte_size"]) + 2) // 3) - elif part.get("type") == "image_url": - image_url = part.get("image_url") - url = ( - image_url.get("url") - if isinstance(image_url, dict) - else image_url - ) - if isinstance(url, str) and url.startswith("data:image/"): - comma = url.find(",") - if comma >= 0 and ";base64" in url[:comma]: - # Do not slice the full Base64 suffix merely to count it. - size = len(url) - comma - 1 - if size > max_encoded_bytes: - raise ImagePayloadTooLargeError( - f"A selected image uses {size} Base64 bytes, exceeding the " - f"{max_encoded_bytes}-byte limit. Historical images were not changed." - ) - total += size - return total diff --git a/astrbot/core/agent/runners/coze/coze_agent_runner.py b/astrbot/core/agent/runners/coze/coze_agent_runner.py index a4054d9ce..59e5aacff 100644 --- a/astrbot/core/agent/runners/coze/coze_agent_runner.py +++ b/astrbot/core/agent/runners/coze/coze_agent_runner.py @@ -1,13 +1,11 @@ import json import sys import typing as T -from dataclasses import replace from pathlib import Path import astrbot.core.message.components as Comp from astrbot import logger from astrbot.core import sp -from astrbot.core.agent.context.image_budget import validate_context_image_bytes from astrbot.core.message.message_event_result import MessageChain from astrbot.core.provider.entities import ( LLMResponse, @@ -169,18 +167,12 @@ class CozeAgentRunner(BaseAgentRunner[TContext]): image_options = ( self.req.image_preparation_options or ImagePreparationOptions() ) - image_limit = image_options.max_encoded_bytes - if image_limit is None: - image_limit = ImagePreparationOptions().max_encoded_bytes - validate_context_image_bytes( - contexts, - image_limit, - ) contexts = await materialize_image_media_refs( contexts, ImageMediaStore(Path(get_astrbot_data_path()) / "media"), + options=image_options, ) - history_image_options = replace(image_options, enabled=False) + history_image_options = image_options for ctx in contexts: if is_checkpoint_message(ctx): continue @@ -255,6 +247,10 @@ class CozeAgentRunner(BaseAgentRunner[TContext]): file_id = await self._download_and_upload_image( url, session_id, + image_options=( + self.req.image_preparation_options + or ImagePreparationOptions() + ), ) object_string_content.append( { diff --git a/astrbot/core/agent/runners/deerflow/deerflow_agent_runner.py b/astrbot/core/agent/runners/deerflow/deerflow_agent_runner.py index 9ef9c9cb2..37174f8f2 100644 --- a/astrbot/core/agent/runners/deerflow/deerflow_agent_runner.py +++ b/astrbot/core/agent/runners/deerflow/deerflow_agent_runner.py @@ -512,8 +512,9 @@ class DeerFlowAgentRunner(BaseAgentRunner[TContext]): """ runtime_configurable = self._build_runtime_configurable(thread_id) - image_options = getattr( - getattr(self, "req", None), "image_preparation_options", None + image_options = ( + getattr(getattr(self, "req", None), "image_preparation_options", None) + or ImagePreparationOptions() ) return { "assistant_id": self.assistant_id, diff --git a/astrbot/core/agent/runners/dify/dify_agent_runner.py b/astrbot/core/agent/runners/dify/dify_agent_runner.py index c4af5dd90..d07257798 100644 --- a/astrbot/core/agent/runners/dify/dify_agent_runner.py +++ b/astrbot/core/agent/runners/dify/dify_agent_runner.py @@ -10,6 +10,7 @@ from astrbot.core.provider.entities import ( ) from astrbot.core.utils.media_utils import ( ImagePayloadTooLargeError, + ImagePreparationOptions, MediaResolver, resolve_media_ref_to_base64_data, ) @@ -117,8 +118,9 @@ class DifyAgentRunner(BaseAgentRunner[TContext]): image_url: str, session_id: str, ) -> dict[str, str] | None: - image_options = getattr( - getattr(self, "req", None), "image_preparation_options", None + image_options = ( + getattr(getattr(self, "req", None), "image_preparation_options", None) + or ImagePreparationOptions() ) image_data = await resolve_media_ref_to_base64_data( image_url, diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index 5b5b28637..99b9b2c4f 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -61,10 +61,6 @@ from astrbot.core.utils.media_utils import ( from ..context.compressor import ContextCompressor from ..context.config import ContextConfig -from ..context.image_budget import ( - get_image_encoded_byte_limit, - validate_context_image_bytes, -) from ..context.manager import ContextManager from ..context.token_counter import EstimateTokenCounter, TokenCounter from ..hooks import BaseAgentRunHooks @@ -550,18 +546,21 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): Provider-compatible messages without changing persisted history. Raises: - ImagePayloadTooLargeError: An image exceeds the encoded-byte limit. + ImagePayloadTooLargeError: An image exceeds the model input limit. MemoryError: Image materialization exhausts available memory. """ selected_contexts = self._sanitize_contexts_for_provider(contexts) - limit = get_image_encoded_byte_limit( - getattr(self.provider, "provider_settings", {}) + image_options = ( + self.req.image_preparation_options + or get_image_preparation_options( + getattr(self.provider, "provider_settings", {}) + ) ) - validate_context_image_bytes(selected_contexts, limit) return await materialize_image_media_refs( selected_contexts, self.image_media_store or ImageMediaStore(Path(get_astrbot_data_path()) / "media"), + options=image_options, ) async def _iter_llm_responses( @@ -1200,13 +1199,9 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): else "tool_image" ), ), - options=replace( - self.req.image_preparation_options - or get_image_preparation_options( - getattr(self.provider, "provider_settings", {}) - ), - preserve_dimensions=cached_img.tool_name - == "astrbot_cua_screenshot", + options=self.req.image_preparation_options + or get_image_preparation_options( + getattr(self.provider, "provider_settings", {}) ), ) if img_data: diff --git a/astrbot/core/astr_main_agent.py b/astrbot/core/astr_main_agent.py index 91eb13f4e..dea2ea577 100644 --- a/astrbot/core/astr_main_agent.py +++ b/astrbot/core/astr_main_agent.py @@ -32,6 +32,7 @@ from astrbot.core.astr_main_agent_resources import ( from astrbot.core.computer.booters.local import resolve_windows_shell from astrbot.core.conversation_mgr import Conversation from astrbot.core.db import BaseDatabase +from astrbot.core.exceptions import ProviderRequestTooLargeError from astrbot.core.message.components import File, Image, Record, Reply, Video from astrbot.core.persona_error_reply import ( extract_persona_custom_error_message_from_persona, @@ -110,6 +111,7 @@ 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 ( + ImagePayloadTooLargeError, is_file_uri, is_recoverable_image_error, normalize_model_image_max_size, @@ -789,6 +791,11 @@ async def _ensure_img_caption( ) return image_refs except Exception as exc: # noqa: BLE001 + if isinstance( + exc, + (MemoryError, ImagePayloadTooLargeError, ProviderRequestTooLargeError), + ): + raise logger.error("处理图片描述失败: %s", exc) req.extra_user_content_parts.append(TextPart(text="[Image Captioning Failed]")) finally: diff --git a/astrbot/core/provider/sources/anthropic_source.py b/astrbot/core/provider/sources/anthropic_source.py index 03f33867e..88a92baea 100644 --- a/astrbot/core/provider/sources/anthropic_source.py +++ b/astrbot/core/provider/sources/anthropic_source.py @@ -13,10 +13,6 @@ from anthropic.types.usage import Usage from astrbot import logger from astrbot.api.provider import Provider -from astrbot.core.agent.context.image_budget import ( - get_image_encoded_byte_limit, - validate_context_image_bytes, -) from astrbot.core.agent.message import AudioURLPart, ContentPart, ImageURLPart, TextPart from astrbot.core.exceptions import EmptyModelOutputError from astrbot.core.provider.entities import LLMResponse, TokenUsage @@ -798,11 +794,10 @@ class ProviderAnthropic(Provider): contexts, _ = sanitize_contexts_by_modalities( contexts, self.provider_config.get("modalities") ) - validate_context_image_bytes( - contexts, get_image_encoded_byte_limit(self.provider_settings) - ) contexts = await materialize_image_media_refs( - contexts, ImageMediaStore(Path(get_astrbot_data_path()) / "media") + contexts, + ImageMediaStore(Path(get_astrbot_data_path()) / "media"), + options=get_image_preparation_options(self.provider_settings), ) new_record = None if prompt is not None: @@ -880,11 +875,10 @@ class ProviderAnthropic(Provider): contexts, _ = sanitize_contexts_by_modalities( contexts, self.provider_config.get("modalities") ) - validate_context_image_bytes( - contexts, get_image_encoded_byte_limit(self.provider_settings) - ) contexts = await materialize_image_media_refs( - contexts, ImageMediaStore(Path(get_astrbot_data_path()) / "media") + contexts, + ImageMediaStore(Path(get_astrbot_data_path()) / "media"), + options=get_image_preparation_options(self.provider_settings), ) new_record = None if prompt is not None: diff --git a/astrbot/core/provider/sources/gemini_source.py b/astrbot/core/provider/sources/gemini_source.py index a244748d7..992f24694 100644 --- a/astrbot/core/provider/sources/gemini_source.py +++ b/astrbot/core/provider/sources/gemini_source.py @@ -4,7 +4,6 @@ import json import logging import random from collections.abc import AsyncGenerator -from dataclasses import replace from pathlib import Path from typing import Literal, cast @@ -16,10 +15,6 @@ from google.genai.errors import APIError import astrbot.core.message.components as Comp from astrbot import logger from astrbot.api.provider import Provider -from astrbot.core.agent.context.image_budget import ( - get_image_encoded_byte_limit, - validate_context_image_bytes, -) from astrbot.core.agent.message import AudioURLPart, ContentPart, ImageURLPart, TextPart from astrbot.core.exceptions import EmptyModelOutputError, ProviderRequestTooLargeError from astrbot.core.message.message_event_result import MessageChain @@ -323,13 +318,11 @@ class ProviderGoogleGenAI(Provider): messages, _ = sanitize_contexts_by_modalities( payloads.get("messages", []), provider_config.get("modalities") ) - validate_context_image_bytes( - messages, get_image_encoded_byte_limit(provider_settings) - ) payloads["messages"] = messages payloads["messages"] = await materialize_image_media_refs( payloads.get("messages", []), ImageMediaStore(Path(get_astrbot_data_path()) / "media"), + options=get_image_preparation_options(provider_settings), ) def create_text_part(text: str) -> types.Part: @@ -360,10 +353,7 @@ class ProviderGoogleGenAI(Provider): url, media_type="image", strict=True, - image_options=replace( - get_image_preparation_options(provider_settings), - enabled=False, - ), + image_options=get_image_preparation_options(provider_settings), ) if image_data is None: raise ValueError( diff --git a/astrbot/core/provider/sources/openai_responses_source.py b/astrbot/core/provider/sources/openai_responses_source.py index 76a19040e..d37174c16 100644 --- a/astrbot/core/provider/sources/openai_responses_source.py +++ b/astrbot/core/provider/sources/openai_responses_source.py @@ -9,10 +9,6 @@ from openai.types.responses import Response import astrbot.core.message.components as Comp from astrbot import logger -from astrbot.core.agent.context.image_budget import ( - get_image_encoded_byte_limit, - validate_context_image_bytes, -) from astrbot.core.agent.message import ContentPart, Message from astrbot.core.agent.tool import ToolSet from astrbot.core.exceptions import EmptyModelOutputError @@ -24,6 +20,7 @@ from astrbot.core.utils.image_media_store import ( ImageMediaStore, materialize_image_media_refs, ) +from astrbot.core.utils.media_utils import get_image_preparation_options from ..register import register_provider_adapter from .openai_source import ProviderOpenAIOfficial @@ -274,12 +271,10 @@ class ProviderOpenAIResponses(ProviderOpenAIOfficial): contexts, _ = sanitize_contexts_by_modalities( contexts, self.provider_config.get("modalities") ) - validate_context_image_bytes( - contexts, get_image_encoded_byte_limit(self.provider_settings) - ) contexts = await materialize_image_media_refs( contexts, ImageMediaStore(Path(get_astrbot_data_path()) / "media"), + options=get_image_preparation_options(self.provider_settings), ) context_query = copy.deepcopy(self._ensure_message_to_dicts(contexts)) if prompt is not None: diff --git a/astrbot/core/provider/sources/openai_source.py b/astrbot/core/provider/sources/openai_source.py index b46d8b95f..68769a872 100644 --- a/astrbot/core/provider/sources/openai_source.py +++ b/astrbot/core/provider/sources/openai_source.py @@ -5,7 +5,6 @@ import json import random import re from collections.abc import AsyncGenerator -from dataclasses import replace from pathlib import Path from typing import Any, Literal @@ -20,10 +19,6 @@ from openai.types.completion_usage import CompletionUsage import astrbot.core.message.components as Comp from astrbot import logger from astrbot.api.provider import Provider -from astrbot.core.agent.context.image_budget import ( - get_image_encoded_byte_limit, - validate_context_image_bytes, -) from astrbot.core.agent.message import ( AudioURLPart, ContentPart, @@ -355,9 +350,7 @@ class ProviderOpenAIOfficial(Provider): async def _materialize_context_image_parts( self, context_query: list[dict] ) -> list[dict]: - options = replace( - get_image_preparation_options(self.provider_settings), enabled=False - ) + options = get_image_preparation_options(self.provider_settings) return [ await self._materialize_message_image_parts(message, options=options) for message in context_query @@ -1015,11 +1008,10 @@ class ProviderOpenAIOfficial(Provider): contexts, _ = sanitize_contexts_by_modalities( contexts, self.provider_config.get("modalities") ) - validate_context_image_bytes( - contexts, get_image_encoded_byte_limit(self.provider_settings) - ) contexts = await materialize_image_media_refs( - contexts, ImageMediaStore(Path(get_astrbot_data_path()) / "media") + contexts, + ImageMediaStore(Path(get_astrbot_data_path()) / "media"), + options=get_image_preparation_options(self.provider_settings), ) new_record = None if prompt is not None: diff --git a/astrbot/core/utils/image_input.py b/astrbot/core/utils/image_input.py index f3f9558b5..c92cadd28 100644 --- a/astrbot/core/utils/image_input.py +++ b/astrbot/core/utils/image_input.py @@ -75,7 +75,7 @@ async def prepare_request_images( ) except ImageInputTooLargeError as exc: image = None - original_path = str(exc) + original_path = exc.path status = "oversized" if image: path, is_montage, needs_cleanup, original_path = image diff --git a/astrbot/core/utils/image_media_store.py b/astrbot/core/utils/image_media_store.py index e33c4ac46..d5eb2c098 100644 --- a/astrbot/core/utils/image_media_store.py +++ b/astrbot/core/utils/image_media_store.py @@ -13,7 +13,12 @@ from pathlib import Path from PIL import Image -from astrbot.core.utils.media_utils import validate_image_input_size +from astrbot.core.utils.media_utils import ( + ImagePayloadTooLargeError, + ImagePreparationOptions, + prepare_image_source, + validate_image_input_size, +) @dataclass(frozen=True, slots=True) @@ -91,6 +96,7 @@ class ImageMediaStore: OSError: The object or metadata cannot be committed atomically. ValueError: The stored bytes are not a readable image. """ + validate_image_input_size(data) media_id = hashlib.sha256(data).hexdigest() self.root.mkdir(parents=True, exist_ok=True) if self.root.is_symlink() or not self.root.is_dir(): @@ -301,6 +307,7 @@ async def materialize_image_media_refs( store: ImageMediaStore, *, strict: bool = False, + options: ImagePreparationOptions | None = None, ) -> list: """Resolve only the selected references into a request-local view. @@ -308,6 +315,8 @@ async def materialize_image_media_refs( contexts: Messages selected by the context manager. store: Store for the application's configured data root. strict: Fail on missing media when preparing a rollback export. + options: Mainline model-image preparation options. When omitted, return + the exact stored bytes for exports and dashboard previews. Returns: Request-local messages with images, or the original list if no refs exist. @@ -326,14 +335,54 @@ async def materialize_image_media_refs( serialized = message.model_dump() if isinstance(message, Message) else message parts = serialized.get("content") if not isinstance(parts, list) or not any( - isinstance(part, dict) and part.get("type") == "image_media_ref" + isinstance(part, dict) + and ( + part.get("type") == "image_media_ref" + or (options is not None and part.get("type") == "image_url") + ) for part in parts ): output.append(message) continue resolved = [] for part in parts: - if not isinstance(part, dict) or part.get("type") != "image_media_ref": + if not isinstance(part, dict): + resolved.append(part) + continue + if part.get("type") == "image_url" and options is not None: + image_url = part.get("image_url") + url = image_url.get("url") if isinstance(image_url, dict) else None + if not isinstance(url, str) or not url: + resolved.append(part) + continue + try: + image_data = await prepare_image_source( + url, + options=options, + default_mime_type=( + image_url.get("mime_type") + if isinstance(image_url, dict) + else None + ), + ) + resolved.append( + { + **part, + "image_url": { + **image_url, + "url": image_data.to_data_url(), + }, + } + ) + except (ImagePayloadTooLargeError, MemoryError): + raise + except (TypeError, ValueError, OSError): + if strict: + raise + logger.warning("A selected conversation image is unavailable") + resolved.append({"type": "text", "text": "[Image unavailable]"}) + continue + if part.get("type") != "image_media_ref": resolved.append(part) continue try: @@ -347,16 +396,27 @@ async def materialize_image_media_refs( part.get("version", 1), part.get("image_id"), ) + validate_image_input_size(ref.byte_size) payload = await asyncio.to_thread(store.read, ref, {ref.media_id}) - image_url = { - "url": f"data:{ref.mime_type};base64,{base64.b64encode(payload).decode('ascii')}", - } + if options is None: + image_url = { + "url": f"data:{ref.mime_type};base64,{base64.b64encode(payload).decode('ascii')}", + } + else: + image_data = await prepare_image_source( + payload, + options=options, + default_mime_type=ref.mime_type, + ) + image_url = {"url": image_data.to_data_url()} del payload if ref.detail is not None: image_url["detail"] = ref.detail if ref.image_id is not None: image_url["id"] = ref.image_id resolved.append({"type": "image_url", "image_url": image_url}) + except (ImagePayloadTooLargeError, MemoryError): + raise except (KeyError, TypeError, ValueError, OSError): if strict: raise diff --git a/astrbot/core/utils/media_utils.py b/astrbot/core/utils/media_utils.py index 321a6d17e..ea7c0925f 100644 --- a/astrbot/core/utils/media_utils.py +++ b/astrbot/core/utils/media_utils.py @@ -44,6 +44,44 @@ IMAGE_COMPRESS_DEFAULT_MIN_FILE_SIZE_MB = 1.0 # Model image inputs larger than this are skipped before decoding. MODEL_IMAGE_MAX_INPUT_BYTES = 64 * 1024 * 1024 + +class ImagePayloadTooLargeError(ValueError): + """Raised when an image exceeds a safe input or output byte budget.""" + + +@dataclass(slots=True) +class ImagePreparationOptions: + """Options for the shared provider-facing image preparation boundary. + + Args: + max_size: Longest edge of the prepared image in pixels. + """ + + max_size: int = IMAGE_COMPRESS_DEFAULT_MAX_SIZE + + +def get_image_preparation_options( + provider_settings: dict | None, +) -> ImagePreparationOptions: + """Build preparation options from the current provider settings. + + Args: + provider_settings: Provider-level settings. Only the mainline image + dimension option is read; removed legacy byte-budget settings are + intentionally ignored. + + Returns: + Validated options for the shared image preparation path. + """ + raw_options = ( + provider_settings.get("image_compress_options", {}) + if isinstance(provider_settings, dict) + else {} + ) + max_size = raw_options.get("max_size") if isinstance(raw_options, dict) else None + return ImagePreparationOptions(max_size=normalize_model_image_max_size(max_size)) + + MEDIA_MIME_EXTENSIONS = { "audio/wav": ".wav", "audio/wave": ".wav", @@ -120,6 +158,21 @@ Examples: """ +@dataclass(frozen=True, slots=True) +class ImagePreparationInput: + """Describe an image entering the shared preparation boundary. + + Args: + value: Image path, URL, data URI, base64 reference, or raw bytes. + source_kind: Logical producer name used for diagnostics. + cleanup_paths: Temporary paths owned by the caller. + """ + + value: MediaRefStr | bytes + source_kind: str = "unknown" + cleanup_paths: tuple[Path, ...] = () + + @dataclass(slots=True) class ResolvedMediaData: """Base64 media bytes plus the metadata needed by provider payloads. @@ -133,6 +186,7 @@ class ResolvedMediaData: base64_data: str mime_type: str format: str | None = None + byte_size: int | None = None def to_bytes(self) -> bytes: """Decode the base64 payload, accepting missing padding.""" @@ -322,6 +376,133 @@ def _decode_base64_payload( raise ValueError(error_message) from exc +def _estimate_base64_decoded_size( + payload: str, + *, + start: int = 0, + require_valid_chars: bool = False, +) -> int | None: + """Estimate decoded bytes without creating a compact payload copy. + + Args: + payload: Base64 text, possibly containing whitespace. + start: Index at which the payload begins. + require_valid_chars: Whether to reject non-base64 characters. + + Returns: + Estimated decoded byte count, or ``None`` when the payload is not a + complete standard Base64 string. + """ + encoded_size = 0 + padding_size = 0 + for index, char in enumerate(payload): + if index < start or char.isspace(): + continue + if require_valid_chars and char not in ( + "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=" + ): + return None + encoded_size += 1 + if char == "=": + padding_size += 1 + if encoded_size % 4 == 1: + return None + return max(0, encoded_size * 3 // 4 - padding_size) + + +def validate_image_input_size(image_source: bytes | str | Path | int) -> int | None: + """Reject known oversized image inputs before decoding them. + + Args: + image_source: Raw bytes, a local path, a media reference, or a known + byte count. Remote URLs return ``None`` until downloaded. + + Returns: + Known encoded byte count, or ``None`` when it cannot be determined + without reading the source. + + Raises: + ImagePayloadTooLargeError: The known source exceeds the 64 MiB input cap. + TypeError: The source type is unsupported. + ValueError: The source size is negative. + """ + source_size: int | None = None + if isinstance(image_source, bool): + raise TypeError("Image source size must not be a boolean") + if isinstance(image_source, int): + source_size = image_source + elif isinstance(image_source, bytes): + source_size = len(image_source) + elif isinstance(image_source, Path): + try: + source_size = image_source.stat().st_size + except FileNotFoundError: + return None + elif isinstance(image_source, str): + if image_source.startswith(("http://", "https://")): + return None + if image_source.startswith("data:"): + comma_index = image_source.find(",") + if comma_index < 0: + return None + header_parts = image_source[5:comma_index].split(";") + if any(part.lower() == "base64" for part in header_parts[1:]): + source_size = _estimate_base64_decoded_size( + image_source, + start=comma_index + 1, + ) + elif image_source.startswith("base64://"): + source_size = _estimate_base64_decoded_size( + image_source, + start=len("base64://"), + ) + else: + is_uri = is_file_uri(image_source) + path = ( + Path(file_uri_to_path(image_source)) if is_uri else Path(image_source) + ) + try: + source_size = path.stat().st_size + except (FileNotFoundError, OSError, ValueError): + source_size = None + if source_size is None and not is_uri: + source_size = _estimate_base64_decoded_size( + image_source, + require_valid_chars=True, + ) + else: + raise TypeError(f"Unsupported image source type: {type(image_source).__name__}") + + if source_size is None: + return None + if source_size < 0: + raise ValueError("Image source size must not be negative") + if source_size > MODEL_IMAGE_MAX_INPUT_BYTES: + if isinstance(image_source, Path): + raise ImageInputTooLargeError(str(image_source)) + raise ImageInputTooLargeError( + "Image input exceeds the " + f"{MODEL_IMAGE_MAX_INPUT_BYTES}-byte limit ({source_size} bytes)" + ) + return source_size + + +def _encode_file_to_base64(path: Path) -> str: + """Encode a local file in chunks without retaining raw and encoded copies.""" + encoded = io.StringIO() + remainder = b"" + with path.open("rb") as source: + while chunk := source.read(1023 * 1024): + block = remainder + chunk if remainder else chunk + complete_size = len(block) - len(block) % 3 + if complete_size: + encoded.write(base64.b64encode(block[:complete_size]).decode("ascii")) + remainder = block[complete_size:] + if remainder: + encoded.write(base64.b64encode(remainder).decode("ascii")) + return encoded.getvalue() + + def describe_media_ref(media_ref: object | None) -> str: """Return a log-safe description of a media reference. @@ -486,6 +667,11 @@ async def _materialize_media_ref( cleanup_paths.append(target_path) try: await download_file(media_ref, str(target_path)) + if media_type == "image": + validate_image_input_size(target_path) + except ImageInputTooLargeError as exc: + exc.path = str(target_path) + raise except Exception: _cleanup_paths(cleanup_paths) raise @@ -514,6 +700,8 @@ async def _materialize_media_ref( return _LocalMediaFile(path=path, mime_type=_guess_mime_type(path)) if media_ref.startswith("data:"): + if media_type == "image": + validate_image_input_size(media_ref) mime_type, media_bytes = _parse_base64_data_uri(media_ref) target_suffix = _extension_from_mime_type(mime_type) or suffix if media_type == "image" and target_suffix == suffix: @@ -538,6 +726,8 @@ async def _materialize_media_ref( ) if media_ref.startswith("base64://"): + if media_type == "image": + validate_image_input_size(media_ref) media_bytes = _decode_base64_payload( media_ref.removeprefix("base64://"), error_message="invalid base64 media payload", @@ -572,6 +762,8 @@ async def _materialize_media_ref( if path_exists: return _LocalMediaFile(path=path, mime_type=_guess_mime_type(path)) + if media_type == "image": + validate_image_input_size(media_ref) compact_media_ref = "".join(media_ref.split()) if compact_media_ref: try: @@ -858,6 +1050,7 @@ class MediaResolver: return ResolvedMediaData( base64_data=base64.b64encode(media_bytes).decode("utf-8"), mime_type=mime_type, + byte_size=len(media_bytes), ) async with self.as_path( @@ -913,29 +1106,43 @@ class MediaResolver: async def resolve_image_ref_to_base64_data( - image_ref: MediaRefStr, + image_ref: MediaRefStr | bytes, *, strict: bool = False, default_mime_type: str | None = "image/jpeg", + options: ImagePreparationOptions | None = None, ) -> ResolvedMediaData | None: - """Resolve an image reference to losslessly encoded base64 data. + """Resolve and prepare an image reference for a provider request. - Only materializes the source and detects its MIME type; no - provider-specific format conversion, frame extraction, or montage is - performed here. Platform senders and generic request assembly rely on - this to encode image bytes without transforming their content. + When ``options`` is provided, the mainline preparation path bounds a + model-facing image while keeping historical media references outside the + in-memory request until selected. Without options, preserve the resolver's + original byte-for-byte behavior for platform and compatibility callers. ``strict=False`` returns ``None`` for invalid images so payload assembly can skip bad image refs without failing the whole request. """ - return await MediaResolver( - image_ref, - media_type="image", - default_suffix=".bin", - ).to_base64_data( - strict=strict, - default_mime_type=default_mime_type, - ) + try: + if options is None: + return await MediaResolver( + image_ref, + media_type="image", + default_suffix=".bin", + ).to_base64_data( + strict=strict, + default_mime_type=default_mime_type, + ) + return await prepare_image_source( + image_ref, + options=options, + default_mime_type=default_mime_type, + ) + except (ImagePayloadTooLargeError, MemoryError): + raise + except Exception: + if strict: + raise + return None def is_recoverable_image_error(error: Exception) -> bool: @@ -1198,9 +1405,14 @@ def _extract_animation_montage_sync(image: PILImage.Image, max_size: int) -> byt return encoded -class ImageInputTooLargeError(ValueError): +class ImageInputTooLargeError(ImagePayloadTooLargeError): """Raised with the retained source path when an image exceeds the input cap.""" + def __init__(self, path: str) -> None: + """Keep the source path for the file-reading fallback.""" + super().__init__(path) + self.path = path + async def prepare_model_image( image_ref: str, @@ -1233,17 +1445,19 @@ async def prepare_model_image( input_size, source.path, ) + original_path = str(source.path) source.detach() - raise ImageInputTooLargeError(str(source.path)) + raise ImageInputTooLargeError(original_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) + needs_cleanup = bool(source.cleanup_paths) # Final image labels expose this original for later file-tool access. source.detach() if converted_bytes is image_bytes: - return original_path, is_montage, False, original_path + return original_path, is_montage, needs_cleanup, 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) @@ -1259,7 +1473,7 @@ async def prepare_model_image( output_path.unlink(missing_ok=True) raise return str(output_path), is_montage, True, original_path - except ImageInputTooLargeError: + except (ImageInputTooLargeError, ImagePayloadTooLargeError): raise except Exception as exc: if not is_recoverable_image_error(exc): @@ -1297,31 +1511,6 @@ async def resolve_audio_ref_to_base64_data( return audio_data -async def resolve_media_ref_to_base64_data( - media_ref: MediaRefStr, - *, - media_type: str, - strict: bool = False, -) -> ResolvedMediaData | None: - """Resolve a media reference to base64 data through one shared entrypoint. - - This helper keeps provider sources from knowing whether a reference is local, - HTTP(S), ``base64://``, a data URI, or a legacy bare base64 payload. - """ - - if media_type == "image": - return await resolve_image_ref_to_base64_data(media_ref, strict=strict) - if media_type == "audio": - return await resolve_audio_ref_to_base64_data(media_ref) - - return await MediaResolver( - media_ref, - media_type=media_type, - ).to_base64_data( - strict=strict, - ) - - async def get_media_duration(file_path: str) -> int | None: """Probe media duration with ffprobe. @@ -1920,6 +2109,8 @@ async def compress_image( url_or_path: str, max_size: int = IMAGE_COMPRESS_DEFAULT_MAX_SIZE, quality: int = IMAGE_COMPRESS_DEFAULT_QUALITY, + *, + optimize: bool = IMAGE_COMPRESS_DEFAULT_OPTIMIZE, ) -> str: """Compress large user-uploaded images. @@ -1927,6 +2118,7 @@ async def compress_image( url_or_path: Image path or URL. max_size: Longest edge of the compressed image in pixels. quality: JPEG output quality in the range 1-100. + optimize: Whether Pillow should optimize the encoded output. Returns: The compressed image path. Returns the original path if compression @@ -1934,7 +2126,6 @@ async def compress_image( """ max_size = max(int(max_size), 1) quality = min(max(int(quality), 1), 100) - optimize = IMAGE_COMPRESS_DEFAULT_OPTIMIZE min_file_size_bytes = int(IMAGE_COMPRESS_DEFAULT_MIN_FILE_SIZE_MB * 1024 * 1024) image_source: bytes | Path | None = None @@ -1950,6 +2141,7 @@ async def compress_image( if url_or_path.startswith("http"): return url_or_path elif url_or_path.startswith("data:image"): + validate_image_input_size(url_or_path) _header, encoded = url_or_path.split(",", 1) image_source = _decode_base64_payload( encoded, @@ -1963,9 +2155,9 @@ async def compress_image( local_path = Path(url_or_path) if not local_path.exists(): return url_or_path - if local_path.stat().st_size < min_file_size_bytes and not _exceeds_max_size( - local_path - ): + validate_image_input_size(local_path) + source_size = local_path.stat().st_size + if source_size < min_file_size_bytes and not _exceeds_max_size(local_path): return url_or_path image_source = local_path @@ -1985,3 +2177,137 @@ async def compress_image( optimize, ) return compressed_path or url_or_path + + +async def prepare_image_source( + image_ref: MediaRefStr | bytes | ImagePreparationInput, + *, + options: ImagePreparationOptions | None = None, + default_mime_type: str | None = "image/jpeg", +) -> ResolvedMediaData: + """Prepare one image through the mainline model-image pipeline. + + Args: + image_ref: Local path, URL, data URI, base64 reference, or raw bytes. + options: Optional dimension settings for the request. + default_mime_type: Fallback MIME type when detection is unavailable. + + Returns: + Provider-ready Base64 data and its MIME type. + + Raises: + ImagePayloadTooLargeError: The input or prepared image exceeds a safe + byte budget. + ValueError: The source is not a readable image. + """ + selected = options or ImagePreparationOptions() + preparation_input = ( + image_ref + if isinstance(image_ref, ImagePreparationInput) + else ImagePreparationInput(image_ref) + ) + source_ref = preparation_input.value + try: + if isinstance(source_ref, bytes): + validate_image_input_size(source_ref) + prepared_bytes, _ = await asyncio.to_thread( + _prepare_model_image_sync, + source_ref, + max(selected.max_size, 1), + ) + mime_type = detect_image_mime_type( + prepared_bytes, + default_mime_type=default_mime_type, + ) + if not mime_type: + raise ValueError("image content could not be identified") + return ResolvedMediaData( + base64_data=base64.b64encode(prepared_bytes).decode("ascii"), + mime_type=mime_type, + byte_size=len(prepared_bytes), + ) + + if not isinstance(source_ref, str): + raise TypeError("image reference must be a string or bytes") + + prepared = await prepare_model_image( + source_ref, + max_size=max(selected.max_size, 1), + output_dir=Path(get_astrbot_temp_path()), + ) + if prepared is None: + # Preserve the historical opaque-base64 fallback for callers that + # intentionally pass bytes Pillow cannot identify as an image. + resolved = await MediaResolver( + source_ref, + media_type="image", + default_suffix=".bin", + ).to_base64_data( + strict=True, + default_mime_type=default_mime_type, + ) + if resolved is not None: + if ( + resolved.byte_size is not None + and resolved.byte_size >= MODEL_IMAGE_MAX_BYTES + ): + raise ImagePayloadTooLargeError( + "Prepared image exceeds the model input byte limit" + ) + return resolved + raise ValueError(f"Invalid image file: {describe_media_ref(source_ref)}") + prepared_path, _is_montage, needs_cleanup, _original_path = prepared + path = Path(prepared_path) + try: + mime_type = await detect_image_mime_type_async( + path, + default_mime_type=default_mime_type, + ) + if not mime_type: + raise ValueError("image content could not be identified") + encoded_data = await asyncio.to_thread(_encode_file_to_base64, path) + byte_size = path.stat().st_size + return ResolvedMediaData( + base64_data=encoded_data, + mime_type=mime_type, + byte_size=byte_size, + ) + finally: + if needs_cleanup: + path.unlink(missing_ok=True) + finally: + for cleanup_path in preparation_input.cleanup_paths: + cleanup_path.unlink(missing_ok=True) + + +async def resolve_media_ref_to_base64_data( + media_ref: MediaRefStr | bytes, + *, + media_type: str, + strict: bool = False, + image_options: ImagePreparationOptions | None = None, + default_mime_type: str | None = "image/jpeg", +) -> ResolvedMediaData | None: + """Resolve a media reference through the shared provider boundary. + + Args: + media_ref: Media path, URL, data URI, Base64 reference, or raw bytes. + media_type: Logical media family. + strict: Whether ordinary resolution failures should propagate. + image_options: Optional image preparation settings. + default_mime_type: MIME fallback for image data. + + Returns: + Resolved media data, or ``None`` for a recoverable non-strict failure. + """ + if media_type == "image": + return await resolve_image_ref_to_base64_data( + media_ref, + strict=strict, + default_mime_type=default_mime_type, + options=image_options, + ) + return await MediaResolver( + media_ref, + media_type=media_type, + ).to_base64_data(strict=strict, default_mime_type=default_mime_type) diff --git a/tests/test_agent_runner_media_resolver.py b/tests/test_agent_runner_media_resolver.py index 4b5bbf289..76cab73be 100644 --- a/tests/test_agent_runner_media_resolver.py +++ b/tests/test_agent_runner_media_resolver.py @@ -134,7 +134,7 @@ async def test_coze_history_reference_is_materialized_before_upload( assert len(responses) == 1 assert len(uploaded) == 1 - assert upload_options[0].enabled is False + assert upload_options[0].max_size == 1280 assert base64.b64decode(uploaded[0].split(",", 1)[1]) == image_bytes assert captured["additional_messages"][0]["content"][0] == { "type": "file", diff --git a/tests/test_media_utils.py b/tests/test_media_utils.py index 3c0094062..c4901bdca 100644 --- a/tests/test_media_utils.py +++ b/tests/test_media_utils.py @@ -359,30 +359,6 @@ async def test_compress_image_preserves_alpha_png(tmp_path, monkeypatch): compressed_path.unlink(missing_ok=True) -@pytest.mark.asyncio -async def test_compress_image_rejects_oversized_encoded_payload(tmp_path, monkeypatch): - from PIL import Image as PILImage - - temp_dir = tmp_path / "temp" - monkeypatch.setattr(media_utils, "get_astrbot_temp_path", lambda: str(temp_dir)) - image_path = tmp_path / "high_entropy.png" - image = PILImage.new("RGBA", (32, 32)) - image.putdata( - [ - (index % 256, (index * 7) % 256, (index * 13) % 256, 255) - for index in range(1024) - ] - ) - image.save(image_path, format="PNG") - - with pytest.raises(media_utils.ImagePayloadTooLargeError, match="encoding limit"): - await media_utils.compress_image( - str(image_path), max_size=128, max_encoded_bytes=1 - ) - - assert not temp_dir.exists() or not list(temp_dir.iterdir()) - - @pytest.mark.asyncio async def test_compress_image_keeps_animated_gif(tmp_path, monkeypatch): from PIL import Image as PILImage diff --git a/tests/test_openai_source.py b/tests/test_openai_source.py index 62dc182d6..ce3493f18 100644 --- a/tests/test_openai_source.py +++ b/tests/test_openai_source.py @@ -774,23 +774,44 @@ async def test_prepare_chat_payload_materializes_context_http_image_urls(monkeyp provider = _make_provider() try: - async def fake_resolve_media_ref_to_base64_data( + async def fake_prepare_image_source( media_ref: str, *, - media_type: str, - strict: bool = False, image_options=None, + default_mime_type=None, ) -> ResolvedMediaData: assert media_ref == "https://example.com/quoted.png" - assert media_type == "image" - assert strict is False assert image_options is not None return ResolvedMediaData(base64_data="abcd", mime_type="image/png") + async def fake_materialize_image_media_refs(contexts, _store, *, options=None): + assert options is not None + return [ + { + **contexts[0], + "content": [ + contexts[0]["content"][0], + { + "type": "image_url", + "image_url": { + "url": ( + await fake_prepare_image_source( + "https://example.com/quoted.png", + image_options=options, + ) + ).to_data_url(), + "id": "ctx-img", + "detail": "high", + }, + }, + ], + } + ] + monkeypatch.setattr( openai_source_module, - "resolve_media_ref_to_base64_data", - fake_resolve_media_ref_to_base64_data, + "materialize_image_media_refs", + fake_materialize_image_media_refs, ) contexts = [ @@ -822,11 +843,12 @@ async def test_prepare_chat_payload_materializes_context_http_image_urls(monkeyp "type": "image_url", "image_url": { "url": "data:image/png;base64,abcd", + "id": "ctx-img", "detail": "high", }, }, ] - assert payloads["messages"][0]["content"][1]["image_url"].get("id") is None + assert payloads["messages"][0]["content"][1]["image_url"]["id"] == "ctx-img" assert contexts[0]["content"][1]["image_url"] == { "url": "https://example.com/quoted.png", "id": "ctx-img", @@ -1320,27 +1342,31 @@ async def test_audio_preprocess_failure_does_not_log_media_ref(monkeypatch): @pytest.mark.asyncio -async def test_prepare_chat_payload_keeps_original_context_image_when_materialization_fails( +async def test_prepare_chat_payload_replaces_unavailable_context_image( monkeypatch, ): provider = _make_provider() try: - async def fake_resolve_media_ref_to_base64_data( - media_ref: str, - *, - media_type: str, - strict: bool = False, - ) -> None: - assert media_ref == "https://example.com/expired.png" - assert media_type == "image" - assert strict is False - return None + async def fake_prepare_image_source(*_args, **_kwargs): + raise ValueError("expired image") + + async def fake_materialize_image_media_refs(contexts, _store, *, options=None): + assert options is not None + return [ + { + **contexts[0], + "content": [ + contexts[0]["content"][0], + {"type": "text", "text": "[Image unavailable]"}, + ], + } + ] monkeypatch.setattr( openai_source_module, - "resolve_media_ref_to_base64_data", - fake_resolve_media_ref_to_base64_data, + "materialize_image_media_refs", + fake_materialize_image_media_refs, ) payloads, _ = await provider._prepare_chat_payload( @@ -1363,12 +1389,7 @@ async def test_prepare_chat_payload_keeps_original_context_image_when_materializ assert payloads["messages"][0]["content"] == [ {"type": "text", "text": "look"}, - { - "type": "image_url", - "image_url": { - "url": "https://example.com/expired.png", - }, - }, + {"type": "text", "text": "[Image unavailable]"}, ] finally: await provider.terminate() diff --git a/tests/test_process_stage_images.py b/tests/test_process_stage_images.py index 5f1aafaaa..365f74d79 100644 --- a/tests/test_process_stage_images.py +++ b/tests/test_process_stage_images.py @@ -199,18 +199,13 @@ async def process_event(harness, event, *, preprocess_first=False, stage=None): @pytest.mark.asyncio -@pytest.mark.parametrize("enabled", [None, True, False]) @pytest.mark.parametrize("fmt", ["JPEG", "PNG", "BMP", "WEBP", "GIF"]) -async def test_legacy_toggle_does_not_disable_preparation( - harness, tmp_path, monkeypatch, enabled, fmt +async def test_image_formats_are_prepared_for_provider( + harness, tmp_path, monkeypatch, fmt ): source = source_image(tmp_path, fmt) original = source.read_bytes() event = make_event([Image(file=str(source))], text="") - if enabled is None: - harness.config["provider_settings"].pop("image_compress_enabled", None) - else: - harness.config["provider_settings"]["image_compress_enabled"] = enabled await process_event(harness, event, preprocess_first=True) assert len(harness.captured) == 1 req = harness.captured[0].req @@ -331,14 +326,7 @@ async def test_animation_montage_notice_reaches_model(harness, tmp_path): still = tmp_path / "still.png" PILImage.new("RGB", (60, 30), "red").save(still) - # Legacy settings cannot disable animation preparation. - for source, enabled, expected in ( - (animated, True, True), - (animated, True, True), - (still, True, False), - (animated, False, True), - ): - harness.config["provider_settings"]["image_compress_enabled"] = enabled + for source, expected in ((animated, True), (still, False)): await process_event(harness, make_event([Image(file=str(source))])) notices = [ part @@ -480,9 +468,7 @@ async def test_plugin_request_extra_metadata_hook_and_history( materialized = await materialize_image_media_refs( [saved[-1]], ImageMediaStore(tmp_path / "data" / "media") ) - images = [ - p for p in materialized[0]["content"] if p["type"] == "image_url" - ] + images = [p for p in materialized[0]["content"] if p["type"] == "image_url"] assert images and all( p["image_url"]["url"].startswith("data:image/jpeg;base64,") for p in images ) @@ -510,6 +496,11 @@ async def test_plugin_request_extra_metadata_hook_and_history( async def test_provider_and_history_keep_bounded_previews_and_original_paths( harness, tmp_path, monkeypatch, entry, fmt ): + monkeypatch.setattr( + tool_loop_agent_runner, + "get_astrbot_data_path", + lambda: str(tmp_path / "data"), + ) source = tmp_path / f"original.{fmt.lower()}" pixels = random.Random(9703).randbytes(1280 * 960 * 3) image = PILImage.frombytes("RGB", (1280, 960), pixels) @@ -528,9 +519,7 @@ async def test_provider_and_history_keep_bounded_previews_and_original_paths( assert len(original) >= 512 * 1024 harness.config["provider_settings"]["image_compress_options"] = { "max_size": 1280, - "quality": 100, } - harness.config["provider_settings"]["image_compress_enabled"] = False part = Image(file=str(source)) event = make_event( [Reply(id="quoted", chain=[part])] if entry == "quote" else [part] @@ -575,6 +564,18 @@ async def test_provider_and_history_keep_bounded_previews_and_original_paths( ] ) for messages in (provider_messages, history): + if any( + part.get("type") == "image_media_ref" + for message in messages + if isinstance(message.get("content"), list) + for part in message["content"] + if isinstance(part, dict) + ): + messages = await materialize_image_media_refs( + messages, + ImageMediaStore(tmp_path / "data" / "media"), + options=media.ImagePreparationOptions(max_size=1280), + ) images = [ part for message in messages @@ -881,10 +882,9 @@ def test_image_preparation_stays_outside_agent_runner_and_providers(): @pytest.mark.asyncio -@pytest.mark.parametrize("enabled", [True, False]) @pytest.mark.parametrize("reference", ["data", "base64", "http", "file"]) async def test_localized_reference_lifetime_and_ownership( - harness, tmp_path, monkeypatch, enabled, reference + harness, tmp_path, monkeypatch, reference ): source = source_image(tmp_path) encoded = base64.b64encode(source.read_bytes()).decode() @@ -899,7 +899,6 @@ async def test_localized_reference_lifetime_and_ownership( Path(target).write_bytes(source.read_bytes()) monkeypatch.setattr(media, "download_file", download) - harness.config["provider_settings"]["image_compress_enabled"] = enabled event = make_event() event.set_extra( "provider_request", ProviderRequest(prompt="", image_urls=[refs[reference]]) @@ -947,10 +946,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 == [] + assert event._temporary_local_files == ([] if borrowed else [str(path)]) assert set(harness.work.rglob("*")) == (set() if borrowed else {path}) event.cleanup_temporary_local_files() - assert path.exists() + assert path.exists() == borrowed assert source.read_bytes() == original @@ -1049,12 +1048,11 @@ async def test_plugin_and_hook_image_list_normalization(harness, tmp_path, monke @pytest.mark.asyncio @pytest.mark.parametrize("quoted", [False, True]) -@pytest.mark.parametrize("enabled", [False, True]) @pytest.mark.parametrize("late", [False, True]) @pytest.mark.parametrize("reference", ["file", "base64", "data", "http"]) @pytest.mark.parametrize("fmt", ["JPEG", "GIF"]) async def test_attachment_sources_survive_event_cleanup( - harness, monkeypatch, quoted, enabled, late, reference, fmt + harness, monkeypatch, quoted, late, reference, fmt ): harness.work.mkdir() source = source_image(harness.work, fmt) @@ -1074,7 +1072,6 @@ async def test_attachment_sources_survive_event_cleanup( unrelated = harness.work / "other.tmp" unrelated.write_bytes(b"temporary") event.track_temporary_local_file(str(unrelated)) - harness.config["provider_settings"]["image_compress_enabled"] = enabled harness.config["provider_settings"]["image_compress_options"]["max_size"] = 12 async def download(url, target): diff --git a/tests/test_tool_loop_agent_runner.py b/tests/test_tool_loop_agent_runner.py index d2a140762..b32f118d5 100644 --- a/tests/test_tool_loop_agent_runner.py +++ b/tests/test_tool_loop_agent_runner.py @@ -1401,8 +1401,6 @@ async def test_real_tool_image_uses_preparation_and_reference_storage( provider.provider_settings = { "image_compress_options": { "max_size": 512, - "quality": 42, - "max_encoded_bytes": 2 * 1024 * 1024, } } provider.max_calls_before_normal_response = 1 @@ -1424,11 +1422,10 @@ async def test_real_tool_image_uses_preparation_and_reference_storage( if isinstance(part, ImageMediaRefPart) ] assert len(images) == 1 - assert images[0].mime_type == "image/webp" + # The mainline preparation keeps alpha-safe images losslessly as PNG. + assert images[0].mime_type == "image/png" assert images[0].image_id.endswith("call_123_0.webp") assert captured_options[0].max_size == 512 - assert captured_options[0].quality == 42 - assert captured_options[0].max_encoded_bytes == 2 * 1024 * 1024 assert provider.call_count == 2 diff --git a/tests/unit/test_astr_main_agent.py b/tests/unit/test_astr_main_agent.py index 864f80c74..22eae5999 100644 --- a/tests/unit/test_astr_main_agent.py +++ b/tests/unit/test_astr_main_agent.py @@ -1907,7 +1907,6 @@ class TestBuildMainAgent: @pytest.mark.asyncio @pytest.mark.parametrize("quoted", [False, True]) - @pytest.mark.parametrize("compression_enabled", [False, True]) async def test_build_main_agent_with_images( self, mock_event, @@ -1916,7 +1915,6 @@ class TestBuildMainAgent: tmp_path, monkeypatch, quoted, - compression_enabled, ): """Direct builders prepare both ordinary and quoted images before reset.""" @@ -1955,7 +1953,6 @@ class TestBuildMainAgent: config=module.MainAgentBuildConfig( tool_call_timeout=60, provider_settings={ - "image_compress_enabled": compression_enabled, "image_compress_options": {"max_size": 4}, }, ), diff --git a/tests/unit/test_context_image_budget.py b/tests/unit/test_context_image_budget.py deleted file mode 100644 index 809001b51..000000000 --- a/tests/unit/test_context_image_budget.py +++ /dev/null @@ -1,153 +0,0 @@ -"""Image-byte budgets must not change history or token accounting.""" - -import base64 -import copy -from types import SimpleNamespace -from unittest.mock import AsyncMock - -import pytest -from PIL import Image - -from astrbot.core.agent.context.compressor import LLMSummaryCompressor -from astrbot.core.agent.context.image_budget import ( - get_image_encoded_byte_limit, - validate_context_image_bytes, -) -from astrbot.core.agent.message import Message -from astrbot.core.exceptions import ProviderRequestTooLargeError -from astrbot.core.utils.image_media_store import ImageMediaStore -from astrbot.core.utils.media_utils import ImagePayloadTooLargeError - - -@pytest.mark.parametrize("limit", [None, True, False, 0, -1, "100"]) -def test_invalid_provider_image_limit_uses_default(limit): - assert ( - get_image_encoded_byte_limit( - {"image_compress_options": {"max_encoded_bytes": limit}} - ) - == 4 * 1024 * 1024 - ) - - -def test_provider_image_limit_is_respected(): - assert ( - get_image_encoded_byte_limit( - {"image_compress_options": {"max_encoded_bytes": 16 * 1024 * 1024}} - ) - == 16 * 1024 * 1024 - ) - - -@pytest.mark.parametrize("as_models", [False, True]) -def test_byte_budget_does_not_rewrite_history(as_models): - history = [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": {"url": "data:image/png;base64,AAAA"}, - } - ], - } - ] - messages = ( - [Message.model_validate(item) for item in history] if as_models else history - ) - before = copy.deepcopy(messages) - assert validate_context_image_bytes(messages, 4) == 4 - with pytest.raises(ImagePayloadTooLargeError): - validate_context_image_bytes(messages, 3) - assert messages == before - - -def test_reference_size_is_checked_without_opening_media(): - history = [ - { - "role": "user", - "content": [ - { - "type": "image_media_ref", - "media_id": "a" * 64, - "byte_size": 12, - "mime_type": "image/png", - "version": 1, - } - ], - } - ] - assert validate_context_image_bytes(history, 16) == 16 - with pytest.raises(ImagePayloadTooLargeError): - validate_context_image_bytes(history, 15) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("vision", [True, False]) -async def test_summary_loads_only_its_selected_images(tmp_path, monkeypatch, vision): - image_path = tmp_path / "old.png" - Image.new("RGB", (8, 8), "blue").save(image_path) - original = image_path.read_bytes() - store = ImageMediaStore(tmp_path / "media") - old_ref = store.put(original, detail="high", image_id="old-image") - messages = [ - Message.model_validate({"role": "user", "content": [old_ref.model_dump()]}), - Message(role="assistant", content="old answer"), - Message.model_validate( - { - "role": "user", - "content": [{**old_ref.model_dump(), "media_id": "f" * 64}], - } - ), - ] - before = [message.model_dump() for message in messages] - opened = [] - original_read = store.read - - def read(ref, allowed_ids): - opened.append(ref.media_id) - return original_read(ref, allowed_ids) - - monkeypatch.setattr(store, "read", read) - provider = SimpleNamespace( - provider_config={"modalities": ["text", "image"] if vision else ["text"]}, - provider_settings={}, - text_chat=AsyncMock(return_value=SimpleNamespace(completion_text="summary")), - ) - result = await LLMSummaryCompressor( - provider, keep_recent_ratio=0, image_media_store=store - )(messages) - assert opened == ([old_ref.media_id] if vision else []) - assert result[-1] is messages[-1] - assert [message.model_dump() for message in messages] == before - if vision: - sent = provider.text_chat.call_args.kwargs["contexts"][0] - image = sent["content"][0]["image_url"] - assert base64.b64decode(image["url"].split(",", 1)[1]) == original - assert image["detail"] == "high" - assert image["id"] == "old-image" - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "error", - [ - MemoryError(), - ImagePayloadTooLargeError("oversize"), - ProviderRequestTooLargeError("413"), - ], -) -async def test_summary_resource_failure_is_not_swallowed(error): - provider = SimpleNamespace( - provider_config={"modalities": ["text"]}, - provider_settings={}, - text_chat=AsyncMock(side_effect=error), - ) - messages = [ - Message(role="user", content="old question"), - Message(role="assistant", content="old answer"), - Message(role="user", content="current question"), - ] - with pytest.raises(type(error)) as caught: - await LLMSummaryCompressor(provider, keep_recent_ratio=0)(messages) - assert caught.value is error - assert provider.text_chat.await_count == 1 diff --git a/tests/unit/test_image_media_store.py b/tests/unit/test_image_media_store.py index 937987c1a..14436dfe8 100644 --- a/tests/unit/test_image_media_store.py +++ b/tests/unit/test_image_media_store.py @@ -4,6 +4,7 @@ import base64 import io import json import os +import random from concurrent.futures import ThreadPoolExecutor import pytest @@ -20,6 +21,14 @@ def _png() -> bytes: return output.getvalue() +def _noisy_png() -> bytes: + output = io.BytesIO() + pixels = random.Random(9703).randbytes(1280 * 960 * 3) + with Image.frombytes("RGB", (1280, 960), pixels) as image: + image.save(output, "PNG") + return output.getvalue() + + def test_store_deduplicates_and_round_trips_exact_bytes(tmp_path): store = ImageMediaStore(tmp_path / "media") data = _png() @@ -126,6 +135,69 @@ async def test_reference_materialization_preserves_history_and_bytes( assert messages[0].content[0].type == "image_media_ref" +@pytest.mark.asyncio +async def test_model_materialization_bounds_selected_reference_without_rewriting_history( + tmp_path, +): + from astrbot.core.utils.image_media_store import materialize_image_media_refs + from astrbot.core.utils.media_utils import ImagePreparationOptions + + data = _noisy_png() + assert len(data) > 512 * 1024 + store = ImageMediaStore(tmp_path / "media") + ref = store.put(data) + history = [{"role": "user", "content": [ref.model_dump()]}] + + result = await materialize_image_media_refs( + history, + store, + options=ImagePreparationOptions(max_size=1280), + ) + + image_url = result[0]["content"][0]["image_url"]["url"] + prepared = base64.b64decode(image_url.split(",", 1)[1]) + assert len(prepared) < 512 * 1024 + assert image_url.startswith("data:image/jpeg;base64,") + assert history[0]["content"][0]["type"] == "image_media_ref" + assert store.read(ref, {ref.media_id}) == data + + +@pytest.mark.asyncio +async def test_model_materialization_bounds_legacy_inline_image_without_rewriting_history( + tmp_path, +): + from astrbot.core.utils.image_media_store import materialize_image_media_refs + from astrbot.core.utils.media_utils import ImagePreparationOptions + + data = _noisy_png() + encoded = base64.b64encode(data).decode("ascii") + history = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{encoded}"}, + } + ], + } + ] + + result = await materialize_image_media_refs( + history, + ImageMediaStore(tmp_path / "media"), + options=ImagePreparationOptions(max_size=1280), + ) + + image_url = result[0]["content"][0]["image_url"]["url"] + prepared = base64.b64decode(image_url.split(",", 1)[1]) + assert len(prepared) < 512 * 1024 + assert image_url.startswith("data:image/jpeg;base64,") + assert history[0]["content"][0]["image_url"]["url"] == ( + f"data:image/png;base64,{encoded}" + ) + + @pytest.mark.asyncio async def test_unselected_reference_is_not_read(tmp_path, monkeypatch): from astrbot.core.utils.image_media_store import materialize_image_media_refs diff --git a/tests/unit/test_image_preparation_budget.py b/tests/unit/test_image_preparation_budget.py deleted file mode 100644 index 86aafe219..000000000 --- a/tests/unit/test_image_preparation_budget.py +++ /dev/null @@ -1,377 +0,0 @@ -"""Regression tests for bounded image preparation and output ownership.""" - -import asyncio -import base64 -import io -import random -import threading -from pathlib import Path - -import pytest -from PIL import Image - -from astrbot.core.utils import media_utils - - -def test_file_base64_encoding_stream_matches_standard_library(tmp_path): - source = tmp_path / "payload.bin" - payload = random.Random(31).randbytes(1023 * 1024 + 5) - source.write_bytes(payload) - - assert media_utils._encode_file_to_base64(source) == base64.b64encode( - payload - ).decode("ascii") - - -@pytest.mark.parametrize( - "size,target", [((1280, 1280), (960, 960)), ((131, 197), (71, 103))] -) -def test_strip_resize_matches_composited_lanczos(tmp_path, size, target): - from PIL import ImageChops - - with Image.frombytes( - "RGBA", size, random.Random(43).randbytes(size[0] * size[1] * 4) - ) as source: - with ( - source.resize(target, Image.Resampling.LANCZOS) as expected, - media_utils._resize_alpha_in_strips(source, target) as actual, - ): - assert actual.size == expected.size - for color in ("black", "white"): - with ( - Image.new("RGBA", target, color) as background, - Image.alpha_composite(background, expected) as expected_view, - Image.alpha_composite(background, actual) as actual_view, - ImageChops.difference(expected_view, actual_view) as delta, - ): - assert max(high for _, high in delta.getextrema()) <= 2 - - -@pytest.mark.parametrize( - "image_format,mode", [("PNG", "RGBA"), ("JPEG", "RGB"), ("WEBP", "RGBA")] -) -def test_compliant_bytes_are_preserved(tmp_path, image_format, mode): - path = tmp_path / "original" - with Image.new(mode, (32, 20)) as image: - image.save(path, image_format) - original = path.read_bytes() - result = media_utils._compress_image_sync(path, tmp_path, 1280, 95, True) - assert result is None - assert path.read_bytes() == original - assert list(tmp_path.iterdir()) == [path] - - -@pytest.mark.parametrize( - "mode,save_options,expected_format", - [("RGB", {}, "JPEG"), ("RGBA", {"lossless": True}, "PNG")], -) -def test_oversized_static_webp_avoids_pillow_decoder( - tmp_path, monkeypatch, mode, save_options, expected_format -): - if media_utils._get_webp_decoder() is None: - pytest.skip("Pillow does not expose the bundled WebP decoder") - - source = tmp_path / "source.webp" - pixel_count = 512 * 512 - channels = 4 if mode == "RGBA" else 3 - with Image.frombytes( - mode, - (512, 512), - random.Random(17).randbytes(pixel_count * channels), - ) as image: - image.save(source, "WEBP", **save_options) - - def unexpected_pillow_open(*args, **kwargs): - pytest.fail("Static WebP should use the preallocated decoder path") - - monkeypatch.setattr(media_utils.PILImage, "open", unexpected_pillow_open) - output = media_utils._compress_image_sync(source, tmp_path, 64, 95, True, 20_000) - monkeypatch.undo() - - assert output is not None - with Image.open(output) as prepared: - assert prepared.format == expected_format - assert max(prepared.size) <= 64 - - -def test_pixel_compliant_noise_fits_byte_budget(tmp_path): - source = tmp_path / "noise.png" - with Image.frombytes( - "RGBA", (1280, 1280), random.Random(7).randbytes(1280 * 1280 * 4) - ) as image: - image.save(source) - budget = 4 * 1024 * 1024 - assert 4 * ((source.stat().st_size + 2) // 3) > budget - result = media_utils._compress_image_sync(source, tmp_path, 1280, 95, True, budget) - assert result is not None - assert 4 * ((Path(result).stat().st_size + 2) // 3) <= budget - with Image.open(result) as image: - assert image.mode == "RGBA" - assert image.width < 1280 - - -def test_animation_is_preserved_when_compliant_and_rejected_when_oversized(tmp_path): - source = tmp_path / "animated.gif" - frames = [Image.new("RGB", (16, 12), color) for color in ("red", "blue")] - try: - frames[0].save( - source, - format="GIF", - save_all=True, - append_images=frames[1:], - duration=40, - loop=0, - ) - finally: - for frame in frames: - frame.close() - - original = source.read_bytes() - assert media_utils._compress_image_sync(source, tmp_path, 4, 95, True) is None - assert source.read_bytes() == original - - with pytest.raises(media_utils.ImagePayloadTooLargeError): - media_utils._compress_image_sync(source, tmp_path, 4, 95, True, 1) - assert list(tmp_path.glob("compressed_*")) == [] - - -def test_webp_animation_is_preserved_without_flattening(tmp_path): - source = tmp_path / "animated.webp" - frames = [Image.new("RGBA", (16, 12), color) for color in ("red", "blue")] - try: - frames[0].save( - source, - format="WEBP", - save_all=True, - append_images=frames[1:], - duration=40, - loop=0, - ) - finally: - for frame in frames: - frame.close() - - original = source.read_bytes() - assert media_utils._compress_image_sync(source, tmp_path, 4, 95, True) is None - assert source.read_bytes() == original - - with pytest.raises(media_utils.ImagePayloadTooLargeError): - media_utils._compress_image_sync(source, tmp_path, 4, 95, True, 1) - - -def test_screenshot_preserves_oriented_coordinates(tmp_path): - path = tmp_path / "oriented.jpg" - with Image.frombytes( - "RGB", (120, 80), random.Random(1).randbytes(120 * 80 * 3) - ) as image: - exif = image.getexif() - exif[274] = 6 - image.save(path, quality=100, exif=exif) - output = media_utils._compress_image_sync( - path, tmp_path, 10, 80, True, 10_000, preserve_dimensions=True - ) - assert output is not None - with Image.open(output) as image: - assert image.size == (80, 120) - assert image.getexif().get(274, 1) == 1 - - -@pytest.mark.parametrize("orientation", range(1, 9)) -def test_thumbnail_preserves_exif_corner_placement(tmp_path, orientation): - from PIL import ImageOps - - source = tmp_path / "corners.jpg" - with Image.new("RGB", (800, 600), "black") as image: - image.paste("red", (0, 0, 400, 300)) - image.paste("green", (400, 0, 800, 300)) - image.paste("blue", (0, 300, 400, 600)) - image.paste("white", (400, 300, 800, 600)) - exif = image.getexif() - exif[274] = orientation - image.save(source, quality=95, exif=exif) - with Image.open(source) as image: - expected = ImageOps.exif_transpose(image) - expected.thumbnail((160, 160)) - try: - output = media_utils._compress_image_sync(source, tmp_path, 160, 95, True) - assert output is not None - with Image.open(output) as actual: - assert actual.size == expected.size - assert actual.getexif().get(274, 1) == 1 - for x in (actual.width // 4, actual.width * 3 // 4): - for y in (actual.height // 4, actual.height * 3 // 4): - assert ( - max( - abs(a - b) - for a, b in zip( - actual.getpixel((x, y)), expected.getpixel((x, y)) - ) - ) - <= 5 - ) - finally: - expected.close() - - -def test_jpeg_stops_at_first_fitting_quality(tmp_path, monkeypatch): - path = tmp_path / "photo.jpg" - with Image.new("RGB", (200, 100), (20, 40, 80)) as image: - image.save(path, "JPEG", quality=95) - qualities = [] - original_save = Image.Image.save - - def save(image, target, fmt=None, **kwargs): - if fmt == "JPEG": - qualities.append(kwargs.get("quality")) - return original_save(image, target, fmt, **kwargs) - - monkeypatch.setattr(Image.Image, "save", save) - output = media_utils._compress_image_sync( - path, tmp_path, 100, 95, True, 4 * 1024 * 1024 - ) - assert output is not None - assert qualities == [95] - - -def test_cua_tries_next_jpeg_quality_without_resizing(tmp_path, monkeypatch): - path = tmp_path / "opaque.png" - with Image.frombytes( - "RGB", (120, 80), random.Random(23).randbytes(120 * 80 * 3) - ) as image: - image.save(path, "PNG") - qualities = [] - original_save = Image.Image.save - - def save(image, target, fmt=None, **kwargs): - if fmt == "JPEG": - quality = kwargs.get("quality") - qualities.append(quality) - if quality == 95: - original_save(image, target, fmt, **kwargs) - with Path(target).open("ab") as candidate: - candidate.write(b"x" * 20_000) - return - return original_save(image, target, fmt, **kwargs) - - monkeypatch.setattr(Image.Image, "save", save) - output = media_utils._compress_image_sync( - path, tmp_path, 1, 95, True, 20_000, preserve_dimensions=True - ) - assert output is not None - assert qualities[:2] == [95, 85] - with Image.open(output) as image: - assert image.size == (120, 80) - - -def test_impossible_screenshot_budget_rejects_without_resizing(tmp_path): - path = tmp_path / "screenshot.png" - with Image.new("RGBA", (100, 60), (20, 100, 50, 127)) as image: - image.save(path) - with pytest.raises(media_utils.ImagePayloadTooLargeError): - media_utils._compress_image_sync( - path, tmp_path, 1, 95, True, 1, preserve_dimensions=True - ) - assert list(tmp_path.iterdir()) == [path] - - -def test_write_failure_removes_partial_candidate(tmp_path, monkeypatch): - source = io.BytesIO() - with Image.new("RGB", (100, 60)) as image: - image.save(source, "PNG") - - def fail_write(image, path, *args, **kwargs): - Path(path).write_bytes(b"partial") - raise OSError("disk full") - - monkeypatch.setattr(Image.Image, "save", fail_write) - with pytest.raises(OSError, match="disk full"): - media_utils._compress_image_sync(source.getvalue(), tmp_path, 10, 95, True) - assert not list(tmp_path.iterdir()) - - -def test_decode_failure_leaves_no_output(tmp_path): - with pytest.raises(OSError): - media_utils._compress_image_sync(b"invalid", tmp_path, 10, 95, True) - assert not list(tmp_path.iterdir()) - - -@pytest.mark.asyncio -async def test_disabled_compression_rejects_oversize_before_read(tmp_path, monkeypatch): - source = tmp_path / "oversized.png" - with Image.new("RGB", (32, 32), "red") as image: - image.save(source) - - def unexpected_read(path): - pytest.fail("Oversized source was read before checking the byte limit") - - monkeypatch.setattr(Path, "read_bytes", unexpected_read) - with pytest.raises(media_utils.ImagePayloadTooLargeError): - await media_utils.prepare_image_source( - str(source), - options=media_utils.ImagePreparationOptions( - enabled=False, max_encoded_bytes=1 - ), - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "failure", [media_utils.ImagePayloadTooLargeError, MemoryError, OSError, ValueError] -) -async def test_agent_does_not_fall_back_to_original_on_resource_error( - tmp_path, monkeypatch, failure -): - source = tmp_path / "image.png" - with Image.new("RGB", (8, 8), "red") as image: - image.save(source) - - async def fail(*args, **kwargs): - raise failure("bounded failure") - - monkeypatch.setattr(media_utils, "compress_image", fail) - with pytest.raises(failure): - await media_utils.prepare_image_source(str(source)) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("timeout", [False, True]) -async def test_cancelled_worker_cleans_after_exit(tmp_path, monkeypatch, timeout): - source = tmp_path / "source" - source.write_bytes(b"source") - monkeypatch.setattr(media_utils, "IMAGE_COMPRESS_DEFAULT_MIN_FILE_SIZE_BYTES", 1) - output = tmp_path / "worker-output" - entered = threading.Event() - release = threading.Event() - cleaned = asyncio.Event() - original_unlink = Path.unlink - - def unlink(path, *args, **kwargs): - original_unlink(path, *args, **kwargs) - if path == output: - cleaned.set() - - monkeypatch.setattr(Path, "unlink", unlink) - - def blocked_worker(*args, **kwargs): - entered.set() - release.wait(5) - output.write_bytes(b"prepared") - return str(output) - - monkeypatch.setattr(media_utils, "_compress_image_sync", blocked_worker) - task = asyncio.create_task(media_utils.compress_image(str(source))) - try: - assert await asyncio.to_thread(entered.wait, 2) - if timeout: - with pytest.raises(asyncio.TimeoutError): - await asyncio.wait_for(task, timeout=0.01) - else: - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - finally: - release.set() - # A finished task can still have its cleanup callback queued on the loop. - await asyncio.wait_for(cleaned.wait(), timeout=5) - assert not output.exists() - assert source.read_bytes() == b"source" diff --git a/tests/unit/test_image_screenshot_encoding.py b/tests/unit/test_image_screenshot_encoding.py index cec79c5b0..06ced3a8c 100644 --- a/tests/unit/test_image_screenshot_encoding.py +++ b/tests/unit/test_image_screenshot_encoding.py @@ -1,11 +1,13 @@ from pathlib import Path +import pytest from PIL import Image, ImageDraw from astrbot.core.utils import media_utils -def test_prefers_smaller_lossless_encoding_for_ui_screenshot(tmp_path): +@pytest.mark.asyncio +async def test_model_preparation_bounds_ui_screenshot(tmp_path): source = tmp_path / "screenshot-ordinary.png" with Image.new("RGB", (1920, 1080), "white") as image: draw = ImageDraw.Draw(image) @@ -21,23 +23,15 @@ def test_prefers_smaller_lossless_encoding_for_ui_screenshot(tmp_path): ) image.save(source, "PNG", optimize=True) - output = media_utils._compress_image_sync( - source, tmp_path, 1280, 95, True, 4 * 1024 * 1024 + prepared = await media_utils.prepare_model_image( + str(source), max_size=1280, output_dir=tmp_path ) - assert output is not None - with Image.open(source) as original: - resized = original.copy() - resized.thumbnail((1280, 1280), Image.Resampling.LANCZOS) - png_candidate = tmp_path / "expected.png" - jpeg_candidate = tmp_path / "expected.jpg" - resized.save(png_candidate, "PNG", optimize=True) - resized.save(jpeg_candidate, "JPEG", quality=95, optimize=True) - # The old implementation always selected JPEG for opaque PNG input, - # even when the same resized pixels have a smaller PNG encoding. - assert png_candidate.stat().st_size < jpeg_candidate.stat().st_size - assert max(resized.size) <= 1280 - assert Path(output).suffix == ".png" - assert Path(output).stat().st_size == png_candidate.stat().st_size + assert prepared is not None + output, is_montage, needs_cleanup, _ = prepared + assert not is_montage and needs_cleanup + assert Path(output).suffix == ".jpg" + assert Path(output).stat().st_size < 512 * 1024 with Image.open(output) as image: assert max(image.size) <= 1280 + Path(output).unlink() diff --git a/tests/unit/test_image_source_integration.py b/tests/unit/test_image_source_integration.py index 41623a78d..110af9c8d 100644 --- a/tests/unit/test_image_source_integration.py +++ b/tests/unit/test_image_source_integration.py @@ -11,7 +11,6 @@ from astrbot.core.utils.image_media_store import ( persist_inline_image_refs, ) from astrbot.core.utils.media_utils import ( - ImagePayloadTooLargeError, ImagePreparationInput, ImagePreparationOptions, prepare_image_source, @@ -47,31 +46,15 @@ async def test_current_plugin_image_part_is_prepared_without_mutating_caller(tmp @pytest.mark.asyncio -async def test_quoted_user_path_uses_configured_encoded_limit(tmp_path): - source = tmp_path / "quoted.png" - source.write_bytes(_image()) - options = ImagePreparationOptions(max_encoded_bytes=1) - - with pytest.raises(ImagePayloadTooLargeError): - await prepare_image_source(str(source), options=options) - - -@pytest.mark.asyncio -async def test_cua_dimensions_are_explicit_and_normal_tool_can_resize(tmp_path): +async def test_tool_images_use_the_mainline_dimension_limit(tmp_path): source = tmp_path / "tool.png" source.write_bytes(_image((80, 40))) normal = await prepare_image_source( str(source), options=ImagePreparationOptions(max_size=20) ) - cua = await prepare_image_source( - str(source), - options=ImagePreparationOptions(max_size=20, preserve_dimensions=True), - ) with Image.open(io.BytesIO(normal.to_bytes())) as image: assert max(image.size) <= 20 - with Image.open(io.BytesIO(cua.to_bytes())) as image: - assert image.size == (80, 40) def test_temporary_image_part_is_not_persisted(tmp_path): diff --git a/tests/unit/test_image_source_preparation.py b/tests/unit/test_image_source_preparation.py index 70de889f4..93a54e244 100644 --- a/tests/unit/test_image_source_preparation.py +++ b/tests/unit/test_image_source_preparation.py @@ -1,9 +1,7 @@ -import asyncio import base64 import io import threading from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer -from pathlib import Path import pytest from PIL import Image @@ -52,70 +50,17 @@ async def test_prepare_image_source_accepts_all_reference_forms(tmp_path, monkey server.server_close() -@pytest.mark.asyncio -async def test_prepare_cancel_keeps_resolver_source_until_worker_exits( - tmp_path, monkeypatch -): - entered = threading.Event() - release = threading.Event() - source_seen = {} - original = media_utils._compress_image_sync - - def blocked(source, *args, **kwargs): - source_seen["path"] = Path(source) if isinstance(source, (str, Path)) else None - entered.set() - release.wait(5) - return original(source, *args, **kwargs) - - monkeypatch.setattr(media_utils, "_compress_image_sync", blocked) - encoded = base64.b64encode(_png()).decode() - task = asyncio.create_task( - media_utils.prepare_image_source( - f"data:image/png;base64,{encoded}", - options=media_utils.ImagePreparationOptions(max_size=2), - ) - ) - assert await asyncio.to_thread(entered.wait, 2) - assert source_seen["path"] is not None - assert source_seen["path"].exists() - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - assert source_seen["path"].exists() - release.set() - await asyncio.sleep(0.2) - assert not source_seen["path"].exists() - - -@pytest.mark.asyncio -async def test_prepare_bytes_write_failure_cleans_owned_source(monkeypatch, tmp_path): - owned = tmp_path / "owned.bin" - - monkeypatch.setattr(media_utils, "_temp_media_path", lambda *_args: owned) - - original_write_bytes = Path.write_bytes - - def fail_write(path, data): - original_write_bytes(path, b"partial") - raise OSError("disk full") - - monkeypatch.setattr(Path, "write_bytes", fail_write) - with pytest.raises(OSError, match="disk full"): - await media_utils.prepare_image_source(_png()) - assert not owned.exists() - - @pytest.mark.asyncio async def test_provider_request_passes_preparation_options(): data = base64.b64encode(_png()).decode() request = ProviderRequest( image_urls=[f"data:image/png;base64,{data}"], - image_preparation_options=media_utils.ImagePreparationOptions( - enabled=False, max_encoded_bytes=None - ), + image_preparation_options=media_utils.ImagePreparationOptions(max_size=2), ) context = await request.assemble_context() - assert context["content"][1]["image_url"]["url"].endswith(data) + assert context["content"][1]["image_url"]["url"].startswith( + "data:image/png;base64," + ) @pytest.mark.asyncio diff --git a/tests/unit/test_provider_image_references.py b/tests/unit/test_provider_image_references.py index 0089dc9f3..4d5f612d3 100644 --- a/tests/unit/test_provider_image_references.py +++ b/tests/unit/test_provider_image_references.py @@ -8,7 +8,10 @@ from astrbot.core.provider.sources.anthropic_source import ProviderAnthropic from astrbot.core.provider.sources.gemini_source import ProviderGoogleGenAI from astrbot.core.provider.sources.openai_source import ProviderOpenAIOfficial from astrbot.core.utils.image_media_store import ImageMediaStore -from astrbot.core.utils.media_utils import ImagePayloadTooLargeError +from astrbot.core.utils.media_utils import ( + MODEL_IMAGE_MAX_INPUT_BYTES, + ImagePayloadTooLargeError, +) @pytest.fixture @@ -50,34 +53,12 @@ async def test_openai_payload_materializes_reference_without_mutating(image_cont @pytest.mark.asyncio async def test_oversized_reference_fails_before_materialization(image_context): oversized = deepcopy(image_context) - oversized[0]["content"][0]["byte_size"] = 10_000_000 + oversized[0]["content"][0]["byte_size"] = MODEL_IMAGE_MAX_INPUT_BYTES + 1 adapter = _bare(ProviderOpenAIOfficial) with pytest.raises(ImagePayloadTooLargeError): await adapter._prepare_chat_payload(None, contexts=oversized) -@pytest.mark.asyncio -@pytest.mark.parametrize("adapter_type", [ProviderOpenAIOfficial, ProviderGoogleGenAI]) -async def test_provider_settings_control_reference_budget(adapter_type, image_context): - adapter = _bare(adapter_type) - adapter.provider_settings = {"image_compress_options": {"max_encoded_bytes": 1}} - with pytest.raises(ImagePayloadTooLargeError): - if adapter_type is ProviderOpenAIOfficial: - await adapter._prepare_chat_payload(None, contexts=image_context) - else: - await adapter._prepare_conversation({"messages": image_context}) - - adapter.provider_settings = { - "image_compress_options": {"max_encoded_bytes": 1024 * 1024} - } - if adapter_type is ProviderOpenAIOfficial: - payload, _ = await adapter._prepare_chat_payload(None, contexts=image_context) - assert payload["messages"][0]["content"][0]["image_url"] - else: - contents = await adapter._prepare_conversation({"messages": image_context}) - assert contents[0].parts[0].inline_data is not None - - @pytest.mark.asyncio async def test_text_only_modality_scrubs_missing_oversized_reference_without_read( image_context, @@ -110,14 +91,6 @@ async def test_anthropic_text_chat_materializes_reference_without_mutating( assert image_context[0]["content"][0]["type"] == "image_media_ref" -@pytest.mark.asyncio -async def test_anthropic_provider_settings_control_reference_budget(image_context): - adapter = _bare(ProviderAnthropic) - adapter.provider_settings = {"image_compress_options": {"max_encoded_bytes": 1}} - with pytest.raises(ImagePayloadTooLargeError): - await adapter.text_chat(contexts=image_context) - - @pytest.mark.asyncio async def test_gemini_conversation_materializes_reference(image_context): adapter = _bare(ProviderGoogleGenAI)