@@ -4,10 +4,19 @@ from datetime import timedelta from decimal import Decimal from io import BytesIO from typing import Any, Dict, Iterator +from pathlib import Path import filetype from django.db.models.fields.files import FieldFile from PIL import Image +from ml_model.services.FileService import FileProcessingService +from ml_model.exceptions import FileExtensionNotSupported, CorruptedFileError +from ml_model.services.EmbeddingService import EmbeddingService +from payments.exceptions.insufficient_balance import InsufficientBalance +from payments.selectors.payment_plan_selector import PaymentPlanSelector + + +from poller.models import Proxy from messages.models import Message from ml_model.services.base import SimpleService @@ -35,8 +44,13 @@ class Grok(SimpleService): }, # 1M tokens } + TOOLS_TOKEN_COSTS = {'text-embedding-3-small': {'output': Decimal('0.00001')}} + + SUPPORTED_EXTENSIONS = ['PDF', 'DOC', 'DOCX', 'XLSX', 'JPG', 'JPEG', 'PNG', 'WEBP'] + + def calculate_price( - self, version: str, input_tokens: int, output_tokens: int, image: FieldFile + self, version: str, input_tokens: int, output_tokens: int, image: FieldFile, embedding_tokens: int ) -> Decimal: price_map = self.TOKENS_COST[version.split('/')[1]] price = ( @@ -44,6 +58,8 @@ class Grok(SimpleService): ) if image: price += price_map['input_imgs'] / 1_000 + if embedding_tokens: + price += embedding_tokens * self.TOOLS_TOKEN_COSTS['text-embedding-3-small']['output'] return price.quantize(Decimal('0.1'), rounding='ROUND_UP') def save_results(self, content: Iterator[Any], t: timedelta, save: bool = True) -> list[Message]: @@ -64,31 +80,89 @@ class Grok(SimpleService): callback_data = {**input_message.info} messages = self.get_chat_history() 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}}, - ] - result = openrouter_run(version, messages, callback_data, 'Grok') - process_time = timedelta(seconds=(time.time() - start_time)) - self.handle_invoice( - input_message.content_object.model, - version=version, - input_tokens=result[1], - output_tokens=result[2], - image=image, - ) - msgs = self.save_results(result[0], process_time) - return msgs + embedding_tokens = 0 + + image = None + if input_message.file: + file_service = FileProcessingService + file_bytes = input_message.file.read() + kind = filetype.guess(file_bytes[:550]) + + if not kind: + if Path(input_message.file.name).suffix[1:].upper() not in self.SUPPORTED_EXTENSIONS: + raise FileExtensionNotSupported(self.SUPPORTED_EXTENSIONS) + 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) + chunks = EmbeddingService.split_text_to_chunks(text) + approx_tokens = sum([len(message['content']) for message in messages]) / 3 + predict_price = ( + Decimal(approx_tokens) + * self.TOKENS_COST[version.split('/')[1]]['input'] + / Decimal('1000000') + + len(chunks) * 2100 * self.TOOLS_TOKEN_COSTS['text-embedding-3-small']['output'] + ).quantize(Decimal('0.1'), rounding='ROUND_UP') + if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < predict_price: + raise InsufficientBalance(balance, predict_price) + + if len(text) > 20_000: + 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, + 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}' + ) + + + elif file_extension in ('jpg', 'jpeg', 'png', 'webp'): + image = input_message.file + + mime = kind.mime if kind else 'application/octet-stream' + normalized_image = Image.open(input_message.file) + 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 FileExtensionNotSupported(self.SUPPORTED_EXTENSIONS) + + result = openrouter_run(version, messages, callback_data, 'Grok') + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice( + input_message.content_object.model, + version=version, + input_tokens=result[1], + output_tokens=result[2], + embedding_tokens=embedding_tokens, + image=image, + ) + msgs = self.save_results(result[0], process_time) + return msgs + + + def get_chat_history(self, message_limit: int = 10, max_character_limit: int = 1500) -> list[dict[str, str | list]]: if isinstance(self.store, Chat): @@ -39,7 +39,6 @@ class UnsupportedSize(Exception): self.current_size | self.required_size ) - class ModelTimeoutError(Exception): def __str__(self): return _('The model is not responding')