"tests/kernels/attention/test_attention.py" did not exist on "56f738ae9b631189e67795b397258afbed59b042"
serving.py 83.5 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
import asyncio
5
import json
6
import time
7
8
from collections.abc import AsyncGenerator, AsyncIterator
from collections.abc import Sequence as GenericSequence
9
from typing import Any, Final
10

11
import jinja2
12
import partial_json_parser
13
import regex as re
14
from fastapi import Request
15
from openai_harmony import Message as OpenAIMessage
16
from partial_json_parser.core.options import Allow
17

18
from vllm.engine.protocol import EngineClient
19
20
21
22
23
24
from vllm.entrypoints.chat_utils import (
    ChatTemplateContentFormatOption,
    ConversationMessage,
    get_history_tool_calls_cnt,
    make_tool_call_id,
)
25
from vllm.entrypoints.logger import RequestLogger
26
from vllm.entrypoints.openai.chat_completion.protocol import (
27
28
29
30
31
32
33
34
35
36
    ChatCompletionLogProb,
    ChatCompletionLogProbs,
    ChatCompletionLogProbsContent,
    ChatCompletionNamedToolChoiceParam,
    ChatCompletionRequest,
    ChatCompletionResponse,
    ChatCompletionResponseChoice,
    ChatCompletionResponseStreamChoice,
    ChatCompletionStreamResponse,
    ChatMessage,
37
38
)
from vllm.entrypoints.openai.chat_completion.stream_harmony import (
39
    TokenState,
40
41
42
    extract_harmony_streaming_delta,
)
from vllm.entrypoints.openai.engine.protocol import (
43
44
45
46
47
48
49
50
51
    DeltaFunctionCall,
    DeltaMessage,
    DeltaToolCall,
    ErrorResponse,
    PromptTokenUsageInfo,
    RequestResponseMetadata,
    ToolCall,
    UsageInfo,
)
52
from vllm.entrypoints.openai.engine.serving import (
53
54
55
56
    GenerationError,
    OpenAIServing,
    clamp_prompt_logprobs,
)
57
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
58
59
60
61
62
63
64
65
66
from vllm.entrypoints.openai.parser.harmony_utils import (
    get_developer_message,
    get_stop_tokens_for_assistant_actions,
    get_streamable_parser_for_assistant,
    get_system_message,
    parse_chat_inputs_to_harmony_messages,
    parse_chat_output,
    render_for_completion,
)
67
from vllm.entrypoints.openai.utils import maybe_filter_parallel_tool_calls
68
from vllm.entrypoints.utils import get_max_tokens, should_include_usage
69
from vllm.inputs.data import TokensPrompt
70
from vllm.logger import init_logger
71
from vllm.logprobs import Logprob
72
from vllm.outputs import CompletionOutput, RequestOutput
73
from vllm.sampling_params import BeamSearchParams, SamplingParams
74
75
76
from vllm.tokenizers import TokenizerLike
from vllm.tokenizers.mistral import (
    MistralTokenizer,
77
78
79
80
    maybe_serialize_tool_calls,
    truncate_tool_call_ids,
    validate_request_params,
)
81
82
from vllm.tool_parsers import ToolParser
from vllm.tool_parsers.mistral_tool_parser import MistralToolCall
83
from vllm.tool_parsers.utils import partial_json_loads
84
from vllm.utils.collection_utils import as_list
85
from vllm.v1.sample.logits_processor import validate_logits_processors_parameters
86
87
88
89
90

logger = init_logger(__name__)


class OpenAIServingChat(OpenAIServing):
91
92
93
    def __init__(
        self,
        engine_client: EngineClient,
94
        models: OpenAIServingModels,
95
96
        response_role: str,
        *,
97
98
        request_logger: RequestLogger | None,
        chat_template: str | None,
99
        chat_template_content_format: ChatTemplateContentFormatOption,
100
        trust_request_chat_template: bool = False,
101
        return_tokens_as_token_ids: bool = False,
102
        reasoning_parser: str = "",
103
        enable_auto_tools: bool = False,
104
        exclude_tools_when_tool_choice_none: bool = False,
105
        tool_parser: str | None = None,
106
        enable_prompt_tokens_details: bool = False,
107
        enable_force_include_usage: bool = False,
108
        enable_log_outputs: bool = False,
109
        enable_log_deltas: bool = True,
110
        log_error_stack: bool = False,
111
        default_chat_template_kwargs: dict[str, Any] | None = None,
112
    ) -> None:
113
114
115
116
117
118
119
        super().__init__(
            engine_client=engine_client,
            models=models,
            request_logger=request_logger,
            return_tokens_as_token_ids=return_tokens_as_token_ids,
            log_error_stack=log_error_stack,
        )
120

121
        self.response_role = response_role
122
123
        self.chat_template = chat_template
        self.chat_template_content_format: Final = chat_template_content_format
124
        self.trust_request_chat_template = trust_request_chat_template
125
        self.default_chat_template_kwargs = default_chat_template_kwargs or {}
126
        self.enable_log_outputs = enable_log_outputs
127
        self.enable_log_deltas = enable_log_deltas
128

129
130
131
        # set up logits processors
        self.logits_processors = self.model_config.logits_processors

132
133
134
135
        # set up reasoning parser
        self.reasoning_parser = self._get_reasoning_parser(
            reasoning_parser_name=reasoning_parser
        )
136
137
        # set up tool use
        self.enable_auto_tools: bool = enable_auto_tools
138
139
        self.tool_parser = self._get_tool_parser(
            tool_parser_name=tool_parser, enable_auto_tools=enable_auto_tools
140
141
        )
        self.exclude_tools_when_tool_choice_none = exclude_tools_when_tool_choice_none
142

143
        self.enable_prompt_tokens_details = enable_prompt_tokens_details
144
        self.enable_force_include_usage = enable_force_include_usage
145
        self.default_sampling_params = self.model_config.get_diff_sampling_param()
zhuwenwen's avatar
zhuwenwen committed
146
147
148
149
150
        if self.model_config.hf_config.model_type == "kimi_k2":
            self.tool_call_id_type = "kimi_k2"
        else:
            self.tool_call_id_type = "random"

151
        self.use_harmony = self.model_config.hf_config.model_type == "gpt_oss"
152
153
154
155
        if self.use_harmony:
            if "stop_token_ids" not in self.default_sampling_params:
                self.default_sampling_params["stop_token_ids"] = []
            self.default_sampling_params["stop_token_ids"].extend(
156
157
                get_stop_tokens_for_assistant_actions()
            )
158
159
160
161
162
163
164
165
166
167
168

        # NOTE(woosuk): While OpenAI's chat completion API supports browsing
        # for some models, currently vLLM doesn't support it. Please use the
        # Responses API instead.
        self.supports_browsing = False
        self.browser_tool = None
        # NOTE(woosuk): Chat completion API does not support code interpreter.
        # Please use the Responses API instead.
        self.supports_code_interpreter = False
        self.python_tool = None

169
170
171
172
173
174
175
176
177
178
179
180
    async def warmup(self) -> None:
        """
        Warm up the chat template processing to avoid first-request latency.

        This method triggers Jinja2 template compilation and content format
        detection that would otherwise happen on the first real request,
        causing increased latency on the first request.
        """
        logger.info("Warming up chat template processing...")
        start_time = time.perf_counter()

        try:
181
            renderer = self.engine_client.renderer
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196

            # Create a minimal dummy request
            dummy_request = ChatCompletionRequest(
                messages=[{"role": "user", "content": "warmup"}],
                model=None,
                max_completion_tokens=1,
            )

            # Call _preprocess_chat to trigger template compilation
            # This forces:
            # 1. Chat template content format detection
            # 2. Jinja2 template compilation
            # 3. Tokenizer initialization for chat
            await self._preprocess_chat(
                dummy_request,
197
                renderer,
198
199
200
201
202
203
204
205
                dummy_request.messages,
                chat_template=self.chat_template,
                chat_template_content_format=self.chat_template_content_format,
                add_generation_prompt=True,
                continue_final_message=False,
                tool_dicts=None,
                documents=None,
                chat_template_kwargs=None,
206
                default_chat_template_kwargs=self.default_chat_template_kwargs,
207
208
209
210
211
212
213
214
215
216
217
                tool_parser=None,
                add_special_tokens=False,
            )

            elapsed = (time.perf_counter() - start_time) * 1000
            logger.info("Chat template warmup completed in %.1fms", elapsed)

        except Exception:
            # Log but don't fail server startup if warmup fails
            logger.exception("Chat template warmup failed")

218
    async def render_chat_request(
219
220
        self,
        request: ChatCompletionRequest,
221
    ) -> tuple[list[ConversationMessage], list[Any]] | ErrorResponse:
222
        """
223
        render chat request by validating and preprocessing inputs.
224

225
226
227
        Returns:
            A tuple of (conversation, engine_prompts) on success,
            or an ErrorResponse on failure.
228
229
230
        """
        error_check_ret = await self._check_model(request)
        if error_check_ret is not None:
231
            logger.error("Error with model %s", error_check_ret)
232
233
            return error_check_ret

