Convert images to jpeg (#2737)

* convert to jpeg

* k

* typing
This commit is contained in:
pablodanswer 2024-10-09 14:56:45 -07:00 committed by GitHub
parent 804de3248e
commit 30d17ef9ee
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -3,6 +3,7 @@ import io
import uuid
from collections.abc import Callable
from collections.abc import Generator
from typing import Tuple
from fastapi import APIRouter
from fastapi import Depends
@ -11,6 +12,7 @@ from fastapi import Request
from fastapi import Response
from fastapi import UploadFile
from fastapi.responses import StreamingResponse
from PIL import Image
from pydantic import BaseModel
from sqlalchemy.orm import Session
@ -508,6 +510,21 @@ def seed_chat(
"""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")
def upload_files_for_chat(
files: list[UploadFile],
@ -570,19 +587,25 @@ def upload_files_for_chat(
for file in files:
if file.content_type in image_content_types:
file_type = ChatFileType.IMAGE
# Convert image to JPEG
file_content, new_content_type = convert_to_jpeg(file)
elif file.content_type in document_content_types:
file_type = ChatFileType.DOC
file_content = io.BytesIO(file.file.read())
new_content_type = file.content_type or ""
else:
file_type = ChatFileType.PLAIN_TEXT
file_content = io.BytesIO(file.file.read())
new_content_type = file.content_type or ""
# store the raw file
# store the file (now JPEG for images)
file_id = str(uuid.uuid4())
file_store.save_file(
file_name=file_id,
content=file.file,
content=file_content,
display_name=file.filename,
file_origin=FileOrigin.CHAT_UPLOAD,
file_type=file.content_type or file_type.value,
file_type=new_content_type or file_type.value,
)
# if the file is a doc, extract text and store that so we don't need