@@ -4,10 +4,11 @@ from datetime import timedelta from decimal import Decimal from io import BytesIO from typing import Any, Iterator +from pathlib import Path import filetype from django.db.models.fields.files import FieldFile -from PIL import Image +from PIL import Image, UnidentifiedImageError from messages.models import Message from ml_model.services.base import SimpleService @@ -66,18 +67,22 @@ class Claude(SimpleService): messages.append({'role': 'user', 'content': input_message.content}) image = input_message.file if image: - kind = filetype.guess(image.read(20)) - mime = kind.mime if kind else 'application/octet-stream' - normalized_image = Image.open(image) - format = 'jpeg' if kind.extension == 'jpg' else kind.extension - buf = BytesIO() - normalized_image.save(buf, format=format) - image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' - buf.close() - messages[-1]['content'] = [ - {'type': 'text', 'text': input_message.content}, - {'type': 'image_url', 'image_url': {'url': image_url}}, - ] + image_extension = Path(image.name).suffix + if image_extension in ('.png', '.jpg', '.jpeg'): + kind = filetype.guess(image.read(20)) + mime = kind.mime if kind else 'application/octet-stream' + normalized_image = Image.open(image) + format = 'jpeg' if kind.extension == 'jpg' else kind.extension + buf = BytesIO() + normalized_image.save(buf, format=format) + image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' + buf.close() + messages[-1]['content'] = [ + {'type': 'text', 'text': input_message.content}, + {'type': 'image_url', 'image_url': {'url': image_url}}, + ] + else: + raise UnidentifiedImageError result = openrouter_run(version, messages, callback_data, 'Claude') process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice( @@ -3,10 +3,11 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from pathlib import Path import filetype from django.db.models.fields.files import FieldFile -from PIL import Image +from PIL import Image, UnidentifiedImageError from messages.models import Message from ml_model.services.base import SimpleService @@ -64,18 +65,22 @@ class Gemini(SimpleService): messages.append({'role': 'user', 'content': input_message.content}) image = input_message.file if image: - kind = filetype.guess(image.read(20)) - mime = kind.mime if kind else 'application/octet-stream' - normalized_image = Image.open(image) - format = 'jpeg' if kind.extension == 'jpg' else kind.extension - buf = BytesIO() - normalized_image.save(buf, format=format) - image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' - buf.close() - messages[-1]['content'] = [ - {'type': 'text', 'text': input_message.content}, - {'type': 'image_url', 'image_url': {'url': image_url}}, - ] + image_extension = Path(image.name).suffix + if image_extension in ('.png', '.jpg', '.jpeg'): + kind = filetype.guess(image.read(20)) + mime = kind.mime if kind else 'application/octet-stream' + normalized_image = Image.open(image) + format = 'jpeg' if kind.extension == 'jpg' else kind.extension + buf = BytesIO() + normalized_image.save(buf, format=format) + image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' + buf.close() + messages[-1]['content'] = [ + {'type': 'text', 'text': input_message.content}, + {'type': 'image_url', 'image_url': {'url': image_url}}, + ] + else: + raise UnidentifiedImageError start_time = time.time() result = openrouter_run(version, messages, callback_data, 'Gemini') process_time = timedelta(seconds=(time.time() - start_time)) @@ -3,11 +3,12 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO -from typing import Any, Dict, Iterator +from typing import Any, Iterator +from pathlib import Path import filetype from django.db.models.fields.files import FieldFile -from PIL import Image +from PIL import Image, UnidentifiedImageError from messages.models import Message from ml_model.services.base import SimpleService @@ -62,18 +63,22 @@ class Grok(SimpleService): messages.append({'role': 'user', 'content': input_message.content}) image = input_message.file if image: - kind = filetype.guess(image.read(20)) - mime = kind.mime if kind else 'application/octet-stream' - normalized_image = Image.open(image) - format = 'jpeg' if kind.extension == 'jpg' else kind.extension - buf = BytesIO() - normalized_image.save(buf, format=format) - image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' - buf.close() - messages[-1]['content'] = [ - {'type': 'text', 'text': input_message.content}, - {'type': 'image_url', 'image_url': {'url': image_url}}, - ] + image_extension = Path(image.name).suffix + if image_extension in ('.png', '.jpg', '.jpeg'): + kind = filetype.guess(image.read(20)) + mime = kind.mime if kind else 'application/octet-stream' + normalized_image = Image.open(image) + format = 'jpeg' if kind.extension == 'jpg' else kind.extension + buf = BytesIO() + normalized_image.save(buf, format=format) + image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' + buf.close() + messages[-1]['content'] = [ + {'type': 'text', 'text': input_message.content}, + {'type': 'image_url', 'image_url': {'url': image_url}}, + ] + else: + raise UnidentifiedImageError result = openrouter_run(version, messages, callback_data, 'Grok') process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice( @@ -1,6 +1,7 @@ import base64 import time from io import BytesIO +from pathlib import Path import filetype @@ -15,7 +16,7 @@ from ml_model.tasks import openrouter_run from tools.chats.models import Chat from tools.copywrite.models import Copywrite from tools.public_api.models import APIStore -from PIL import Image +from PIL import Image, UnidentifiedImageError class Llama(SimpleService): @@ -67,18 +68,22 @@ class Llama(SimpleService): messages.append({'role': 'user', 'content': input_message.content}) image = input_message.file if image: - kind = filetype.guess(image.read(20)) - mime = kind.mime if kind else 'application/octet-stream' - normalized_image = Image.open(image) - format = 'jpeg' if kind.extension == 'jpg' else kind.extension - buf = BytesIO() - normalized_image.save(buf, format=format) - image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' - buf.close() - messages[-1]['content'] = [ - {'type': 'text', 'text': input_message.content}, - {'type': 'image_url', 'image_url': {'url': image_url}}, - ] + image_extension = Path(image.name).suffix + if image_extension in ('.png', '.jpg', '.jpeg'): + kind = filetype.guess(image.read(20)) + mime = kind.mime if kind else 'application/octet-stream' + normalized_image = Image.open(image) + format = 'jpeg' if kind.extension == 'jpg' else kind.extension + buf = BytesIO() + normalized_image.save(buf, format=format) + image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' + buf.close() + messages[-1]['content'] = [ + {'type': 'text', 'text': input_message.content}, + {'type': 'image_url', 'image_url': {'url': image_url}}, + ] + else: + raise UnidentifiedImageError result = openrouter_run(version, messages, callback_data, 'LLaMA') process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice( @@ -3,10 +3,11 @@ import time from _decimal import Decimal from datetime import timedelta from io import BytesIO +from pathlib import Path import filetype from django.db.models.fields.files import FieldFile -from PIL import Image +from PIL import Image, UnidentifiedImageError from messages.models import Message from ml_model.services.base import SimpleService @@ -61,18 +62,22 @@ class Mistral(SimpleService): messages.append({'role': 'user', 'content': input_message.content}) image = input_message.file if image: - kind = filetype.guess(image.read(20)) - mime = kind.mime if kind else 'application/octet-stream' - normalized_image = Image.open(image) - format = 'jpeg' if kind and kind.extension == 'jpg' else kind.extension - buf = BytesIO() - normalized_image.save(buf, format=format) - image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' - buf.close() - messages[-1]['content'] = [ - {'type': 'text', 'text': input_message.content}, - {'type': 'image_url', 'image_url': {'url': image_url}}, - ] + image_extension = Path(image.name).suffix + if image_extension in ('.png', '.jpg', '.jpeg'): + kind = filetype.guess(image.read(20)) + mime = kind.mime if kind else 'application/octet-stream' + normalized_image = Image.open(image) + format = 'jpeg' if kind.extension == 'jpg' else kind.extension + buf = BytesIO() + normalized_image.save(buf, format=format) + image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' + buf.close() + messages[-1]['content'] = [ + {'type': 'text', 'text': input_message.content}, + {'type': 'image_url', 'image_url': {'url': image_url}}, + ] + else: + raise UnidentifiedImageError start_time = time.time() result = openrouter_run(version, messages, callback_data, 'Mistral') process_time = timedelta(seconds=(time.time() - start_time)) @@ -11,9 +11,10 @@ from rest_framework.permissions import IsAuthenticated from rest_framework.response import Response from rest_framework.status import ( HTTP_402_PAYMENT_REQUIRED, - HTTP_500_INTERNAL_SERVER_ERROR, + HTTP_500_INTERNAL_SERVER_ERROR, HTTP_400_BAD_REQUEST, ) from rest_framework.views import APIView +from PIL import UnidentifiedImageError from messages.models import Message from messages.serializers import MessageSerializer @@ -148,6 +149,13 @@ class MessagesAPIView(APIView): ) try: output_messages = service(chat).make(input_message) + except UnidentifiedImageError: + return Response({ + 'detail': _( + 'Unable to recognize the image. (Supported formats are PNG, JPG, JPEG)' + ) + }, + status=HTTP_400_BAD_REQUEST) except Exception as exc: input_message.is_sent = False input_message.save()