234
235
236
237
238
239
        # If the engine is dead, raise the engine's DEAD_ERROR.
        # This is required for the streaming case, where we return a
        # success status before we actually start generating text :).
        if self.engine_client.errored:
            raise self.engine_client.dead_error

240
        try:
241
242
            renderer = self.engine_client.renderer
            tokenizer = renderer.tokenizer
243

244
245
            tool_parser = self.tool_parser

246
            if isinstance(tokenizer, MistralTokenizer):
247
248
249
                # because of issues with pydantic we need to potentially
                # re-serialize the tool_calls field of the request
                # for more info: see comment in `maybe_serialize_tool_calls`
zhuwenwen's avatar
zhuwenwen committed
250
251
                maybe_serialize_tool_calls(request)
                truncate_tool_call_ids(request)
252
                validate_request_params(request)
253

254
255
256
            # Check if tool parsing is unavailable (common condition)
            tool_parsing_unavailable = (
                tool_parser is None
257
258
                and not isinstance(tokenizer, MistralTokenizer)
                and not self.use_harmony
259
260
261
262
263
264
            )

            # Validate tool_choice when tool parsing is required but unavailable
            if tool_parsing_unavailable and request.tool_choice not in (
                None,
                "none",
265
            ):
266
267
268
269
270
271
272
273
274
275
276
277
278
                if request.tool_choice == "auto" and not self.enable_auto_tools:
                    # for hf tokenizers, "auto" tools requires
                    # --enable-auto-tool-choice and --tool-call-parser
                    return self.create_error_response(
                        '"auto" tool choice requires '
                        "--enable-auto-tool-choice and --tool-call-parser to be set"
                    )
                elif request.tool_choice != "auto":
                    # "required" or named tool requires tool parser
                    return self.create_error_response(
                        f'tool_choice="{request.tool_choice}" requires '
                        "--tool-call-parser to be set"
                    )
279

280
281
282
283
            if request.tools is None or (
                request.tool_choice == "none"
                and self.exclude_tools_when_tool_choice_none
            ):
284
285
286
                tool_dicts = None
            else:
                tool_dicts = [tool.model_dump() for tool in request.tools]
287

288
289
            if not self.use_harmony:
                # Common case.
290
291
292
                error_check_ret = self._validate_chat_template(
                    request_chat_template=request.chat_template,
                    chat_template_kwargs=request.chat_template_kwargs,
293
                    trust_request_chat_template=self.trust_request_chat_template,
294
295
296
                )
                if error_check_ret is not None:
                    return error_check_ret
297
298
299
300

                chat_template_kwargs = request.chat_template_kwargs or {}
                chat_template_kwargs.update(reasoning_effort=request.reasoning_effort)

301
                conversation, engine_prompts = await self._preprocess_chat(
302
                    request,
303
                    renderer,
304
                    request.messages,
305
                    chat_template=request.chat_template or self.chat_template,
306
                    chat_template_content_format=self.chat_template_content_format,
307
308
309
310
                    add_generation_prompt=request.add_generation_prompt,
                    continue_final_message=request.continue_final_message,
                    tool_dicts=tool_dicts,
                    documents=request.documents,
311
                    chat_template_kwargs=chat_template_kwargs,
312
                    default_chat_template_kwargs=self.default_chat_template_kwargs,
313
314
315
316
317
                    tool_parser=tool_parser,
                    add_special_tokens=request.add_special_tokens,
                )
            else:
                # For GPT-OSS.
318
319
320
321
                should_include_tools = tool_dicts is not None
                conversation, engine_prompts = self._make_request_with_harmony(
                    request, should_include_tools
                )
322
        except (ValueError, TypeError, RuntimeError, jinja2.TemplateError) as e:
323
            logger.exception("Error in preprocessing prompt inputs")
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
            return self.create_error_response(e)

        return conversation, engine_prompts

    async def create_chat_completion(
        self,
        request: ChatCompletionRequest,
        raw_request: Request | None = None,
    ) -> AsyncGenerator[str, None] | ChatCompletionResponse | ErrorResponse:
        """
        Chat Completion API similar to OpenAI's API.

        See https://platform.openai.com/docs/api-reference/chat/create
        for the API specification. This API mimics the OpenAI
        Chat Completion API.
        """
        result = await self.render_chat_request(request)
        if isinstance(result, ErrorResponse):
            return result

        conversation, engine_prompts = result
345

346
347
348
        request_id = (
            f"chatcmpl-{self._base_request_id(raw_request, request.request_id)}"
        )
349
350
351
352
353

        request_metadata = RequestResponseMetadata(request_id=request_id)
        if raw_request:
            raw_request.state.request_metadata = request_metadata

354
355
356
357
358
359
360
361
362
363
        try:
            lora_request = self._maybe_get_adapters(
                request, supports_default_mm_loras=True
            )

            model_name = self.models.model_name(lora_request)
        except (ValueError, TypeError, RuntimeError) as e:
            logger.exception("Error preparing request components")
            return self.create_error_response(e)

364
365
366
        # Extract data_parallel_rank from header (router can inject it)
        data_parallel_rank = self._get_data_parallel_rank(raw_request)

367
        # Schedule the request and get the result generator.
368
        generators: list[AsyncGenerator[RequestOutput, None]] = []
369
        try:
370
            for i, engine_prompt in enumerate(engine_prompts):
zhuwenwen's avatar
zhuwenwen committed
371
                prompt_text, _, _ = self._get_prompt_components(engine_prompt)
372
373
374
375
376
                # If we are creating sub requests for multiple prompts, ensure that they
                # have unique request ids.
                sub_request_id = (
                    request_id if len(engine_prompts) == 1 else f"{request_id}_{i}"
                )
377

zhuwenwen's avatar
zhuwenwen committed
378
379
380
                if self.default_sampling_params is None:
                    self.default_sampling_params = {}

381
382
383
                max_tokens = get_max_tokens(
                    max_model_len=self.max_model_len,
                    request=request,
zhuwenwen's avatar
zhuwenwen committed
384
                    input_length=len(engine_prompt["prompt_token_ids"]),
385
386
                    default_sampling_params=self.default_sampling_params,
                )
387

388
                sampling_params: SamplingParams | BeamSearchParams
389
390
                if request.use_beam_search:
                    sampling_params = request.to_beam_search_params(
391
392
                        max_tokens, self.default_sampling_params
                    )
393
394
                else:
                    sampling_params = request.to_sampling_params(
395
396
397
398
                        max_tokens,
                        self.model_config.logits_processor_pattern,
                        self.default_sampling_params,
                    )
399
400
401
402
                    validate_logits_processors_parameters(
                        self.logits_processors,
                        sampling_params,
                    )
403

404
                self._log_inputs(
405
                    sub_request_id,
406
                    engine_prompt,
407
408
409
                    params=sampling_params,
                    lora_request=lora_request,
                )
410

411
412
413
414
415
                trace_headers = (
                    None
                    if raw_request is None
                    else await self._get_trace_headers(raw_request.headers)
                )
416
417

                if isinstance(sampling_params, BeamSearchParams):
418
                    generator = self.beam_search(
419
                        prompt=engine_prompt,
420
                        request_id=sub_request_id,
421
                        params=sampling_params,
422
                        lora_request=lora_request,
423
                        trace_headers=trace_headers,
424
425
                    )
                else:
426
                    engine_request, tokenization_kwargs = await self._process_inputs(
427
                        sub_request_id,
428
429
430
431
432
                        engine_prompt,
                        sampling_params,
                        lora_request=lora_request,
                        trace_headers=trace_headers,
                        priority=request.priority,
433
                        data_parallel_rank=data_parallel_rank,
434
                    )
435

436
                    generator = self.engine_client.generate(
437
                        engine_request,
438
                        sampling_params,
439
                        sub_request_id,
440
441
442
                        lora_request=lora_request,
                        trace_headers=trace_headers,
                        priority=request.priority,
443
444
                        prompt_text=prompt_text,
                        tokenization_kwargs=tokenization_kwargs,
445
                        data_parallel_rank=data_parallel_rank,
446
447
448
                    )

                generators.append(generator)
449
        except ValueError as e:
450
            return self.create_error_response(e)
451

452
        assert len(generators) == 1
453
        (result_generator,) = generators
454

455
        # Streaming response
456
457
        tokenizer = self.renderer.tokenizer

458
459
        if request.stream:
            return self.chat_completion_stream_generator(
460
461
462
463
464
465
466
                request,
                result_generator,
                request_id,
                model_name,
                conversation,
                tokenizer,
                request_metadata,
467
            )
468

469
470
        try:
            return await self.chat_completion_full_generator(
471
472
473
474
475
476
477
478
                request,
                result_generator,
                request_id,
                model_name,
                conversation,
                tokenizer,
                request_metadata,
            )
479
480
        except GenerationError as e:
            return self._convert_generation_error_to_response(e)
481
        except ValueError as e:
482
            return self.create_error_response(e)
483
484
485
486

    def get_chat_request_role(self, request: ChatCompletionRequest) -> str:
        if request.add_generation_prompt:
            return self.response_role
487
        return request.messages[-1]["role"]
488

489
    @staticmethod
490
    def _bracket_level(s: str, opening="{", closing="}") -> int:
