@@ -1,6 +1,7 @@ import logging.config from pathlib import Path +from PIL import ImageFile from celery.schedules import crontab from environs import Env @@ -336,11 +337,12 @@ YANDEX_CLOUD_ID = env.str('YANDEX_CLOUD_ID', 'defaultapikey') OPENAI_PROXY_HOST = env.str('OPENAI_PROXY_HOST', 'neuron-proxy:8080') UPSCALE_MULTIPLIER_HOST = env.str('UPSCALE_MULTIPLIER_HOST', 'packet:8080') - +# FILES MAX_UPLOAD_SIZE_PER_MODEL = { 'raifgpt': 50, 'default': 8, } +ImageFile.LOAD_TRUNCATED_IMAGES = True # Payments @@ -4,6 +4,7 @@ import logging import re import subprocess import time +import zipfile from concurrent.futures import ThreadPoolExecutor, as_completed @@ -17,14 +18,12 @@ from datetime import timedelta from decimal import Decimal from io import BufferedReader, BytesIO from math import ceil -from pathlib import Path from typing import Generator, List, Optional, Dict, Any, Tuple import docx2txt import filetype import httpx import tiktoken -from django.core.files.uploadedfile import UploadedFile from langchain.chains import ConversationChain from langchain_core.chat_history import InMemoryChatMessageHistory from langchain_core.messages import ( @@ -37,14 +36,14 @@ from langchain_core.prompts.prompt import PromptTemplate from langchain_core.runnables import RunnableWithMessageHistory from langchain_openai.chat_models import ChatOpenAI from langchain_text_splitters import RecursiveCharacterTextSplitter -from PIL import Image, UnidentifiedImageError +from PIL import Image from redis.commands.search.document import Document from redis.commands.search.query import Query from backend import settings from messages.models import BaseStore, Message from ml_model.constants import TEMPORARY_TEST_TEXT -from ml_model.exceptions import GenerationException +from ml_model.exceptions import FileExtensionNotSupported from ml_model.models import ( ModelConfiguration, NeuronModel @@ -157,34 +156,48 @@ class Chatgpt(SimpleService): normalized_image = None embedding_tokens = 0 if file: - file_extension = Path(file.name).suffix - if file_extension == '.pdf': - raw_text = self.get_pdf_data(file) - text = re.sub(r'\n{2,}', '\n', raw_text) - chunks = self.split_text_to_chunks(text) - elif file_extension in ('.doc', '.docx'): - raw_text = self.get_word_data(file_extension, file) - text = re.sub(r'\n{2,}', '\n', raw_text) - chunks = self.split_text_to_chunks(text) - elif file_extension == '.xlsx': - chunks = self.split_text_to_chunks(self.get_xlsx_data(file)) - else: - image = file - if image: try: - kind = filetype.guess(input_message.file.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() - image_size = normalized_image.size - image_data = {'type': 'image_url', 'image_url': {'url': image_url}} - input_content.append(image_data) - except UnidentifiedImageError: - raise Exception(_('Unable to recognize the image. (Supported formats are PNG, JPG, JPEG)')) + file_bytes = input_message.file.read() + kind = filetype.guess(file_bytes[:20]) + file_extension = kind.extension + if file_extension == 'zip': + file_extension = None + signatures = { + 'xlsx': 'xl/workbook.xml', + 'docx': 'word/document.xml' + } + with zipfile.ZipFile(BytesIO(file_bytes), 'r') as zip_file: + namelist = zip_file.namelist() + for format_name, required_file in signatures.items(): + if required_file in namelist: + file_extension = format_name + break + if not file_extension: + raise + if file_extension == 'pdf': + raw_text = self.get_pdf_data(file_bytes) + text = re.sub(r'\n{2,}', '\n', raw_text) + chunks = self.split_text_to_chunks(text) + elif file_extension in ('doc', 'docx'): + raw_text = self.get_word_data(file_extension, file_bytes) + text = re.sub(r'\n{2,}', '\n', raw_text) + chunks = self.split_text_to_chunks(text) + elif file_extension == 'xlsx': + chunks = self.split_text_to_chunks(self.get_xlsx_data(file_bytes)) + else: + image = file + mime = kind.mime if kind else 'application/octet-stream' + normalized_image = Image.open(BytesIO(file_bytes)) + format = 'jpeg' if file_extension == 'jpg' else file_extension + buf = BytesIO() + normalized_image.save(buf, format=format) + image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' + buf.close() + image_size = normalized_image.size + image_data = {'type': 'image_url', 'image_url': {'url': image_url}} + input_content.append(image_data) + except: + raise FileExtensionNotSupported(['PDF', 'DOC', 'DOCX', 'XLSX', 'JPG', 'JPEG', 'PNG']) for proxy in Proxy.objects.all(): self.llm = ChatOpenAI( model=model_name, @@ -666,14 +679,13 @@ class Chatgpt(SimpleService): } return search_context_size, json_data - def get_pdf_data(self, pdf_file: UploadedFile) -> str: + def get_pdf_data(self, pdf_data: bytes) -> str: """ Extracting text from pdf-file :param pdf_file: uploaded pdf file :return: pdf-file content """ try: - pdf_data = pdf_file.read() doc = fitz.open(stream=pdf_data, filetype="pdf") raw_text = '' for page_number, page in enumerate(doc, start=1): @@ -686,14 +698,14 @@ class Chatgpt(SimpleService): return f"Ошибка: Файл поврежден или не может быть прочитан." return f'Содержимое файла: {raw_text.strip()}' - def get_xlsx_data(self, xlsx_file: UploadedFile) -> str: + def get_xlsx_data(self, xlsx_data: bytes) -> str: """ Extracting text from xlsx-file :param xlsx_file: uploaded xlsx file :return: xlsx_file content """ try: - xlsx_content = BytesIO(xlsx_file.read()) + xlsx_content = BytesIO(xlsx_data) workbook = openpyxl.load_workbook(xlsx_content) raw_text = '' for sheet_name in workbook.sheetnames: @@ -704,7 +716,7 @@ class Chatgpt(SimpleService): raw_text = 'Произошла ошибка во время чтения файла' return f'Содержимое файла: {raw_text}' - def get_word_data(self, extension: str, word_file: UploadedFile) -> str: + def get_word_data(self, extension: str, word_data: bytes) -> str: """ Extracting text from word-file :param extension: extension of uploaded word file @@ -712,17 +724,16 @@ class Chatgpt(SimpleService): :return: word-file content """ try: - file_content = word_file.read() - if extension == '.docx': - text = docx2txt.process(BytesIO(file_content)) - elif extension == '.doc': + if extension == 'docx': + text = docx2txt.process(BytesIO(word_data)) + elif extension == 'doc': process = subprocess.Popen( ['antiword', '-w', '0', '-'], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, ) - text, _ = process.communicate(input=file_content) + text, _ = process.communicate(input=word_data) text = text.decode('utf-8') else: text = ''