Fix image wonkiness (#3735)

* fix images

* quick nit

* quick nit

* update

* update for clarity
This commit is contained in:
pablonyx 2025-01-22 18:38:51 -08:00 committed by GitHub
parent 69c60feda4
commit e94ffbc2a1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -5,7 +5,6 @@ import os
import uuid import uuid
from collections.abc import Callable from collections.abc import Callable
from collections.abc import Generator from collections.abc import Generator
from typing import Tuple
from uuid import UUID from uuid import UUID
from fastapi import APIRouter from fastapi import APIRouter
@ -15,7 +14,6 @@ from fastapi import Request
from fastapi import Response from fastapi import Response
from fastapi import UploadFile from fastapi import UploadFile
from fastapi.responses import StreamingResponse from fastapi.responses import StreamingResponse
from PIL import Image
from pydantic import BaseModel from pydantic import BaseModel
from sqlalchemy.orm import Session from sqlalchemy.orm import Session
@ -595,21 +593,6 @@ def seed_chat_from_slack(
"""File upload""" """File upload"""
def convert_to_jpeg(file: UploadFile) -> Tuple[io.BytesIO, str]:
try:
with Image.open(file.file) as img:
if img.mode != "RGB":
img = img.convert("RGB")
jpeg_io = io.BytesIO()
img.save(jpeg_io, format="JPEG", quality=85)
jpeg_io.seek(0)
return jpeg_io, "image/jpeg"
except Exception as e:
raise HTTPException(
status_code=400, detail=f"Failed to convert image: {str(e)}"
)
@router.post("/file") @router.post("/file")
def upload_files_for_chat( def upload_files_for_chat(
files: list[UploadFile], files: list[UploadFile],
@ -645,6 +628,9 @@ def upload_files_for_chat(
) )
for file in files: for file in files:
if not file.content_type:
raise HTTPException(status_code=400, detail="File content type is required")
if file.content_type not in allowed_content_types: if file.content_type not in allowed_content_types:
if file.content_type in image_content_types: if file.content_type in image_content_types:
error_detail = "Unsupported image file type. Supported image types include .jpg, .jpeg, .png, .webp." error_detail = "Unsupported image file type. Supported image types include .jpg, .jpeg, .png, .webp."
@ -676,22 +662,27 @@ def upload_files_for_chat(
file_info: list[tuple[str, str | None, ChatFileType]] = [] file_info: list[tuple[str, str | None, ChatFileType]] = []
for file in files: for file in files:
if file.content_type in image_content_types: file_type = (
file_type = ChatFileType.IMAGE ChatFileType.IMAGE
# Convert image to JPEG if file.content_type in image_content_types
file_content, new_content_type = convert_to_jpeg(file) else ChatFileType.CSV
elif file.content_type in csv_content_types: if file.content_type in csv_content_types
file_type = ChatFileType.CSV else ChatFileType.DOC
file_content = io.BytesIO(file.file.read()) if file.content_type in document_content_types
new_content_type = file.content_type or "" else ChatFileType.PLAIN_TEXT
elif file.content_type in document_content_types: )
file_type = ChatFileType.DOC
file_content = io.BytesIO(file.file.read()) if file_type == ChatFileType.IMAGE:
new_content_type = file.content_type or "" file_content = file.file
# NOTE: Image conversion to JPEG used to be enforced here.
# This was removed to:
# 1. Preserve original file content for downloads
# 2. Maintain transparency in formats like PNG
# 3. Ameliorate issue with file conversion
else: else:
file_type = ChatFileType.PLAIN_TEXT
file_content = io.BytesIO(file.file.read()) file_content = io.BytesIO(file.file.read())
new_content_type = file.content_type or ""
new_content_type = file.content_type
# store the file (now JPEG for images) # store the file (now JPEG for images)
file_id = str(uuid.uuid4()) file_id = str(uuid.uuid4())