491
492
493
494
495
496
497
498
499
500
501
502
        """
        Calculate the current level of nested brackets in a given string.
        """
        level = 0
        for char in s:
            if char == opening:
                level += 1
            elif char == closing:
                level -= 1
        return level

    @staticmethod
503
    def _filter_delta_text(delta_text: str, previous_text: str) -> tuple[str, bool]:
504
505
506
507
508
509
510
511
512
        # remove last '},' of the tool definition stemming from the
        # "name"/"parameters" outer object or closing ']' of the tool list
        # count occurrences of opening and closing curly braces and
        # once level 0 is reached stop outputting text
        # if 0 is reached while parsing the delta_text we know the current
        # tool will finish in this current iteration
        bracket_level = OpenAIServingChat._bracket_level(previous_text)
        updated_delta, passed_zero = "", False
        for c in delta_text:
513
            if c == "{":
514
515
                bracket_level += 1
                passed_zero = bracket_level == 0
516
            elif c == "}":
517
518
519
520
521
522
523
                bracket_level -= 1
                passed_zero = bracket_level == 0

            if bracket_level != 0:
                updated_delta += c
            else:
                # if a comma is reached at level 0 we can stop
524
                if c == ",":
525
526
527
528
529
530
                    break
        return updated_delta, passed_zero

    def extract_tool_call_required_streaming(
        self,
        previous_text: str,
531
        current_text: str | None,
532
533
        delta_text: str,
        function_name_returned: bool,
534
535
        tool_call_idx: int | None = None,
    ) -> tuple[DeltaMessage | None, bool]:
536
537
538
        if current_text is None or current_text == "":
            # if the current text is empty, we cannot parse it
            return None, function_name_returned
539
        try:
540
541
542
543
544
545
            flags = Allow.ALL
            obj, _ = partial_json_loads(current_text, flags)
        except (
            partial_json_parser.core.exceptions.MalformedJSON,
            json.JSONDecodeError,
        ):
546
            logger.debug("not enough tokens to parse into JSON yet")
547
548
549
550
551
552
553
554
555
556
            obj = None

        # check if the current text is a valid array
        # containing a partial tool calling object
        # if not repeat
        if obj is None or not isinstance(obj, list) or not len(obj) > 0:
            function_name_returned = False
            delta_message = None
        else:
            _, finishes_previous_tool = OpenAIServingChat._filter_delta_text(
557
558
                delta_text, previous_text
            )
559
560
561
562
            # take the last tool call from the generated list
            current_tool_call = obj[-1]

            # once parameters have been generated the name is complete as well
563
564
565
            if not finishes_previous_tool and (
                "name" not in current_tool_call or "parameters" not in current_tool_call
            ):
566
567
568
569
570
                function_name_returned = False
                delta_message = None
            else:
                if not function_name_returned:
                    # get partly generated arguments from the latest tool call
571
572
573
                    param_match = re.search(
                        r'.*"parameters":\s*(.*)', current_text, re.DOTALL
                    )
574
575
                    arguments = param_match.group(1) if param_match else ""
                    arguments, _ = OpenAIServingChat._filter_delta_text(
576
577
                        arguments, previous_text
                    )
578
579
580
581

                    # if this iteration finishes a previous tool call but a
                    # new incomplete tool is already generated, take the
                    # previous from the list
582
                    if finishes_previous_tool and "parameters" not in current_tool_call:
583
584
585
                        current_tool_call = obj[-2]

                    function_name_returned = True
586
587
588
                    tool_call_id = make_tool_call_id(
                        id_type=self.tool_call_id_type,
                        func_name=current_tool_call["name"],
589
590
591
592
593
594
595
596
597
598
599
600
601
602
                        idx=tool_call_idx,
                    )
                    delta_message = DeltaMessage(
                        tool_calls=[
                            DeltaToolCall(
                                id=tool_call_id,
                                function=DeltaFunctionCall(
                                    name=current_tool_call["name"], arguments=arguments
                                ),
                                index=len(obj) - 1,
                                type="function",
                            )
                        ]
                    )
603
604
605

                else:
                    delta_text, _ = OpenAIServingChat._filter_delta_text(
606
607
                        delta_text, previous_text
                    )
608
609

                    if delta_text != "":
610
611
612
613
614
615
616
617
618
619
620
621
622
                        delta_message = DeltaMessage(
                            tool_calls=[
                                DeltaToolCall(
                                    function=DeltaFunctionCall(
                                        # OpenAI API returns None
                                        # instead of name every time
                                        name=None,
                                        arguments=delta_text,
                                    ),
                                    index=len(obj) - 1,
                                )
                            ]
                        )
623
624
625
626
627
                    else:
                        delta_message = None

        return delta_message, function_name_returned

628
    async def chat_completion_stream_generator(
629
630
631
632
        self,
        request: ChatCompletionRequest,
        result_generator: AsyncIterator[RequestOutput],
        request_id: str,
633
        model_name: str,
634
        conversation: list[ConversationMessage],
zhuwenwen's avatar
zhuwenwen committed
635
        tokenizer: TokenizerLike | None,
636
        request_metadata: RequestResponseMetadata,
637
    ) -> AsyncGenerator[str, None]:
638
        created_time = int(time.time())
639
        chunk_object_type: Final = "chat.completion.chunk"
640
        first_iteration = True
641
642

        # Send response for each token for each request.n (index)
643
644
645
        num_choices = 1 if request.n is None else request.n
        previous_num_tokens = [0] * num_choices
        finish_reason_sent = [False] * num_choices
646
        num_prompt_tokens = 0
647
        num_cached_tokens = None
648
649
        if self.use_harmony:
            harmony_parsers = [
650
                get_streamable_parser_for_assistant() for _ in range(num_choices)
651
            ]
652
653
            harmony_tools_streamed = [False] * num_choices
        tools_streamed = [False] * num_choices
654
655
656
657
658
659
660
661
662

        if isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam):
            tool_choice_function_name = request.tool_choice.function.name
        else:
            tool_choice_function_name = None

        # Determine whether tools are in use with "auto" tool choice
        tool_choice_auto = (
            not tool_choice_function_name
663
664
            and self._should_stream_with_auto_tool_parsing(request)
        )
665

666
        all_previous_token_ids: list[list[int]] | None
667
        function_name_returned = [False] * num_choices
668
        if self.tool_call_id_type == "kimi_k2":
669
670
671
            history_tool_call_cnt = get_history_tool_calls_cnt(conversation)
        else:
            history_tool_call_cnt = 0
672

673
674
675
        # Always track previous_texts for comprehensive output logging
        previous_texts = [""] * num_choices

676
677
        # Only one of these will be used, thus previous_texts and
        # all_previous_token_ids will not be used twice in the same iteration.
678
        if tool_choice_auto or self.reasoning_parser:
679
680
            # These are only required in "auto" tool choice case
            all_previous_token_ids = [[]] * num_choices
681
682
683
            # For reasoning parser and tool call all enabled
            added_content_delta_arr = [False] * num_choices
            reasoning_end_arr = [False] * num_choices
684
        else:
685
            all_previous_token_ids = None
686

687
        try:
688
            if self.reasoning_parser:
689
690
691
692
693
                if tokenizer is None:
                    raise ValueError(
                        "Tokenizer not available when `skip_tokenizer_init=True`"
                    )

694
695
696
697
698
                # Pass the same chat template kwargs as used in tokenization
                chat_template_kwargs = self._prepare_extra_chat_template_kwargs(
                    request.chat_template_kwargs,
                    self.default_chat_template_kwargs,
                )
699
700
                reasoning_parser = self.reasoning_parser(
                    tokenizer,
zhuwenwen's avatar
zhuwenwen committed
701
                    chat_template_kwargs=chat_template_kwargs,  # type: ignore[call-arg]
702
                )
703
704
705
706
707
708
        except RuntimeError as e:
            logger.exception("Error in reasoning parser creation.")
            data = self.create_streaming_error_response(str(e))
            yield f"data: {data}\n\n"
            yield "data: [DONE]\n\n"
            return
709
710
711
        # Prepare the tool parser if it's needed
        try:
            if tool_choice_auto and self.tool_parser:
712
713
714
715
716
                if tokenizer is None:
                    raise ValueError(
                        "Tokenizer not available when `skip_tokenizer_init=True`"
                    )

717
                tool_parsers: list[ToolParser | None] = [
718
719
720
721
                    self.tool_parser(tokenizer)
                ] * num_choices
            else:
                tool_parsers = [None] * num_choices
722
        except Exception as e:
723
            logger.exception("Error in tool parser creation.")
724
            data = self.create_streaming_error_response(e)
725
726
727
728
            yield f"data: {data}\n\n"
            yield "data: [DONE]\n\n"
            return

729
        stream_options = request.stream_options
730
731
732
        include_usage, include_continuous_usage = should_include_usage(
            stream_options, self.enable_force_include_usage
        )
733

734
735
        try:
            async for res in result_generator:
736
737
                if res.prompt_token_ids is not None:
                    num_prompt_tokens = len(res.prompt_token_ids)
738
739
                    if res.encoder_prompt_token_ids is not None:
                        num_prompt_tokens += len(res.encoder_prompt_token_ids)
740

