maga_transformer/openai/renderers/custom_renderer.py

from typing import Optional, List, Dict, Any, Union, Tuple, Callable, AsyncGenerator
import functools
import os
import json
import torch
import asyncio
import logging
from dataclasses import dataclass, field
from PIL import Image
from concurrent.futures import Future

from transformers import PreTrainedTokenizerBase

from maga_transformer.models.base_model import GenerateOutput, BaseModel, GenerateInput, GenerateOutputs, AuxInfo
from maga_transformer.config.generate_config import GenerateConfig
from maga_transformer.config.gpt_init_model_parameters import TemplateType
from maga_transformer.utils.mm_process_engine import MMProcessEngine
from maga_transformer.openai.api_datatype import ChatMessage, GPTFunctionDefinition, UsageInfo, \
    ChatCompletionRequest, ChatCompletionResponseStreamChoice, DeltaMessage, FinisheReason, \
    RoleEnum, RendererInfo, ChatCompletionStreamResponse, CompletionTokensDetails, PromptTokensDetails, \
    ChatCompletionTokenLogprob, TopLogprob, ChoiceLogprobs, \
    ChatCompletionResponseChoice, ChatCompletionResponse, DebugInfo
from maga_transformer.async_decoder_engine.async_model import AsyncModel
from maga_transformer.utils.word_util import get_stop_word_slices, truncate_response_with_stop_words, is_truncated
from maga_transformer.utils.util import has_overlap, has_overlap_kmp
from maga_transformer.utils.multimodal_util import MMUrlType, MultimodalInput, MMPreprocessConfig
from maga_transformer.async_decoder_engine.backend_rpc_server_visitor import BackendRPCServerVisitor

think_mode = bool(int(os.environ.get("THINK_MODE", 0)))
think_start_tag = os.environ.get("THINK_START_TAG", "<think>\n").encode('utf-8').decode('unicode_escape')
think_end_tag = os.environ.get("THINK_END_TAG", "</think>\n\n").encode('utf-8').decode('unicode_escape')

class StreamStatus:
    index: int = 0
    request: ChatCompletionRequest
    output: Optional[GenerateOutput] = None
    origin_output_ids: torch.Tensor = torch.empty(0, dtype=torch.int32)
    output_ids: torch.Tensor = torch.empty(0, dtype=torch.int32)
    last_output_ids: List[int] = []
    last_token_length: int = 0
    finish_reason = None
    tokenizer = None
    responded_string = ""
    delta_output_string = ""

    def __init__(self, request: ChatCompletionRequest):
        self.request = request

    def update_output(self,
                      output: GenerateOutput,
                      clean_output_func,
                      check_finish_func,
                      remove_stop_word_ids_func):
        self.index += 1
        self.output = output
        self.origin_output_ids = torch.cat((self.origin_output_ids, output.output_ids), dim=1)
        self.output_ids = clean_output_func(self.origin_output_ids)
        self.finish_reason = check_finish_func(self.output_ids, self.input_token_length)
        self.output_ids = remove_stop_word_ids_func(self.output_ids)

    def update_result(self):
        self.last_token_length = len(self.output_ids) - len(self.last_output_ids)
        self.last_output_ids = self.output_ids
        self.responded_string += self.delta_output_string

    @property
    def output_token_length(self):
        return len(self.output_ids)

    @property
    def input_token_length(self):
        return self.output.aux_info.input_len

    @property
    def reuse_length(self):
        return self.output.aux_info.reuse_len

    @property
    def prev_token_id(self):
        return self.last_output_ids[-self.last_token_length:]

    @property
    def tokens_to_decode(self):
        return self.prev_token_id + self.output_ids[len(self.last_output_ids):]

class StreamStatusSync:
    index: int = 0
    request: ChatCompletionRequest
    origin_output_ids: torch.Tensor = torch.empty(0, dtype=torch.int32)
    output_ids: torch.Tensor = torch.empty(0, dtype=torch.int32)
    last_output_ids: List[int] = []
    last_token_length: int = 0
    finish_reason = None
    tokenizer = None
    responded_string = ""
    delta_output_string = ""

    def __init__(self, request: ChatCompletionRequest):
        self.request = request

    def update_output_sync(self,
                           output_ids,
                           input_len,
                           clean_output_func,
                           check_finish_func,
                           remove_stop_word_ids_func):
        self.index += 1
        self.origin_output_ids = torch.cat((self.origin_output_ids, output_ids), dim=1)
        self.output_ids = clean_output_func(self.origin_output_ids)
        self.finish_reason = check_finish_func(self.output_ids, input_len)
        self.output_ids = remove_stop_word_ids_func(self.output_ids)

    def update_result(self):
        self.last_token_length = len(self.output_ids) - len(self.last_output_ids)
        self.last_output_ids = self.output_ids
        self.responded_string += self.delta_output_string

    @property
    def prev_token_id(self):
        return self.last_output_ids[-self.last_token_length:]

    @property
    def tokens_to_decode(self):
        return self.prev_token_id + self.output_ids[len(self.last_output_ids):]

