@@ -9,6 +9,7 @@ from ninja import Router from authentication.security import AsyncAuthBearer from ml_model.schemas import PredictPriceSchema, PredictPriceInputSchema from ml_model.services.base import SimpleService +from tools.chats.models import Chat router = Router(auth=AsyncAuthBearer(), tags=['ml_model']) @@ -16,17 +17,38 @@ router = Router(auth=AsyncAuthBearer(), tags=['ml_model']) @router.post('predict-price/', tags=['ml_model/predict-price'], response=PredictPriceSchema) def calculate_predict_price(request, body: PredictPriceInputSchema): payload = body.dict() - content = payload.pop('content') + + chat_id = payload.pop('chat_id', None) + content = payload.pop('content') if not chat_id else payload.get('content') + version = body.info.get("version", "base") + model_slug = body.model_slug.title() + json_str = json.dumps(payload, sort_keys=True, separators=(',', ':')) signature = hashlib.sha256(json_str.encode('utf-8')).hexdigest() cache_key = f'predict_price:{signature}' predicted_price = cache.get(cache_key) + history_price = None + if predicted_price is None: service: type[SimpleService] = getattr( - sys.modules['ml_model.services'], f'{body.model_slug.title()}' + sys.modules['ml_model.services'], f'{model_slug}' ) + + if body.chat_id is not None: + history_key = f'history_price:{model_slug}:{chat_id}:{version}' + history_price = cache.get(history_key) + + if history_price is None: + history_price = service.calculate_history_price(chat_id=chat_id, version=version) + if history_price is not None: + cache.set(history_key, history_price) + predicted_price = service.predict_price(content=content, file_exists=body.file_exists, info=body.info) + + if history_price: + predicted_price += history_price + if predicted_price: cache.set(cache_key, predicted_price) @@ -1,6 +1,7 @@ from abc import ABC, abstractmethod from decimal import Decimal from typing import Never, Any +from uuid import UUID from asgiref.sync import async_to_sync from googletrans import Translator @@ -61,6 +62,10 @@ class SimpleService(ABC): raise PromptLengthExceeded return async_to_sync(self.translator.translate)(prompt, dest=to).text + @classmethod + def calculate_history_price(cls, chat_id: UUID, version: str) -> Decimal | None: + return None + @abstractmethod def calculate_price(self, *args, **kwargs) -> Decimal: ... @@ -2,8 +2,9 @@ import base64 import itertools import logging import time +from uuid import UUID - +from django.core.cache import cache from django.utils.translation import gettext_lazy as _ from datetime import timedelta from decimal import Decimal @@ -90,6 +91,8 @@ class Chatgpt(SimpleService): 'gpt-oss-120b': 65_500, } + ENCODING_NAME = 'o200k_base' + def __init__(self, store: BaseStore) -> None: super().__init__(store) self.logger = logging.getLogger(self.__class__.__name__) @@ -376,6 +379,13 @@ class Chatgpt(SimpleService): embedding_tokens ) msgs = self.save_results([response], process_time, save) + if save: + model_slug = self.__class__.__name__ + chat_id = self.store.uid + history_price = self.calculate_history_price(chat_id=chat_id, version=model_name) + + history_key = f'history_price:{model_slug}:{chat_id}:{model_name}' + cache.set(history_key, history_price) return msgs def get_chat_history(self, model_name: str) -> InMemoryChatMessageHistory: @@ -498,8 +508,26 @@ class Chatgpt(SimpleService): + extra_tokens[model_version]['tile_tokens'] * tiles_size ) + @classmethod + def calculate_history_price(cls, chat_id: UUID, version: str) -> Decimal | None: + service = cls(store=Chat.objects.get(pk=chat_id)) + + history = service.get_chat_history(model_name=version) + + input_tokens = service.count_text_tokens(history.messages) + return Decimal(input_tokens * service.TOKENS_COST[version]['input']) + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict) -> Decimal | None: + model_name = info.get('version') + message = HumanMessage(content=content.strip()) + + input_tokens = cls.count_text_tokens(self=cls, messages=[message]) + return Decimal(input_tokens * cls.TOKENS_COST[model_name]['input']) + def count_text_tokens(self, messages: list[BaseMessage]) -> int: - encoding = tiktoken.get_encoding('o200k_base') + encoding = tiktoken.get_encoding(self.ENCODING_NAME) + total_tokens = 0 for message in messages: if isinstance(message.content, str): @@ -508,6 +536,7 @@ class Chatgpt(SimpleService): total_tokens += len( encoding.encode(''.join([input_data.get('text', '') for input_data in message.content])) ) + else: total_tokens += len(encoding.encode(''.join(message.content))) @@ -2,6 +2,7 @@ import time from datetime import timedelta from decimal import Decimal +from django.core.cache import cache import filetype from langchain_core.messages import HumanMessage, SystemMessage @@ -86,6 +87,8 @@ class Chatgpt_5(Chatgpt): TOKEN_LIMITS = {key: 200_000 for key in TOKENS_COST.keys()} + ENCODING_NAME = 'o200k_base' + @property def neuron_model(self): return NeuronModel.objects.get(slug='chatgpt_5') @@ -241,4 +244,11 @@ class Chatgpt_5(Chatgpt): self.neuron_model, input_tokens, output_tokens, model_name, info, embedding_tokens ) msgs = self.save_results([response], process_time, save) + if save: + model_slug = self.__class__.__name__ + chat_id = self.store.uid + history_price = self.calculate_history_price(chat_id=chat_id, version=model_name) + + history_key = f'history_price:{model_slug}:{chat_id}:{model_name}' + cache.set(history_key, history_price) return msgs @@ -5,6 +5,7 @@ from decimal import Decimal from io import BytesIO import filetype +from django.core.cache import cache from django.core.files import File from django.utils.translation import gettext_lazy from langchain_core.messages import HumanMessage, SystemMessage, BaseMessage @@ -48,6 +49,8 @@ class Chatgpt_5_4(Chatgpt): 'gpt-5.4-pro': 1_050_000 // 2, } + ENCODING_NAME = 'o200k_base' + @property def neuron_model(self): return NeuronModel.objects.get(slug='chatgpt_5_4') @@ -283,4 +286,11 @@ class Chatgpt_5_4(Chatgpt): self.neuron_model, input_tokens, output_tokens, model_name, info, embedding_tokens, generated_image ) msgs = self.save_results([response], process_time, generated_image, save) + if save: + model_slug = self.__class__.__name__ + chat_id = self.store.uid + history_price = self.calculate_history_price(chat_id=chat_id, version=model_name) + + history_key = f'history_price:{model_slug}:{chat_id}:{model_name}' + cache.set(history_key, history_price) return msgs @@ -1,4 +1,5 @@ from typing import List, Optional, Any +from uuid import UUID from ninja import ModelSchema, Schema from pydantic import condecimal @@ -39,6 +40,7 @@ class PredictPriceInputSchema(Schema): model_slug: str content: str file_exists: bool + chat_id: UUID | None = None info: dict[str, Any]