741
742
743
744
                # We need to do it here, because if there are exceptions in
                # the result_generator, it needs to be sent as the FIRST
                # response (by the try...catch).
                if first_iteration:
745
                    num_cached_tokens = res.num_cached_tokens
746
747
                    # Send first response for each request.n (index) with
                    # the role
748
                    role = self.get_chat_request_role(request)
749
750
751

                    # NOTE num_choices defaults to 1 so this usually executes
                    # once per request
752
                    for i in range(num_choices):
753
754
                        choice_data = ChatCompletionResponseStreamChoice(
                            index=i,
755
756
757
758
                            delta=DeltaMessage(
                                role=role,
                                content="",
                            ),
759
                            logprobs=None,
760
761
                            finish_reason=None,
                        )
762
763

                        # return prompt_token_ids at the first chunk ever
764
765
766
767
768
                        chunk = ChatCompletionStreamResponse(
                            id=request_id,
                            object=chunk_object_type,
                            created=created_time,
                            choices=[choice_data],
769
                            model=model_name,
770
771
772
773
774
775
                            prompt_token_ids=(
                                res.prompt_token_ids
                                if request.return_token_ids
                                else None
                            ),
                        )
776

777
778
779
780
781
                        # if continuous usage stats are requested, add it
                        if include_continuous_usage:
                            chunk.usage = UsageInfo(
                                prompt_tokens=num_prompt_tokens,
                                completion_tokens=0,
782
783
                                total_tokens=num_prompt_tokens,
                            )
784

785
786
787
                        data = chunk.model_dump_json(exclude_unset=True)
                        yield f"data: {data}\n\n"

788
789
                    # Send response to echo the input portion of the
                    # last message
790
                    if request.echo:
791
                        last_msg_content: str | list[dict[str, str]] = ""
792
793
794
795
796
                        if (
                            conversation
                            and "content" in conversation[-1]
                            and conversation[-1].get("role") == role
                        ):
797
                            last_msg_content = conversation[-1]["content"] or ""
798
799

                        if last_msg_content:
800
                            for i in range(num_choices):
801
802
803
804
805
806
                                choice_data = ChatCompletionResponseStreamChoice(
                                    index=i,
                                    delta=DeltaMessage(content=last_msg_content),
                                    logprobs=None,
                                    finish_reason=None,
                                )
807
808
809
810
811
                                chunk = ChatCompletionStreamResponse(
                                    id=request_id,
                                    object=chunk_object_type,
                                    created=created_time,
                                    choices=[choice_data],
812
813
                                    model=model_name,
                                )
814
815
816
817
                                if include_continuous_usage:
                                    chunk.usage = UsageInfo(
                                        prompt_tokens=num_prompt_tokens,
                                        completion_tokens=0,
818
819
                                        total_tokens=num_prompt_tokens,
                                    )
820

821
                                data = chunk.model_dump_json(exclude_unset=True)
822
823
824
                                yield f"data: {data}\n\n"
                    first_iteration = False

825
826
827
828
829
830
831
832
833
                prompt_is_reasoning_end_arr = [False] * num_choices
                if self.reasoning_parser and res.prompt_token_ids:
                    prompt_is_reasoning_end = reasoning_parser.is_reasoning_end(
                        res.prompt_token_ids
                    )
                    prompt_is_reasoning_end_arr = [
                        prompt_is_reasoning_end
                    ] * num_choices

834
835
                for output in res.outputs:
                    i = output.index
836
                    tool_parser = tool_parsers[i]
837
838
839
840

                    if finish_reason_sent[i]:
                        continue

841
                    if request.logprobs and request.top_logprobs is not None:
842
                        assert output.logprobs is not None, "Did not output logprobs"
843
                        logprobs = self._create_chat_logprobs(
844
845
                            token_ids=output.token_ids,
                            top_logprobs=output.logprobs,
846
                            tokenizer=tokenizer,
847
                            num_output_top_logprobs=request.top_logprobs,
848
                            return_as_token_id=request.return_tokens_as_token_ids,
849
850
851
852
                        )
                    else:
                        logprobs = None

853
854
                    if self.use_harmony:
                        harmony_parser = harmony_parsers[i]
855
                        prev_recipient = harmony_parser.current_recipient
856
857
858

                        # Track accumulated content per token with their state
                        token_states: list[TokenState] = []
859
860
                        for token_id in output.token_ids:
                            harmony_parser.process(token_id)
861
862
863
864
865
866
867
868
869
                            token_delta = harmony_parser.last_content_delta or ""
                            token_states.append(
                                TokenState(
                                    harmony_parser.current_channel,
                                    harmony_parser.current_recipient,
                                    token_delta,
                                )
                            )
                        delta_text = "".join(delta for _, _, delta in token_states)
870
                        cur_channel = harmony_parser.current_channel
871

872
873
874
875
876
                        # handle the case where several tokens where generated at once
                        # including the final token, leading to a delta in the text
                        # but the current channel to be empty (start state)
                        if not cur_channel and delta_text:
                            cur_channel = "final"
877
878
                    else:
                        delta_text = output.text
879

880
881
882
883
884
                    if (
                        not delta_text
                        and not output.token_ids
                        and not previous_num_tokens[i]
                    ):
885
886
887
                        # Chunked prefill case, don't return empty chunks
                        continue

888
                    delta_message: DeltaMessage | None
889

890
                    # just update previous_texts and previous_token_ids
891
                    if tool_choice_auto or self.reasoning_parser:
892
893
894
895
896
                        assert previous_texts is not None
                        assert all_previous_token_ids is not None
                        previous_text = previous_texts[i]
                        previous_token_ids = all_previous_token_ids[i]
                        current_text = previous_text + delta_text
897
898
                        # avoid the None + list error.
                        if previous_token_ids:
899
                            current_token_ids = previous_token_ids + as_list(
900
901
                                output.token_ids
                            )
902
                        else:
903
                            current_token_ids = as_list(output.token_ids)
904

905
                    if self.use_harmony:
906
907
908
                        delta_message, tools_streamed_flag = (
                            extract_harmony_streaming_delta(
                                harmony_parser=harmony_parser,
909
                                token_states=token_states,
910
911
912
913
914
                                prev_recipient=prev_recipient,
                                include_reasoning=request.include_reasoning,
                            )
                        )
                        harmony_tools_streamed[i] |= tools_streamed_flag
915
                    # handle streaming deltas for tools with named tool_choice
916
                    elif tool_choice_function_name:
917
918
919
920
921
922
923
                        if (
                            reasoning_parser
                            and not reasoning_end_arr[i]
                            and prompt_is_reasoning_end_arr[i]
                        ):
                            reasoning_end_arr[i] = True

924
925
926
927
928
929
930
                        if (
                            self.reasoning_parser
                            and not reasoning_end_arr[i]
                            and not reasoning_parser.is_reasoning_end(
                                previous_token_ids
                            )
                        ):
931
932
                            assert reasoning_parser is not None
                            delta_message = (
933
                                reasoning_parser.extract_reasoning_streaming(
934
935
936
937
938
939
                                    previous_text,
                                    current_text,
                                    delta_text,
                                    previous_token_ids,
                                    current_token_ids,
                                    output.token_ids,
940
941
                                )
                            )
942
                            # When encountering think end id in delta_token_ids,
943
                            # set reasoning status to end.
944
                            # Only keep 'content', remove 'reasoning'.
945
                            if reasoning_parser.is_reasoning_end(
946
947
                                as_list(output.token_ids)
                            ):
948
                                reasoning_end_arr[i] = True
949
950
951
952
953
954
955
956
                                if delta_message and delta_message.content:
                                    # This need to be added to next `delta_text`
                                    current_text = delta_message.content
                                    delta_message.content = None
                                else:
                                    current_text = ""
                        else:
                            # Just to add remaining `content`
957
                            if self.reasoning_parser:
958
959
960
                                delta_text = previous_text + delta_text
                                current_text = ""

961
962
                            if function_name_returned[i]:
                                delta_tool_call = DeltaToolCall(
963
964
965
                                    function=DeltaFunctionCall(arguments=delta_text),
                                    index=i,
                                )
966
967
                            else:
                                delta_tool_call = DeltaToolCall(
zhuwenwen's avatar
zhuwenwen committed
968
                                    id=make_tool_call_id(),
969
970
971
                                    type="function",
                                    function=DeltaFunctionCall(
                                        name=tool_choice_function_name,
972
973
974
975
                                        arguments=delta_text,
                                    ),
                                    index=i,
                                )
976
977
                                function_name_returned[i] = True

978
979
980
981
982
                            delta_message = DeltaMessage(
                                tool_calls=[
                                    delta_tool_call,
                                ]
                            )
983
                            tools_streamed[i] = True
984

985
986
987
988
989
                    elif request.tool_choice == "required":
                        assert previous_texts is not None
                        previous_text = previous_texts[i]
                        current_text = previous_text + delta_text
                        fn_name_returned = function_name_returned[i]
990
991
992
993
994
995
996
997
998
                        output_token_ids = as_list(output.token_ids)

                        if (
                            self.reasoning_parser is not None
                            and not reasoning_end_arr[i]
                            and res.prompt_token_ids
                            and reasoning_parser.is_reasoning_end(res.prompt_token_ids)
                        ):
                            reasoning_end_arr[i] = True
