@@ -0,0 +1,10 @@ +class MessageService: + @classmethod + def prepare_output_message(cls, reasoning_text: str, output_text: str) -> str: + reasoning = reasoning_text.strip() + output = output_text.strip() + if reasoning and output: + return f'{reasoning}\n{output}' + if reasoning: + return f'{reasoning}' + return output @@ -7,6 +7,7 @@ from typing import Any, Generator, TypeAlias, TypedDict import httpx from backend import settings +from messages.services.message_service import MessageService from ml_model.exceptions import ( FileExtensionNotSupported, GenerationException, @@ -15,6 +16,7 @@ from ml_model.exceptions import ( RequestBlocked, ) from poller.models import Proxy +from tools.chats.domain import RawSSEChunk logger = logging.getLogger(__name__) @@ -70,7 +72,7 @@ class BytedanceVideoTaskResponse(TypedDict, total=False): RunChatResult: TypeAlias = tuple[str, int, int] -RunStreamChatResult: TypeAlias = Generator[str, None, BytedanceUsage] +RunStreamChatResult: TypeAlias = Generator[RawSSEChunk, None, BytedanceUsage] RunImageResult: TypeAlias = list[str] RunVideoResult: TypeAlias = tuple[str, int] BytedanceRunResult: TypeAlias = RunChatResult | RunImageResult | RunVideoResult @@ -131,7 +133,6 @@ class BytedanceModelArkAdapter: return str(error.get('code', '')) == 'InputTextSensitiveContentDetected' - @classmethod def _extract_chat_answer( cls, @@ -144,11 +145,14 @@ class BytedanceModelArkAdapter: for choice in choices if choice.get('finish_reason') in (BytedanceFinishReason.STOP, BytedanceFinishReason.LENGTH) and choice.get('message') - and choice['message'].get('content') is not None + and ( + choice['message'].get('content') is not None + or choice['message'].get('reasoning_content') is not None + ) ] if stop_choices: - content = ','.join(str(choice['message']['content']) for choice in stop_choices) - # TODO: включить после разделения reasoning и content в хранении + content = ','.join(choice['message'].get('content') or '' for choice in stop_choices) + reasoning = '' if include_reasoning: reasoning = ','.join( str(reasoning) @@ -156,11 +160,7 @@ class BytedanceModelArkAdapter: if choice.get('message') and (reasoning := choice['message'].get('reasoning_content')) is not None ) - if reasoning and content: - return f'**Рассуждение:**\n\n{reasoning}\n\n**Основная мысль:**\n\n{content}' - if reasoning: - return reasoning - return content + return MessageService.prepare_output_message(reasoning, content) cls._raise_by_error_payload(data, choices) @@ -286,8 +286,6 @@ class BytedanceModelArkAdapter: ) raise GenerationException usage: BytedanceUsage = {} - reasoning_started = False - content_started = False for line in resp.iter_lines(): if not line: continue @@ -305,16 +303,10 @@ class BytedanceModelArkAdapter: 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 + yield RawSSEChunk(event='think', data={'content': 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 + yield RawSSEChunk(event='token', data={'content': chunk}) if ( fr := choices[0].get('finish_reason') ) and fr not in ( @@ -7,7 +7,9 @@ import httpx import tiktoken from backend import settings +from messages.services.message_service import MessageService from poller.models import Proxy +from tools.chats.domain import RawSSEChunk from .models import ModelResponse @@ -35,7 +37,7 @@ class OpenrouterAdapter: @classmethod def run_streaming_api( cls, version: str, messages: list, callback_data: dict, model_name: str - ) -> Iterator[str]: + ) -> Iterator[RawSSEChunk]: for proxy in Proxy.objects.all(): with httpx.Client( base_url=cls.BASE_URL, @@ -55,8 +57,6 @@ class OpenrouterAdapter: }, ) as resp: content = '' - # reasoning используем только для фоллбэк-подсчёта токенизатора - # в ответ не кладём, заполняет буфер истории сообщений reasoning = '' input_tokens = output_tokens = cost = 0 for line in resp.iter_lines(): @@ -70,11 +70,14 @@ class OpenrouterAdapter: try: data_obj = json.loads(data) - chunk = data_obj['choices'][0]['delta'].get('content') or '' - reasoning += data_obj['choices'][0]['delta'].get('reasoning') or '' - if chunk: - content += chunk - yield chunk + content_chunk = data_obj['choices'][0]['delta'].get('content') or '' + reasoning_chunk = data_obj['choices'][0]['delta'].get('reasoning') or '' + if content_chunk: + content += content_chunk + yield RawSSEChunk(event='token', data={'content': content_chunk}) + if reasoning_chunk: + reasoning += reasoning_chunk + yield RawSSEChunk(event='think', data={'content': reasoning_chunk}) if data_obj.get('usage'): input_tokens = data_obj['usage']['prompt_tokens'] output_tokens = data_obj['usage']['completion_tokens'] @@ -100,15 +103,24 @@ class OpenrouterAdapter: cls, version: str, messages: list, callback_data: dict, model_name: str ) -> ModelResponse: stream = cls.run_streaming_api(version, messages, callback_data, model_name) - content_parts: list[str] = [] + reasoning = '' + content = '' try: while True: chunk = next(stream) if chunk: - content_parts.append(chunk) + if chunk.event == 'think': + reasoning += chunk.data['content'] + else: + content += chunk.data['content'] except StopIteration as exc: input_tokens, output_tokens, cost = exc.value - return ModelResponse(''.join(content_parts), input_tokens, output_tokens, cost) + return ModelResponse( + MessageService.prepare_output_message(reasoning, content), + input_tokens, + output_tokens, + cost, + ) @classmethod def _fallback_tokenize(cls, model_name: str, messages: list, content: str) -> tuple[int, int]: @@ -1,6 +1,6 @@ from abc import ABC, abstractmethod from decimal import Decimal -from typing import Any, Generator, Never +from typing import Any, Iterator, Never from asgiref.sync import async_to_sync from googletrans import Translator @@ -86,4 +86,4 @@ class SimpleService(ABC): class StreamSimpleService(SimpleService): @abstractmethod - def make_stream(self, input_message: Message, save: bool = True) -> Generator: ... + def make_stream(self, input_message: Message, save: bool = True) -> Iterator: ... @@ -37,6 +37,7 @@ from payments.exceptions.insufficient_balance import InsufficientBalance from payments.selectors.payment_plan_selector import PaymentPlanSelector from poller.models import Proxy +from tools.chats.domain import RawSSEChunk class Chatgpt(Chatgpt_4, StreamSimpleService, OpenAIStreamMixin): @@ -265,7 +266,7 @@ class Chatgpt(Chatgpt_4, StreamSimpleService, OpenAIStreamMixin): ) return self.save_results([response], process_time, generated_image, save) - def make_stream(self, input_message: Message, save: bool = True) -> Iterator[str]: + def make_stream(self, input_message: Message, save: bool = True) -> Iterator[RawSSEChunk]: ctx: dict[str, Any] = {} content_parts: list[str] = [] input_tokens = output_tokens = 0 @@ -284,7 +285,7 @@ class Chatgpt(Chatgpt_4, StreamSimpleService, OpenAIStreamMixin): chunk = next(stream) if chunk: content_parts.append(chunk) - yield chunk + yield RawSSEChunk(event='token', data={'content': chunk}) except StopIteration as exc: input_tokens, output_tokens = exc.value or (0, 0) break @@ -11,6 +11,7 @@ from messages.models import Message from ml_model.exceptions import ModelVersionNotAvailable, PaidPlanRequiredError from ml_model.services.chatgpt import Chatgpt from poller.models import Proxy +from tools.chats.domain import RawSSEChunk class Chatgpt_5_4(Chatgpt): @@ -94,7 +95,7 @@ class Chatgpt_5_4(Chatgpt): price += self.TOKENS_COST[model]['generated_image'] return price.quantize(Decimal('0.1'), rounding='ROUND_UP') - def make_stream(self, input_message: Message, save: bool = True) -> Iterator[str]: + def make_stream(self, input_message: Message, save: bool = True) -> Iterator[RawSSEChunk]: return (yield from super().make_stream(input_message, save)) def _build_payload( @@ -11,6 +11,7 @@ from PIL import Image from django.utils.translation import gettext from messages.models import Message +from messages.services.message_service import MessageService from ml_model.adapters.openrouter import OpenrouterAdapter from ml_model.exceptions import ( CorruptedFileError, @@ -25,6 +26,7 @@ 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 +from tools.chats.domain import RawSSEChunk from tools.chats.models import Chat from tools.copywrite.models import Copywrite from tools.public_api.models import APIStore @@ -132,7 +134,7 @@ class Claude(SerperMixin, StreamSimpleService): ) return self.save_results(result.content, process_time, save) - def make_stream(self, input_message: Message, save: bool = True) -> Iterator[str]: + def make_stream(self, input_message: Message, save: bool = True) -> Iterator[RawSSEChunk]: start_time = time.time() version_slug = input_message.info.get('version') if version_slug is None or version_slug not in self.TOKENS_COST: @@ -140,9 +142,10 @@ class Claude(SerperMixin, StreamSimpleService): model_slug = f'anthropic/{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 + reasoning = '' + content = '' result = '' try: @@ -151,13 +154,16 @@ class Claude(SerperMixin, StreamSimpleService): while True: chunk = next(stream) if chunk: - content_parts.append(chunk) + if chunk.event == 'think': + reasoning += chunk.data['content'] + else: + content += chunk.data['content'] yield chunk except StopIteration as exc: input_tokens, output_tokens, cost = exc.value finally: - if content_parts: - result = ''.join(content_parts) + if reasoning or content: + result = MessageService.prepare_output_message(reasoning, content) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice( input_message.content_object.model, @@ -12,6 +12,7 @@ import filetype from PIL import Image from messages.models import Message +from messages.services.message_service import MessageService from ml_model.adapters.bytedance_model_ark import BytedanceContentType, BytedanceModelArkAdapter from ml_model.services.FileService import FileProcessingService from ml_model.exceptions import GenerationException, ModelVersionNotAvailable @@ -19,6 +20,7 @@ from ml_model.services.base import SimpleService 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.domain import RawSSEChunk from tools.chats.models import Chat from tools.copywrite.models import Copywrite from tools.public_api.models import APIStore @@ -284,12 +286,13 @@ class Dola_Seed(SimpleService): msgs = self.save_results(result[0], process_time, save) return msgs - def make_stream(self, input_message: Message, save: bool = True) -> Iterator[str]: + def make_stream(self, input_message: Message, save: bool = True) -> Iterator[RawSSEChunk]: 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 + reasoning = '' + content = '' result = '' try: @@ -303,15 +306,18 @@ class Dola_Seed(SimpleService): while True: chunk = next(stream) if chunk: - content_parts.append(chunk) + if chunk.event == 'think': + reasoning += chunk.data['content'] + else: + content += chunk.data['content'] 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 reasoning or content: + result = MessageService.prepare_output_message(reasoning, content) if not (input_tokens + output_tokens): input_text_parts = [] for message in messages: @@ -10,6 +10,7 @@ import filetype from PIL import Image from messages.models import Message +from messages.services.message_service import MessageService from ml_model.adapters.openrouter import OpenrouterAdapter from ml_model.exceptions import CorruptedFileError, FileExtensionNotSupported, ModelVersionNotAvailable from ml_model.services.EmbeddingService import EmbeddingService @@ -17,6 +18,7 @@ from ml_model.services.FileService import FileProcessingService from ml_model.services.base import StreamSimpleService from ml_model.tasks import openrouter_run from poller.models import Proxy +from tools.chats.domain import RawSSEChunk from tools.chats.models import Chat from tools.copywrite.models import Copywrite from tools.public_api.models import APIStore @@ -94,7 +96,7 @@ class Gemini_3_1(StreamSimpleService): ) return self.save_results(result[0], process_time, save) - def make_stream(self, input_message: Message, save: bool = True) -> Iterator[str]: + def make_stream(self, input_message: Message, save: bool = True) -> Iterator[RawSSEChunk]: start_time = time.time() version_slug = input_message.info.get('version') if version_slug is None or version_slug not in self.TOKENS_COST: @@ -105,8 +107,9 @@ class Gemini_3_1(StreamSimpleService): **input_message.info, } messages, embedding_tokens = self._prepare_messages(input_message) - content_parts: list[str] = [] input_tokens = output_tokens = 0 + reasoning = '' + content = '' result = '' try: @@ -115,13 +118,16 @@ class Gemini_3_1(StreamSimpleService): while True: chunk = next(stream) if chunk: - content_parts.append(chunk) + if chunk.event == 'think': + reasoning += chunk.data['content'] + else: + content += chunk.data['content'] yield chunk except StopIteration as exc: input_tokens, output_tokens, _ = exc.value finally: - if content_parts: - result = ''.join(content_parts) + if reasoning or content: + result = MessageService.prepare_output_message(reasoning, content) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice( input_message.content_object.model, @@ -5,12 +5,14 @@ from decimal import Decimal from typing import Any, Iterator from messages.models import Message +from messages.services.message_service import MessageService 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, stream_bytedance_model_ark_run from payments.exceptions.insufficient_balance import InsufficientBalance from payments.selectors.payment_plan_selector import PaymentPlanSelector +from tools.chats.domain import RawSSEChunk from tools.chats.models import Chat from tools.copywrite.models import Copywrite from tools.public_api.models import APIStore @@ -107,11 +109,12 @@ class Glm_4_7(SimpleService): msgs = self.save_results(result[0], process_time, save) return msgs - def make_stream(self, input_message: Message, save: bool = True) -> Iterator[str]: + def make_stream(self, input_message: Message, save: bool = True) -> Iterator[RawSSEChunk]: version, callback_data, messages = self._prepare_data(input_message) start_time = time.time() - content_parts: list[str] = [] input_tokens = output_tokens = 0 + reasoning = '' + content = '' result = '' try: @@ -125,15 +128,18 @@ class Glm_4_7(SimpleService): while True: chunk = next(stream) if chunk: - content_parts.append(chunk) + if chunk.event == 'think': + reasoning += chunk.data['content'] + else: + content += chunk.data['content'] 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 reasoning or content: + result = MessageService.prepare_output_message(reasoning, content) if not (input_tokens + output_tokens): input_tokens = BytedanceModelArkAdapter.tokenize( version, ''.join([m['content'] for m in messages]) @@ -11,6 +11,7 @@ from PIL import Image from django.utils.translation import gettext from messages.models import Message +from messages.services.message_service import MessageService from ml_model.adapters.openrouter import OpenrouterAdapter from ml_model.exceptions import CorruptedFileError, FileExtensionNotSupported, PaidPlanRequiredError from ml_model.services.base import StreamSimpleService @@ -20,6 +21,7 @@ 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 +from tools.chats.domain import RawSSEChunk from tools.chats.models import Chat from tools.copywrite.models import Copywrite from tools.public_api.models import APIStore @@ -81,14 +83,15 @@ class Grok(SerperMixin, StreamSimpleService): ) return self.save_results(result.content, process_time, save) - def make_stream(self, input_message: Message, save: bool = True) -> Iterator[str]: + def make_stream(self, input_message: Message, save: bool = True) -> Iterator[RawSSEChunk]: start_time = time.time() 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 + reasoning = '' + content = '' result = '' try: @@ -97,13 +100,16 @@ class Grok(SerperMixin, StreamSimpleService): while True: chunk = next(stream) if chunk: - content_parts.append(chunk) + if chunk.event == 'think': + reasoning += chunk.data['content'] + else: + content += chunk.data['content'] yield chunk except StopIteration as exc: input_tokens, output_tokens, cost = exc.value finally: - if content_parts: - result = ''.join(content_parts) + if reasoning or content: + result = MessageService.prepare_output_message(reasoning, content) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice( input_message.content_object.model, @@ -1,4 +1,3 @@ -import json import time from datetime import timedelta from decimal import Decimal @@ -6,10 +5,9 @@ from pathlib import Path from typing import Iterator import filetype -import httpx -from backend import settings from messages.models import Message +from messages.services.message_service import MessageService from ml_model.adapters.openrouter import OpenrouterAdapter from ml_model.exceptions import CorruptedFileError, FileExtensionNotSupported from ml_model.services.EmbeddingService import EmbeddingService @@ -18,6 +16,7 @@ from ml_model.exceptions import ModelVersionNotAvailable from ml_model.services.base import StreamSimpleService from ml_model.tasks import openrouter_run from poller.models import Proxy +from tools.chats.domain import RawSSEChunk from tools.chats.models import Chat from tools.copywrite.models import Copywrite from tools.public_api.models import APIStore @@ -92,32 +91,34 @@ class Qwen_3_7(StreamSimpleService): msgs = self.save_results(result[0], process_time) return msgs - def make_stream(self, input_message: Message, save: bool = True) -> Iterator[str]: + def make_stream(self, input_message: Message, save: bool = True) -> Iterator[RawSSEChunk]: start_time = time.time() version_slug, model_slug, callback_data, messages, embedding_tokens = self._prepare_data( input_message ) - content_parts: list[str] = [] + input_tokens = output_tokens = 0 cost = 0 + reasoning = '' + content = '' result = '' try: - stream = self._run_streaming_api(model_slug, messages, callback_data) + stream = OpenrouterAdapter.run_streaming_api(model_slug, messages, callback_data, 'Qwen') try: while True: chunk = next(stream) if chunk: - content_parts.append(chunk) + if chunk.event == 'think': + reasoning += chunk.data['content'] + else: + content += chunk.data['content'] yield chunk except StopIteration as exc: - cost = exc.value + input_tokens, output_tokens, cost = exc.value finally: - if content_parts: - result = ''.join(content_parts) + if reasoning or content: + result = MessageService.prepare_output_message(reasoning, content) if not cost: - input_tokens, output_tokens = OpenrouterAdapter.count_tokens_fallback( - 'Qwen', messages, result - ) cost = self._estimate_cost(version_slug, input_tokens, output_tokens) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice( @@ -240,46 +241,3 @@ class Qwen_3_7(StreamSimpleService): + output_tokens * price_map['output'] / 1_000_000 ) return float(price / self.COEFFICIENT) - - def _run_streaming_api( - self, version: str, messages: list, callback_data: dict - ) -> Iterator[str]: - for proxy in Proxy.objects.all(): - with httpx.Client( - base_url='https://openrouter.ai/api/v1', - headers={'Authorization': f'Bearer {settings.OPENROUTER_API_KEY}'}, - proxy=f'{proxy.protocol}://{proxy.address}', - timeout=600, - ) as client: - with client.stream( - 'POST', - 'chat/completions', - json={ - 'model': version, - 'stream': True, - 'messages': messages, - 'transforms': ['middle-out'], - **callback_data, - }, - ) as resp: - cost = 0 - for line in resp.iter_lines(): - line = line.strip() - if not line or not line.startswith('data: '): - continue - - data = line[6:] - if data == '[DONE]': - break - - try: - data_obj = json.loads(data) - chunk = data_obj['choices'][0]['delta'].get('content') or '' - if chunk: - yield chunk - if data_obj.get('usage'): - cost = data_obj['usage'].get('cost') or 0 - except json.JSONDecodeError: - continue - - return cost @@ -9,10 +9,14 @@ from decimal import Decimal from django.core.cache import cache from messages.models import Message +from messages.services.message_service import MessageService +from ml_model.exceptions import GenerationException from ml_model.services.base import StreamSimpleService from typing import Iterator +from tools.chats.domain import RawSSEChunk + type TestTextTokens = list[str] @@ -20,6 +24,7 @@ type TestTextTokens = list[str] class PrepareTestData: input_tokens: TestTextTokens output_tokens: TestTextTokens + reasoning_tokens: TestTextTokens ttft: float tbt: float start_time: float @@ -39,6 +44,13 @@ class Text_Test_Model(StreamSimpleService): fermentum lorem sit amet tortor ultricies, id pulvinar nibh pulvinar. """ + BASE_REASONING_MESSAGE = """ + In a dapibus nulla. Aenean erat orci, egestas non orci at, varius tempus risus. Ut suscipit lorem magna, + quis auctor leo molestie ac. Integer ut efficitur neque. Curabitur sollicitudin ipsum dolor, et tempus massa + lacinia a. Donec efficitur egestas facilisis. Aliquam feugiat convallis arcu quis sollicitudin. + Nullam eleifend iaculis sapien id scelerisque. + """ + def calculate_price(self, input_tokens: int, output_tokens: int) -> Decimal: price = ( input_tokens * self.TOKENS_COST['input'] / 1_000_000 @@ -58,23 +70,49 @@ class Text_Test_Model(StreamSimpleService): def make(self, input_message: Message, save: bool = True) -> list[Message]: prepare = self._prepare(input_message) - result = ''.join(self._stream(prepare)) + reasoning = '' + output = '' + for token in self._stream(prepare): + if token.event == 'think': + reasoning += token.data['content'] + else: + output += token.data['content'] + result = MessageService.prepare_output_message(reasoning, output) + if not result: + raise GenerationException return self._finalize( - input_message, prepare, result, save, output_token_count=len(prepare.output_tokens) + input_message, + prepare, + result, + save, + output_token_count=len(prepare.output_tokens) + len(prepare.reasoning_tokens), ) - def make_stream(self, input_message: Message, save: bool = True) -> Iterator[str]: + def make_stream(self, input_message: Message, save: bool = True) -> Iterator[RawSSEChunk]: prepare = self._prepare(input_message) - result = '' + output = '' + reasoning = '' output_token_count = 0 try: for token in self._stream(prepare): - result += token + if token.event == 'think': + reasoning += token.data['content'] + else: + output += token.data['content'] output_token_count += 1 yield token finally: + result = MessageService.prepare_output_message(reasoning, output) if result: - self._finalize(input_message, prepare, result, save, output_token_count=output_token_count) + self._finalize( + input_message, + prepare, + result, + save, + output_token_count=output_token_count, + ) + if not result: + raise GenerationException return result def _prepare(self, input_message: Message) -> PrepareTestData: @@ -82,16 +120,21 @@ class Text_Test_Model(StreamSimpleService): info = input_message.info.copy() input_tokens = self._get_cached_tokens(input_message.content) output_tokens = self._get_cached_tokens(info.get('cm') or self.BASE_OUTPUT_MESSAGE) + reasoning_tokens = [] + if info.get('reasoning'): + reasoning_tokens = self._get_cached_tokens(info.get('rm') or self.BASE_REASONING_MESSAGE) ttft = info.get('ttft', 0.5) tbt = info.get('tbt', 0.35) - return PrepareTestData(input_tokens, output_tokens, ttft, tbt, start_time) + return PrepareTestData(input_tokens, output_tokens, reasoning_tokens, ttft, tbt, start_time) - def _stream(self, prepare: PrepareTestData) -> Iterator[str]: + def _stream(self, prepare: PrepareTestData) -> Iterator[RawSSEChunk]: time.sleep(prepare.ttft) - for i, token in enumerate(prepare.output_tokens, start=1): - yield token - if i < len(prepare.output_tokens): - time.sleep(prepare.tbt) + streaming_data = {'think': prepare.reasoning_tokens, 'token': prepare.output_tokens} + for k, v in streaming_data.items(): + for i, token in enumerate(v, start=1): + yield RawSSEChunk(event=k, data={'content': token}) + if i < len(v): + time.sleep(prepare.tbt) def _finalize( self, @@ -20,6 +20,7 @@ from replicate.exceptions import ModelError from requests import Response from backend import settings +from messages.services.message_service import MessageService from ml_model.adapters.bytedance_model_ark import BytedanceContentType, BytedanceModelArkAdapter from ml_model.adapters.openrouter import OpenrouterAdapter from ml_model.exceptions import ( @@ -167,7 +168,12 @@ def openrouter_run(version: str, messages: list, callback_data: dict, model_name for c in data.get('choices', []) if c.get('message') and c['message'].get('content') is not None ] - if not raw_content: + raw_reasoning = ','.join( + reasoning + for choice in data.get('choices', []) + if (reasoning := choice['message'].get('reasoning')) is not None + ) + if not raw_content and not raw_reasoning: if any( c.get('native_finish_reason') == 'SAFETY_CHECK_TYPE_CSAM' for c in (data.get('choices') or []) @@ -182,22 +188,8 @@ def openrouter_run(version: str, messages: list, callback_data: dict, model_name ) raise GenerationException content = ','.join(raw_content) - reasoning = ','.join( - reasoning - for choice in data.get('choices', []) - if (reasoning := choice['message'].get('reasoning')) is not None - ) - reasoning = re.sub(r'Вывод:|Основная мысль:|Рассуждение:|\*\*', '', reasoning) - answer = reasoning - if any(m in data['model'] for m in ('google/gemini', 'x-ai/grok-4.3')) or re.match( - r'^qwen/qwen3\.(?:5|6|7)-.*$', data['model'] - ): - answer = content - elif reasoning and content: - # TODO: переделать рендеринг сообщения на Jinja 2 - answer = f'**Рассуждение:**\n\n{reasoning}\n\n**Основная мысль:**\n\n{content}' - elif content: - answer = content + reasoning = re.sub(r'Вывод:|Основная мысль:|Рассуждение:|\*\*', '', raw_reasoning) + answer = MessageService.prepare_output_message(reasoning, content) if int(data.get('choices')[0].get('error', {}).get('code', 0)) == 502: error_type = re.sub(r'["\']', '', str(data['choices'][0]['error']['message'])) if error_type == 'Overloaded': @@ -15,6 +15,10 @@ class SSEChunkService: def token(cls, event_id: int, content: str) -> SSEChunk: return cls._chunk(event_id, 'token', {'content': content}) + @classmethod + def think(cls, event_id: int, content: str) -> SSEChunk: + return cls._chunk(event_id, 'think', {'content': content}) + @classmethod def error(cls, event_id: int, content: str) -> SSEChunk: return cls._chunk(event_id, 'error', {'detail': content}) @@ -5,12 +5,16 @@ from dataclasses import asdict, dataclass from tools.chats.typing import SSEData, SSEEvent -@dataclass -class SSEChunk: - event_id: int +@dataclass(frozen=True, slots=True) +class RawSSEChunk: event: SSEEvent data: SSEData + +@dataclass(frozen=True, slots=True) +class SSEChunk(RawSSEChunk): + event_id: int + def encode(self): payload = orjson.dumps(self.data, default=str).decode() return f'id: {self.event_id}\nevent: {self.event}\ndata: {payload}\n\n' @@ -8,6 +8,7 @@ from django.db.models.functions import Greatest from messages.models import Message from ml_model.models import NeuronModel from ml_model.services.base import StreamSimpleService +from tools.chats.domain import RawSSEChunk from tools.chats.models import Chat from tools.chats.services.sse_chunk_service import SSEChunkService from tools.chats.services.sse_store import PublicSSEStoreService, SSEStoreService @@ -35,10 +36,17 @@ def _run_stream( stream = service.make_stream(message) while True: token = next(stream) - if not token: + if not isinstance(token, RawSSEChunk): + continue + token_data = token.data.get('content') + if not token_data or not isinstance(token_data, str): continue event_id += 1 - store.push(SSEChunkService.token(event_id, token)) + store.push( + SSEChunkService.think(event_id, token_data) + if token.event == 'think' + else SSEChunkService.token(event_id, token_data) + ) except StopIteration as exc: event_id += 1 store.push(SSEChunkService.done(event_id, exc.value or ''), ttl=settings.SSE_DONE_STREAM_TTL) @@ -1,4 +1,4 @@ from typing import Any, Literal -type SSEEvent = Literal['pending', 'start', 'token', 'error', 'done'] +type SSEEvent = Literal['pending', 'start', 'token', 'think', 'error', 'done'] type SSEData = dict[str, Any]