# vllm/v1/outputs.py

@dataclass
class SamplerOutput:
    """采样器的输出"""

    # 采样的 token IDs [num_requests, 1]
    sampled_token_ids: torch.Tensor

    # Logprobs 信息（可选）
    logprobs_tensors: LogprobsTensors | None = None

@dataclass
class LogprobsTensors:
    """Logprobs 张量"""

    # Top-k token indices [num_tokens, num_logprobs + 1]
    indices: torch.Tensor

    # Top-k logprobs [num_tokens, num_logprobs + 1]
    logprobs: torch.Tensor

    # 采样 token 的排名 [num_tokens]
    ranks: torch.Tensor

# vllm/v1/outputs.py

@dataclass
class ModelRunnerOutput:
    """ModelRunner 的输出"""

    # 每个请求的采样结果
    # Dict[request_id, SamplerOutput]
    sampler_output: dict[str, SamplerOutput]

    # 模型特定的输出（如 pooling embeddings）
    model_output: Any | None = None

# vllm/v1/engine/__init__.py

@dataclass
class EngineCoreOutput:
    """EngineCore 的单个请求输出"""

    request_id: str

    # 新生成的 token IDs
    new_token_ids: list[int]

    # 完成原因（如果完成）
    finish_reason: FinishReason | None = None

    # 停止字符串（如果因停止字符串完成）
    stop_str: str | None = None

    # Logprobs 信息
    new_logprobs: list[dict[int, Logprob]] | None = None

@dataclass
class EngineCoreOutputs:
    """EngineCore 的批量输出"""

    outputs: list[EngineCoreOutput]

    # 完成的请求 IDs
    finished_req_ids: set[str] | None = None

# vllm/outputs.py

@dataclass
class RequestOutput:
    """用户可见的最终输出"""

    request_id: str

    # 原始 prompt
    prompt: str | None
    prompt_token_ids: list[int]

    # 生成结果（可能多个，如 beam search）
    outputs: list[CompletionOutput]

    # 是否完成
    finished: bool

    # 指标
    metrics: RequestMetrics | None = None

@dataclass
class CompletionOutput:
    """单个生成序列的输出"""

    index: int                    # 序列索引
    text: str                     # 生成的文本
    token_ids: list[int]          # 生成的 token IDs
    cumulative_logprob: float | None  # 累积对数概率
    logprobs: list[dict] | None   # 每个 token 的 logprobs
    finish_reason: str | None     # 完成原因
    stop_reason: str | int | None # 停止原因

# vllm/v1/core/sched/scheduler.py

def update_from_output(
    self,
    model_runner_output: ModelRunnerOutput,
    sampler_output: SamplerOutput | None,
    scheduler_output: SchedulerOutput,
) -> EngineCoreOutputs:
    """根据模型输出更新请求状态"""

    outputs: list[EngineCoreOutput] = []

    for req_id, req_output in model_runner_output.items():
        request = self.requests[req_id]

        # 1. 获取新生成的 token IDs
        new_token_ids = req_output.sampled_token_ids.tolist()

        # 2. 追加到请求
        request.append_output_token_ids(new_token_ids)

        # 3. 检查停止条件
        finish_reason, stop_str = check_stop(request, self.max_model_len)

        # 4. 处理完成的请求
        if finish_reason is not None:
            self._finish_request(request, finish_reason)

        # 5. 构建输出
        output = EngineCoreOutput(
            request_id=req_id,
            new_token_ids=new_token_ids,
            finish_reason=finish_reason,
            stop_str=stop_str,
            new_logprobs=self._process_logprobs(req_output),
        )
        outputs.append(output)

    return EngineCoreOutputs(outputs=outputs)

# vllm/v1/core/sched/utils.py

def check_stop(
    request: Request,
    max_model_len: int,
) -> tuple[FinishReason | None, str | None]:
    """检查请求是否应该停止"""

    # 1. 检查 EOS token
    last_token_id = request.all_token_ids[-1]
    if last_token_id == request.eos_token_id:
        return FinishReason.STOP, None

    # 2. 检查最大输出长度
    if request.num_output_tokens >= request.max_tokens:
        return FinishReason.LENGTH, None

    # 3. 检查模型最大长度
    if len(request.all_token_ids) >= max_model_len:
        return FinishReason.LENGTH, None

    # 4. 检查停止 token IDs
    if request.stop_token_ids:
        if last_token_id in request.stop_token_ids:
            return FinishReason.STOP, None

    # 5. 检查停止字符串
    if request.stop_strings:
        output_text = request.get_output_text()
        for stop_str in request.stop_strings:
            if stop_str in output_text:
                return FinishReason.STOP, stop_str

    return None, None  # 继续生成