999

1000
1001
                        if self.reasoning_parser and not reasoning_end_arr[i]:
                            delta_message = (
1002
                                reasoning_parser.extract_reasoning_streaming(
1003
1004
1005
1006
1007
1008
1009
                                    previous_text,
                                    current_text,
                                    delta_text,
                                    previous_token_ids,
                                    current_token_ids,
                                    output_token_ids,
                                )
1010
                            )
1011
1012
1013
1014
1015
1016
1017
1018
1019
                            if reasoning_parser.is_reasoning_end(output_token_ids):
                                reasoning_end_arr[i] = True
                                if delta_message and delta_message.content:
                                    current_text = delta_message.content
                                    delta_message.content = None
                                else:
                                    # reasoning ended
                                    current_text = ""

1020
                        else:
1021
                            # either finished reasoning or no reasoning at all
1022
                            content = current_text
1023
1024
1025
1026
1027
1028
1029
1030
1031

                            delta_message, function_name_returned[i] = (
                                self.extract_tool_call_required_streaming(
                                    previous_text=previous_text,
                                    current_text=content,
                                    delta_text=delta_text,
                                    function_name_returned=fn_name_returned,
                                    tool_call_idx=history_tool_call_cnt,
                                )
1032
                            )
1033
1034
1035
1036
1037
1038
1039
                            if (
                                delta_message
                                and delta_message.tool_calls
                                and delta_message.tool_calls[0].id is not None
                            ):
                                history_tool_call_cnt += 1
                                tools_streamed[i] = True
1040

1041
1042
                    # handle streaming deltas for tools with "auto" tool choice
                    # and reasoning parser
1043
                    elif tool_choice_auto and self.reasoning_parser:
1044
1045
1046
1047
                        assert tool_parser is not None
                        assert reasoning_parser is not None
                        assert added_content_delta_arr is not None
                        assert reasoning_end_arr is not None
1048
                        output_token_ids = as_list(output.token_ids)
1049
                        if not reasoning_end_arr[i]:
1050
1051
1052
                            # When encountering think end id in prompt_token_ids
                            # i.e {"enable_thinking": False},
                            # set reasoning status to end.
1053
1054
1055
1056
1057
1058
                            if (
                                res.prompt_token_ids
                                and reasoning_parser.is_reasoning_end(
                                    res.prompt_token_ids
                                )
                            ):
1059
                                reasoning_end_arr[i] = True
1060
                                current_token_ids = output_token_ids
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
                                # Don't update current_text, keep it as is from delta
                            else:
                                delta_message = (
                                    reasoning_parser.extract_reasoning_streaming(
                                        previous_text,
                                        current_text,
                                        delta_text,
                                        previous_token_ids,
                                        current_token_ids,
                                        output_token_ids,
1071
1072
                                    )
                                )
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089

                                # When encountering think end id in delta_token_ids,
                                # set reasoning status to end.
                                # Remove the text and token ids related
                                # to 'reasoning'.
                                if reasoning_parser.is_reasoning_end(output_token_ids):
                                    reasoning_end_arr[i] = True
                                    current_token_ids = (
                                        reasoning_parser.extract_content_ids(
                                            output_token_ids
                                        )
                                    )
                                    if delta_message and delta_message.content:
                                        current_text = delta_message.content
                                        delta_message.content = None
                                    else:
                                        current_text = ""
1090
1091

                        # handle tool calls only after reasoning is done,
1092
                        if reasoning_end_arr[i]:
1093
                            delta_token_ids = output_token_ids
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
                            # First time to tool call,
                            # add the remaining text and token ids
                            # to delta from previous
                            if not added_content_delta_arr[i]:
                                added_content_delta_arr[i] = True
                                previous_text = ""
                                previous_token_ids = []
                                delta_text = current_text
                                delta_token_ids = current_token_ids

1104
                            delta_message = tool_parser.extract_tool_calls_streaming(
1105
1106
                                previous_text=previous_text,
                                current_text=current_text,
1107
                                delta_text=delta_text,
1108
1109
                                previous_token_ids=previous_token_ids,
                                current_token_ids=current_token_ids,
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
                                delta_token_ids=delta_token_ids,
                                request=request,
                            )
                            if delta_message and delta_message.tool_calls:
                                tools_streamed[i] = True
                    # when only tool calls
                    elif tool_choice_auto:
                        assert tool_parser is not None
                        delta_message = tool_parser.extract_tool_calls_streaming(
                            previous_text=previous_text,
                            current_text=current_text,
                            delta_text=delta_text,
                            previous_token_ids=previous_token_ids,
                            current_token_ids=current_token_ids,
                            delta_token_ids=output.token_ids,
                            request=request,
                        )
1127
1128
                        if delta_message and delta_message.tool_calls:
                            tools_streamed[i] = True
1129

1130
                    # when only reasoning
1131
                    elif self.reasoning_parser:
1132
1133
1134
                        if prompt_is_reasoning_end_arr[i]:
                            delta_message = DeltaMessage(content=delta_text)
                        else:
1135
1136
1137
1138
1139
1140
1141
                            delta_message = reasoning_parser.extract_reasoning_streaming(
                                previous_text,
                                current_text,
                                delta_text,
                                previous_token_ids,
                                current_token_ids,
                                output.token_ids,
1142
                            )
1143
                    # handle streaming just a content delta
1144
1145
1146
                    else:
                        delta_message = DeltaMessage(content=delta_text)

1147
                    # update the previous values for the next iteration
1148
1149
1150
                    if (
                        tool_choice_auto or self.reasoning_parser
                    ) and not self.use_harmony:
1151
1152
1153
1154
                        assert previous_texts is not None
                        assert all_previous_token_ids is not None
                        previous_texts[i] = current_text
                        all_previous_token_ids[i] = current_token_ids
1155
1156
1157
1158
                    else:
                        # Update for comprehensive logging even in simple case
                        assert previous_texts is not None
                        previous_texts[i] += delta_text
1159

1160
                    # set the previous values for the next iteration
1161
                    previous_num_tokens[i] += len(output.token_ids)
1162
1163
1164
1165
1166
1167

                    # if the message delta is None (e.g. because it was a
                    # "control token" for tool calls or the parser otherwise
                    # wasn't ready to send a token, then
                    #   get the next token without streaming a chunk
                    if delta_message is None:
1168
1169
1170
1171
1172
1173
1174
                        # NOTE: If return_token_ids is enabled, we still need to
                        # send a chunk with token_ids even if delta_message is None
                        # to ensure all tokens are included in the response
                        if (
                            output.finish_reason is None
                            and not request.return_token_ids
                        ):
1175
                            continue
1176
                        delta_message = DeltaMessage()
1177

1178
1179
                    # Log streaming delta if output logging is enabled
                    if self.enable_log_outputs and self.request_logger:
1180
                        delta_content_parts = []
1181
                        if delta_message.content:
1182
1183
1184
1185
1186
1187
                            delta_content_parts.append(delta_message.content)
                        if delta_message.reasoning_content:
                            reasoning = delta_message.reasoning_content
                            delta_content_parts.append(f"[reasoning: {reasoning}]")
                        if delta_message.tool_calls:
                            tool_args = "".join(
1188
1189
                                tc.function.arguments
                                for tc in delta_message.tool_calls
1190
1191
                                if tc.function and tc.function.arguments
                            )
1192
1193
                            if tool_args:
                                delta_content_parts.append(f"[tool_calls: {tool_args}]")
1194

1195
1196
                        if delta_content_parts and self.enable_log_deltas:
                            delta_content = " ".join(delta_content_parts)
1197
1198
1199
                            self.request_logger.log_outputs(
                                request_id=request_id,
                                outputs=delta_content,
1200
                                output_token_ids=as_list(output.token_ids),
1201
1202
1203
1204
1205
                                finish_reason=output.finish_reason,
                                is_streaming=True,
                                delta=True,
                            )

1206
1207
1208
1209
                    if output.finish_reason is None:
                        # Send token-by-token response for each request.n
                        choice_data = ChatCompletionResponseStreamChoice(
                            index=i,
1210
                            delta=delta_message,
1211
                            logprobs=logprobs,
1212
                            finish_reason=None,
1213
1214
1215
1216
1217
1218
                            token_ids=(
                                as_list(output.token_ids)
                                if request.return_token_ids
                                else None
                            ),
                        )
1219
1220

                    # if the model is finished generating
1221
                    else:
1222
1223
1224
1225
                        # check for error finish reason and abort streaming
                        # finish_reason='error' indicates a retryable error
                        self._raise_if_error(output.finish_reason, request_id)

1226
1227
1228
                        # check to make sure we haven't "forgotten" to stream
                        #   any tokens that were generated but previously
                        #   matched by partial json parsing
1229
                        # only happens if we are NOT using structured outputs
1230
                        auto_tools_called = False
1231
                        if tool_parser:
1232
1233
1234
1235
1236
1237
                            auto_tools_called = len(tool_parser.prev_tool_call_arr) > 0
                            index = (
                                len(tool_parser.prev_tool_call_arr) - 1
                                if auto_tools_called
                                else 0
                            )
