diff --git a/backend/app/api/upload.py b/backend/app/api/upload.py index 28c7867c9..9774f527d 100644 --- a/backend/app/api/upload.py +++ b/backend/app/api/upload.py @@ -1,14 +1,17 @@ """File upload API for chat — saves files to agent workspace and extracts text.""" -import base64 import os import uuid from pathlib import Path -from fastapi import APIRouter, Depends, File, HTTPException, UploadFile, Form -from loguru import logger +from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile + from app.core.security import get_current_user from app.models.user import User +from app.services.chat_image_preview import ( + VisionImagePreviewError, + build_vision_image_data_url, +) from app.services.storage import ensure_local_path, get_storage_backend, guess_content_type, normalize_storage_key router = APIRouter(prefix="/chat", tags=["chat"]) @@ -23,10 +26,7 @@ IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"} EXTRACTABLE = TEXT_EXTENSIONS | OFFICE_EXTENSIONS -MIME_MAP = { - ".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", - ".gif": "image/gif", ".webp": "image/webp", ".bmp": "image/bmp", -} +MAX_CHAT_IMAGE_BYTES = 10 * 1024 * 1024 def extract_text(file_path: Path, extension: str) -> str: @@ -147,12 +147,16 @@ async def upload_file( is_image = ext in IMAGE_EXTENSIONS image_data_url = "" if is_image: - # For images: generate base64 data URL for vision models - if len(content) > 10 * 1024 * 1024: # 10MB limit + # Store the original, but send a bounded preview to vision models so a + # batch of images still fits in a single WebSocket message. + if len(content) > MAX_CHAT_IMAGE_BYTES: raise HTTPException(status_code=400, detail="Image too large (max 10MB)") - mime = MIME_MAP.get(ext, "image/png") - b64 = base64.b64encode(content).decode("ascii") - image_data_url = f"data:{mime};base64,{b64}" + try: + image_data_url = build_vision_image_data_url(content, ext) + except VisionImagePreviewError as exc: + raise HTTPException( + status_code=400, detail="Invalid or unsupported image" + ) from exc extracted = f"[图片文件: {file.filename}，需要视觉模型分析]" elif ext in EXTRACTABLE: extracted = extract_text(save_path, ext) diff --git a/backend/app/services/chat_image_preview.py b/backend/app/services/chat_image_preview.py new file mode 100644 index 000000000..d24a47272 --- /dev/null +++ b/backend/app/services/chat_image_preview.py @@ -0,0 +1,73 @@ +"""Bounded previews for image-bearing chat messages.""" + +import base64 +from io import BytesIO + +from PIL import Image, ImageOps, UnidentifiedImageError + +MAX_VISION_IMAGE_BYTES = 512 * 1024 +MAX_VISION_IMAGE_EDGE = 1568 + +_MIME_BY_EXTENSION = { + ".png": "image/png", + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".gif": "image/gif", + ".webp": "image/webp", + ".bmp": "image/bmp", +} + + +class VisionImagePreviewError(ValueError): + """The uploaded image cannot be converted into a safe preview.""" + + +def build_vision_image_data_url(content: bytes, extension: str) -> str: + """Return a bounded data URL while leaving the stored original intact.""" + try: + with Image.open(BytesIO(content)) as opened: + opened.load() + image = ImageOps.exif_transpose(opened) + if ( + len(content) <= MAX_VISION_IMAGE_BYTES + and max(image.size) <= MAX_VISION_IMAGE_EDGE + ): + mime = _MIME_BY_EXTENSION.get(extension, "image/png") + encoded = base64.b64encode(content).decode("ascii") + return f"data:{mime};base64,{encoded}" + + if image.mode in {"RGBA", "LA"} or ( + image.mode == "P" and "transparency" in image.info + ): + rgba = image.convert("RGBA") + background = Image.new("RGB", rgba.size, "white") + background.paste(rgba, mask=rgba.getchannel("A")) + image = background + else: + image = image.convert("RGB") + + image.thumbnail( + (MAX_VISION_IMAGE_EDGE, MAX_VISION_IMAGE_EDGE), + Image.Resampling.LANCZOS, + ) + quality = 85 + for _ in range(16): + output = BytesIO() + image.save(output, format="JPEG", quality=quality, optimize=True) + preview = output.getvalue() + if len(preview) <= MAX_VISION_IMAGE_BYTES: + encoded = base64.b64encode(preview).decode("ascii") + return f"data:image/jpeg;base64,{encoded}" + if quality > 55: + quality -= 10 + else: + width, height = image.size + image = image.resize( + (max(1, int(width * 0.8)), max(1, int(height * 0.8))), + Image.Resampling.LANCZOS, + ) + quality = 75 + except (OSError, UnidentifiedImageError, ValueError) as exc: + raise VisionImagePreviewError("Invalid or unsupported image") from exc + + raise VisionImagePreviewError("Could not create a bounded image preview") diff --git a/backend/tests/test_chat_upload_images.py b/backend/tests/test_chat_upload_images.py new file mode 100644 index 000000000..ad07921eb --- /dev/null +++ b/backend/tests/test_chat_upload_images.py @@ -0,0 +1,43 @@ +"""Regression tests for bounded chat image payloads.""" + +import base64 +import os +from io import BytesIO + +from PIL import Image + +from app.services.chat_image_preview import ( + MAX_VISION_IMAGE_BYTES, + MAX_VISION_IMAGE_EDGE, + build_vision_image_data_url, +) + + +def _decode_data_url(data_url: str) -> bytes: + return base64.b64decode(data_url.split(",", 1)[1]) + + +def test_small_chat_image_keeps_original_encoding() -> None: + output = BytesIO() + Image.new("RGB", (32, 24), "red").save(output, format="PNG") + original = output.getvalue() + + data_url = build_vision_image_data_url(original, ".png") + + assert data_url.startswith("data:image/png;base64,") + assert _decode_data_url(data_url) == original + + +def test_large_chat_image_is_bounded_for_multi_image_messages() -> None: + image = Image.frombytes("RGB", (2400, 1800), os.urandom(2400 * 1800 * 3)) + original = BytesIO() + image.save(original, format="JPEG", quality=100) + assert len(original.getvalue()) > MAX_VISION_IMAGE_BYTES + + data_url = build_vision_image_data_url(original.getvalue(), ".jpg") + preview = _decode_data_url(data_url) + + assert data_url.startswith("data:image/jpeg;base64,") + assert len(preview) <= MAX_VISION_IMAGE_BYTES + with Image.open(BytesIO(preview)) as bounded: + assert max(bounded.size) <= MAX_VISION_IMAGE_EDGE