@@ -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 @@ -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( @@ -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):