1238
1239
1240
                        else:
                            index = 0

1241
                        if (
1242
1243
                            self._should_check_for_unstreamed_tool_arg_tokens(
                                delta_message, output
1244
                            )
1245
                            and tool_parser
1246
                        ):
1247
                            latest_delta_len = 0
1248
1249
                            if (
                                isinstance(
1250
                                    delta_message.tool_calls[0].function,
1251
1252
1253
1254
1255
                                    DeltaFunctionCall,
                                )
                            ) and isinstance(
                                delta_message.tool_calls[0].function.arguments, str
                            ):
1256
                                latest_delta_len = len(
1257
1258
                                    delta_message.tool_calls[0].function.arguments
                                )
1259

1260
                            # get the expected call based on partial JSON
1261
1262
1263
1264
1265
1266
                            # parsing which "autocompletes" the JSON
                            expected_call = json.dumps(
                                tool_parser.prev_tool_call_arr[index].get(
                                    "arguments", {}
                                ),
                                ensure_ascii=False,
1267
                            )
1268

1269
                            # get what we've streamed so far for arguments
1270
                            # for the current tool
1271
1272
                            actual_call = tool_parser.streamed_args_for_tool[index]
                            if latest_delta_len > 0:
1273
                                actual_call = actual_call[:-latest_delta_len]
1274
1275

                            # check to see if there's anything left to stream
1276
                            remaining_call = expected_call.replace(actual_call, "", 1)
1277
                            # set that as a delta message
1278
1279
                            delta_message = self._create_remaining_args_delta(
                                delta_message, remaining_call, index
1280
                            )
1281

1282
                        # Send the finish response for each request.n only once
1283
1284
1285
1286
                        # In OpenAI's API, when a tool is called, the
                        # finish_reason is:
                        # "tool_calls" for "auto" or "required" tool calls,
                        # and "stop" for named tool calls.
1287
1288
                        if (
                            auto_tools_called
1289
                            or (tools_streamed[i] and not tool_choice_function_name)
1290
1291
                            or (self.use_harmony and harmony_tools_streamed[i])
                        ):
1292
1293
                            finish_reason_ = "tool_calls"
                        else:
1294
1295
1296
                            finish_reason_ = (
                                output.finish_reason if output.finish_reason else "stop"
                            )
1297
1298
                        choice_data = ChatCompletionResponseStreamChoice(
                            index=i,
1299
                            delta=delta_message,
1300
                            logprobs=logprobs,
1301
                            finish_reason=finish_reason_,
1302
                            stop_reason=output.stop_reason,
1303
1304
1305
1306
1307
1308
                            token_ids=(
                                as_list(output.token_ids)
                                if request.return_token_ids
                                else None
                            ),
                        )
1309

1310
                        finish_reason_sent[i] = True
1311

1312
                    choice_data = maybe_filter_parallel_tool_calls(choice_data, request)
1313
1314
1315
1316
1317
                    chunk = ChatCompletionStreamResponse(
                        id=request_id,
                        object=chunk_object_type,
                        created=created_time,
                        choices=[choice_data],
1318
1319
                        model=model_name,
                    )
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329

                    # handle usage stats if requested & if continuous
                    if include_continuous_usage:
                        completion_tokens = previous_num_tokens[i]
                        chunk.usage = UsageInfo(
                            prompt_tokens=num_prompt_tokens,
                            completion_tokens=completion_tokens,
                            total_tokens=num_prompt_tokens + completion_tokens,
                        )

1330
                    data = chunk.model_dump_json(exclude_unset=True)
1331
1332
                    yield f"data: {data}\n\n"

1333
1334
            # once the final token is handled, if stream_options.include_usage
            # is sent, send the usage
1335
1336
            if include_usage:
                completion_tokens = sum(previous_num_tokens)
1337
1338
1339
1340
1341
                final_usage = UsageInfo(
                    prompt_tokens=num_prompt_tokens,
                    completion_tokens=completion_tokens,
                    total_tokens=num_prompt_tokens + completion_tokens,
                )
1342
1343
                if self.enable_prompt_tokens_details and num_cached_tokens:
                    final_usage.prompt_tokens_details = PromptTokenUsageInfo(
1344
1345
                        cached_tokens=num_cached_tokens
                    )
1346
1347
1348
1349
1350
1351
1352

                final_usage_chunk = ChatCompletionStreamResponse(
                    id=request_id,
                    object=chunk_object_type,
                    created=created_time,
                    choices=[],
                    model=model_name,
1353
1354
1355
1356
1357
                    usage=final_usage,
                )
                final_usage_data = final_usage_chunk.model_dump_json(
                    exclude_unset=True, exclude_none=True
                )
1358
                yield f"data: {final_usage_data}\n\n"
1359

1360
1361
1362
1363
1364
            # report to FastAPI middleware aggregate usage across all choices
            num_completion_tokens = sum(previous_num_tokens)
            request_metadata.final_usage_info = UsageInfo(
                prompt_tokens=num_prompt_tokens,
                completion_tokens=num_completion_tokens,
1365
1366
1367
1368
1369
1370
1371
1372
1373
                total_tokens=num_prompt_tokens + num_completion_tokens,
            )

            # Log complete streaming response if output logging is enabled
            if self.enable_log_outputs and self.request_logger:
                # Log the complete response for each choice
                for i in range(num_choices):
                    full_text = (
                        previous_texts[i]
1374
1375
                        if previous_texts and i < len(previous_texts)
                        else f"<streaming_complete: {previous_num_tokens[i]} tokens>"
1376
1377
1378
1379
                    )
                    self.request_logger.log_outputs(
                        request_id=request_id,
                        outputs=full_text,
1380
                        output_token_ids=None,  # Consider also logging all token IDs
1381
1382
1383
1384
                        finish_reason="streaming_complete",
                        is_streaming=True,
                        delta=False,
                    )
1385

1386
1387
        except GenerationError as e:
            yield f"data: {self._convert_generation_error_to_streaming_response(e)}\n\n"
1388
        except Exception as e:
1389
            logger.exception("Error in chat completion stream generator.")
1390
            data = self.create_streaming_error_response(e)
1391
            yield f"data: {data}\n\n"
1392
1393
1394
1395
        # Send the final done message after all response.n are finished
        yield "data: [DONE]\n\n"

    async def chat_completion_full_generator(
1396
1397
1398
1399
        self,
        request: ChatCompletionRequest,
        result_generator: AsyncIterator[RequestOutput],
        request_id: str,
1400
        model_name: str,
1401
        conversation: list[ConversationMessage],
zhuwenwen's avatar
zhuwenwen committed
1402
        tokenizer: TokenizerLike | None,
1403
        request_metadata: RequestResponseMetadata,
1404
    ) -> ErrorResponse | ChatCompletionResponse:
1405
        created_time = int(time.time())
1406
        final_res: RequestOutput | None = None
1407

1408
1409
1410
1411
1412
        try:
            async for res in result_generator:
                final_res = res
        except asyncio.CancelledError:
            return self.create_error_response("Client disconnected")
1413
        except ValueError as e:
1414
            return self.create_error_response(e)
1415

1416
1417
        assert final_res is not None

1418
        choices: list[ChatCompletionResponseChoice] = []
1419
        if self.tool_call_id_type == "kimi_k2":
1420
1421
1422
            history_tool_call_cnt = get_history_tool_calls_cnt(conversation)
        else:
            history_tool_call_cnt = 0
1423

1424
1425
        role = self.get_chat_request_role(request)
        for output in final_res.outputs:
1426
1427
1428
            # check for error finish reason and raise GenerationError
            # finish_reason='error' indicates a retryable request-level internal error
            self._raise_if_error(output.finish_reason, request_id)
1429
            token_ids = output.token_ids
1430
            out_logprobs = output.logprobs
1431
            tool_call_info = None
1432

1433
1434
            if request.logprobs and request.top_logprobs is not None:
                assert out_logprobs is not None, "Did not output logprobs"
1435
                logprobs = self._create_chat_logprobs(
1436
                    token_ids=token_ids,
1437
                    top_logprobs=out_logprobs,
1438
                    num_output_top_logprobs=request.top_logprobs,
1439
                    tokenizer=tokenizer,
1440
                    return_as_token_id=request.return_tokens_as_token_ids,
1441
1442
1443
                )
            else:
                logprobs = None
1444
1445

            if self.use_harmony:
1446
                reasoning, content, _ = parse_chat_output(token_ids)
1447
                if not request.include_reasoning:
1448
                    reasoning = None
1449

1450
                if self.tool_parser is not None:
1451
1452
1453
1454
1455
                    if tokenizer is None:
                        raise ValueError(
                            "Tokenizer not available when `skip_tokenizer_init=True`"
                        )

1456
1457
1458
1459
1460
1461
1462
                    tool_parser = self.tool_parser(tokenizer)
                    # NOTE: We use token_ids for openai tool parser
                    tool_call_info = tool_parser.extract_tool_calls(
                        "",
                        request=request,
                        token_ids=token_ids,  # type: ignore
                    )
1463
                    content = tool_call_info.content
