@@ -1,7 +1,8 @@ +import json from enum import StrEnum import logging import time -from typing import Any, TypeAlias, TypedDict +from typing import Any, Generator, TypeAlias, TypedDict import httpx @@ -69,6 +70,7 @@ class BytedanceVideoTaskResponse(TypedDict, total=False): RunChatResult: TypeAlias = tuple[str, int, int] +RunStreamChatResult: TypeAlias = Generator[str, None, BytedanceUsage] RunImageResult: TypeAlias = list[str] RunVideoResult: TypeAlias = tuple[str, int] BytedanceRunResult: TypeAlias = RunChatResult | RunImageResult | RunVideoResult @@ -247,6 +249,74 @@ class BytedanceModelArkAdapter: raise GenerationException + @classmethod + def _stream_chat( + cls, + model: str, + callback_data: dict[str, Any], + messages: list[dict[str, Any]] | None = None, + include_reasoning: bool = False, + ) -> RunStreamChatResult: + payload = dict( + **callback_data, + model=model, + messages=messages, + stream=True, + stream_options={'include_usage': True}, + ) + 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: + with client.stream( + 'POST', cls.CONTENT_TYPE_TO_ENDPOINT[BytedanceContentType.CHAT], json=payload + ) as resp: + if resp.status_code >= 400: + logger.error( + 'Bytedance request failed status=%s route=%s body=%s', + resp.status_code, + cls.CONTENT_TYPE_TO_ENDPOINT[BytedanceContentType.CHAT], + resp.text, + ) + raise GenerationException + usage: BytedanceUsage = {} + for line in resp.iter_lines(): + if not line: + continue + if isinstance(line, bytes): + line = line.decode('utf-8') + line = line.strip() + if not line.startswith('data:'): + continue + data = line[5:].lstrip() + if data == '[DONE]': + break + try: + data_obj = json.loads(data) + if choices := data_obj.get('choices'): + chunk = choices[0].get('delta', {}).get('content') or '' + if chunk: + yield chunk + if ( + fr := choices[0].get('finish_reason') + ) and fr != BytedanceFinishReason.STOP: + cls._raise_by_error_payload(data_obj, choices) + if raw_usage := data_obj.get('usage'): + usage = { + 'prompt_tokens': int(raw_usage.get('prompt_tokens') or 0), + 'completion_tokens': int(raw_usage.get('completion_tokens') or 0), + } + except json.JSONDecodeError: + continue + return usage + raise GenerationException + @classmethod def _generate_video( cls, @@ -288,7 +358,7 @@ class BytedanceModelArkAdapter: if error_code := data.get('error', {}).get('code', ''): if error_code == 'OutputImageSensitiveContentDetected': raise NSFWDetectedException - + if image_data := data.get('data'): urls = [item.get('url') for item in image_data if isinstance(item, dict) and item.get('url')] if urls: @@ -398,3 +468,21 @@ class BytedanceModelArkAdapter: raise GenerationException time.sleep(cls.VIDEO_POLL_DELAY_SECONDS) raise GenerationException + + @classmethod + def tokenize(cls, model: str, text: 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: + try: + response = client.post('tokenization', json={'model': model, 'text': [text]}).json() + return response['data'][0]['total_tokens'] + except: + return 0 @@ -6,8 +6,9 @@ from typing import Any, Iterator from messages.models import Message from ml_model.adapters.bytedance_model_ark import BytedanceContentType, BytedanceModelArkAdapter +from ml_model.exceptions import GenerationException from ml_model.services.base import SimpleService -from ml_model.tasks import bytedance_model_ark_run +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 @@ -31,7 +32,7 @@ class Glm_4_7(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, @@ -43,7 +44,7 @@ class Glm_4_7(SimpleService): return Message.objects.bulk_create(msgs) return msgs - def make(self, input_message: Message, save: bool = True) -> list[Message]: + 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", @@ -53,7 +54,10 @@ class Glm_4_7(SimpleService): 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') + return version, callback_data, messages + 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( @@ -72,7 +76,52 @@ class Glm_4_7(SimpleService): 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) + start_time = time.time() + content_parts: list[str] = [] + input_tokens = output_tokens = 0 + result = '' + + try: + stream = stream_bytedance_model_ark_run( + model=version, + callback_data=callback_data, + messages=messages, + ) + 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_tokens = BytedanceModelArkAdapter.tokenize( + version, ''.join([m['content'] for m in messages]) + ) + output_tokens = BytedanceModelArkAdapter.tokenize(version, 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( @@ -281,6 +281,18 @@ def bytedance_model_ark_run( ) +def stream_bytedance_model_ark_run( + model: str, + callback_data: dict, + messages: list | None = None, + include_reasoning: bool = False, +): + usage = yield from BytedanceModelArkAdapter._stream_chat( + model, callback_data, messages, include_reasoning + ) + return usage + + @shared_task def drop_redis_vectors(message_uid: str) -> None: redis_client = redis.Redis(host=settings.REDIS_HOST, port=settings.REDIS_PORT, db=0)