diff --git a/backend/app/api/upload.py b/backend/app/api/upload.py index e0fc14b84..243933955 100644 --- a/backend/app/api/upload.py +++ b/backend/app/api/upload.py @@ -1,14 +1,23 @@ """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 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.storage import ensure_local_path, get_storage_backend, guess_content_type, normalize_storage_key +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, +) from app.services.text_extractor import extract_text as extract_document_text router = APIRouter(prefix="/chat", tags=["chat"]) @@ -23,10 +32,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: @@ -60,7 +66,7 @@ async def upload_file( file: UploadFile = File(...), agent_id: str = Form(""), current_user: User = Depends(get_current_user), -): +) -> dict[str, object]: """Upload a file for chat context. Saves to agent workspace/uploads/ and returns extracted text.""" if not file.filename: raise HTTPException(status_code=400, detail="No filename") @@ -97,12 +103,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..2afca1167 --- /dev/null +++ b/backend/app/services/chat_image_preview.py @@ -0,0 +1,87 @@ +"""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 +MAX_VISION_IMAGE_PIXELS = 40_000_000 + +_MIME_BY_EXTENSION = { + ".png": "image/png", + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".gif": "image/gif", + ".webp": "image/webp", + ".bmp": "image/bmp", +} + +_MIME_BY_FORMAT = { + "BMP": "image/bmp", + "GIF": "image/gif", + "JPEG": "image/jpeg", + "PNG": "image/png", + "WEBP": "image/webp", +} + + +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: + if opened.width * opened.height > MAX_VISION_IMAGE_PIXELS: + raise VisionImagePreviewError("Image dimensions are too large") + source_format = opened.format + 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_FORMAT.get( + source_format, _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..3017b7d02 --- /dev/null +++ b/backend/tests/test_chat_upload_images.py @@ -0,0 +1,63 @@ +"""Regression tests for bounded chat image payloads.""" + +import base64 +import os +from io import BytesIO + +import pytest +from PIL import Image + +from app.services import chat_image_preview +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_small_chat_image_uses_detected_format_for_mime_type() -> None: + output = BytesIO() + Image.new("RGB", (32, 24), "red").save(output, format="PNG") + + data_url = build_vision_image_data_url(output.getvalue(), ".jpg") + + assert data_url.startswith("data:image/png;base64,") + + +def test_image_pixel_limit_is_checked_before_decode(monkeypatch) -> None: + output = BytesIO() + Image.new("RGB", (32, 24), "red").save(output, format="PNG") + monkeypatch.setattr(chat_image_preview, "MAX_VISION_IMAGE_PIXELS", 100) + + with pytest.raises(chat_image_preview.VisionImagePreviewError): + build_vision_image_data_url(output.getvalue(), ".png") + + +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