@@ -1,11 +1,16 @@ import itertools +import base64 import time +import logging + from datetime import timedelta from decimal import Decimal +from math import ceil from io import BufferedReader, BytesIO -from typing import Generator, List, Optional +from typing import Any, Dict, Generator, List, Optional -from django.core.files.uploadedfile import InMemoryUploadedFile +import filetype +import tiktoken from langchain import hub from langchain.agents import AgentExecutor, create_structured_chat_agent from langchain.chains import ConversationChain @@ -17,8 +22,9 @@ from langchain_openai.chat_models import ChatOpenAI from PIL import Image from backend import settings -from messages.models import Message from ml_model.constants import TEMPORARY_TEST_TEXT +from messages.models import Message, BaseStore +from ml_model.models import ModelCategory, ModelInput, ModelParameter, ModelVersion from ml_model.models import ( ModelCategory, ModelConfiguration, @@ -49,9 +55,6 @@ class Chatgpt(SimpleService): ModelVersion(name='GPT-4o1 Mini', slug='o1-mini'), ModelVersion(name='GPT-4omni Mini', slug='gpt-4o-mini'), ModelVersion(name='GPT-4omni', slug='gpt-4o'), - ModelVersion(name='GPT-4 Turbo', slug='gpt-4-turbo'), - ModelVersion(name='GPT-4', slug='gpt-4'), - ModelVersion(name='GPT-3.5', slug='gpt-3.5-turbo'), ] inputs = [ @@ -88,15 +91,16 @@ class Chatgpt(SimpleService): category = ModelCategory(title='Чат-боты', slug='chat-bots') TOKENS_COST = { - 'gpt-3.5-turbo': Decimal('0.0085'), - 'gpt-4': Decimal('0.08925'), - 'gpt-4-turbo': Decimal('0.25'), - 'gpt-4o': Decimal('0.075'), - 'o1-preview': Decimal('0.09'), - 'o1-mini': Decimal('0.02'), - 'gpt-4o-mini': Decimal('0.03'), + 'o1-preview': {'input': Decimal('0.008250'), 'output': Decimal('0.033000')}, + 'o1-mini': {'input': Decimal('0.001650'), 'output': Decimal('0.006600')}, + 'gpt-4o-mini': {'input': Decimal('0.000083'), 'output': Decimal('0.000330')}, + 'gpt-4o': {'input': Decimal('0.001375'), 'output': Decimal('0.005500')}, } + def __init__(self, store: BaseStore) -> None: + super().__init__(store) + self.logger = logging.getLogger(self.__class__.__name__) + def make( self, input_message: Message, @@ -107,8 +111,19 @@ class Chatgpt(SimpleService): model_name = info.pop('version', 'gpt-4o') input_content = [{'type': 'text', 'text': input_message.content or ''}] image = input_message.file + image_size = None if image: - input_content.append({'type': 'image_url', 'image_url': {'url': image.url}}) + kind = filetype.guess(input_message.file.read(20)) + mime = kind.mime if kind else 'application/octet-stream' + normalized_image = Image.open(image) + buf = BytesIO() + normalized_image.save(buf, format=kind.extension.upper()) + image_url = ( + f"data:{mime},base64,{base64.b64encode(buf.getvalue()).decode('utf-8')}" + ) + buf.close() + image_size = normalized_image.size + input_content.append({'type': 'image_url', 'image_url': {'url': image_url}}) self.llm = ChatOpenAI( model=model_name, temperature=info.pop('temperature', 0.5), @@ -154,8 +169,9 @@ class Chatgpt(SimpleService): ), ) llm_input = HumanMessage(content=input_content) + input_tokens = self.count_text_tokens([*chat_history.buffer_as_messages, llm_input]) self.assert_enough_balance( - [*chat_history.buffer_as_messages, llm_input], model=self.llm.model_name + input_tokens, image_size, model=self.llm.model_name ) if image: response = conversation.llm.invoke([llm_input]) @@ -190,12 +206,22 @@ class Chatgpt(SimpleService): chat_history.chat_memory.add_ai_message(response) process_time = timedelta(seconds=time.time() - start_time) - total_tokens = self.llm.get_num_tokens_from_messages( - chat_history.chat_memory.messages - ) + output_tokens = self.count_text_tokens([response]) + if image: - total_tokens += self.count_image_tokens(image) - self.handle_invoice(self.neuron_model, total_tokens, self.llm.model_name) + self.logger.info(f'Input количество токенов БЕЗ картинки {model_name} - {input_tokens}') + input_tokens += self.count_image_tokens(normalized_image.size, model_name) + + self.logger.info(f'Input количество токенов для {model_name} - {input_tokens}') + self.logger.info(f'Output количество токенов для {model_name} - {output_tokens}') + self.logger.info(f'Общее количество токенов для {model_name} - {input_tokens + output_tokens}') + + self.handle_invoice( + self.neuron_model, + input_tokens, + output_tokens, + self.llm.model_name, + ) msgs = self.save_results([response], process_time, save) return msgs @@ -243,32 +269,65 @@ class Chatgpt(SimpleService): return memory def assert_enough_balance( - self, for_input: list[BaseMessage], model: str = 'gpt-3.5-turbo' + self, input_tokens: int, image_size: tuple, model: str = 'gpt-3.5-turbo' ): balance = PaymentPlanSelector(self.store.user).get_current_balance() - input_cost = self.TOKENS_COST[ - self.llm.model_name - ] * self.llm.get_num_tokens_from_messages(for_input) - if getattr(for_input[-1], 'image', None): - input_cost += self.count_image_tokens(for_input[-1].image) - input_cost *= self.TOKENS_COST[model] + total_tokens = input_tokens + + if image_size: + total_tokens += self.count_image_tokens(image_size) + + input_cost = self.TOKENS_COST[model]['input'] * total_tokens if input_cost > balance: raise InsufficientBalance(balance, input_cost) - def calculate_price(self, usage: int, model: str, *args, **kwargs) -> Decimal: - price = usage * self.TOKENS_COST[model] + def calculate_price( + self, input_tokens: int, output_tokens: int, model: str, *args, **kwargs + ) -> Decimal: + price = ( + input_tokens * self.TOKENS_COST[model]['input'] + + output_tokens * self.TOKENS_COST[model]['output'] + ) return price.quantize(Decimal('0.1'), rounding='ROUND_UP') def count_image_tokens( - self, image: InMemoryUploadedFile, high_resolution: bool = True + self, image_size: tuple, model_version: str = 'gpt-4o' ) -> int: - TILE_SIZE = 512 * 512 - RESOLUTION_MULTIPLIER = 170 if high_resolution else 85 - minio_image = Image.open(BytesIO(image.read())) - width, height = minio_image.size - area = width * height - num_of_tiles = (area // TILE_SIZE) + 1 - return num_of_tiles * RESOLUTION_MULTIPLIER + 85 + extra_tokens = { + 'gpt-4o': { + 'tile_tokens': 170, + 'base_tokens': 85, + }, + 'gpt-4o-mini': { + 'tile_tokens': 5667, + 'base_tokens': 2833, + } + } + width, height = image_size + + if max(width, height) > 2048: + a_ratio = width / height + width, height = (2048, int(2048 / a_ratio)) if a_ratio > 1 else (int(2048 * a_ratio), 2048) + if width >= height and height > 768: + width, height = int((768 / height) * width), 768 + elif height > width and width > 768: + width, height = 768, int((768 / width) * height) + tiles_size = ceil(width / 512) * ceil(height / 512) + + return extra_tokens[model_version]['base_tokens'] + extra_tokens[model_version]['tile_tokens'] * tiles_size + + def count_text_tokens(self, messages: list[BaseMessage]) -> int: + encoding = tiktoken.get_encoding("cl100k_base") + total_tokens = 0 + + for message in messages: + total_tokens += len( + encoding.encode( + message.content if isinstance(message.content, str) + else message.content[0]['text'] + ) + ) + return total_tokens def save_results( self, results: list[BaseMessage], elapsed_time: timedelta, save: bool = True @@ -8,7 +8,6 @@ from django.core.files import File from backend import settings from messages.models import BaseStore, Message -from ml_model.exceptions.external_api import ExternalAPIException from ml_model.models import ( ModelCategory, ModelInput, @@ -123,7 +122,7 @@ class Recraft(SimpleService): } data = {'input': payload} response = requests.post( - url=f'{self.generate_url}{payload['version']}/predictions', + url=f'{self.generate_url}{payload.get('version', 'recraft-v3')}/predictions', headers=headers, json=data, ) @@ -186,7 +185,7 @@ class Recraft(SimpleService): 'линогравюра': 'linocut' } start_time = time.time() - extension = '.svg' if input_message.info['version'] == self.versions[1].slug else '.png' + extension = '.svg' if input_message.info.get('version', 'recraft-v3') == self.versions[1].slug else '.png' size = self._get_size(input_message.info.pop('width', 1024), input_message.info.pop('height', 1024)) callback_data = dict( { @@ -195,7 +194,7 @@ class Recraft(SimpleService): f'Input description: {self.translate_prompt(input_message.content)}' ), 'size': size, - 'style': styles.get(input_message.info.pop('style'), 'любой'), + 'style': styles.get(input_message.info.get('style', 'любой')), **input_message.info } ) @@ -17,6 +17,7 @@ djangorestframework-simplejwt = "^5.2.2" djangorestframework = "^3.14.0" yookassa = "^2.4.0" environs = "^9.5.0" +gunicorn = {extras = ["gevent"], version = "^23.0.0"} dj-rest-auth = "^4.0.1" django-celery-beat = "^2.5.0" minio = "^7.1.15" @@ -56,6 +57,7 @@ httptools = "^0.6.4" uvloop = "^0.21.0" wsproto = "^1.2.0" channels-redis = "^4.2.1" +tiktoken = "<0.6.0" [tool.poetry.group.test.dependencies]