1464
1465
                    message = ChatMessage(
                        role=role,
1466
                        reasoning=reasoning,
1467
1468
1469
1470
1471
1472
                        content=content,
                        tool_calls=tool_call_info.tool_calls,
                    )
                else:
                    message = ChatMessage(
                        role=role,
1473
                        reasoning=reasoning,
1474
1475
                        content=content,
                    )
1476
1477
1478
1479
1480

                choice_data = ChatCompletionResponseChoice(
                    index=output.index,
                    message=message,
                    logprobs=logprobs,
1481
1482
1483
1484
1485
1486
1487
                    finish_reason=(
                        "tool_calls"
                        if (tool_call_info is not None and tool_call_info.tools_called)
                        else output.finish_reason
                        if output.finish_reason
                        else "stop"
                    ),
1488
                    stop_reason=output.stop_reason,
1489
1490
1491
                    token_ids=(
                        as_list(output.token_ids) if request.return_token_ids else None
                    ),
1492
1493
1494
                )
                choices.append(choice_data)
                continue
1495

1496
            if self.reasoning_parser:
1497
                try:
1498
1499
1500
1501
1502
                    if tokenizer is None:
                        raise ValueError(
                            "Tokenizer not available when `skip_tokenizer_init=True`"
                        )

1503
1504
1505
1506
1507
                    # Pass the same chat template kwargs as used in tokenization
                    chat_template_kwargs = self._prepare_extra_chat_template_kwargs(
                        request.chat_template_kwargs,
                        self.default_chat_template_kwargs,
                    )
1508
1509
                    reasoning_parser = self.reasoning_parser(
                        tokenizer,
1510
                        chat_template_kwargs=chat_template_kwargs,  # type: ignore[call-arg]
1511
                    )
1512
1513
1514
                except RuntimeError as e:
                    logger.exception("Error in reasoning parser creation.")
                    return self.create_error_response(str(e))
1515
1516
                # If the reasoning parser is enabled,
                # tool calls are extracted exclusively from the content.
1517
                reasoning, content = reasoning_parser.extract_reasoning(
1518
1519
                    output.text, request=request
                )
1520
                if not request.include_reasoning:
1521
                    reasoning = None
1522
            else:
1523
                reasoning = None
1524
                content = output.text
1525

1526
            auto_tools_called = False
1527
1528
            # if auto tools are not enabled, and a named tool choice using
            #   outlines is not being used
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
            tool_calls, content = self._parse_tool_calls_from_content(
                request=request,
                tokenizer=tokenizer,
                content=content,
                enable_auto_tools=self.enable_auto_tools,
                tool_parser_cls=self.tool_parser,
            )
            tool_call_class = (
                MistralToolCall if isinstance(tokenizer, MistralTokenizer) else ToolCall
            )
zhuwenwen's avatar
zhuwenwen committed
1539
            if (not self.enable_auto_tools or not self.tool_parser) and (
1540
1541
1542
                not isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam)
                and request.tool_choice != "required"
            ):
1543
                message = ChatMessage(role=role, reasoning=reasoning, content=content)
1544

zhuwenwen's avatar
zhuwenwen committed
1545
            # if the request uses tools and specified a tool choice
1546
1547
1548
1549
            elif (
                request.tool_choice
                and type(request.tool_choice) is ChatCompletionNamedToolChoiceParam
            ):
1550
                assert tool_calls is not None and len(tool_calls) > 0
1551
1552
                message = ChatMessage(
                    role=role,
1553
                    reasoning=reasoning,
1554
                    content="",
zhuwenwen's avatar
zhuwenwen committed
1555
                    tool_calls=[tool_call_class(function=tc) for tc in tool_calls],
1556
                )
1557

1558
            elif request.tool_choice and request.tool_choice == "required":
1559
1560
                tool_call_class_items = []
                assert tool_calls is not None and len(tool_calls) > 0
zhuwenwen's avatar
zhuwenwen committed
1561
1562
1563
1564
                for tool_call in tool_calls:
                    tool_call_class_items.append(
                        tool_call_class(
                            id=make_tool_call_id(
1565
1566
                                id_type=self.tool_call_id_type,
                                func_name=tool_call.name,
zhuwenwen's avatar
zhuwenwen committed
1567
1568
1569
1570
1571
                                idx=history_tool_call_cnt,
                            ),
                            function=tool_call,
                        )
                    )
1572
                    history_tool_call_cnt += 1
1573
1574
1575
                message = ChatMessage(
                    role=role,
                    content="",
1576
                    tool_calls=tool_call_class_items,
1577
                    reasoning=reasoning,
1578
                )
1579

1580
1581
            # if the request doesn't use tool choice
            # OR specifies to not use a tool
1582
            elif not request.tool_choice or request.tool_choice == "none":
1583
                message = ChatMessage(role=role, reasoning=reasoning, content=content)
1584
1585

            # handle when there are tools and tool choice is auto
1586
1587
1588
1589
1590
1591
            elif (
                request.tools
                and (request.tool_choice == "auto" or request.tool_choice is None)
                and self.enable_auto_tools
                and self.tool_parser
            ):
1592
1593
1594
                # In the OpenAI API the finish_reason is "tools_called"
                # if the tool choice is auto and the model produced a tool
                # call. The same is not true for named function calls
1595
1596
                auto_tools_called = tool_calls is not None and len(tool_calls) > 0
                if tool_calls:
1597
1598
                    message = ChatMessage(
                        role=role,
1599
                        reasoning=reasoning,
1600
                        content=content,
zhuwenwen's avatar
zhuwenwen committed
1601
1602
1603
1604
1605
1606
1607
                        tool_calls=[
                            ToolCall(
                                function=tc,
                                type="function",
                            )
                            for tc in tool_calls
                        ],
1608
                    )
1609
1610
1611
1612

                else:
                    # FOR NOW make it a chat message; we will have to detect
                    # the type to make it later.
1613
1614
1615
1616
                    ret_content = content

                    # try to use content return from tool parser first,
                    # tool parser may do some modify for the content.
1617
1618
                    if content and len(content) > 0:
                        ret_content = content
1619
1620
                    message = ChatMessage(
                        role=role,
1621
                        reasoning=reasoning,
1622
1623
                        content=ret_content,
                    )
1624
1625
1626
1627
1628
1629

            # undetermined case that is still important to handle
            else:
                logger.error(
                    "Error in chat_completion_full_generator - cannot determine"
                    " if tools should be extracted. Returning a standard chat "
1630
1631
                    "completion."
                )
1632
                message = ChatMessage(role=role, reasoning=reasoning, content=content)
1633
1634
1635
1636
1637
1638
1639
1640
            # In OpenAI's API, when a tool is called, the finish_reason is:
            # "tool_calls" for "auto" or "required" tool calls,
            # and "stop" for named tool calls.
            is_finish_reason_tool_calls = auto_tools_called or (
                request.tool_choice
                and request.tool_choice == "required"
                and output.finish_reason == "stop"
            )
1641

1642
1643
            choice_data = ChatCompletionResponseChoice(
                index=output.index,
1644
                message=message,
1645
                logprobs=logprobs,
1646
1647
1648
1649
1650
                finish_reason="tool_calls"
                if is_finish_reason_tool_calls
                else output.finish_reason
                if output.finish_reason
                else "stop",
1651
                stop_reason=output.stop_reason,
1652
1653
1654
                token_ids=(
                    as_list(output.token_ids) if request.return_token_ids else None
                ),
1655
            )
1656
            choice_data = maybe_filter_parallel_tool_calls(choice_data, request)
1657

1658
1659
            choices.append(choice_data)

1660
        if request.echo:
1661
            last_msg_content: str | list[dict[str, str]] = ""
1662
1663
1664
1665
1666
            if (
                conversation
                and "content" in conversation[-1]
                and conversation[-1].get("role") == role
            ):
1667
                last_msg_content = conversation[-1]["content"] or ""
1668
            if isinstance(last_msg_content, list):
1669
                last_msg_content = "\n".join(msg["text"] for msg in last_msg_content)
1670
1671

            for choice in choices:
1672
                full_message = last_msg_content + (choice.message.content or "")
1673
1674
                choice.message.content = full_message

1675
        assert final_res.prompt_token_ids is not None
1676
        num_prompt_tokens = len(final_res.prompt_token_ids)
1677
1678
        if final_res.encoder_prompt_token_ids is not None:
            num_prompt_tokens += len(final_res.encoder_prompt_token_ids)
1679
        num_generated_tokens = sum(
1680
1681
1682
1683
1684
1685
1686
            len(output.token_ids) for output in final_res.outputs
        )
        usage = UsageInfo(
            prompt_tokens=num_prompt_tokens,
            completion_tokens=num_generated_tokens,
            total_tokens=num_prompt_tokens + num_generated_tokens,
        )
1687
1688
        if self.enable_prompt_tokens_details and final_res.num_cached_tokens:
            usage.prompt_tokens_details = PromptTokenUsageInfo(
1689
1690
                cached_tokens=final_res.num_cached_tokens
            )
1691
1692
1693

        request_metadata.final_usage_info = usage

