@@ -112,6 +112,7 @@ class EmbeddingService: user_content, model: str = 'text-embedding-3-large', index_name: str = 'ml_model-index', + top_k: int = 10, ): redis_client = redis.Redis(host=settings.REDIS_HOST, port=settings.REDIS_PORT, db=0) embedding_tokens = 0 @@ -148,6 +149,7 @@ class EmbeddingService: redis_client=redis_client, message_uid=message_uid, user_query_embeddings=query_embedding, + top_k=top_k, index_name=index_name, ) ] @@ -1,17 +1,8 @@ import base64 import itertools import logging -import re -import subprocess import time -import zipfile -from concurrent.futures import ThreadPoolExecutor, as_completed - -import numpy as np -import openpyxl -import fitz -import redis from django.utils.translation import gettext_lazy as _ from datetime import timedelta @@ -20,7 +11,6 @@ from io import BufferedReader, BytesIO from math import ceil from typing import Generator, List, Optional, Dict, Any, Tuple -import docx2txt import filetype import httpx import tiktoken @@ -35,21 +25,16 @@ from langchain_core.messages import ( 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 -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 FileExtensionNotSupported -from ml_model.models import ( - ModelConfiguration, - NeuronModel -) +from ml_model.exceptions import FileExtensionNotSupported, CorruptedFileError +from ml_model.models import ModelConfiguration, NeuronModel +from ml_model.services.EmbeddingService import EmbeddingService +from ml_model.services.FileService import FileProcessingService from ml_model.services.base import SimpleService -from ml_model.tasks import drop_redis_vectors from payments.exceptions.insufficient_balance import InsufficientBalance from payments.selectors.payment_plan_selector import PaymentPlanSelector from poller.models import Proxy @@ -73,32 +58,30 @@ class Chatgpt(SimpleService): 'input': Decimal('0.0003'), 'output': Decimal('0.0003'), 'web_search': { - 'low': Decimal('12.5'), # 1 call - 'medium': Decimal('13.75'), # 1 call - 'high': Decimal('15') # 1 call - } + 'low': Decimal('12.5'), # 1 call + 'medium': Decimal('13.75'), # 1 call + 'high': Decimal('15'), # 1 call + }, }, 'gpt-4o': { 'input': Decimal('0.005'), 'output': Decimal('0.005'), 'web_search': { - 'low': Decimal('15'), # 1 call - 'medium': Decimal('17.5'), # 1 call - 'high': Decimal('25') # 1 call - } + 'low': Decimal('15'), # 1 call + 'medium': Decimal('17.5'), # 1 call + 'high': Decimal('25'), # 1 call + }, }, - 'gpt-oss-120b': { - 'input': Decimal('0.0002'), - 'output': Decimal('0.0002') - } + 'gpt-oss-120b': {'input': Decimal('0.0002'), 'output': Decimal('0.0002')}, } TOOLS_TOKEN_COSTS = { - 'text-embedding-3-large': { - 'output': Decimal('0.000065') - } + 'text-embedding-3-small': {'output': Decimal('0.00001')}, + 'text-embedding-3-large': {'output': Decimal('0.000065')}, } + EMBEDDING_MODEL_FOR_BILLING = 'text-embedding-3-small' + TOKEN_LIMITS = { 'o3-mini': 100_000, 'gpt-4o-mini': 64_000, @@ -131,19 +114,24 @@ class Chatgpt(SimpleService): normalized_image = None embedding_tokens = 0 chunks = [] + text_chunks: list[str] = [] if file: - try: - file_bytes = input_message.file.read() - kind = filetype.guess(file_bytes[:20]) - raw_file_extension = kind.extension - file_extension = self._get_file_extension(raw_file_extension, file_bytes) - if file_extension in ('pdf', 'doc', 'docx', 'xlsx'): - chunks = self._get_file_data(file_extension, file_bytes) - else: - image = file - normalized_image, image_size, image_data = self._get_image_data(file_bytes, file_extension) - input_content.append(image_data) - except: + file_service = FileProcessingService + file_bytes = input_message.file.read() + kind = filetype.guess(file_bytes[:20]) + if not kind: + raise CorruptedFileError + raw_file_extension = kind.extension + file_extension = file_service.get_file_extension(raw_file_extension, file_bytes) + if file_extension in ('pdf', 'doc', 'docx', 'xlsx'): + text = file_service.get_file_data(file_extension, file_bytes) + text_chunks = EmbeddingService.split_text_to_chunks(text) + chunks = [HumanMessage(content=chunk_text) for chunk_text in text_chunks] + elif file_extension in ('jpg', 'jpeg', 'png', 'webp'): + image = file + normalized_image, image_size, image_data = self._get_image_data(file_bytes, file_extension) + input_content.append(image_data) + else: raise FileExtensionNotSupported(['PDF', 'DOC', 'DOCX', 'XLSX', 'JPG', 'JPEG', 'PNG', 'WEBP']) for proxy in Proxy.objects.all(): self.llm = ChatOpenAI( @@ -173,28 +161,44 @@ class Chatgpt(SimpleService): if model_name == 'gpt-oss-120b': system = chat_history.messages.pop(0) messages = [ - {'role': 'user' if isinstance(msg, HumanMessage) else 'assistant', 'content': msg.content} + { + 'role': 'user' if isinstance(msg, HumanMessage) else 'assistant', + 'content': msg.content, + } for msg in chat_history.messages ] messages.insert(0, {'role': 'system', 'content': system.content}) messages.insert(0, {'role': 'system', 'content': user_system_prompt}) if file and not image: if sum([len(chunk.content) for chunk in chunks]) > 20_000: - document_name = chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] - embedding_tokens, file_data = self.get_large_file_data(chunks, proxy, input_message.content) - messages[-1]['content'] = self.make_embeddings_prompt( - document_name=document_name, section_texts=file_data, question=input_message.content + document_name = ( + chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] + ) + embedding_tokens, file_data = EmbeddingService.get_large_file_data( + self.store.messages.first().pk, + text_chunks, + proxy, + input_message.content, + model='text-embedding-3-small', + index_name='ml_model-index-1536', + ) + messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( + document_name=document_name, + section_texts=file_data, + question=input_message.content, ) else: - messages[-1]['content'] = (f'Используй системный промпт. Содержание файла: ' - f'{chunks}. Вопрос: {input_message.content}') - json_data = { - 'model': f'openai/{model_name}', - 'messages': messages - } + messages[-1]['content'] = ( + 'Используй системный промпт. Содержание файла: ' + f'{"".join(text_chunks)}. Вопрос: {input_message.content}' + ) + json_data = {'model': f'openai/{model_name}', 'messages': messages} response = httpx.post( - url='https://openrouter.ai/api/v1/chat/completions', proxy=f'{proxy.protocol}://{proxy.address}', - headers={'Authorization': f'Bearer {settings.OPENROUTER_API_KEY}'}, timeout=600, json=json_data + url='https://openrouter.ai/api/v1/chat/completions', + proxy=f'{proxy.protocol}://{proxy.address}', + headers={'Authorization': f'Bearer {settings.OPENROUTER_API_KEY}'}, + timeout=600, + json=json_data, ) if ( (data := response.json()) @@ -212,7 +216,10 @@ class Chatgpt(SimpleService): elif model_name == 'o3-mini': system = chat_history.messages.pop(0) messages = [ - {'role': 'user' if isinstance(msg, HumanMessage) else 'assistant', 'content': msg.content} + { + 'role': 'user' if isinstance(msg, HumanMessage) else 'assistant', + 'content': msg.content, + } for msg in chat_history.messages ] messages.insert(0, {'role': 'system', 'content': system.content}) @@ -225,13 +232,24 @@ class Chatgpt(SimpleService): elif file: if sum([len(chunk.content) for chunk in chunks]) > 20_000: document_name = chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] - embedding_tokens, file_data = self.get_large_file_data(chunks, proxy, input_message.content) - messages[-1]['content'] = self.make_embeddings_prompt( - document_name=document_name, section_texts=file_data, question=input_message.content + embedding_tokens, file_data = EmbeddingService.get_large_file_data( + self.store.messages.first().pk, + text_chunks, + proxy, + input_message.content, + model='text-embedding-3-small', + index_name='ml_model-index-1536', + ) + messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( + document_name=document_name, + section_texts=file_data, + question=input_message.content, ) else: - messages[-1]['content'] = (f'Используй системный промпт. Содержание файла: ' - f'{chunks}. Вопрос: {input_message.content}') + messages[-1]['content'] = ( + 'Используй системный промпт. Содержание файла: ' + f'{"".join(text_chunks)}. Вопрос: {input_message.content}' + ) json_data = { 'model': model_name, 'messages': messages @@ -246,27 +264,38 @@ class Chatgpt(SimpleService): messages.insert(0, {'role': 'system', 'content': system.content}) messages.insert(0, {'role': 'system', 'content': user_system_prompt}) search_context_size, json_data = self.get_web_search_data( - info.get('web_search', 'Средний контекст'), - model_name, - messages + info.get('web_search', 'Средний контекст'), model_name, messages ) info['web_search'] = search_context_size if image: messages[-1]['content'] = [ {'type': 'input_text', 'text': input_message.content}, - {'type': 'input_image', 'image_url': image_data['image_url']['url']} + {'type': 'input_image', 'image_url': image_data['image_url']['url']}, ] elif file: if sum([len(chunk.content) for chunk in chunks]) > 20_000: document_name = chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] - embedding_tokens, file_data = self.get_large_file_data(chunks, proxy, input_message.content) - messages[-1]['content'] = self.make_embeddings_prompt( - document_name=document_name, section_texts=file_data, question=input_message.content + embedding_tokens, file_data = EmbeddingService.get_large_file_data( + self.store.messages.first().pk, + text_chunks, + proxy, + input_message.content, + model='text-embedding-3-small', + index_name='ml_model-index-1536', + ) + messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( + document_name=document_name, + section_texts=file_data, + question=input_message.content, ) else: - messages[-1]['content'] = (f'Используй системный промпт. Содержание файла: ' - f'{chunks}. Вопрос: {input_message.content}') - input_tokens, output_tokens, response = self.call_openai_api(proxy=proxy, endpoint='responses',json_data=json_data) + messages[-1]['content'] = ( + 'Используй системный промпт. Содержание файла: ' + f'{"".join(text_chunks)}. Вопрос: {input_message.content}' + ) + input_tokens, output_tokens, response = self.call_openai_api( + proxy=proxy, endpoint='responses', json_data=json_data + ) elif image: response = self.llm.invoke(llm_input) chat_history.add_ai_message(response) @@ -274,12 +303,23 @@ class Chatgpt(SimpleService): input_tokens = self.count_text_tokens([*chat_history.messages]) if sum([len(chunk.content) for chunk in chunks]) > 20_000: document_name = chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] - embedding_tokens, file_data = self.get_large_file_data(chunks, proxy, input_message.content) + embedding_tokens, file_data = EmbeddingService.get_large_file_data( + self.store.messages.first().pk, + text_chunks, + proxy, + input_message.content, + model='text-embedding-3-small', + index_name='ml_model-index-1536', + ) user_input = [ SystemMessage(content=user_system_prompt), - HumanMessage(self.make_embeddings_prompt( - document_name=document_name, section_texts=file_data, question=input_message.content - )) + HumanMessage( + EmbeddingService.make_embeddings_prompt( + document_name=document_name, + section_texts=file_data, + question=input_message.content, + ) + ), ] input_tokens += self.count_text_tokens(user_input) response = conversation.invoke( @@ -290,9 +330,11 @@ class Chatgpt(SimpleService): input = [ SystemMessage(content=user_system_prompt), HumanMessage( - content=f'Используй системный промпт. Содержание файла: ' - f'{chunks}. Вопрос: {input_message.content}' - ) + content=( + 'Используй системный промпт. Содержание файла: ' + f'{"".join(text_chunks)}. Вопрос: {input_message.content}' + ) + ), ] input_tokens += self.count_text_tokens(input) response = conversation.invoke( @@ -393,7 +435,10 @@ class Chatgpt(SimpleService): total_tokens += self.count_image_tokens(image_size) input_cost = self.TOKENS_COST[model]['input'] * total_tokens if embedding_tokens > 0: - input_cost += embedding_tokens * self.TOOLS_TOKEN_COSTS['text-embedding-3-large']['output'] + input_cost += ( + embedding_tokens + * self.TOOLS_TOKEN_COSTS[self.EMBEDDING_MODEL_FOR_BILLING]['output'] + ) if input_cost > balance: raise InsufficientBalance(balance, input_cost) @@ -416,7 +461,10 @@ class Chatgpt(SimpleService): if info.get('code_interpreter', False): price += self.TOKENS_COST[model]['code_interpreter'] if embedding_tokens > 0: - price += self.TOOLS_TOKEN_COSTS['text-embedding-3-large']['output'] * embedding_tokens + price += ( + self.TOOLS_TOKEN_COSTS[self.EMBEDDING_MODEL_FOR_BILLING]['output'] + * embedding_tokens + ) return price.quantize(Decimal('0.1'), rounding='ROUND_UP') def count_image_tokens(self, image_size: tuple, model_version: str = 'gpt-4o') -> int: @@ -461,20 +509,6 @@ class Chatgpt(SimpleService): return total_tokens - def _get_file_extension(self, raw_file_extension: str, file_bytes: bytes) -> str: - if raw_file_extension == 'zip': - 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: - return format_name - raise - return raw_file_extension - def _get_image_data(self, file_bytes: bytes, file_extension: str) -> Tuple: normalized_image = Image.open(BytesIO(file_bytes)).convert('RGB') buf = BytesIO() @@ -486,15 +520,6 @@ class Chatgpt(SimpleService): image_data = {'type': 'image_url', 'image_url': {'url': image_url}} return normalized_image, image_size, image_data - def _get_file_data(self, file_extension: str, file_bytes: bytes) -> list[HumanMessage]: - is_word = file_extension in ('doc', 'docx') - method_name = 'word' if is_word else file_extension - operation = getattr(self, f'get_{method_name}_data') - text = operation(file_extension, file_bytes) if is_word else operation(file_bytes) - if file_extension != 'xlsx': - text = re.sub(r'\n{2,}', '\n', text) - return self.split_text_to_chunks(text) - def _get_input_tokens(self, file, image, chunks, chat_history, llm_input): input_embedding_tokens = 0 if file and not image: @@ -509,43 +534,11 @@ class Chatgpt(SimpleService): input_tokens = self.count_text_tokens([*chat_history.messages, *llm_input]) return input_tokens, input_embedding_tokens - def get_large_file_data(self, chunks, proxy, user_content): - redis_client = redis.Redis(host=settings.REDIS_HOST, port=settings.REDIS_PORT, db=0) - embedding_tokens = 0 - message_uid = str(self.store.messages.first().pk).replace('-', '_') - with httpx.Client( - base_url='https://api.openai.com/v1/', - proxy=f'{proxy.protocol}://{proxy.address}', - headers={'Authorization': f'Bearer {settings.OPENAI_API_KEY}'}, - timeout=600, - ) as client: - threads = [] - with ThreadPoolExecutor(max_workers=settings.MAX_THREADS) as executor: - for chunk_id, chunk in enumerate(chunks): - threads.append( - executor.submit(self.process_chunk, client, chunk, redis_client, message_uid, chunk_id) - ) - for thread in as_completed(threads): - embedding_tokens += thread.result() - query_embedding, e_total_tokens = self.get_embedding(client=client, content=user_content) - embedding_tokens += e_total_tokens - result = [ - s['section_text'] - for s in self.search_via_embeddings( - redis_client=redis_client, - message_uid=message_uid, - user_query_embeddings=query_embedding - ) - ] - drop_redis_vectors.delay(message_uid) - redis_client.close() - return embedding_tokens, result - def get_web_search_data(self, search_size: str, model_name: str, messages: List[Dict[str, any]]): search_context_sizes = { 'Малый контекст': 'low', 'Средний контекст': 'medium', - 'Большой контекст': 'high' + 'Большой контекст': 'high', } search_context_size = search_context_sizes.get(search_size) json_data = { @@ -555,202 +548,23 @@ class Chatgpt(SimpleService): { 'type': 'web_search_preview', 'search_context_size': search_context_size, - 'user_location': {'type': 'approximate', 'country': 'RU'} + 'user_location': {'type': 'approximate', 'country': 'RU'}, } - ] + ], } return search_context_size, json_data - 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: - doc = fitz.open(stream=pdf_data, filetype="pdf") - raw_text = '' - for page_number, page in enumerate(doc, start=1): - content = page.get_text("text") - if content: - raw_text += content - doc.close() - fitz.TOOLS.store_shrink(100) - except Exception: - return f"Ошибка: Файл поврежден или не может быть прочитан." - return f'Содержимое файла: {raw_text.strip()}' - - 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_data) - workbook = openpyxl.load_workbook(xlsx_content) - raw_text = '' - for sheet_name in workbook.sheetnames: - sheet = workbook[sheet_name] - for row in sheet.iter_rows(values_only=True): - raw_text += f'Данные ряда: {row}\n' - except Exception: - raw_text = 'Произошла ошибка во время чтения файла' - return f'Содержимое файла: {raw_text}' - - def get_word_data(self, extension: str, word_data: bytes) -> str: - """ - Extracting text from word-file - :param extension: extension of uploaded word file - :param word_file: uploaded word file - :return: word-file content - """ - try: - 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=word_data) - text = text.decode('utf-8') - else: - text = '' - except Exception: - text = 'Файл поврежден или не может быть прочитан.' - if text.strip(): - return f'Это текст, извлечённый из загруженного WORD-файла:\n{text}' - else: - return 'Файл пуст или содержит изображения, из которых невозможно извлечь текст.' - - def split_text_to_chunks( - self, raw_text: str, chunk_size: int = 4000, overlap: int = 200 - ) -> list[HumanMessage]: - """ - Splitting file raw text to chunks - :param raw_text: full text which file includes - :param chunk_size: еhe maximum size of each chunk - :param overlap: еhe number of overlapping characters between chunks - :return: list of chunks - """ - text_splitter = RecursiveCharacterTextSplitter( - chunk_size=chunk_size, chunk_overlap=overlap, length_function=len, separators=["\n\n", "\n", ".", " ", ""] - ) - chunks = text_splitter.split_text(raw_text) - return [HumanMessage(chunk) for chunk in chunks] - - def process_chunk( - self, client: httpx.Client, chunk: HumanMessage, redis_client: redis.Redis, message_uid:str, chunk_id: int - ) -> int: - ''' - A method for getting and saving embeddings from a single chunk - :param client: Httpx client - :param chunk: a HumanMessage object with a content as a part of a full text - :param redis_client: Redis client - :param message_uid: UID of user's message - :param chunk_id: a sequence number of a chunk - ''' - embedding, e_total_tokens = self.get_embedding(client=client, content=chunk.content) - self.save_embeddings( - redis_client=redis_client, - message_uid=message_uid, - chunk_id=chunk_id, - text=chunk.content, - embeddings=embedding - ) - return e_total_tokens - - def get_embedding(self, client: httpx.Client, content: str) -> Tuple[List[float], int]: - ''' - A method for converting raw text (content) into embeddings - using OpenAI API request - :param client: Httpx client - :param content: raw text of a chunk - ''' - response = client.post( - url="embeddings", - json={ - 'model': 'text-embedding-3-large', - 'input': content - } - ) - response.raise_for_status() - data = response.json() - return data['data'][0]['embedding'], data['usage']['total_tokens'] - - def save_embeddings( - self, redis_client: redis.Redis, message_uid: str, chunk_id: int, text: str, embeddings: List[float] - ) -> None: - ''' - A method for saving embeddings in Redis - :param redis_client: Redis client - :param message_uid: UID of user's message - :param chunk_id: a sequence number of a chunk - :param text: a chunk content - :param embeddings: a list of embeddings getting from a chunk - ''' - embeddings_bytes = np.array(embeddings).astype(dtype=np.float32).tobytes() - redis_client.hset( - f'ml_model:messages:{message_uid}:vectors:{chunk_id}', - mapping={ - 'message_uid': message_uid, - 'section_text': text, - 'section_embeddings': embeddings_bytes - } - ) - - def search_via_embeddings( - self, redis_client: redis.Redis, message_uid: str, user_query_embeddings: List[float], top_k: int = 10 - ) -> List[Document]: - ''' - A method for searching similar vectors to user's query - :param redis_client: Redis client - :param message_uid: UID of user's message - :param user_query_embeddings: a list of embeddings getting from user's query - :param top_k: a number of max return documents - ''' - base_query = f'@message_uid:{{{message_uid}}}=>[KNN {top_k} @section_embeddings $vector AS vector_score]' - query = ( - Query(base_query) - .return_fields('section_text') - .sort_by("vector_score") - .paging(0, top_k) - .dialect(2) - ) - params_dict = {"vector": np.array(user_query_embeddings).astype(dtype=np.float32).tobytes()} - results = redis_client.ft('ml_model-index').search(query, params_dict) - return results.docs - - def make_embeddings_prompt(self, document_name: str, section_texts: List[str], question: str) -> str: - ''' - A method for making a prompt using found embeddings - :param document_name: name of the loaded document - :param section_texts: list of sections' contents - :param question: user question - ''' - return f"""Ты — аналитик данных. Отвечай только на основе предоставленного контекста. - Название файла: {document_name} - Фрагменты: - { - '\n'.join(section_texts) - } - Вопрос: {question} - """ - def call_openai_api( - self, proxy: Proxy, endpoint: str, json_data: Dict[str, Any] + self, proxy: Proxy, endpoint: str, json_data: Dict[str, Any] ) -> Tuple[Any, Any, AIMessage] | Tuple[List[float], int]: - ''' + """ A method for sending a request to official openai API :param proxy: Proxy settings object with protocol and address. :param endpoint: Str URL part for the OpenAI API request :param json_data: Payload for the OpenAI API request :return: Tuple of (input_tokens, output_tokens, AIMessage instance with response content) :raises: Exception: If the response is invalid or incomplete - ''' + """ with httpx.Client( base_url='https://api.openai.com/v1', proxy=f'{proxy.protocol}://{proxy.address}', @@ -6,9 +6,11 @@ import filetype from langchain_core.messages import HumanMessage, SystemMessage from messages.models import Message -from ml_model.exceptions import FileExtensionNotSupported +from ml_model.exceptions import FileExtensionNotSupported, CorruptedFileError from ml_model.models import NeuronModel from ml_model.services import Chatgpt +from ml_model.services.EmbeddingService import EmbeddingService +from ml_model.services.FileService import FileProcessingService from poller.models import Proxy @@ -103,18 +105,23 @@ class Chatgpt_5(Chatgpt): image_size = None embedding_tokens = 0 chunks = [] + text_chunks: list[str] = [] if file: - try: - file_bytes = input_message.file.read() - kind = filetype.guess(file_bytes[:20]) - raw_file_extension = kind.extension - file_extension = self._get_file_extension(raw_file_extension, file_bytes) - if file_extension in ('pdf', 'doc', 'docx', 'xlsx'): - chunks = self._get_file_data(file_extension, file_bytes) - else: - image = file - _, image_size, image_data = self._get_image_data(file_bytes, file_extension) - except Exception: + file_service = FileProcessingService + file_bytes = input_message.file.read() + kind = filetype.guess(file_bytes[:20]) + if not kind: + raise CorruptedFileError + raw_file_extension = kind.extension + file_extension = file_service.get_file_extension(raw_file_extension, file_bytes) + if file_extension in ('pdf', 'doc', 'docx', 'xlsx'): + text = file_service.get_file_data(file_extension, file_bytes) + text_chunks = EmbeddingService.split_text_to_chunks(text) + chunks = [HumanMessage(content=chunk_text) for chunk_text in text_chunks] + elif file_extension in ('jpg', 'jpeg', 'png', 'webp'): + image = file + _, image_size, image_data = self._get_image_data(file_bytes, file_extension) + else: raise FileExtensionNotSupported(['PDF', 'DOC', 'DOCX', 'XLSX', 'JPG', 'JPEG', 'PNG', 'WEBP']) chat_history = self.get_chat_history(model_name=model_name) chat_history.add_message(HumanMessage(content=input_message.content)) @@ -143,16 +150,23 @@ class Chatgpt_5(Chatgpt): document_name = ( chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] ) - embedding_tokens, file_data = self.get_large_file_data( - chunks, proxy, input_message.content + embedding_tokens, file_data = EmbeddingService.get_large_file_data( + self.store.messages.first().pk, + text_chunks, + proxy, + input_message.content, + model='text-embedding-3-small', + index_name='ml_model-index-1536', ) - messages[-1]['content'] = self.make_embeddings_prompt( - document_name=document_name, section_texts=file_data, question=input_message.content + messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( + document_name=document_name, + section_texts=file_data, + question=input_message.content, ) else: messages[-1]['content'] = ( - f'Используй системный промпт. Содержание файла: ' - f'{chunks}. Вопрос: {input_message.content}' + 'Используй системный промпт. Содержание файла: ' + f'{"".join(text_chunks)}. Вопрос: {input_message.content}' ) json_data = { 'model': model_name, @@ -45,7 +45,7 @@ class Claude(SimpleService): }, # 1M tokens } - TOOLS_TOKEN_COSTS = {'text-embedding-3-large': {'output': Decimal('0.000065')}} + TOOLS_TOKEN_COSTS = {'text-embedding-3-small': {'output': Decimal('0.00001')}} def calculate_price( self, version: str, input_tokens: int, output_tokens: int, embedding_tokens: int @@ -55,7 +55,7 @@ class Claude(SimpleService): input_tokens * price_map['input'] / 1_000_000 + output_tokens * price_map['output'] / 1_000_000 ) if embedding_tokens > 0: - price += self.TOOLS_TOKEN_COSTS['text-embedding-3-large']['output'] * embedding_tokens + price += self.TOOLS_TOKEN_COSTS['text-embedding-3-small']['output'] * embedding_tokens return price.quantize(Decimal('0.1'), rounding='ROUND_UP') def save_results(self, content: Iterator[Any], t: timedelta, save: bool = True) -> list[Message]: @@ -98,7 +98,12 @@ class Claude(SimpleService): chunks[0].partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] ) embedding_tokens, file_data = EmbeddingService.get_large_file_data( - self.store.messages.first().pk, chunks, proxy, input_message.content + self.store.messages.first().pk, + chunks, + proxy, + input_message.content, + model='text-embedding-3-small', + index_name='ml_model-index-1536', ) messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( document_name=document_name, @@ -9,6 +9,7 @@ from django.db.models.fields.files import FieldFile from PIL import Image from messages.models import Message +from ml_model.exceptions import FileExtensionNotSupported from ml_model.services.EmbeddingService import EmbeddingService from ml_model.services.FileService import FileProcessingService from ml_model.services.base import SimpleService @@ -57,7 +58,7 @@ class Gemini(SimpleService): }, } - TOOLS_TOKEN_COSTS = {'text-embedding-3-large': {'output': Decimal('0.000065')}} + TOOLS_TOKEN_COSTS = {'text-embedding-3-small': {'output': Decimal('0.00001')}} def calculate_price( self, version: str, input_tokens: int, output_tokens: int, image: FieldFile, embedding_tokens: int @@ -76,7 +77,7 @@ class Gemini(SimpleService): if image: price += price_map['input_imgs'] / 1_000 if embedding_tokens > 0: - price += self.TOOLS_TOKEN_COSTS['text-embedding-3-large']['output'] * embedding_tokens + price += self.TOOLS_TOKEN_COSTS['text-embedding-3-small']['output'] * embedding_tokens return price.quantize(Decimal('0.1'), rounding='ROUND_UP') def save_results(self, content: str, t: timedelta, save: bool = True) -> list[Message]: @@ -145,7 +146,12 @@ class Gemini(SimpleService): for proxy in Proxy.objects.all(): document_name = chunks[0].partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] embedding_tokens, file_data = EmbeddingService.get_large_file_data( - self.store.messages.first().pk, chunks, proxy, input_message.content + self.store.messages.first().pk, + chunks, + proxy, + input_message.content, + model='text-embedding-3-small', + index_name='ml_model-index-1536', ) messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( document_name=document_name, @@ -157,8 +163,7 @@ class Gemini(SimpleService): f'Используй системный промпт. Содержание файла: ' f'{chunks}. Вопрос: {input_message.content}' ) - else: - kind = filetype.guess(file_bytes[:20]) + elif file_extension in ('jpg', 'jpeg', 'png', 'webp'): mime = kind.mime if kind else 'application/octet-stream' normalized_image = Image.open(file) format = 'jpeg' if kind.extension == 'jpg' else kind.extension @@ -171,6 +176,8 @@ class Gemini(SimpleService): {'type': 'image_url', 'image_url': {'url': image_url}}, ] image = file + else: + raise FileExtensionNotSupported(['PDF', 'DOC', 'DOCX', 'XLSX', 'JPG', 'JPEG', 'PNG', 'WEBP']) start_time = time.time() result = openrouter_run(version, messages, callback_data, 'Gemini') process_time = timedelta(seconds=(time.time() - start_time)) @@ -1,9 +1,12 @@ import re +import subprocess import time from concurrent.futures import ThreadPoolExecutor from concurrent.futures._base import as_completed from typing import List, Tuple +import docx2txt +import filetype import httpx import base64 @@ -20,17 +23,20 @@ from django.template.loader import get_template from openai import BadRequestError from backend import settings -from ml_model.exceptions import TemplateNotFound, TemplateUnknownException, FileExtensionNotSupported, \ - ExceededContextLengthError +from ml_model.exceptions import ( + TemplateNotFound, + TemplateUnknownException, + FileExtensionNotSupported, + ExceededContextLengthError, + CorruptedFileError, +) from ml_model.models import NeuronModel from ml_model.services import Chatgpt from django.core.files.uploadedfile import UploadedFile -from django.utils.translation import gettext_lazy as _ from datetime import timedelta -from pathlib import Path from langchain_core.messages import ( HumanMessage, @@ -40,14 +46,18 @@ from langchain_core.runnables import RunnableWithMessageHistory from langchain_openai.chat_models import ChatOpenAI from messages.models import Message -from ml_model.exceptions import GenerationException +from ml_model.services.EmbeddingService import EmbeddingService +from ml_model.services.FileService import FileProcessingService from poller.models import Proxy from ml_model.tasks import drop_redis_vectors from ml_model.constants import ANCHORS + class Raifgpt(Chatgpt): + EMBEDDING_MODEL_FOR_BILLING = 'text-embedding-3-large' + @property def neuron_model(self): return NeuronModel.objects.get(slug='raifgpt') @@ -68,15 +78,22 @@ class Raifgpt(Chatgpt): embedding_tokens = 0 file = input_message.file if file: - file_extension = Path(file.name).suffix - if file_extension == '.pdf': - raw_text = self.get_pdf_data(file) - elif file_extension in ('.doc', '.docx'): - raw_text = self.get_word_data(file_extension[1:], file.read()) + file_service = FileProcessingService + file_bytes = file.read() + kind = filetype.guess(file_bytes[:20]) + if not kind: + raise CorruptedFileError + raw_file_extension = kind.extension + file_extension = file_service.get_file_extension(raw_file_extension, file_bytes) + if file_extension == 'pdf': + raw_text = self.get_pdf_data(file_bytes) + elif file_extension in ('doc', 'docx'): + raw_text = self.get_word_data(file_extension, file_bytes) else: raise FileExtensionNotSupported(['PDF', 'DOC', 'DOCX']) text = re.sub(r'\n{2,}', '\n', raw_text) - chunks = self.split_text_to_chunks(text, chunk_size=1000) + text_chunks = EmbeddingService.split_text_to_chunks(text, chunk_size=1000) + chunks = [HumanMessage(content=chunk_text) for chunk_text in text_chunks] for proxy in Proxy.objects.all(): self.llm = ChatOpenAI( model='gpt-4o', @@ -107,7 +124,9 @@ class Raifgpt(Chatgpt): input_tokens = self.count_text_tokens([*chat_history.messages]) if sum([len(chunk.content) for chunk in chunks]) > 40_000: redis_client = redis.Redis(host=settings.REDIS_HOST, port=settings.REDIS_PORT, db=0) - document_name = chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] + document_name = ( + chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] + ) message_uid = str(self.store.messages.first().pk).replace('-', '_') with httpx.Client( base_url='https://api.openai.com/v1/', @@ -120,7 +139,12 @@ class Raifgpt(Chatgpt): for chunk_id, chunk in enumerate(chunks): threads.append( executor.submit( - self.process_chunk, client, chunk, redis_client, message_uid, chunk_id + EmbeddingService.process_chunk, + client, + chunk.content, + redis_client, + message_uid, + chunk_id, ) ) for thread in as_completed(threads): @@ -130,7 +154,11 @@ class Raifgpt(Chatgpt): anchor_embeddings = {} with ThreadPoolExecutor(max_workers=settings.MAX_THREADS) as executor: for identify, value in ANCHORS.items(): - threads.append(executor.submit(self.get_anchor_embedding, client, value[0], identify)) + threads.append( + executor.submit( + self.get_anchor_embedding, client, value[0], identify + ) + ) for thread in as_completed(threads): thread_result = thread.result() embedding_tokens += thread_result[1] @@ -140,31 +168,41 @@ class Raifgpt(Chatgpt): for identify, embeddings in anchor_embeddings.items(): threads.append( executor.submit( - self.search_via_embeddings, + EmbeddingService.search_via_embeddings, redis_client, message_uid, - embeddings, - top_k=ANCHORS[identify][1] + user_query_embeddings=embeddings, + top_k=ANCHORS[identify][1], ) ) - result = [s['section_text'] for thread in as_completed(threads) for s in thread.result()] + result = [ + s['section_text'] + for thread in as_completed(threads) + for s in thread.result() + ] else: - query_embedding, e_total_tokens = self.get_embedding(client=client, content=input_message.content) + query_embedding, e_total_tokens = EmbeddingService._get_embedding( + client=client, content=input_message.content + ) embedding_tokens += e_total_tokens result = [ s['section_text'] - for s in self.search_via_embeddings( + for s in EmbeddingService.search_via_embeddings( redis_client=redis_client, message_uid=message_uid, user_query_embeddings=query_embedding, - top_k=25 + top_k=25, ) ] user_input = [ SystemMessage(content=user_system_prompt), - HumanMessage(self.make_embeddings_prompt( - document_name=document_name, section_texts=result, question=input_message.content - )) + HumanMessage( + self.make_embeddings_prompt( + document_name=document_name, + section_texts=result, + question=input_message.content, + ) + ), ] input_tokens += self.count_text_tokens(user_input) response = conversation.invoke( @@ -177,9 +215,11 @@ class Raifgpt(Chatgpt): input = [ SystemMessage(content=user_system_prompt), HumanMessage( - content=f'Используй системный промпт. Содержание файла: ' - f'{chunks}. Вопрос: {input_message.content}' + content=( + 'Используй системный промпт. Содержание файла: ' + f'{"".join(chunk.content for chunk in chunks)}. Вопрос: {input_message.content}' ) + ), ] input_tokens += self.count_text_tokens(input) response = conversation.invoke( @@ -202,23 +242,19 @@ class Raifgpt(Chatgpt): self.logger.info(f'Input количество токенов для raifgpt - {input_tokens}') self.logger.info(f'Output количество токенов для raifgpt - {output_tokens}') self.logger.info(f'Embedding количество токенов для raifgpt - {embedding_tokens}') - self.logger.info(f'Общее количество токенов для raifgpt - {input_tokens + output_tokens + embedding_tokens}') + self.logger.info( + f'Общее количество токенов для raifgpt - {input_tokens + output_tokens + embedding_tokens}' + ) process_time = timedelta(seconds=time.time() - start_time) self.handle_invoice( - self.neuron_model, - input_tokens, - output_tokens, - self.llm.model_name, - {}, - embedding_tokens + self.neuron_model, input_tokens, output_tokens, self.llm.model_name, {}, embedding_tokens ) msgs = self.save_results([response], process_time, save) return msgs - def get_pdf_data(self, pdf_file: UploadedFile) -> str: + def get_pdf_data(self, pdf_data: bytes) -> str: max_batch_size = 3.9 * 1024 * 1024 - pdf_data = pdf_file.read() image_count = 0 try: doc = fitz.open(stream=pdf_data, filetype="pdf") @@ -346,20 +382,43 @@ class Raifgpt(Chatgpt): """ def get_anchor_embedding(self, client: httpx.Client, content: str, anchor: str) -> Tuple[List[float], int, str]: - ''' + """ A method for converting raw text (anchor content) into embeddings using OpenAI API request :param client: Httpx client :param content: raw text of a chunk :param anchor: anchor identifier - ''' - response = client.post( - url="embeddings", - json={ - 'model': 'text-embedding-3-large', - 'input': content - } - ) + """ + response = client.post(url='embeddings', json={'model': 'text-embedding-3-large', 'input': content}) response.raise_for_status() data = response.json() - return data['data'][0]['embedding'], data['usage']['total_tokens'], anchor \ No newline at end of file + return data['data'][0]['embedding'], data['usage']['total_tokens'], anchor + + def get_word_data(self, extension: str, word_data: bytes) -> str: + """ + Extracting text from word-file + :param extension: extension of uploaded word file + :param word_file: uploaded word file + :return: word-file content + """ + try: + 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=word_data) + text = text.decode('utf-8') + else: + text = '' + except Exception: + text = 'Файл поврежден или не может быть прочитан.' + if text.strip(): + return f'Это текст, извлечённый из загруженного WORD-файла:\n{text}' + else: + return 'Файл пуст или содержит изображения, из которых невозможно извлечь текст.' +