# vllm/v1/engine/detokenizer.py

class IncrementalDetokenizer:
    """增量反向分词器"""

    def __init__(self, tokenizer: TokenizerLike, request: Request):
        self.tokenizer = tokenizer
        self.request = request

        # 已解码的文本
        self.output_text = ""

        # 待解码的 token 缓冲区
        self.token_buffer: list[int] = []

        # 上一次的偏移量（用于流式输出）
        self.prev_output_len = 0

    def decode(self, new_token_ids: list[int]) -> str:
        """解码新的 token，返回新增的文本"""

        # 添加到缓冲区
        self.token_buffer.extend(new_token_ids)

        # 尝试解码
        text = self.tokenizer.decode(
            self.token_buffer,
            skip_special_tokens=True,
        )

        # 检查是否有完整的字符
        # （某些 token 可能是部分字符，需要等待后续 token）
        if self._is_valid_utf8(text):
            self.output_text = text
            new_text = text[self.prev_output_len:]
            self.prev_output_len = len(text)
            return new_text

        return ""

# 流式响应示例
async def stream_response():
    async for output in engine.generate_stream(prompt, params):
        yield f"data: {json.dumps(output)}\n\n"

    yield "data: [DONE]\n\n"

@dataclass
class Logprob:
    """单个 token 的 logprob 信息"""

    logprob: float           # 对数概率
    rank: int | None         # 在词表中的排名
    decoded_token: str       # 解码后的文本

# 每个位置的 logprobs
# Dict[token_id, Logprob]

def _process_logprobs(
    self,
    sampler_output: SamplerOutput,
) -> list[dict[int, Logprob]] | None:
    """处理 logprobs 输出"""

    if sampler_output.logprobs_tensors is None:
        return None

    tensors = sampler_output.logprobs_tensors
    result = []

    for i in range(tensors.indices.shape[0]):
        token_logprobs = {}

        # 获取 top-k token IDs 和 logprobs
        for j in range(tensors.indices.shape[1]):
            token_id = tensors.indices[i, j].item()
            logprob = tensors.logprobs[i, j].item()

            # 解码 token
            decoded = self.tokenizer.decode([token_id])

            token_logprobs[token_id] = Logprob(
                logprob=logprob,
                rank=j if j > 0 else tensors.ranks[i].item(),
                decoded_token=decoded,
            )

        result.append(token_logprobs)

    return result

# vllm/entrypoints/openai/protocol.py

class ChatCompletionResponse(BaseModel):
    """OpenAI Chat Completion 响应格式"""

    id: str
    object: str = "chat.completion"
    created: int
    model: str
    choices: list[ChatCompletionResponseChoice]
    usage: UsageInfo

class ChatCompletionResponseChoice(BaseModel):
    index: int
    message: ChatMessage
    finish_reason: str | None
    logprobs: ChoiceLogprobs | None

class ChatMessage(BaseModel):
    role: str
    content: str

def create_chat_completion_response(
    request_output: RequestOutput,
    model_name: str,
) -> ChatCompletionResponse:
    """将 RequestOutput 转换为 OpenAI 格式"""

    choices = []
    for i, output in enumerate(request_output.outputs):
        choice = ChatCompletionResponseChoice(
            index=i,
            message=ChatMessage(
                role="assistant",
                content=output.text,
            ),
            finish_reason=output.finish_reason,
            logprobs=_convert_logprobs(output.logprobs),
        )
        choices.append(choice)

    return ChatCompletionResponse(
        id=f"chatcmpl-{request_output.request_id}",
        created=int(time.time()),
        model=model_name,
        choices=choices,
        usage=UsageInfo(
            prompt_tokens=len(request_output.prompt_token_ids),
            completion_tokens=sum(len(o.token_ids) for o in request_output.outputs),
            total_tokens=...,
        ),
    )