@@ -142,7 +142,7 @@ class BytedanceModelArkAdapter: stop_choices = [ choice for choice in choices - if choice.get('finish_reason') == BytedanceFinishReason.STOP + if choice.get('finish_reason') in (BytedanceFinishReason.STOP, BytedanceFinishReason.LENGTH) and choice.get('message') and choice['message'].get('content') is not None ] @@ -286,6 +286,8 @@ class BytedanceModelArkAdapter: ) raise GenerationException usage: BytedanceUsage = {} + reasoning_started = False + content_started = False for line in resp.iter_lines(): if not line: continue @@ -300,12 +302,25 @@ class BytedanceModelArkAdapter: try: data_obj = json.loads(data) if choices := data_obj.get('choices'): - chunk = choices[0].get('delta', {}).get('content') or '' + delta = choices[0].get('delta', {}) + reasoning_chunk = delta.get('reasoning_content') or '' + if include_reasoning and reasoning_chunk: + if not reasoning_started: + yield '**Рассуждение:**\n\n' + reasoning_started = True + yield reasoning_chunk + chunk = delta.get('content') or '' if chunk: + if include_reasoning and reasoning_started and not content_started: + yield '\n\n**Основная мысль:**\n\n' + content_started = True yield chunk if ( fr := choices[0].get('finish_reason') - ) and fr != BytedanceFinishReason.STOP: + ) and fr not in ( + BytedanceFinishReason.STOP, + BytedanceFinishReason.LENGTH, + ): cls._raise_by_error_payload(data_obj, choices) if raw_usage := data_obj.get('usage'): usage = { @@ -486,3 +501,25 @@ class BytedanceModelArkAdapter: return response['data'][0]['total_tokens'] except: return 0 + + @classmethod + def batch_tokenize(cls, model: str, texts: list[str]): + for proxy in Proxy.objects.all(): + with httpx.Client( + base_url=cls.BASE_URL, + headers={ + 'Authorization': f'Bearer {settings.BYTEDANCE_MODEL_ARK_API_KEY}', + 'Content-Type': 'application/json', + }, + proxy=f'{proxy.protocol}://{proxy.address}', + timeout=600, + ) as client: + total_tokens = 0 + try: + response = client.post('tokenization', json={'model': model, 'text': texts}).json() + data = response['data'] or [] + for token_info in data: + total_tokens += token_info['total_tokens'] + return total_tokens + except: + return 0 \ No newline at end of file @@ -6,3 +6,4 @@ class ModelResponse: content: str input_tokens: int output_tokens: int + cost: float @@ -58,7 +58,7 @@ class OpenrouterAdapter: # reasoning используем только для фоллбэк-подсчёта токенизатора # в ответ не кладём, заполняет буфер истории сообщений reasoning = '' - input_tokens = output_tokens = 0 + input_tokens = output_tokens = cost = 0 for line in resp.iter_lines(): line = line.strip() if not line or not line.startswith('data: '): @@ -78,6 +78,7 @@ class OpenrouterAdapter: if data_obj.get('usage'): input_tokens = data_obj['usage']['prompt_tokens'] output_tokens = data_obj['usage']['completion_tokens'] + cost = data_obj['usage']['cost'] except json.JSONDecodeError: logger.warning( @@ -85,13 +86,14 @@ class OpenrouterAdapter: ) continue - if not input_tokens and not output_tokens: + if not input_tokens and not output_tokens and not cost: logger.error(f'Opernrouter failed get data about tokens for model {model_name}') input_tokens, output_tokens = cls._fallback_tokenize( model_name, messages, content + reasoning ) + cost = 0 - return input_tokens, output_tokens + return input_tokens, output_tokens, cost @classmethod def collect_streaming_api( @@ -105,8 +107,8 @@ class OpenrouterAdapter: if chunk: content_parts.append(chunk) except StopIteration as exc: - input_tokens, output_tokens = exc.value - return ModelResponse(''.join(content_parts), input_tokens, output_tokens) + input_tokens, output_tokens, cost = exc.value + return ModelResponse(''.join(content_parts), input_tokens, output_tokens, cost) @classmethod def _fallback_tokenize(cls, model_name: str, messages: list, content: str) -> tuple[int, int]: @@ -1,58 +1,45 @@ -import base64 -import json -import time +from copy import copy from datetime import timedelta from decimal import Decimal from io import BytesIO -from pathlib import Path +from typing import Any, Iterator -import filetype -import httpx from django.core.files import File -from django.utils.translation import gettext, gettext_lazy -from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, BaseMessage +from langchain_core.messages import BaseMessage -from backend import settings from messages.models import Message -from ml_model.exceptions import ( - CorruptedFileError, - FileExtensionNotSupported, - InvalidParameterError, - PaidPlanRequiredError, -) -from ml_model.models import NeuronModel -from ml_model.services.chatgpt_4 import Chatgpt_4 -from ml_model.exceptions import ModelVersionNotAvailable -from ml_model.services.EmbeddingService import EmbeddingService -from ml_model.services.FileService import FileProcessingService +from ml_model.exceptions import ModelVersionNotAvailable, PaidPlanRequiredError +from ml_model.services.chatgpt import Chatgpt from poller.models import Proxy -class Chatgpt_5_4(Chatgpt_4): +class Chatgpt_5_4(Chatgpt): TOKENS_COST = { 'gpt-5.4': { 'input': Decimal('0.00125'), 'output': Decimal('0.0075'), 'web_search': { - 'low': Decimal('5'), # 1 call - 'medium': Decimal('5'), # 1 call - 'high': Decimal('5'), # 1 call + 'low': Decimal('5'), + 'medium': Decimal('5'), + 'high': Decimal('5'), }, - 'code_interpreter': Decimal('15'), # 1 call + 'code_interpreter': Decimal('15'), 'generated_image': Decimal('10.2'), }, 'gpt-5.4-pro': { 'input': Decimal('0.0075'), 'output': Decimal('0.045'), 'web_search': { - 'low': Decimal('5'), # 1 call - 'medium': Decimal('5'), # 1 call - 'high': Decimal('5'), # 1 call + 'low': Decimal('5'), + 'medium': Decimal('5'), + 'high': Decimal('5'), }, 'generated_image': Decimal('10.2'), }, } + BASE_VERSION = 'gpt-5.4' + TOKEN_LIMITS = { 'gpt-5.4': 1_050_000 // 2, 'gpt-5.4-pro': 1_050_000 // 2, @@ -107,233 +94,41 @@ class Chatgpt_5_4(Chatgpt_4): price += self.TOKENS_COST[model]['generated_image'] return price.quantize(Decimal('0.1'), rounding='ROUND_UP') - def make( + def make_stream(self, input_message: Message, save: bool = True) -> Iterator[str]: + return (yield from super().make_stream(input_message, save)) + + def _build_payload( self, + proxy: Proxy, input_message: Message, - save: bool = True, - ) -> list[Message]: - start_time = time.time() - info = input_message.info.copy() - model_name = info.pop('version', None) - if model_name is None or model_name not in self.TOKENS_COST: - raise ModelVersionNotAvailable(model_name, self.TOKENS_COST) - user_system_prompt = info.pop('system_prompt', '') - is_free_plan = self.store.user.plan.price <= 0 - if is_free_plan and model_name == 'gpt-5.4-pro': - raise PaidPlanRequiredError('ChatGPT 5.4 PRO') - if is_free_plan: - info.pop('web_search', None) - info.pop('code_interpreter', None) - info.pop('verbosity', None) - input_content = [{'type': 'text', 'text': input_message.content or ''}] - file = input_message.file - image = None - image_size = None - embedding_tokens = 0 - chunks = [] - text_chunks = [] - if file: - supported_extensions = ['PDF', 'DOC', 'DOCX', 'XLSX', 'JPG', 'JPEG', 'PNG', 'WEBP'] - 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 supported_extensions: - raise FileExtensionNotSupported(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) - 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(supported_extensions) - if image and info.get('code_interpreter'): - raise InvalidParameterError( - gettext('The "Use code" option cannot be used together with an attached image.') - ) - chat_history = self.get_chat_history(model_name=model_name) - chat_history.add_message(HumanMessage(content=input_message.content)) - llm_input = [SystemMessage(content=user_system_prompt), HumanMessage(content=input_content)] - input_tokens, input_embedding_tokens = self._get_input_tokens( - file, image, chunks, chat_history, llm_input, model_name - ) - self.assert_enough_balance( - input_tokens, - image_size, - model=model_name, - embedding_tokens=input_embedding_tokens, - output_tokens=500 if is_free_plan else 4000, + ctx: dict[str, Any], + *, + include_image_tool: bool = True, + ) -> tuple[dict[str, Any], int]: + payload_message = input_message + if not ctx: + model_name = input_message.info.get('version', self.BASE_VERSION) + if model_name is None or model_name not in self.TOKENS_COST: + raise ModelVersionNotAvailable(model_name, self.TOKENS_COST) + if self.store.user.plan.price <= 0 and model_name == 'gpt-5.4-pro': + raise PaidPlanRequiredError('ChatGPT 5.4 PRO') + if model_name == 'gpt-5.4-pro' and input_message.info.get('code_interpreter'): + payload_message = copy(input_message) + payload_message.info = input_message.info.copy() + payload_message.info.pop('code_interpreter', None) + + json_data, predicted_input_tokens = super()._build_payload( + proxy, + payload_message, + ctx, + include_image_tool=include_image_tool, ) - for proxy in Proxy.objects.all(): - system = chat_history.messages.pop(0) - messages = [ - {'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': ''' - You are a senior analytical assistant optimized for GPT-5.4 Pro, with strong emphasis on factual accuracy, grounded reasoning, and high-quality web retrieval. - Your goal is to produce correct, evidence-based, and practical answers using reasoning, provided context, and web search when necessary. - Core principles: - - Prioritize correctness and evidence over speed or completeness - - Never fabricate facts, sources, statistics, or citations - - Clearly separate: facts, assumptions, and interpretations - - Prefer primary, official, and high-authority sources over secondary or SEO content - - If evidence is weak or conflicting, explicitly state uncertainty instead of guessing - Web retrieval behavior: - - Use web search only when it materially improves accuracy, freshness, or completeness - - Prefer fewer, higher-quality sources over many low-quality ones - - Actively filter out SEO content, reposts, unverified blogs, and low-authority pages - - Cross-check important facts across multiple reliable sources when possible - - Stop searching once sufficient high-confidence evidence is collected (do not over-search) - Anti-hallucination rules: - - Never guess missing facts - - Never infer specific numbers, dates, names, or capabilities without evidence - - If sources conflict, explicitly report the conflict and prefer the most authoritative source - - If no reliable evidence is found, say so clearly instead of filling gaps - Reasoning behavior: - - Use internal structured reasoning, but respond in a concise and direct way - - Focus on outcome and actionable insight, not process explanation - - Avoid over-analysis once the answer is sufficiently supported - Source selection priority: - 1. Official documentation, standards, and primary publications - 2. Reputable technical or academic sources - 3. Established industry publications - 4. Secondary summaries only if no better sources exist (clearly labeled as such) - Output rules: - - Be high-signal, concise, and structured only when it improves clarity - - Do not include unsupported claims or speculative additions - - Clearly label uncertainty when applicable - - Stop reasoning immediately once the answer is sufficiently supported by evidence - '''}) - messages.insert(0, {'role': 'system', 'content': user_system_prompt}) - if image: - messages[-1]['content'] = [ - {'type': 'input_text', 'text': input_message.content}, - {'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 = 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'{"".join(text_chunks)}. Вопрос: {input_message.content}' - ) - tools = [] - if not is_free_plan: - tools.append( - { - 'type': 'image_generation', - 'size': '1024x1024', - 'quality': 'medium', - 'model': 'gpt-image-1.5', - } - ) - api_model_name = self.API_MODEL_ALIASES.get(model_name, model_name) - json_data = { - 'model': api_model_name, - 'input': messages, - 'tools': tools, - 'instructions': ( - 'Форматирование — обязательное требование. Выполняй строго по правилам:\n\n' - "1) Используй реальные символы новой строки, не выводи '\\n' как текст — вставляй переносы.\n\n" - '2) Абзацы: между абзацами ставь две пустые строки (два символа новой строки подряд).\n\n' - '3) Нумерованные и маркированные списки: каждый пункт на отдельной строке;\n' - ' между списком и текстом оставляй две пустые строки.\n\n' - '4) Блоки кода: любые фрагменты кода выделяй тройными бэктиками (```) с указанием языка программирования;\n' - ' перед и после блока оставляй две пустые строки.\n\n' - "5) Заголовки абзацев: делай крупным, используя Markdown '####' (например, '### Заголовок');\n" - ' выделяй жирным (**Заголовок**); оставляй две пустые строки перед и после заголовка.\n\n' - '6) Используй Markdown для всего форматирования, не используй HTML.\n\n' - '7) Исправление формата: если формат неверный, перепиши ответ и верни исправленный вариант.\n\n' - 'Строго разделяй текст на абзацы с жирными заголовками;\n' - 'нумерованные и маркированные списки выводи с переносами строк;\n' - 'блоки кода — с тройными бэктиками и указанием языка;\n' - "не выводи '\\n' как текст, используйте реальные переносы строк;\n" - 'добавляй две пустые строки между абзацами и блоками для улучшения читаемости.' - ), - } - if is_free_plan: - json_data['max_output_tokens'] = 500 - json_data['reasoning'] = {'effort': 'none', 'summary': 'auto'} - elif reasoning := info.get('reasoning'): - reasoning_data = { - 'Средний': 'low', - 'Высокий': 'medium', - } - json_data['reasoning'] = {'effort': reasoning_data[reasoning], 'summary': 'auto'} - web_search = info.get('web_search', 'Выключен') - if info.get('verbosity', False) and web_search in ('Выключен', 'Низкий'): - json_data['text'] = {'verbosity': 'low'} - search_context_sizes = { - 'Низкий': 'low', - 'Средний': 'medium', - 'Высокий': 'medium', - 'Сверхвысокий': 'medium', - } - if web_search != 'Выключен': - json_data['tools'].append( - { - 'type': 'web_search', - 'search_context_size': search_context_sizes[web_search], - 'user_location': {'type': 'approximate', 'country': 'RU'}, - } - ) - info['web_search'] = search_context_sizes[web_search] - if info.get('code_interpreter') and model_name == 'gpt-5.4': - json_data['tools'].append({'type': 'code_interpreter', 'container': {'type': 'auto'}}) - messages[-1]['content'] += ' the python tool ' - input_tokens, output_tokens, response = self._stream_openai_responses( - proxy=proxy, - json_data=json_data, - model_name=model_name, - ) - generated_image = None - if isinstance(response.content, list): - if isinstance(response.content[0], dict) and response.content[0].get('generate_image'): - generated_image = base64.b64decode(response.content[0]['image']) - response.content = gettext_lazy('Image is ready') - self.logger.info(f'Input количество токенов для {model_name} - {input_tokens}') - self.logger.info(f'Output количество токенов для {model_name} - {output_tokens}') - self.logger.info(f'Embedding количество токенов для {model_name} - {embedding_tokens}') - if generated_image: - self.logger.info( - f'Фиксированная цена за генерацию картинки - ' - f'{self.TOKENS_COST[model_name]["generated_image"]}' - ) - self.logger.info( - f'Общее количество токенов для {model_name} - {input_tokens + output_tokens + embedding_tokens}' - ) - process_time = timedelta(seconds=time.time() - start_time) - self.handle_invoice( - input_message.content_object.model, - input_tokens, - output_tokens, - model_name, - info, - embedding_tokens, - generated_image, - ) - msgs = self.save_results([response], process_time, generated_image, save) - return msgs + json_data['model'] = self.API_MODEL_ALIASES.get(ctx['model_name'], ctx['model_name']) + return json_data, predicted_input_tokens + + def _count_responses_input_tokens(self, proxy: Proxy, json_data: dict) -> int: + payload = { + **json_data, + 'model': self.API_MODEL_ALIASES.get(json_data['model'], json_data['model']), + } + return super()._count_responses_input_tokens(proxy, payload) @@ -8,14 +8,20 @@ from typing import Any, Iterator import filetype from PIL import Image +from django.utils.translation import gettext from messages.models import Message from ml_model.adapters.openrouter import OpenrouterAdapter -from ml_model.exceptions import CorruptedFileError, FileExtensionNotSupported, ModelVersionNotAvailable +from ml_model.exceptions import ( + CorruptedFileError, + FileExtensionNotSupported, + ModelVersionNotAvailable, + PaidPlanRequiredError, +) from ml_model.services.EmbeddingService import EmbeddingService from ml_model.services.FileService import FileProcessingService from ml_model.services.base import StreamSimpleService -from ml_model.tasks import openrouter_run +from ml_model.services.serper_mixin import SerperMixin from payments.exceptions.insufficient_balance import InsufficientBalance from payments.selectors.payment_plan_selector import PaymentPlanSelector from poller.models import Proxy @@ -23,7 +29,7 @@ from tools.chats.models import Chat from tools.copywrite.models import Copywrite from tools.public_api.models import APIStore -class Claude(StreamSimpleService): +class Claude(SerperMixin, StreamSimpleService): """ Claude Service contains abstract method make, which makes a generation @@ -33,14 +39,17 @@ class Claude(StreamSimpleService): 'claude-sonnet-4.6': { 'input': Decimal('900'), 'output': Decimal('4500'), + 'coefficient': Decimal('3'), }, # 1M tokens 'claude-opus-4.6': { 'input': Decimal('1500'), 'output': Decimal('7500'), + 'coefficient': Decimal('3'), }, # 1M tokens 'claude-fable-5': { 'input': Decimal('3000'), 'output': Decimal('15000'), + 'coefficient': Decimal('3'), }, # 1M tokens } @@ -75,15 +84,18 @@ class Claude(StreamSimpleService): SUPPORTED_EXTENSIONS = ['PDF', 'DOC', 'DOCX', 'XLSX', 'JPG', 'JPEG', 'PNG', 'WEBP'] def calculate_price( - self, version: str, input_tokens: int, output_tokens: int, embedding_tokens: int + self, version: str, cost: float, input_tokens: int, output_tokens: int, embedding_tokens: int ) -> Decimal: price_map = self.TOKENS_COST[version] - price = ( - input_tokens * price_map['input'] / 1_000_000 + output_tokens * price_map['output'] / 1_000_000 - ) + if cost: + price = Decimal(str(cost)) * price_map['coefficient'] * Decimal('100') + else: + price = ( + 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-small']['output'] * embedding_tokens - price += Decimal('2') return price.quantize(Decimal('0.1'), rounding='ROUND_UP') def save_results(self, content: str, t: timedelta, save: bool = True) -> list[Message]: @@ -104,20 +116,21 @@ class Claude(StreamSimpleService): if version_slug is None or version_slug not in self.TOKENS_COST: raise ModelVersionNotAvailable(version_slug, self.TOKENS_COST) model_slug = f'anthropic/{version_slug}' - callback_data = self._build_callback_data(input_message, version_slug) - messages, embedding_tokens = self._prepare_messages(input_message, version_slug) + callback_data = self._build_callback_data(input_message) + messages, embedding_tokens = self._prepare_messages(input_message, version_slug, callback_data) - result = openrouter_run(model_slug, messages, callback_data, 'Claude') + result = OpenrouterAdapter.collect_streaming_api(model_slug, messages, callback_data, 'Claude') process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice( input_message.content_object.model, version=version_slug, - input_tokens=result[1], - output_tokens=result[2], + cost=result.cost, + input_tokens=result.input_tokens, + output_tokens=result.output_tokens, embedding_tokens=embedding_tokens, ) - return self.save_results(result[0], process_time, save) + return self.save_results(result.content, process_time, save) def make_stream(self, input_message: Message, save: bool = True) -> Iterator[str]: start_time = time.time() @@ -125,10 +138,11 @@ class Claude(StreamSimpleService): if version_slug is None or version_slug not in self.TOKENS_COST: raise ModelVersionNotAvailable(version_slug, self.TOKENS_COST) model_slug = f'anthropic/{version_slug}' - callback_data = self._build_callback_data(input_message, version_slug) - messages, embedding_tokens = self._prepare_messages(input_message, version_slug) + callback_data = self._build_callback_data(input_message) + messages, embedding_tokens = self._prepare_messages(input_message, version_slug, callback_data) content_parts: list[str] = [] input_tokens = output_tokens = 0 + cost = 0 result = '' try: @@ -140,7 +154,7 @@ class Claude(StreamSimpleService): content_parts.append(chunk) yield chunk except StopIteration as exc: - input_tokens, output_tokens = exc.value + input_tokens, output_tokens, cost = exc.value finally: if content_parts: result = ''.join(content_parts) @@ -148,6 +162,7 @@ class Claude(StreamSimpleService): self.handle_invoice( input_message.content_object.model, version=version_slug, + cost=cost, input_tokens=input_tokens, output_tokens=output_tokens, embedding_tokens=embedding_tokens, @@ -194,32 +209,17 @@ class Claude(StreamSimpleService): return memory - def _build_callback_data(self, input_message: Message, version_slug: str) -> dict[str, Any]: - callback_data = {'provider': {'order': ['anthropic']}, **input_message.info, 'tools': []} - if version_slug == 'claude-fable-5': - if reasoning := input_message.info.get('reasoning'): - reasoning_data = { - 'Средний': 'low', - 'Высокий': 'medium', - } - callback_data['reasoning'] = {'effort': reasoning_data[reasoning]} - callback_data['tools'].append( - { - 'type': 'openrouter:web_search', - 'parameters': { - 'engine': 'parallel', - 'max_results': 1, - 'max_total_results': 3, - 'search_context_size': 'low', - }, - } - ) - return callback_data + def _build_callback_data(self, input_message: Message) -> dict[str, Any]: + return {'provider': {'order': ['anthropic']}, **input_message.info, 'tools': []} def _prepare_messages( - self, input_message: Message, version_slug: str + self, input_message: Message, version_slug: str, callback_data: dict ) -> tuple[list[dict[str, str | list]], int]: - if version_slug == 'claude-fable-5' and input_message.file: + if ( + version_slug == 'claude-fable-5' + and input_message.file + and self.store.user.plan.price > 0 + ): current_user_balance = PaymentPlanSelector(self.store.user).get_current_balance() if current_user_balance < (cost := Decimal('100')): raise InsufficientBalance(current_user_balance, cost) @@ -234,7 +234,17 @@ class Claude(StreamSimpleService): if version_slug == 'claude-fable-5': messages.insert(1, {'role': 'system', 'content': self.FABLE_SYSTEM_PROMPT}) + is_free_plan = self.store.user.plan.price <= 0 + current_user_balance = PaymentPlanSelector(self.store.user).get_current_balance() + is_low_balance = current_user_balance < Decimal('100') embedding_tokens = 0 + predict_embedding_tokens = 0 + predicted_input_price = Decimal(0) + text = '' + chunks = [] + file = None + image = None + image_width = image_height = 0 if input_message.file: file_service = FileProcessingService file_bytes = input_message.file.read() @@ -246,34 +256,35 @@ class Claude(StreamSimpleService): 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'): + if is_free_plan: + raise PaidPlanRequiredError(gettext('File analysis')) text = file_service.get_file_data(file_extension, file_bytes) chunks = EmbeddingService.split_text_to_chunks(text) - 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, + if (chunks_length := sum(len(chunk) for chunk in chunks)) > 20_000: + predict_embedding_tokens = len(chunks) * 2020 + predicted_input_price += ( + ( + Decimal('210') + + Decimal(chunks_length) / Decimal(len(chunks)) * Decimal('10') ) + / Decimal('2.7') + * self.TOKENS_COST[version_slug]['input'] + / Decimal('1_000_000') + ) else: - messages[-1]['content'] = ( - f'Используй системный промпт. Содержание файла: ' - f'{chunks}. Вопрос: {input_message.content}' + predicted_input_price += ( + Decimal(55 + chunks_length + len(input_message.content)) + / Decimal('2.7') + * self.TOKENS_COST[version_slug]['input'] + / Decimal('1_000_000') ) + file = input_message.file elif file_extension in ('jpg', 'jpeg', 'png', 'webp'): kind = filetype.guess(file_bytes[:20]) mime = kind.mime if kind else 'application/octet-stream' format = 'jpeg' if kind.extension == 'jpg' else kind.extension with Image.open(input_message.file) as normalized_image: + image_width, image_height = normalized_image.size with BytesIO() as buf: normalized_image.save(buf, format=format) image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' @@ -281,7 +292,128 @@ class Claude(StreamSimpleService): {'type': 'text', 'text': input_message.content}, {'type': 'image_url', 'image_url': {'url': image_url}}, ] + image = input_message.file else: raise FileExtensionNotSupported(self.SUPPORTED_EXTENSIONS) + serper_sources = 0 + if not is_free_plan and not is_low_balance: + callback_data['tools'].append( + { + 'type': 'openrouter:web_search', + 'parameters': { + 'engine': 'parallel', + 'max_results': 1, + 'max_total_results': 3, + 'search_context_size': 'low', + }, + } + ) + predicted_input_price += Decimal('0.005') * Decimal('100') * Decimal('3') + else: + serper_sources = 5 + predicted_input_price += ( + Decimal(serper_sources * 250) + / Decimal('2.7') + * self.TOKENS_COST[version_slug]['input'] + / Decimal('1_000_000') + ) + + if image: + predicted_image_tokens = min((image_width * image_height + 749) // 750, 1600) + predicted_input_price += ( + Decimal(predicted_image_tokens) + * self.TOKENS_COST[version_slug]['input'] + / Decimal('1_000_000') + ) + + callback_data['reasoning'] = { + 'enabled': True, + 'effort': 'low' if is_free_plan else 'medium', + } + reasoning_effort = callback_data['reasoning']['effort'] + reasoning_input_tokens = ( + {'low': 150, 'medium': 250}.get(reasoning_effort, 0) + if version_slug == 'claude-fable-5' + else 0 + ) + reasoning_output_reserve = ( + {'low': 300, 'medium': 500}.get(reasoning_effort, 0) + if version_slug == 'claude-fable-5' + else 0 + ) + estimated_input_tokens = ( + Decimal( + sum( + len(message['content']) if isinstance(message['content'], str) else 0 + for message in messages + ) + + (len(input_message.content) if image else 0) + ) + / Decimal('2.7') + + reasoning_input_tokens + ) + predicted_input_price += ( + estimated_input_tokens + * self.TOKENS_COST[version_slug]['input'] + / Decimal('1_000_000') + + predict_embedding_tokens * self.TOOLS_TOKEN_COSTS['text-embedding-3-small']['output'] + ) + + balance_buffer = Decimal('0.3') + max_output_tokens = min( + max( + int( + (current_user_balance - predicted_input_price - balance_buffer) + / (self.TOKENS_COST[version_slug]['output'] / Decimal('1_000_000')) + ) + - reasoning_output_reserve, + 0, + ), + 30_000, + ) + min_response_tokens = 150 if is_free_plan else 300 + if max_output_tokens < min_response_tokens: + cost = ( + predicted_input_price + + (min_response_tokens + reasoning_output_reserve) + * self.TOKENS_COST[version_slug]['output'] + / Decimal('1_000_000') + + balance_buffer + ) + raise InsufficientBalance(current_user_balance, cost) + callback_data['max_tokens'] = max_output_tokens + + if file: + 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}' + ) + + if serper_sources: + serp = self.run_serper(input_message.content) + messages.append( + { + 'role': 'user', + 'content': self.serper_to_context(serp, max_sources=serper_sources), + } + ) + return messages, embedding_tokens @@ -1,18 +1,24 @@ import time import logging +import json +import math +import subprocess from datetime import timedelta from decimal import Decimal +from io import BytesIO from typing import Any, Iterator import filetype +from PIL import Image from messages.models import Message from ml_model.adapters.bytedance_model_ark import BytedanceContentType, BytedanceModelArkAdapter -from ml_model.exceptions import FileExtensionNotSupported from ml_model.services.FileService import FileProcessingService -from ml_model.exceptions import ModelVersionNotAvailable +from ml_model.exceptions import GenerationException, ModelVersionNotAvailable from ml_model.services.base import SimpleService -from ml_model.tasks import bytedance_model_ark_run +from payments.exceptions.insufficient_balance import InsufficientBalance +from payments.selectors.payment_plan_selector import PaymentPlanSelector +from ml_model.tasks import bytedance_model_ark_run, stream_bytedance_model_ark_run from tools.chats.models import Chat from tools.copywrite.models import Copywrite from tools.public_api.models import APIStore @@ -48,6 +54,104 @@ class Dola_Seed(SimpleService): 'seed-2-0-pro': 'seed-2-0-pro-260328', 'seed-2-0-mini': 'seed-2-0-mini-260215', } + MAX_IMAGE_TOKENS = 1312 + MIN_VIDEO_FRAME_TOKENS = 128 + MAX_VIDEO_FRAME_TOKENS = 640 + MAX_SAMPLED_FRAMES = 640 + MAX_VIDEO_TOKENS = 80_000 + MIN_VIDEO_FRAME_PIXELS = 100_000 + MAX_VIDEO_FRAME_PIXELS = 500_000 + VIDEO_FRAME_SAMPLING_DIVISOR = 12 + + @classmethod + def count_image_tokens(cls, image_width: int, image_height: int) -> int: + return min( + (image_width * image_height + 783) // 784, + cls.MAX_IMAGE_TOKENS, + ) + + @classmethod + def _count_video_frame_tokens(cls, frame_width: int, frame_height: int) -> int: + pixels = frame_width * frame_height + if pixels <= cls.MIN_VIDEO_FRAME_PIXELS: + return cls.MIN_VIDEO_FRAME_TOKENS + if pixels >= cls.MAX_VIDEO_FRAME_PIXELS: + return cls.MAX_VIDEO_FRAME_TOKENS + return cls.MIN_VIDEO_FRAME_TOKENS + ( + (pixels - cls.MIN_VIDEO_FRAME_PIXELS) + * (cls.MAX_VIDEO_FRAME_TOKENS - cls.MIN_VIDEO_FRAME_TOKENS) + // (cls.MAX_VIDEO_FRAME_PIXELS - cls.MIN_VIDEO_FRAME_PIXELS) + ) + + @classmethod + def count_video_tokens( + cls, + duration_seconds: float, + video_width: int, + video_height: int, + fps: float, + ) -> int: + if duration_seconds <= 0 or fps <= 0: + return 0 + sampled_frames = min( + cls.MAX_SAMPLED_FRAMES, + max(1, math.ceil(duration_seconds * fps / cls.VIDEO_FRAME_SAMPLING_DIVISOR)), + ) + tokens_per_frame = cls._count_video_frame_tokens(video_width, video_height) + return min(cls.MAX_VIDEO_TOKENS, sampled_frames * tokens_per_frame) + + @classmethod + def _get_video_metadata(cls, video_bytes: bytes) -> tuple[int, int, float, float]: + result = subprocess.run( + [ + 'ffprobe', + '-v', + 'quiet', + '-print_format', + 'json', + '-show_streams', + '-show_format', + '-', + ], + input=video_bytes, + capture_output=True, + check=False, + ) + if result.returncode != 0: + raise GenerationException + data = json.loads(result.stdout) + video_stream = next( + ( + stream + for stream in data.get('streams', []) + if stream.get('codec_type') == 'video' + ), + None, + ) + if not video_stream: + raise GenerationException + duration = float(data.get('format', {}).get('duration') or 0) + if duration <= 0: + raise GenerationException + fps = cls._parse_video_fps( + video_stream.get('r_frame_rate') or video_stream.get('avg_frame_rate') + ) + return int(video_stream['width']), int(video_stream['height']), duration, fps + + @staticmethod + def _parse_video_fps(raw_fps: str | None, default: float = 24) -> float: + if not raw_fps: + return default + if '/' in raw_fps: + numerator, denominator = raw_fps.split('/', 1) + denominator_value = float(denominator) + if denominator_value: + return float(numerator) / denominator_value + try: + fps = float(raw_fps) + except ValueError: + return default + return fps if fps > 0 else default def calculate_price(self, version: str, input_tokens: int, output_tokens: int) -> Decimal: price_map = self.TOKENS_COST[version] @@ -59,7 +163,7 @@ class Dola_Seed(SimpleService): return price.quantize(Decimal('0.1'), rounding='ROUND_UP') - def save_results(self, content: Iterator[Any], t: timedelta, save: bool = True) -> list[Message]: + def save_results(self, content: str, t: timedelta, save: bool = True) -> list[Message]: msgs = [ Message( content=content, @@ -72,44 +176,102 @@ class Dola_Seed(SimpleService): return msgs - def make(self, input_message: Message, save: bool = True) -> list[Message]: - version = input_message.info.pop('version', None) + def _prepare_data( + self, input_message: Message + ) -> tuple[str, dict[str, Any], list[dict[str, Any]]]: + info = input_message.info.copy() + version = info.pop('version', None) if version is None or version not in self.TOKENS_COST: raise ModelVersionNotAvailable(version, self.TOKENS_COST) callback_data = { - "reasoning_effort": "minimal", - **input_message.info, + 'reasoning_effort': 'minimal', + **info, } - messages = self.get_chat_history() - content = [{'type': 'text', 'text': input_message.content}] + image_width = image_height = 0 + video_width = video_height = 0 + video_duration = 0.0 + video_fps = 0.0 if file := input_message.file: file_bytes = file.read() kind = filetype.guess(file_bytes[:50]) file_extension = ( FileProcessingService.get_file_extension(kind.extension, file_bytes) if kind else None ) - BytedanceModelArkAdapter.validate_user_attachment_extension(file_extension) - - if file_extension and file_extension.upper() == 'MP4': - content.append( - {'type': 'video_url', 'video_url': {'url': file.url}} + attachment_type = BytedanceModelArkAdapter.get_user_attachment_type(file_extension) + if attachment_type == 'image_url': + with Image.open(BytesIO(file_bytes)) as normalized_image: + image_width, image_height = normalized_image.size + elif attachment_type == 'video_url': + video_width, video_height, video_duration, video_fps = self._get_video_metadata( + file_bytes ) + content.append({'type': attachment_type, attachment_type: {'url': file.url}}) + messages.append({'role': 'user', 'content': content}) + api_model = self.VERSION_MAPPING[version] + texts = [] + for message in messages: + message_content = message['content'] + if isinstance(message_content, str): + texts.append(message_content) else: - content.append( - {'type': 'image_url', 'image_url': {'url': file.url}} + texts.append( + ''.join( + str(item.get('text') or '') + for item in message_content + if isinstance(item, dict) + ) ) + predicted_input_tokens = BytedanceModelArkAdapter.batch_tokenize(api_model, texts) + if not predicted_input_tokens: + raise GenerationException + if image_width and image_height: + predicted_input_tokens += self.count_image_tokens(image_width, image_height) + if video_width and video_height and video_duration and video_fps: + predicted_input_tokens += self.count_video_tokens( + video_duration, + video_width, + video_height, + video_fps, + ) + price_map = self.TOKENS_COST[version] + prompt_type = 'short_prompt' if predicted_input_tokens <= 128_000 else 'long_prompt' + predicted_input_price = ( + predicted_input_tokens * price_map[prompt_type]['input'] / Decimal('1_000_000') + ) + is_free_plan = self.store.user.plan.price <= 0 + current_user_balance = PaymentPlanSelector(self.store.user).get_current_balance() + max_output_tokens = min( + max( + int( + (current_user_balance - predicted_input_price - Decimal('0.2')) + / (price_map[prompt_type]['output'] / Decimal('1_000_000')) + ), + 0, + ), + 30_000, + ) + min_response_tokens = 300 if not is_free_plan else 150 + if max_output_tokens < min_response_tokens: + cost = ( + predicted_input_price + + min_response_tokens * price_map[prompt_type]['output'] / Decimal('1_000_000') + + Decimal('0.2') + ) + raise InsufficientBalance(current_user_balance, cost) + callback_data['max_completion_tokens'] = max_output_tokens + return version, callback_data, messages - messages.append({'role': 'user', 'content': content}) - + def make(self, input_message: Message, save: bool = True) -> list[Message]: + version, callback_data, messages = self._prepare_data(input_message) start_time = time.time() - result = bytedance_model_ark_run( model=self.VERSION_MAPPING[version], callback_data=callback_data, content_type=BytedanceContentType.CHAT, messages=messages, + include_reasoning=True ) process_time = timedelta(seconds=(time.time() - start_time)) @@ -120,10 +282,67 @@ class Dola_Seed(SimpleService): output_tokens=result[2], ) msgs = self.save_results(result[0], process_time, save) - return msgs - def get_chat_history(self, message_limit: int = 10, max_character_limit: int = 1500) -> list[dict[str, str | list]]: + def make_stream(self, input_message: Message, save: bool = True) -> Iterator[str]: + version, callback_data, messages = self._prepare_data(input_message) + model = self.VERSION_MAPPING[version] + start_time = time.time() + content_parts: list[str] = [] + input_tokens = output_tokens = 0 + result = '' + + try: + stream = stream_bytedance_model_ark_run( + model=model, + callback_data=callback_data, + messages=messages, + include_reasoning=True, + ) + try: + while True: + chunk = next(stream) + if chunk: + content_parts.append(chunk) + yield chunk + except StopIteration as exc: + usage = exc.value or {} + input_tokens = int(usage.get('prompt_tokens') or 0) + output_tokens = int(usage.get('completion_tokens') or 0) + finally: + if content_parts: + result = ''.join(content_parts) + if not (input_tokens + output_tokens): + input_text_parts = [] + for message in messages: + content = message['content'] + if isinstance(content, str): + input_text_parts.append(content) + else: + input_text_parts.extend( + str(item.get('text') or '') + for item in content + if isinstance(item, dict) + ) + input_tokens = BytedanceModelArkAdapter.tokenize( + model, ''.join(input_text_parts) + ) + output_tokens = BytedanceModelArkAdapter.tokenize(model, result) + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice( + input_message.content_object.model, + version=version, + input_tokens=input_tokens, + output_tokens=output_tokens, + ) + self.save_results(result, process_time, save) + if result: + return result + raise GenerationException + + def get_chat_history( + self, message_limit: int = 10, max_character_limit: int = 1500 + ) -> list[dict[str, str | list]]: if isinstance(self.store, Chat): air_messages = list( reversed( @@ -118,7 +118,7 @@ class Gemini_3_1(StreamSimpleService): content_parts.append(chunk) yield chunk except StopIteration as exc: - input_tokens, output_tokens = exc.value + input_tokens, output_tokens, _ = exc.value finally: if content_parts: result = ''.join(content_parts) @@ -9,6 +9,8 @@ from ml_model.adapters.bytedance_model_ark import BytedanceContentType, Bytedanc from ml_model.exceptions import GenerationException from ml_model.services.base import SimpleService from ml_model.tasks import bytedance_model_ark_run, stream_bytedance_model_ark_run +from payments.exceptions.insufficient_balance import InsufficientBalance +from payments.selectors.payment_plan_selector import PaymentPlanSelector from tools.chats.models import Chat from tools.copywrite.models import Copywrite from tools.public_api.models import APIStore @@ -27,8 +29,7 @@ class Glm_4_7(SimpleService): def calculate_price(self, version: str, input_tokens: int, output_tokens: int) -> Decimal: price_map = self.TOKENS_COST[version] price = ( - input_tokens * price_map['input'] / 1_000_000 - + output_tokens * price_map['output'] / 1_000_000 + input_tokens * price_map['input'] / 1_000_000 + output_tokens * price_map['output'] / 1_000_000 ) return price.quantize(Decimal('0.1'), rounding='ROUND_UP') @@ -47,13 +48,42 @@ class Glm_4_7(SimpleService): def _prepare_data(self, input_message: Message) -> tuple[str, dict[str, Any], list[dict[str, Any]]]: version = 'glm-4-7-251222' callback_data = { - "reasoning_effort": "minimal", + 'reasoning_effort': 'minimal', **input_message.info, } messages = self.get_chat_history() messages.append({'role': 'user', 'content': input_message.content}) if input_message.file: logger.info('GLM file input ignored: model does not support image/video input') + predicted_input_tokens = BytedanceModelArkAdapter.batch_tokenize( + version, [message['content'] for message in messages] + ) + if not predicted_input_tokens: + raise GenerationException + predicted_input_price = ( + predicted_input_tokens * self.TOKENS_COST[version]['input'] / Decimal('1_000_000') + ) + is_free_plan = self.store.user.plan.price <= 0 + current_user_balance = PaymentPlanSelector(self.store.user).get_current_balance() + max_output_tokens = min( + max( + int( + (current_user_balance - predicted_input_price - Decimal('0.2')) + / (self.TOKENS_COST[version]['output'] / Decimal('1_000_000')) + ), + 0, + ), + 30_000, + ) + min_response_tokens = 300 if not is_free_plan else 150 + if max_output_tokens < min_response_tokens: + cost = ( + predicted_input_price + + min_response_tokens * self.TOKENS_COST[version]['output'] / Decimal('1_000_000') + + Decimal('0.2') + ) + raise InsufficientBalance(current_user_balance, cost) + callback_data['max_completion_tokens'] = max_output_tokens return version, callback_data, messages def make(self, input_message: Message, save: bool = True) -> list[Message]: @@ -65,6 +95,7 @@ class Glm_4_7(SimpleService): callback_data=callback_data, content_type=BytedanceContentType.CHAT, messages=messages, + include_reasoning=True, ) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice( @@ -88,6 +119,7 @@ class Glm_4_7(SimpleService): model=version, callback_data=callback_data, messages=messages, + include_reasoning=True, ) try: while True: @@ -127,7 +159,7 @@ class Glm_4_7(SimpleService): reversed( Message.objects.filter( chats_chats_messages=self.store, is_deleted=False, is_sent=True - ).order_by('-created_at')[1:message_limit + 1] + ).order_by('-created_at')[1 : message_limit + 1] ) ) elif isinstance(self.store, APIStore): @@ -8,13 +8,15 @@ from typing import Iterator import filetype from PIL import Image +from django.utils.translation import gettext from messages.models import Message from ml_model.adapters.openrouter import OpenrouterAdapter -from ml_model.exceptions import CorruptedFileError, FileExtensionNotSupported +from ml_model.exceptions import CorruptedFileError, FileExtensionNotSupported, PaidPlanRequiredError from ml_model.services.base import StreamSimpleService from ml_model.services.EmbeddingService import EmbeddingService from ml_model.services.FileService import FileProcessingService +from ml_model.services.serper_mixin import SerperMixin from payments.exceptions.insufficient_balance import InsufficientBalance from payments.selectors.payment_plan_selector import PaymentPlanSelector from poller.models import Proxy @@ -23,36 +25,27 @@ from tools.copywrite.models import Copywrite from tools.public_api.models import APIStore -class Grok(StreamSimpleService): - +class Grok(SerperMixin, StreamSimpleService): TOKENS_COST = { - 'grok-4.3': { - 'input': Decimal('875'), - 'output': Decimal('1750'), - }, - 'grok-4.5': { - 'input': Decimal('1000'), - 'output': Decimal('3000'), - }, + 'grok-4.3': {'input': Decimal('875'), 'output': Decimal('1750'), 'coefficient': Decimal('5')}, + 'grok-4.5': {'input': Decimal('1000'), 'output': Decimal('3000'), 'coefficient': Decimal('5')}, } TOOLS_TOKEN_COSTS = {'text-embedding-3-small': {'output': Decimal('0.00001')}} SUPPORTED_EXTENSIONS = ['PDF', 'DOC', 'DOCX', 'XLSX', 'JPG', 'JPEG', 'PNG', 'WEBP'] - TTFT = 0.5 - TBT = 0.35 - def calculate_price( - self, version: str, input_tokens: int, output_tokens: int, embedding_tokens: int + self, version: str, cost: float, input_tokens: int, output_tokens: int, embedding_tokens: int ) -> Decimal: - price_map = self.TOKENS_COST[version.split('/')[1]] - price = ( - input_tokens * price_map['input'] / 1_000_000 - + output_tokens - * price_map['output'] - / 1_000_000 - ) + price_map = self.TOKENS_COST[version] + if cost: + price = Decimal(str(cost)) * price_map['coefficient'] * Decimal('100') + else: + price = ( + input_tokens * price_map['input'] / 1_000_000 + + output_tokens * price_map['output'] / 1_000_000 + ) if embedding_tokens: price += embedding_tokens * self.TOOLS_TOKEN_COSTS['text-embedding-3-small']['output'] @@ -71,16 +64,17 @@ class Grok(StreamSimpleService): def make(self, input_message: Message, save: bool = True) -> list[Message]: start_time = time.time() - version = input_message.info.get('version') or 'x-ai/grok-4.5' - callback_data = {**input_message.info} - messages, embedding_tokens = self._prepare_messages(input_message, version) + version = input_message.info.get('version') or 'grok-4.5' + callback_data = {**input_message.info, 'tools': []} + messages, embedding_tokens = self._prepare_messages(input_message, version, callback_data) - result = OpenrouterAdapter.collect_streaming_api(version, messages, callback_data, 'Grok') + result = OpenrouterAdapter.collect_streaming_api(f'x-ai/{version}', messages, callback_data, 'Grok') process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice( input_message.content_object.model, version=version, + cost=result.cost, input_tokens=result.input_tokens, output_tokens=result.output_tokens, embedding_tokens=embedding_tokens, @@ -89,15 +83,16 @@ class Grok(StreamSimpleService): def make_stream(self, input_message: Message, save: bool = True) -> Iterator[str]: start_time = time.time() - version = input_message.info.get('version') or 'x-ai/grok-4.5' - callback_data = {**input_message.info} - messages, embedding_tokens = self._prepare_messages(input_message, version) + version = input_message.info.get('version') or 'grok-4.5' + callback_data = {**input_message.info, 'tools': []} + messages, embedding_tokens = self._prepare_messages(input_message, version, callback_data) content_parts: list[str] = [] input_tokens = output_tokens = 0 + cost = 0 result = '' try: - stream = OpenrouterAdapter.run_streaming_api(version, messages, callback_data, 'Grok') + stream = OpenrouterAdapter.run_streaming_api(f'x-ai/{version}', messages, callback_data, 'Grok') try: while True: chunk = next(stream) @@ -105,7 +100,7 @@ class Grok(StreamSimpleService): content_parts.append(chunk) yield chunk except StopIteration as exc: - input_tokens, output_tokens = exc.value + input_tokens, output_tokens, cost = exc.value finally: if content_parts: result = ''.join(content_parts) @@ -113,6 +108,7 @@ class Grok(StreamSimpleService): self.handle_invoice( input_message.content_object.model, version=version, + cost=cost, input_tokens=input_tokens, output_tokens=output_tokens, embedding_tokens=embedding_tokens, @@ -158,11 +154,21 @@ class Grok(StreamSimpleService): return memory def _prepare_messages( - self, input_message: Message, version: str + self, input_message: Message, version: str, callback_data: dict ) -> tuple[list[dict[str, str | list]], int]: messages = self.get_chat_history() + is_free_plan = self.store.user.plan.price <= 0 + current_user_balance = PaymentPlanSelector(self.store.user).get_current_balance() + is_low_balance = current_user_balance < Decimal('100') messages.append({'role': 'user', 'content': input_message.content}) embedding_tokens = 0 + predict_embedding_tokens = Decimal(0) + predicted_input_price = Decimal(0) + text = '' + chunks = [] + file = None + image = None + image_width = image_height = 0 if input_message.file: file_service = FileProcessingService file_bytes = input_message.file.read() @@ -174,46 +180,35 @@ class Grok(StreamSimpleService): 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'): + if is_free_plan: + raise PaidPlanRequiredError(gettext('File analysis')) 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'][ - 'default' if approx_tokens <= 200_000 else 'high' - ] - / 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, + if (chunks_length := sum(len(chunk) for chunk in chunks)) > 20_000: + predict_embedding_tokens = len(chunks) * 2020 + predicted_input_price += ( + ( + Decimal('210') + + Decimal(chunks_length) / Decimal(len(chunks)) * Decimal('10') ) + / Decimal('2.0') + * self.TOKENS_COST[version]['input'] + / Decimal('1_000_000') + ) else: - messages[-1]['content'] = ( - f'Используй системный промпт. Содержание файла: ' - f'{chunks}. Вопрос: {input_message.content}' + predicted_input_price += ( + Decimal(55 + chunks_length + len(input_message.content)) + / Decimal('2.0') + * self.TOKENS_COST[version]['input'] + / Decimal('1_000_000') ) - + file = input_message.file elif file_extension in ('jpg', 'jpeg', 'png', 'webp'): mime = kind.mime if kind else 'application/octet-stream' input_message.file.seek(0) format = 'jpeg' if kind.extension == 'jpg' else kind.extension with Image.open(input_message.file) as normalized_image: + image_width, image_height = normalized_image.size with BytesIO() as buf: normalized_image.save(buf, format=format) image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' @@ -221,7 +216,110 @@ class Grok(StreamSimpleService): {'type': 'text', 'text': input_message.content}, {'type': 'image_url', 'image_url': {'url': image_url}}, ] + image = input_message.file else: raise FileExtensionNotSupported(self.SUPPORTED_EXTENSIONS) - + serper_sources = 0 + if not is_free_plan and not is_low_balance: + callback_data['tools'].append( + { + 'type': 'openrouter:web_search', + 'parameters': { + 'engine': 'parallel', + 'max_results': 1, + 'max_total_results': 3, + 'search_context_size': 'low', + }, + } + ) + predicted_input_price += Decimal('0.005') * 100 * 5 + else: + serper_sources = 5 + predicted_input_price += ( + Decimal(serper_sources * 250) + / Decimal('2.0') + * self.TOKENS_COST[version]['input'] + / Decimal('1_000_000') + ) + if image: + predicted_image_tokens = min((image_width * image_height + 999) // 1000, 2500) + predicted_input_price += ( + Decimal(predicted_image_tokens) + * self.TOKENS_COST[version]['input'] + / Decimal('1_000_000') + ) + estimated_input_tokens = ( + Decimal( + sum( + len(message['content']) if isinstance(message['content'], str) else 0 + for message in messages + ) + + (len(input_message.content) if image else 0) + ) + / Decimal('2.0') + + (150 if is_free_plan else 250) + ) + predicted_input_price += ( + estimated_input_tokens * self.TOKENS_COST[version]['input'] / Decimal('1_000_000') + + predict_embedding_tokens * self.TOOLS_TOKEN_COSTS['text-embedding-3-small']['output'] + ) + callback_data['reasoning'] = { + 'enabled': True, + 'effort': 'low' if is_free_plan else 'medium', + } + max_output_tokens = min( + max( + int( + ( + current_user_balance + - predicted_input_price + - (Decimal('0.4') if is_free_plan else Decimal('0.6')) + ) + / (self.TOKENS_COST[version]['output'] / Decimal('1_000_000')) + ) + - (300 if is_free_plan else 500), + 0, + ), + 30_000, + ) + min_response_tokens = 300 if not is_free_plan else 150 + if max_output_tokens < min_response_tokens: + cost = ( + predicted_input_price + + min_response_tokens * self.TOKENS_COST[version]['output'] / Decimal('1_000_000') + + (Decimal('0.4') if is_free_plan else Decimal('0.6')) + + (300 if is_free_plan else 500) * self.TOKENS_COST[version]['output'] / Decimal('1_000_000') + ) + raise InsufficientBalance(current_user_balance, cost) + callback_data['max_tokens'] = max_output_tokens + if file: + 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}' + ) + if serper_sources: + serp = self.run_serper(input_message.content) + messages.append( + { + 'role': 'user', + 'content': self.serper_to_context(serp, max_sources=serper_sources), + } + ) return messages, embedding_tokens @@ -0,0 +1,41 @@ +import httpx +from django.conf import settings + + +class SerperMixin: + @staticmethod + def run_serper(query: str, **kwargs): + headers = { + 'X-API-KEY': settings.SERPER_API_KEY, + 'Content-Type': 'application/json', + } + params = { + 'q': query, + **{key: value for key, value in kwargs.items() if value is not None}, + } + response = httpx.post('https://google.serper.dev/search', headers=headers, params=params) + response.raise_for_status() + search_results = response.json() + return search_results + + @staticmethod + def serper_to_context(serp: dict, max_sources: int = 3) -> str: + query = (serp.get('searchParameters') or {}).get('q', '').strip() + organic = (serp.get('organic') or [])[:max_sources] + + lines = [f'Результаты веб-поиска по запросу: {query}' if query else 'Результаты веб-поиска:'] + if not organic: + lines.append('(Совпадений не найдено.)') + return '\n'.join(lines) + + for index, item in enumerate(organic, start=1): + block = f'[{index}] {(item.get("title") or "").strip() or "Без названия"}' + if link := (item.get('link') or '').strip(): + block += f'\nURL: {link}' + if date := (item.get('date') or '').strip(): + block += f'\nДата: {date}' + if snippet := (item.get('snippet') or '').strip(): + block += f'\nОписание: {snippet}' + lines.append(block[:250]) + + return '\n'.join(lines).strip() @@ -8,8 +8,8 @@ class SSEChunkService: return SSEChunk(event_id=event_id, event=event, data=data or {}) @classmethod - def start(cls, event_id: int, message_uuid: str) -> SSEChunk: - return cls._chunk(event_id, 'start', {'message_uuid': message_uuid}) + def start(cls, event_id: int, data: SSEData) -> SSEChunk: + return cls._chunk(event_id, 'start', data) @classmethod def token(cls, event_id: int, content: str) -> SSEChunk: @@ -14,13 +14,23 @@ from tools.chats.services.sse_store import PublicSSEStoreService, SSEStoreServic from tools.public_api.models import APIKey, APIStore -def _run_stream(store: SSEStoreService, message_uuid: str, service: StreamSimpleService, message): +def _run_stream( + store: SSEStoreService, + message_uuid: str, + service: StreamSimpleService, + message, + *, + file_url: str | None = None, +): stream = None event_id = 0 try: event_id += 1 - store.push(SSEChunkService.start(event_id, message_uuid)) + start_data = {'message_uuid': message_uuid} + if file_url: + start_data['file'] = file_url + store.push(SSEChunkService.start(event_id, start_data)) stream = service.make_stream(message) while True: @@ -52,7 +62,13 @@ def event_stream_task(chat_uuid: str, message_uuid: str, user_uuid: str) -> None service = chat.model.service(chat) - _run_stream(store, message_uuid, service, message) + _run_stream( + store, + message_uuid, + service, + message, + file_url=message.file.url if message.file else None, + ) @shared_task(soft_time_limit=570, time_limit=600)