@dataclass
class StreamResponseObject:
    choices: List[ChatCompletionResponseStreamChoice] = field(default_factory=list)
    usage: Optional[UsageInfo] = None
    aux_info: Optional[AuxInfo] = None

@dataclass
class ResponseObject:
    choices: List[ChatCompletionResponseChoice] = field(default_factory=list)
    usage: Optional[UsageInfo] = None
    aux_info: Optional[AuxInfo] = None

@dataclass
class RendererParams:
    model_type: str
    max_seq_len: int
    eos_token_id: int
    stop_word_ids_list: List[List[int]]
    template_type: TemplateType = TemplateType.chat
    ckpt_path: str = ""

@dataclass
class OutputDelta():
    output_str: Union[str, DeltaMessage]
    logprobs: Optional[ChatCompletionTokenLogprob]
    input_length: int
    output_length: int
    reuse_length: int

@dataclass
class ThinkStatus():
    in_think_mode: int = 0
    think_buffer: str = ""
    think_tokens: int = 0
    is_streaming: bool = False

class RenderedInputs:
    input_ids: List[int] = []
    multimodal_inputs: List[MultimodalInput] = []
    rendered_prompt: str = ""

    def __init__(self, input_ids: List[int], rendered_prompt: str = "", input_urls: List[str] = [], input_urls_type: List[MMUrlType] = [], preprocess_configs: List[MMPreprocessConfig] = []):
        self.input_ids = input_ids
        self.rendered_prompt = rendered_prompt
        self.multimodal_inputs = []
        if len(input_urls_type) == 0:
            input_urls_type = [MMUrlType.DEFAULT] * len(input_urls)
        elif len(input_urls_type) != len(input_urls):
            raise Exception(f"the number of multimodal input types must match url, now types {len(input_urls_type)} urls {len(input_urls)}")

        if len(preprocess_configs) == 0:
            preprocess_configs = [MMPreprocessConfig()] * len(input_urls)
        elif len(preprocess_configs) != len(preprocess_configs):
            raise Exception(f"the number of multimodal preprocess config must match url, now types {len(preprocess_configs)} urls {len(input_urls)}")

        for url, type, config in zip(input_urls, input_urls_type, preprocess_configs):
            self.multimodal_inputs.append(MultimodalInput(url, type, config))