1694
1695
1696
1697
1698
1699
        response = ChatCompletionResponse(
            id=request_id,
            created=created_time,
            model=model_name,
            choices=choices,
            usage=usage,
1700
            prompt_logprobs=clamp_prompt_logprobs(final_res.prompt_logprobs),
1701
1702
1703
            prompt_token_ids=(
                final_res.prompt_token_ids if request.return_token_ids else None
            ),
Robert Shaw's avatar
Robert Shaw committed
1704
            kv_transfer_params=final_res.kv_transfer_params,
1705
1706
        )

1707
1708
1709
1710
1711
1712
1713
1714
1715
        # Log complete response if output logging is enabled
        if self.enable_log_outputs and self.request_logger:
            for choice in choices:
                output_text = ""
                if choice.message.content:
                    output_text = choice.message.content
                elif choice.message.tool_calls:
                    # For tool calls, log the function name and arguments
                    tool_call_descriptions = []
zhuwenwen's avatar
zhuwenwen committed
1716
1717
1718
1719
1720
1721
1722
                    for tc in choice.message.tool_calls:
                        if hasattr(tc.function, "name") and hasattr(
                            tc.function, "arguments"
                        ):
                            tool_call_descriptions.append(
                                f"{tc.function.name}({tc.function.arguments})"
                            )
1723
1724
1725
1726
1727
1728
1729
                    tool_calls_str = ", ".join(tool_call_descriptions)
                    output_text = f"[tool_calls: {tool_calls_str}]"

                if output_text:
                    # Get the corresponding output token IDs
                    output_token_ids = None
                    if choice.index < len(final_res.outputs):
1730
                        output_token_ids = final_res.outputs[choice.index].token_ids
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740

                    self.request_logger.log_outputs(
                        request_id=request_id,
                        outputs=output_text,
                        output_token_ids=output_token_ids,
                        finish_reason=choice.finish_reason,
                        is_streaming=False,
                        delta=False,
                    )

1741
        return response
1742
1743

    def _get_top_logprobs(
1744
1745
        self,
        logprobs: dict[int, Logprob],
1746
        top_logprobs: int | None,
1747
        tokenizer: TokenizerLike | None,
1748
1749
        should_return_as_token_id: bool,
    ) -> list[ChatCompletionLogProb]:
1750
        return [
1751
            ChatCompletionLogProb(
1752
1753
1754
1755
1756
1757
1758
1759
                token=(
                    token := self._get_decoded_token(
                        p[1],
                        p[0],
                        tokenizer,
                        return_as_token_id=should_return_as_token_id,
                    )
                ),
1760
1761
                logprob=max(p[1].logprob, -9999.0),
                bytes=list(token.encode("utf-8", errors="replace")),
1762
1763
            )
            for i, p in enumerate(logprobs.items())
1764
            if (top_logprobs and i < top_logprobs or top_logprobs == -1)
1765
1766
1767
1768
1769
        ]

    def _create_chat_logprobs(
        self,
        token_ids: GenericSequence[int],
1770
        top_logprobs: GenericSequence[dict[int, Logprob] | None],
1771
        tokenizer: TokenizerLike | None,
1772
1773
        num_output_top_logprobs: int | None = None,
        return_as_token_id: bool | None = None,
1774
1775
    ) -> ChatCompletionLogProbs:
        """Create OpenAI-style logprobs."""
1776
        logprobs_content: list[ChatCompletionLogProbsContent] = []
1777

1778
1779
1780
1781
1782
        should_return_as_token_id = (
            return_as_token_id
            if return_as_token_id is not None
            else self.return_tokens_as_token_ids
        )
1783
1784
        for i, token_id in enumerate(token_ids):
            step_top_logprobs = top_logprobs[i]
1785
            if step_top_logprobs is None or step_top_logprobs.get(token_id) is None:
1786
                if should_return_as_token_id:
1787
                    token = f"token_id:{token_id}"
1788
                else:
1789
1790
                    if tokenizer is None:
                        raise ValueError(
1791
                            "Unable to get tokenizer because `skip_tokenizer_init=True`"
1792
1793
                        )

1794
                    token = tokenizer.decode(token_id)
1795

1796
1797
                logprobs_content.append(
                    ChatCompletionLogProbsContent(
1798
                        token=token,
1799
                        bytes=list(token.encode("utf-8", errors="replace")),
1800
1801
                    )
                )
1802
            else:
1803
1804
1805
                step_token = step_top_logprobs[token_id]
                step_decoded = step_token.decoded_token

1806
1807
                logprobs_content.append(
                    ChatCompletionLogProbsContent(
1808
                        token=self._get_decoded_token(
1809
1810
1811
                            step_token,
                            token_id,
                            tokenizer,
1812
                            should_return_as_token_id,
1813
1814
                        ),
                        logprob=max(step_token.logprob, -9999.0),
1815
1816
1817
1818
1819
                        bytes=(
                            None
                            if step_decoded is None
                            else list(step_decoded.encode("utf-8", errors="replace"))
                        ),
1820
                        top_logprobs=self._get_top_logprobs(
1821
1822
1823
1824
1825
1826
1827
                            step_top_logprobs,
                            num_output_top_logprobs,
                            tokenizer,
                            should_return_as_token_id,
                        ),
                    )
                )
1828
1829

        return ChatCompletionLogProbs(content=logprobs_content)
1830

1831
    def _should_stream_with_auto_tool_parsing(self, request: ChatCompletionRequest):
1832
1833
1834
1835
1836
1837
1838
1839
        """
        Utility function to check if streamed tokens should go through the tool
        call parser that was configured.

        We only want to do this IF user-provided tools are set, a tool parser
        is configured, "auto" tool choice is enabled, and the request's tool
        choice field indicates that "auto" tool choice should be used.
        """
1840
1841
1842
1843
1844
1845
        return (
            request.tools
            and self.tool_parser
            and self.enable_auto_tools
            and request.tool_choice in ["auto", None]
        )
1846
1847
1848

    def _should_check_for_unstreamed_tool_arg_tokens(
        self,
1849
        delta_message: DeltaMessage | None,
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
        output: CompletionOutput,
    ) -> bool:
        """
        Check to see if we should check for unstreamed tool arguments tokens.
        This is only applicable when auto tool parsing is enabled, the delta
        is a tool call with arguments.
        """

        return bool(
            # if there is a delta message that includes tool calls which
            # include a function that has arguments
1861
            output.finish_reason is not None
1862
1863
1864
1865
1866
            and self.enable_auto_tools
            and self.tool_parser
            and delta_message
            and delta_message.tool_calls
            and delta_message.tool_calls[0]
1867
1868
1869
            and delta_message.tool_calls[0].function
            and delta_message.tool_calls[0].function.arguments is not None
        )
1870

1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
    @staticmethod
    def _create_remaining_args_delta(
        delta_message: DeltaMessage,
        remaining_call: str,
        index: int,
    ) -> DeltaMessage:
        """
        Create a delta message for remaining tool arguments, preserving
        id/type/name from the original delta.
        """
        original_tc = next(
            (tc for tc in delta_message.tool_calls if tc.index == index),
            None,
        )
        original_fn = original_tc.function if original_tc else None
        return DeltaMessage(
            tool_calls=[
                DeltaToolCall(
                    index=index,
                    id=original_tc.id if original_tc else None,
                    type=original_tc.type if original_tc else None,
                    function=DeltaFunctionCall(
                        name=original_fn.name if original_fn else None,
                        arguments=remaining_call,
                    ),
                )
            ]
        )

1900
1901
1902
    def _make_request_with_harmony(
        self,
        request: ChatCompletionRequest,
1903
        should_include_tools: bool = True,
1904
1905
1906
    ):
        messages: list[OpenAIMessage] = []

1907
1908
1909
        # because of issues with pydantic we need to potentially
        # re-serialize the tool_calls field of the request
        # for more info: see comment in `maybe_serialize_tool_calls`
zhuwenwen's avatar
zhuwenwen committed
1910
        maybe_serialize_tool_calls(request)
1911

1912
1913
1914
1915
1916
1917
1918
1919
        # Add system message.
        # NOTE: In Chat Completion API, browsing is enabled by default
        # if the model supports it. TODO: Support browsing.
        assert not self.supports_browsing
        assert not self.supports_code_interpreter
        sys_msg = get_system_message(
            reasoning_effort=request.reasoning_effort,
            browser_description=None,
1920
            python_description=None,
1921
            with_custom_tools=should_include_tools,
1922
        )
1923
1924
1925
        messages.append(sys_msg)

        # Add developer message.
1926
1927
        if request.tools:
            dev_msg = get_developer_message(
zhuwenwen's avatar
zhuwenwen committed
1928
                tools=request.tools if should_include_tools else None
1929
1930
            )
            messages.append(dev_msg)
1931
1932

        # Add user message.
1933
        messages.extend(parse_chat_inputs_to_harmony_messages(request.messages))
1934
1935
1936

        # Render prompt token ids.
        prompt_token_ids = render_for_completion(messages)
1937
        engine_prompt = TokensPrompt(prompt_token_ids=prompt_token_ids)
1938
1939
1940
1941
1942

        # Add cache_salt if provided in the request
        if request.cache_salt is not None:
            engine_prompt["cache_salt"] = request.cache_salt

1943
        return messages, [engine_prompt]