class CustomChatRenderer():
    def __init__(self,
                 tokenizer: PreTrainedTokenizerBase,
                 renderer_params: RendererParams,
    ):
        self.tokenizer = tokenizer
        self.model_type = renderer_params.model_type
        self.max_seq_len = renderer_params.max_seq_len
        self.eos_token_id = renderer_params.eos_token_id
        self.stop_words_id_list = renderer_params.stop_word_ids_list
        self.stop_words_str_list = [
            self.tokenizer.decode(stop_word_ids) for stop_word_ids in self.stop_words_id_list
        ]
        self.ckpt_path = renderer_params.ckpt_path
        # NOTE: stop words or their ids only need to be added to one of these two lists.
        self.extra_stop_words: List[str] = []
        self.extra_stop_word_ids_list: List[List[int]] = []

    def __str__(self) -> str:
        return str(self.get_renderer_info())

    def __repr__(self) -> str:
        return self.__str__()

    def get_renderer_info(self) -> RendererInfo:
        extra_stop_word_ids_list = self.get_all_extra_stop_word_ids_list()
        extra_stop_words_list = [
            self.tokenizer.decode(stop_word_ids) for stop_word_ids in extra_stop_word_ids_list
        ]
        if len(extra_stop_words_list) and isinstance(extra_stop_words_list[0], list):
            extra_stop_words_list = [l[0] for l in extra_stop_words_list]
        return RendererInfo(
            class_name=self.__class__.__name__,
            renderer_model_type=self.model_type,
            extra_stop_word_ids_list=extra_stop_word_ids_list,
            extra_stop_words_list=extra_stop_words_list,
        )

    def add_extra_stop_words(self, extra_stop_words: List[str]):
        self.extra_stop_words.extend(extra_stop_words)

    def add_extra_stop_word_ids(self, extra_stop_word_ids: List[List[int]]):
        self.extra_stop_word_ids_list.extend(extra_stop_word_ids)

    def tokenize_words(self, words: List[str]) -> List[List[int]]:
        ids_list = []
        for word in words:
            if isinstance(self.tokenizer, PreTrainedTokenizerBase):
                token_id = self.tokenizer.convert_tokens_to_ids(word)
                if isinstance(token_id, int):
                    ids_list.append([token_id])
                elif isinstance(token_id, list):
                    ids_list.append(token_id)
                else:
                    ids_list.append(self.tokenizer.encode(word, add_special_tokens=True))
            else:
                ids_list.append(self.tokenizer.encode(word))
        return ids_list

    def get_all_extra_stop_word_ids_list(self) -> List[List[int]]:
        ids_list_from_words = self.tokenize_words(self.extra_stop_words)
        return self.extra_stop_word_ids_list + ids_list_from_words

    def _check_all_finished(self, status_list) -> bool:
        for s in status_list:
            if s.finish_reason == None:
                return False
        return True

    def getRequest(self, request: str) -> ChatCompletionRequest:
        return ChatCompletionRequest(**(json.loads(request)))

    def render_chat(self, request: ChatCompletionRequest) -> RenderedInputs:
        raise NotImplementedError

    async def generate_choice(
            self,
            request_id: int,
            input_ids: List[int],
            mm_inputs: List[MultimodalInput],
            generate_config: GenerateConfig,
            backend_rpc_server_visitor: BackendRPCServerVisitor,
            request: ChatCompletionRequest
    ) -> AsyncGenerator[StreamResponseObject, None]:

        token_type_ids = []
        input_id_tensor = torch.Tensor(input_ids).int().unsqueeze(0)
        output_generator: AsyncGenerator[GenerateOutput, None] = backend_rpc_server_visitor.enqueue(
            GenerateInput(
                request_id=request_id,
                token_ids=input_id_tensor,
                mm_inputs=mm_inputs,
                generate_config=generate_config,
                tokenizer=self.tokenizer,
                token_type_ids=token_type_ids
            )
        )

        async for response in self.render_response_stream(output_generator,
                                                          request,
                                                          generate_config):
            yield response

    async def _create_empty_delta(self, aux_info: AuxInfo):
        return OutputDelta(
            output_str="",
            logprobs=None,
            input_length=aux_info.input_len,
            output_length=aux_info.output_len,
            reuse_length=aux_info.reuse_len
        )

    async def _generate_log_probs(self, status: StreamStatus, output: Optional[GenerateOutput]) -> Optional[ChatCompletionTokenLogprob]:
        assert output is not None
        if not status.request.logprobs:
            return None
        prob_return_num = status.request.top_logprobs or 1
        all_probs = output.all_probs
        output_id = output.output_ids
        if output_id == None:
            return None
        selected_id = output_id[-1].item()
        if (all_probs == None):
            raise Exception("all_probs is None when logprobs is true. There should be a internal bug.")
        all_probs = all_probs.squeeze()
        probs, tokens = all_probs.sort(descending=True)
        non_zero_size = probs.nonzero().shape[0]
        log_values = probs.log()
        prob_return_num = min(prob_return_num, non_zero_size)

        selected_token = self.tokenizer.decode([selected_id])
        chat_logprob = ChatCompletionTokenLogprob(
            token=selected_token,
            bytes=list(selected_token.encode("utf-8", errors="replace")),
            logprob=all_probs[output_id].log().item(),
            top_logprobs=[]
        )
        for i in range(prob_return_num):
            token = self.tokenizer.decode(tokens[i].item())
            chat_logprob.top_logprobs.append(TopLogprob(
                token=token,
                logprob=log_values[i].item(),
                bytes=list(token.encode("utf-8", errors="replace")),
            ))

        logging.debug(f"chat_logprob: {chat_logprob.model_dump_json(indent=4)}")

        return chat_logprob

    async def _update_single_status(self, status: StreamStatus, output: GenerateOutput, max_new_tokens: int, stop_words_str: List[str], stop_word_slice_list: List[str], is_streaming: bool) -> OutputDelta:
        if status.finish_reason != None:
            return await self._create_empty_delta(status.output.aux_info)
        status.update_output(output,
                             self._clean_output_ids,
                             functools.partial(self._check_finish_reason, max_new_tokens=max_new_tokens),
                             self._remove_stop_word_ids)
        decoded_prev_token = self.tokenizer.decode(status.prev_token_id)
        decoded_string = self.tokenizer.decode(status.tokens_to_decode)
        # For some tokenizers (e.g. ChatGLM), decode a single token differs from decode a list of tokens.
        if is_streaming:
            if len(decoded_string) > 0 and u'\uFFFD' == decoded_string[-1]:
                return await self._create_empty_delta(output.aux_info)
        else:
            while (len(decoded_string) > 0) and (u'\uFFFD' == decoded_string[-1]):
                decoded_string = decoded_string[:-1]
        status.delta_output_string = decoded_string[len(decoded_prev_token):]
        if is_truncated(status.delta_output_string, stop_words_str, is_streaming):
            status.finish_reason = FinisheReason.stop
            return await self._create_empty_delta(output.aux_info)
        if not is_truncated(status.delta_output_string, stop_word_slice_list, is_streaming):
            status.update_result()
            delta = OutputDelta(
                output_str=status.delta_output_string,
                logprobs=await self._generate_log_probs(status, output),
                input_length=output.aux_info.input_len,
                output_length=output.aux_info.output_len,
                reuse_length=output.aux_info.reuse_len)
            status.delta_output_string = ""
            return delta
        else:
            return await self._create_empty_delta(output.aux_info)

    async def _generate_first(self, n: int):
        return StreamResponseObject(
                    choices=[ChatCompletionResponseStreamChoice(
                        index=i,
                        delta=DeltaMessage(
                            role=RoleEnum.assistant,
                            content="",
                        ),
                    ) for i in range(n)]
        )
        
    def _split_reasoning_text_and_content(self, item: OutputDelta, think_status: ThinkStatus):
        
        if isinstance(item.output_str, str):
            processing_index, output_len = 0, len(item.output_str)
            if output_len == 0:
                return DeltaMessage(content="")
            
            reasoning_text, content = "", ""
            update_think_tokens = think_status.in_think_mode
            while processing_index < output_len:
                if think_status.in_think_mode:
                    think_status.think_buffer += item.output_str[processing_index]
                    if think_status.think_buffer.startswith(think_start_tag):
                        think_status.think_buffer = think_status.think_buffer[len(think_start_tag):]
                        
                    if think_status.think_buffer.endswith(think_end_tag):
                        reasoning_text = think_status.think_buffer[:-len(think_end_tag)]
                        think_status.think_buffer = ""
                        think_status.in_think_mode = False
                    elif has_overlap_kmp(think_status.think_buffer, think_end_tag) \
                        or think_start_tag.startswith(think_status.think_buffer):
                        pass
                    else:
                        reasoning_text = think_status.think_buffer
                    processing_index += 1
                else:
                    content += item.output_str[processing_index:]
                    processing_index = output_len

            if think_status.in_think_mode:
                if has_overlap_kmp(think_status.think_buffer, think_end_tag) \
                    or think_start_tag.startswith(think_status.think_buffer):
                    reasoning_text = ""
                else:
                    think_status.think_buffer = ""
                    
            if think_mode and update_think_tokens:
                if not think_status.is_streaming:
                    think_status.think_tokens = item.output_length - len(self.tokenizer.tokenize(content or ""))
                else:
                    think_status.think_tokens = item.output_length
            return DeltaMessage(reasoning_content=reasoning_text or "", content=content or "")
        
        elif isinstance(item.output_str, DeltaMessage):
            return item.output_str
        
        else:
            raise Exception(f'undefined output_str type[{type(item.output_str)}]')

    async def _generate_stream_response(self, items: List[OutputDelta], think_status: ThinkStatus) -> StreamResponseObject:
        if len(items) == 0:
            raise Exception("output items length should not be 0")
        input_lengths = items[0].input_length
        output_lengths = sum([x.output_length for x in items])
        reuse_lengths = items[0].reuse_length
        
        all_choices = []
        for i, item in enumerate(items):
            delta = self._split_reasoning_text_and_content(item, think_status)
            all_choices.append(ChatCompletionResponseStreamChoice(
                index=i,
                delta=delta,
                logprobs=ChoiceLogprobs(
                        content=[item.logprobs] if item.logprobs != None else None,
                        refusal=None
                    ) if item.logprobs != None else None
                ))
            
        return StreamResponseObject(
                choices=all_choices,
                usage=UsageInfo(
                    prompt_tokens=input_lengths,
                    total_tokens=input_lengths + output_lengths,
                    completion_tokens=output_lengths,
                    completion_tokens_details=CompletionTokensDetails(reasoning_tokens=think_status.think_tokens) if think_mode > 0 else None,
                    prompt_tokens_details=PromptTokensDetails(cached_tokens=reuse_lengths) if reuse_lengths > 0 else None
                )
        )

    async def _flush_buffer(self, buffer_list: List[StreamStatus], stop_words_str: List[str], is_streaming: bool, think_status: ThinkStatus):
        output_items: List[OutputDelta] = []
        for buffer in buffer_list:
            if buffer.output is None:
                raise Exception("last output should not be None")
            aux_info = buffer.output.aux_info
            trunc_string = truncate_response_with_stop_words(buffer.delta_output_string, stop_words_str, is_streaming)
            output_items.append(OutputDelta(
                trunc_string,
                await self._generate_log_probs(buffer, buffer.output),
                aux_info.input_len,
                aux_info.output_len,
                aux_info.reuse_len))
        return await self._generate_stream_response(output_items, think_status)

    async def _generate_final(self, buffer_list: List[StreamStatus], request: ChatCompletionRequest, think_status: ThinkStatus):
        input_token_length = 0
        output_token_length = 0
        reuse_length = 0
        aux_info = None
        for i, buffer in enumerate(buffer_list):
            if buffer.output is None:
                raise Exception("buffer last output should not be None")
            # 延迟引入, 避免循环import
            from maga_transformer.openai.renderers.qwen_tool_renderer import QwenToolStreamStatus
            # 判断buffer有无generating_tool_call这个属性
            if isinstance(buffer, QwenToolStreamStatus) and buffer.generating_tool_call:
                buffer.finish_reason = FinisheReason.tool_calls

            if buffer.finish_reason == None:
                logging.debug(f"output {i} found no stop reason! use stop as default.")
                buffer.finish_reason = FinisheReason.stop
            if i == 0:
                input_token_length = buffer.output.aux_info.input_len
                reuse_length = buffer.output.aux_info.reuse_len
                aux_info = buffer.output.aux_info if request.aux_info else None
            output_token_length += buffer.output.aux_info.output_len
        return StreamResponseObject(
            choices=[ChatCompletionResponseStreamChoice(
                index=i,
                delta=DeltaMessage(
                    content="",
                ),
                finish_reason=buffer.finish_reason
            ) for i, buffer in enumerate(buffer_list)],
            usage=UsageInfo(
                prompt_tokens=input_token_length,
                total_tokens=input_token_length + output_token_length,
                completion_tokens=output_token_length,
                completion_tokens_details=CompletionTokensDetails(reasoning_tokens=think_status.think_tokens) if think_mode > 0 else None,
                prompt_tokens_details=PromptTokensDetails(cached_tokens=reuse_length) if reuse_length > 0 else None
            ),
            aux_info=aux_info
        )

    async def _create_status_list(self, n: int, request: ChatCompletionRequest) -> List[StreamStatus]:
        return [StreamStatus(request) for _ in range(n)]

    async def render_response_stream(
            self,
            output_generator: AsyncGenerator[GenerateOutputs, None],
            request: ChatCompletionRequest,
            generate_config: GenerateConfig
    ) -> AsyncGenerator[StreamResponseObject, None]:
        stop_word_slice_list = get_stop_word_slices(generate_config.stop_words_str)
        num_return_sequences = request.n if request.n is not None else 1
        status_list = await self._create_status_list(num_return_sequences, request)
        index = 0
        global think_mode
        think_status = ThinkStatus(in_think_mode=think_mode, think_buffer="", think_tokens=0, is_streaming=generate_config.is_streaming)
        async for outputs in output_generator:
            if index == 0:
                yield await self._generate_first(num_return_sequences)
            index += 1
            if len(outputs.generate_outputs) != num_return_sequences:
                raise Exception("output num != num_return_sequences")
            delta_list: List[OutputDelta] = []
            for status, output in zip(status_list, outputs.generate_outputs):
                delta_list.append(await self._update_single_status(
                    status, output, generate_config.max_new_tokens, generate_config.stop_words_str,
                    stop_word_slice_list, generate_config.is_streaming))
            yield await self._generate_stream_response(delta_list, think_status)
            if self._check_all_finished(status_list):
                break
        if index != 0:
            yield await self._flush_buffer(status_list, generate_config.stop_words_str, generate_config.is_streaming, think_status)
            yield await self._generate_final(status_list, request, think_status)

    def _create_empty_delta_sync(self, input_len, output_len, reuse_len):
        return OutputDelta(
            output_str="",
            logprobs=None,
            input_length=input_len,
            output_length=output_len,
            reuse_length=reuse_len
        )

    def _generate_log_probs_sync(self,
                                 status: StreamStatusSync,
                                 all_probs: torch.Tensor,
                                 output_ids: torch.Tensor) -> Optional[ChatCompletionTokenLogprob]:
        if not status.request.logprobs:
            return None
        prob_return_num = status.request.top_logprobs or 1
        all_probs = all_probs
        output_id = output_ids
        if output_id == None:
            return None
        selected_id = output_id[-1].item()
        if (all_probs == None):
            raise Exception("all_probs is None when logprobs is true. There should be a internal bug.")
        all_probs = all_probs.squeeze()
        probs, tokens = all_probs.sort(descending=True)
        non_zero_size = probs.nonzero().shape[0]
        log_values = probs.log()
        prob_return_num = min(prob_return_num, non_zero_size)

        selected_token = self.tokenizer.decode([selected_id])
        chat_logprob = ChatCompletionTokenLogprob(
            token=selected_token,
            bytes=list(selected_token.encode("utf-8", errors="replace")),
            logprob=all_probs[output_id].log().item(),
            top_logprobs=[]
        )
        for i in range(prob_return_num):
            token = self.tokenizer.decode(tokens[i].item())
            chat_logprob.top_logprobs.append(TopLogprob(
                token=token,
                logprob=log_values[i].item(),
                bytes=list(token.encode("utf-8", errors="replace")),
            ))

        logging.debug(f"chat_logprob: {chat_logprob.model_dump_json(indent=4)}")

        return chat_logprob

    def _update_single_status_sync(self,
                              status: StreamStatusSync,
                              input_len, # output.aux_info
                              output_len, # output.aux_info
                              reuse_len, # output.aux_info
                              all_probs: torch.Tensor,
                              output_ids: torch.Tensor,
                              max_new_tokens: int,
                              stop_words_str: List[str],
                              stop_word_slice_list: List[str],
                              is_streaming: bool) -> OutputDelta:
        if status.finish_reason != None:
            return self._create_empty_delta_sync(input_len, output_len, reuse_len)
        status.update_output_sync(output_ids, input_len,
                                  self._clean_output_ids,
                                  functools.partial(self._check_finish_reason, max_new_tokens=max_new_tokens),
                                  self._remove_stop_word_ids)
        decoded_prev_token = self.tokenizer.decode(status.prev_token_id)
        decoded_string = self.tokenizer.decode(status.tokens_to_decode)
        # For some tokenizers (e.g. ChatGLM), decode a single token differs from decode a list of tokens.
        if is_streaming:
            if len(decoded_string) > 0 and u'\uFFFD' == decoded_string[-1]:
                return self._create_empty_delta_sync(input_len, output_len, reuse_len)
        else:
            while (len(decoded_string) > 0) and (u'\uFFFD' == decoded_string[-1]):
                decoded_string = decoded_string[:-1]
        status.delta_output_string = decoded_string[len(decoded_prev_token):]
        if is_truncated(status.delta_output_string, stop_words_str, is_streaming):
            status.finish_reason = FinisheReason.stop
            return self._create_empty_delta_sync(input_len, output_len, reuse_len)
        if not is_truncated(status.delta_output_string, stop_word_slice_list, is_streaming):
            status.update_result()
            delta = OutputDelta(
                output_str=status.delta_output_string,
                logprobs=self._generate_log_probs_sync(status, all_probs, output_ids),
                input_length=input_len,
                output_length=output_len,
                reuse_length=reuse_len)
            status.delta_output_string = ""
            return delta
        else:
            return self._create_empty_delta_sync(input_len, output_len, reuse_len)

    def _generate_first_sync(self, n: int):
        return StreamResponseObject(
                    choices=[ChatCompletionResponseStreamChoice(
                        index=i,
                        delta=DeltaMessage(
                            role=RoleEnum.assistant,
                            content="",
                        ),
                    ) for i in range(n)]
        )

    def _generate_stream_response_sync(self, items: List[OutputDelta]) -> StreamResponseObject:
        if len(items) == 0:
            raise Exception("output items length should not be 0")
        input_lengths = items[0].input_length
        output_lengths = sum([x.output_length for x in items])
        reuse_lengths = items[0].reuse_length
        return StreamResponseObject(
                choices=[ChatCompletionResponseStreamChoice(
                    index=i,
                    delta=DeltaMessage(
                        content=item.output_str,
                    ) if isinstance(item.output_str, str) else item.output_str,
                    logprobs=ChoiceLogprobs(
                        content=[item.logprobs] if item.logprobs != None else None,
                        refusal=None
                    ) if item.logprobs != None else None
                ) for i, item in enumerate(items)],
                usage=UsageInfo(
                    prompt_tokens=input_lengths,
                    total_tokens=input_lengths + output_lengths,
                    completion_tokens=output_lengths,
                    prompt_tokens_details=PromptTokensDetails(cached_tokens=reuse_lengths) if reuse_lengths > 0 else None
                )
        )

    def _flush_buffer_sync(self,
                           buffer_list: List[StreamStatusSync],
                           input_len_list, output_len_list, reuse_len_list,
                           all_probs_list, output_ids_list,
                           stop_words_str: List[str],
                           is_streaming: bool):
        output_items: List[OutputDelta] = []
        for buffer, input_len, output_len, reuse_len, all_probs, output_ids in zip(
                buffer_list,
                input_len_list, output_len_list, reuse_len_list,
                all_probs_list, output_ids_list
                ):
            trunc_string = truncate_response_with_stop_words(buffer.delta_output_string, stop_words_str, is_streaming)
            output_items.append(OutputDelta(
                trunc_string,
                self._generate_log_probs_sync(buffer, all_probs, output_ids),
                input_len,
                output_len,
                reuse_len))
        return self._generate_stream_response_sync(output_items)

    def _generate_final_sync(self,
                             buffer_list: List[StreamStatusSync],
                             input_len_list, output_len_list, reuse_len_list):
        input_token_length = 0
        output_token_length = 0
        reuse_length = 0
        aux_info = None
        for i, (buffer, input_len, output_len, reuse_len) in enumerate(zip(buffer_list,
                                                                           input_len_list,
                                                                           output_len_list,
                                                                           reuse_len_list)):
            if buffer.finish_reason == None:
                logging.debug(f"output {i} found no stop reason! use stop as default.")
                buffer.finish_reason = FinisheReason.stop
            if i == 0:
                input_token_length = input_len
                reuse_length = reuse_len
            output_token_length += output_len
        return StreamResponseObject(
            choices=[ChatCompletionResponseStreamChoice(
                index=i,
                delta=DeltaMessage(
                    content="",
                ),
                finish_reason=buffer.finish_reason
            ) for i, buffer in enumerate(buffer_list)],
            usage=UsageInfo(
                prompt_tokens=input_token_length,
                total_tokens=input_token_length + output_token_length,
                completion_tokens=output_token_length,
                prompt_tokens_details=PromptTokensDetails(cached_tokens=reuse_length) if reuse_length > 0 else None
            ),
            aux_info=aux_info
        )

    def _create_status_list_sync(self, n: int, body: str) -> List[StreamStatusSync]:
        request = self.getRequest(body)
        return [StreamStatusSync(request) for _ in range(n)]

    def render_stream_response_first(self, n: int, debug_info: str):
        stream_response = self._generate_first_sync(n)
        chat_response = ChatCompletionStreamResponse(
                            choices=stream_response.choices,
                            usage=stream_response.usage,
                            aux_info=stream_response.aux_info,
                            debug_info=debug_info
                        )
        return chat_response.model_dump_json(exclude_none=True)

    def render_stream_response_refactor(self,
                                        status_list: StreamStatusSync, # pass in from cpp
                                        input_len_list, # output.aux_info
                                        output_len_list, # output.aux_info
                                        reuse_len_list, # output.aux_info
                                        all_probs_list, # GenerateOutput
                                        output_ids_list, # GenerateOutput
                                        max_new_tokens, # GenerateConfig
                                        stop_words_str, # GenerateConfig
                                        is_streaming):
        stop_word_slice_list = get_stop_word_slices(stop_words_str) # move into cpp, then pass in
        delta_list: List[OutputDelta] = []
        for status, input_len, output_len, reuse_len, all_probs, output_ids in zip(
                status_list,
                input_len_list, output_len_list, reuse_len_list, # AuxInfo
                all_probs_list, output_ids_list # GenerateOutput
                ):
            delta_list.append(self._update_single_status_sync(status,
                                                              input_len, output_len, reuse_len,
                                                              all_probs, output_ids,
                                                              max_new_tokens, stop_words_str,
                                                              stop_word_slice_list,
                                                              is_streaming))
        stream_response =  self._generate_stream_response_sync(delta_list)
        chat_response = ChatCompletionStreamResponse(
                            choices=stream_response.choices,
                            usage=stream_response.usage,
                            aux_info=stream_response.aux_info
                        )
        return chat_response.model_dump_json(exclude_none=True)

    def render_stream_response_flush(self,
                                     status_list,
                                     input_len_list, output_len_list, reuse_len_list,
                                     all_probs_list, output_ids_list,
                                     stop_words_str,
                                     is_streaming):
        stream_response = self._flush_buffer_sync(status_list,
                                                  input_len_list, output_len_list, reuse_len_list,
                                                  all_probs_list, output_ids_list,
                                                  stop_words_str,
                                                  is_streaming)
        chat_response = ChatCompletionStreamResponse(
                            choices=stream_response.choices,
                            usage=stream_response.usage,
                            aux_info=stream_response.aux_info
                        )
        return chat_response.model_dump_json(exclude_none=True)

    def render_stream_response_final(self,
                                     status_list,
                                     input_len_list, output_len_list, reuse_len_list):
        stream_response = self._generate_final_sync(status_list,
                                                    input_len_list, output_len_list, reuse_len_list)
        chat_response = ChatCompletionStreamResponse(
                            choices=stream_response.choices,
                            usage=stream_response.usage,
                            aux_info=stream_response.aux_info
                        )
        return chat_response.model_dump_json(exclude_none=True)

    def render_stream_response_first_blocking(self, n: int):
        stream_response = self._generate_first_sync(n)
        return stream_response

    def render_stream_response_blocking(self,
                                        status_list: StreamStatusSync, # pass in from cpp
                                        input_len_list, # output.aux_info
                                        output_len_list, # output.aux_info
                                        reuse_len_list, # output.aux_info
                                        all_probs_list, # GenerateOutput
                                        output_ids_list, # GenerateOutput
                                        max_new_tokens, # GenerateConfig
                                        stop_words_str, # GenerateConfig
                                        is_streaming
                                        ):
        stop_word_slice_list = get_stop_word_slices(stop_words_str) # move into cpp, then pass in
        delta_list: List[OutputDelta] = []
        for status, input_len, output_len, reuse_len, all_probs, output_ids in zip(
                status_list,
                input_len_list, output_len_list, reuse_len_list, # AuxInfo
                all_probs_list, output_ids_list # GenerateOutput
                ):
            delta_list.append(self._update_single_status_sync(status,
                                                              input_len, output_len, reuse_len,
                                                              all_probs, output_ids,
                                                              max_new_tokens, stop_words_str,
                                                              stop_word_slice_list,
                                                              is_streaming))
        stream_response =  self._generate_stream_response_sync(delta_list)
        return stream_response

    def render_stream_response_flush_blocking(self,
                                     status_list,
                                     input_len_list, output_len_list, reuse_len_list,
                                     all_probs_list, output_ids_list,
                                     stop_words_str,
                                     is_streaming):
        stream_response = self._flush_buffer_sync(status_list,
                                                  input_len_list, output_len_list, reuse_len_list,
                                                  all_probs_list, output_ids_list,
                                                  stop_words_str,
                                                  is_streaming)
        return stream_response

    def render_stream_response_final_blocking(self,
                                     status_list,
                                     input_len_list, output_len_list, reuse_len_list):
        stream_response = self._generate_final_sync(status_list,
                                                    input_len_list, output_len_list, reuse_len_list)
        return stream_response

    def collect_complete_response(self, choice_generator):
        all_choices = []
        usage = None
        aux_info = None
        
        def split_think_tag(text: Optional[str]):
            if text is None:
                return None, None
            text_results = text.split(think_end_tag, 1)
            reasoning_content = text_results[0] if len(text_results) == 2 else None
            content = text_results[1] if len(text_results) == 2 else text
            return content, reasoning_content
        
        for response in choice_generator:
            
            if len(response.choices) != len(all_choices):
                if (all_choices == []):
                    for i, choice in enumerate(response.choices):
                        content, reasoning_content = split_think_tag(choice.delta.content)
                        all_choices.append(ChatCompletionResponseChoice(
                                index=i,
                                message=ChatMessage(
                                    role=choice.delta.role or RoleEnum.assistant,
                                    content=content or None,
                                    reasoning_content=reasoning_content or None,
                                    function_call=choice.delta.function_call or None,
                                ),
                                finish_reason=choice.finish_reason,
                                logprobs=choice.logprobs,
                            )
                        )
                else:
                    raise ValueError(f"response.choices has different length! "
                                     f"[{response.choices}] vs [{all_choices}].")
            else:
                for i in range(len(all_choices)):
                    if all_choices[i].message.content == None:
                        all_choices[i].message.content = (response.choices[i].delta.content or None)
                    else:
                        all_choices[i].message.content += (response.choices[i].delta.content or "")
                    content, reasoning_content = split_think_tag(all_choices[i].message.content)
                    all_choices[i].message.content = content
                    all_choices[i].message.reasoning_content = reasoning_content
                    all_choices[i].message.role = response.choices[i].delta.role or all_choices[i].message.role
                    all_choices[i].message.function_call = response.choices[i].delta.function_call or all_choices[i].message.function_call
                    all_choices[i].finish_reason = response.choices[i].finish_reason or all_choices[i].finish_reason
                    if all_choices[i].logprobs != None:
                        if response.choices[i].logprobs != None:
                            all_choices[i].logprobs.content += response.choices[i].logprobs.content
                    else:
                        all_choices[i].logprobs = response.choices[i].logprobs
            usage = response.usage or usage
            aux_info = response.aux_info or aux_info

        if (usage == None):
            logging.warning(f"No usage returned from stream response. use empty value.")
            usage = UsageInfo(
                prompt_tokens=0,
                total_tokens=0,
                completion_tokens=0
            )
        chat_response = ChatCompletionResponse(
            choices=all_choices,
            usage=usage,
            aux_info=aux_info,
            model="AsyncModel",
        )
        return chat_response.model_dump_json(exclude_none=True)

    def _check_finish_reason(self, token_ids: List[int], input_token_length: int, max_new_tokens: int = -1) -> Optional[FinisheReason]:
        stop_word_ids_list_all = self.get_all_extra_stop_word_ids_list() + self.stop_words_id_list
        if max_new_tokens > 0 and len(token_ids) >= max_new_tokens:
            return FinisheReason.length
        if len(token_ids) + input_token_length >= self.max_seq_len:
            return FinisheReason.length
        if token_ids and token_ids[-1] == self.eos_token_id:
            return FinisheReason.stop
        for stop_word_ids in stop_word_ids_list_all:
            if (len(token_ids) >= len(stop_word_ids)) and (token_ids[-len(stop_word_ids):] == stop_word_ids):
                return FinisheReason.stop
        return None

    def _remove_stop_word_ids(self, output_ids: List[int]) -> List[int]:
        stop_word_ids_list_all = self.get_all_extra_stop_word_ids_list() + self.stop_words_id_list
        for stop_word_ids in stop_word_ids_list_all:
            #  此处应该从最大的范围开始判断
            # 有可能会有stopword_ids 重复的情况，比如[144575, 14098, 144575]
            # 若从1开始判断会导致 去除了最后一个 144575 就退出了
            for i in range(len(stop_word_ids) + 1, 1, -1):
                if output_ids[-i:] == stop_word_ids[:i]:
                    output_ids = output_ids[:-i]
                    break
        return output_ids

    def _clean_output_ids(self, output_ids_tensor: torch.Tensor) -> list[int]:
        output_ids_tensor = output_ids_tensor.cpu().reshape([-1])
        # TODO(wangyin): This slicing shouldn't be done here.
        # model should return output length, ids should be sliced with output length.
        output_ids = output_ids_tensor[output_ids_tensor != self.eos_token_id].tolist()
        return output_ids