serving_chat.py 39.5 KB
Newer Older
1
import asyncio
2
import json
3
import time
4
5
from typing import (AsyncGenerator, AsyncIterator, Callable, Dict, Final, List,
                    Optional)
6
from typing import Sequence as GenericSequence
7
from typing import Union
8

9
from fastapi import Request
10

11
from vllm.config import ModelConfig
12
from vllm.engine.async_llm_engine import AsyncLLMEngine
13
from vllm.engine.protocol import EngineClient
14
from vllm.entrypoints.chat_utils import (ConversationMessage,
15
16
                                         apply_hf_chat_template,
                                         apply_mistral_chat_template,
17
                                         load_chat_template,
18
                                         parse_chat_messages_futures)
19
from vllm.entrypoints.logger import RequestLogger
20
from vllm.entrypoints.openai.protocol import (
21
22
    ChatCompletionLogProb, ChatCompletionLogProbs,
    ChatCompletionLogProbsContent, ChatCompletionNamedToolChoiceParam,
23
    ChatCompletionRequest, ChatCompletionResponse,
24
    ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice,
25
    ChatCompletionStreamResponse, ChatMessage, DeltaFunctionCall, DeltaMessage,
26
27
    DeltaToolCall, ErrorResponse, FunctionCall, RequestResponseMetadata,
    ToolCall, UsageInfo)
28
29
from vllm.entrypoints.openai.serving_engine import (BaseModelPath,
                                                    LoRAModulePath,
30
                                                    OpenAIServing,
31
32
                                                    PromptAdapterPath,
                                                    TextTokensPrompt)
33
from vllm.entrypoints.openai.tool_parsers import ToolParser, ToolParserManager
34
from vllm.inputs import TokensPrompt
35
from vllm.logger import init_logger
36
from vllm.outputs import CompletionOutput, RequestOutput
37
from vllm.sampling_params import BeamSearchParams, SamplingParams
38
from vllm.sequence import Logprob
39
40
from vllm.tracing import (contains_trace_headers, extract_trace_headers,
                          log_tracing_disabled_warning)
41
from vllm.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer
42
from vllm.utils import iterate_with_cancellation, random_uuid
43
44
45
46
47
48

logger = init_logger(__name__)


class OpenAIServingChat(OpenAIServing):

49
    def __init__(self,
50
                 engine_client: EngineClient,
51
                 model_config: ModelConfig,
52
                 base_model_paths: List[BaseModelPath],
53
54
55
56
57
58
59
60
61
                 response_role: str,
                 *,
                 lora_modules: Optional[List[LoRAModulePath]],
                 prompt_adapters: Optional[List[PromptAdapterPath]],
                 request_logger: Optional[RequestLogger],
                 chat_template: Optional[str],
                 return_tokens_as_token_ids: bool = False,
                 enable_auto_tools: bool = False,
                 tool_parser: Optional[str] = None):
62
        super().__init__(engine_client=engine_client,
63
                         model_config=model_config,
64
                         base_model_paths=base_model_paths,
65
66
                         lora_modules=lora_modules,
                         prompt_adapters=prompt_adapters,
67
68
                         request_logger=request_logger,
                         return_tokens_as_token_ids=return_tokens_as_token_ids)
69

70
        self.response_role = response_role
71
        self.use_tool_use_model_template = False
72
        self.chat_template = load_chat_template(chat_template)
73

74
75
76
77
78
79
80
81
82
83
        # set up tool use
        self.enable_auto_tools: bool = enable_auto_tools
        if self.enable_auto_tools:
            logger.info(
                "\"auto\" tool choice has been enabled please note that while"
                " the parallel_tool_calls client option is preset for "
                "compatibility reasons, it will be ignored.")

        self.tool_parser: Optional[Callable[[AnyTokenizer], ToolParser]] = None
        if self.enable_auto_tools:
84
85
86
87
            try:
                self.tool_parser = ToolParserManager.get_tool_parser(
                    tool_parser)
            except Exception as e:
88
                raise TypeError("Error: --enable-auto-tool-choice requires "
89
90
                                f"tool_parser:'{tool_parser}' which has not "
                                "been registered") from e
91

92
    async def create_chat_completion(
93
94
        self,
        request: ChatCompletionRequest,
95
96
97
        raw_request: Optional[Request] = None,
    ) -> Union[AsyncGenerator[str, None], ChatCompletionResponse,
               ErrorResponse]:
98
99
        """Completion API similar to OpenAI's API.

100
101
102
        See https://platform.openai.com/docs/api-reference/chat/create
        for the API specification. This API mimics the OpenAI
        ChatCompletion API.
103
104
105
106

        """
        error_check_ret = await self._check_model(request)
        if error_check_ret is not None:
107
            logger.error("Error with model %s", error_check_ret)
108
109
            return error_check_ret

110
111
112
113
114
115
        # 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

116
        try:
117
118
119
120
121
122
            (
                lora_request,
                prompt_adapter_request,
            ) = self._maybe_get_adapters(request)

            model_config = self.model_config
123
            tokenizer = await self.engine_client.get_tokenizer(lora_request)
124

125
            conversation, mm_data_future = parse_chat_messages_futures(
126
                request.messages, model_config, tokenizer)
127

128
129
130
131
            tool_dicts = None if request.tools is None else [
                tool.model_dump() for tool in request.tools
            ]

132
            prompt: Union[str, List[int]]
133
134
            is_mistral_tokenizer = isinstance(tokenizer, MistralTokenizer)
            if is_mistral_tokenizer:
135
136
137
138
139
                prompt = apply_mistral_chat_template(
                    tokenizer,
                    messages=request.messages,
                    chat_template=request.chat_template or self.chat_template,
                    add_generation_prompt=request.add_generation_prompt,
140
                    continue_final_message=request.continue_final_message,
141
142
143
144
145
146
147
148
149
150
                    tools=tool_dicts,
                    documents=request.documents,
                    **(request.chat_template_kwargs or {}),
                )
            else:
                prompt = apply_hf_chat_template(
                    tokenizer,
                    conversation=conversation,
                    chat_template=request.chat_template or self.chat_template,
                    add_generation_prompt=request.add_generation_prompt,
151
                    continue_final_message=request.continue_final_message,
152
153
154
155
                    tools=tool_dicts,
                    documents=request.documents,
                    **(request.chat_template_kwargs or {}),
                )
156
        except Exception as e:
157
            logger.exception("Error in applying chat template from request")
158
159
            return self.create_error_response(str(e))

160
        try:
161
            mm_data = await mm_data_future
162
        except Exception as e:
163
            logger.exception("Error in loading multi-modal data")
164
165
            return self.create_error_response(str(e))

166
167
168
169
170
171
        # validation for OpenAI tools
        # tool_choice = "required" is not supported
        if request.tool_choice == "required":
            return self.create_error_response(
                "tool_choice = \"required\" is not supported!")

172
        if not is_mistral_tokenizer and request.tool_choice == "auto" and not (
173
                self.enable_auto_tools and self.tool_parser is not None):
174
175
            # for hf tokenizers, "auto" tools requires
            # --enable-auto-tool-choice and --tool-call-parser
176
177
178
179
            return self.create_error_response(
                "\"auto\" tool choice requires "
                "--enable-auto-tool-choice and --tool-call-parser to be set")

180
        request_id = f"chat-{random_uuid()}"
181
182
183
184
185

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

186
        try:
187
188
189
190
            if self.enable_auto_tools and self.tool_parser:
                request = self.tool_parser(tokenizer).adjust_request(
                    request=request)

191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
            if isinstance(prompt, str):
                prompt_inputs = self._tokenize_prompt_input(
                    request,
                    tokenizer,
                    prompt,
                    truncate_prompt_tokens=request.truncate_prompt_tokens,
                    add_special_tokens=request.add_special_tokens,
                )
            else:
                assert isinstance(prompt, list) and isinstance(
                    prompt[0], int
                ), "Prompt has to be either a string or a list of token ids"
                prompt_inputs = TextTokensPrompt(
                    prompt=tokenizer.decode(prompt), prompt_token_ids=prompt)

            assert prompt_inputs is not None
207

208
209
210
211
212
213
214
215
216
            sampling_params: Union[SamplingParams, BeamSearchParams]
            default_max_tokens = self.max_model_len - len(
                prompt_inputs["prompt_token_ids"])
            if request.use_beam_search:
                sampling_params = request.to_beam_search_params(
                    default_max_tokens)
            else:
                sampling_params = request.to_sampling_params(
                    default_max_tokens)
217

218
219
220
221
222
223
            self._log_inputs(request_id,
                             prompt_inputs,
                             params=sampling_params,
                             lora_request=lora_request,
                             prompt_adapter_request=prompt_adapter_request)

224
225
            engine_inputs = TokensPrompt(
                prompt_token_ids=prompt_inputs["prompt_token_ids"])
226
227
228
            if mm_data is not None:
                engine_inputs["multi_modal_data"] = mm_data

229
230
            is_tracing_enabled = (await
                                  self.engine_client.is_tracing_enabled())
231
232
233
234
235
236
237
            trace_headers = None
            if is_tracing_enabled and raw_request:
                trace_headers = extract_trace_headers(raw_request.headers)
            if (not is_tracing_enabled and raw_request
                    and contains_trace_headers(raw_request.headers)):
                log_tracing_disabled_warning()

238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
            if isinstance(sampling_params, BeamSearchParams):
                if not isinstance(self.engine_client, AsyncLLMEngine):
                    raise ValueError(
                        "Beam search in the API server is only supported with"
                        " AsyncLLMEngine. please add "
                        "`--disable-frontend-multiprocessing` to "
                        "use beam search.")
                result_generator = self.engine_client.beam_search(
                    engine_inputs['prompt_token_ids'], request_id,
                    sampling_params)
            else:
                result_generator = self.engine_client.generate(
                    engine_inputs,
                    sampling_params,
                    request_id,
                    lora_request=lora_request,
                    trace_headers=trace_headers,
                    prompt_adapter_request=prompt_adapter_request,
                    priority=request.priority,
                )
258
        except ValueError as e:
259
            # TODO: Use a vllm-specific Validation Error
260
261
            return self.create_error_response(str(e))

262
263
264
265
        if raw_request:
            result_generator = iterate_with_cancellation(
                result_generator, raw_request.is_disconnected)

266
267
268
        # Streaming response
        if request.stream:
            return self.chat_completion_stream_generator(
269
270
                request, result_generator, request_id, conversation, tokenizer,
                request_metadata)
271

272
273
        try:
            return await self.chat_completion_full_generator(
274
275
                request, result_generator, request_id, conversation, tokenizer,
                request_metadata)
276
277
278
        except ValueError as e:
            # TODO: Use a vllm-specific Validation Error
            return self.create_error_response(str(e))
279
280
281
282

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

    async def chat_completion_stream_generator(
286
287
288
289
290
        self,
        request: ChatCompletionRequest,
        result_generator: AsyncIterator[RequestOutput],
        request_id: str,
        conversation: List[ConversationMessage],
291
        tokenizer: AnyTokenizer,
292
        request_metadata: RequestResponseMetadata,
293
    ) -> AsyncGenerator[str, None]:
294
        model_name = self.base_model_paths[0].name
295
        created_time = int(time.time())
296
        chunk_object_type: Final = "chat.completion.chunk"
297
        first_iteration = True
298
299

        # Send response for each token for each request.n (index)
300
301
302
        num_choices = 1 if request.n is None else request.n
        previous_num_tokens = [0] * num_choices
        finish_reason_sent = [False] * num_choices
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
        num_prompt_tokens = 0

        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
            and self._should_stream_with_auto_tool_parsing(request))

        all_previous_token_ids: Optional[List[List[int]]]
        if tool_choice_auto:
            # These are only required in "auto" tool choice case
            previous_texts = [""] * num_choices
            all_previous_token_ids = [[]] * num_choices
        else:
            previous_texts, all_previous_token_ids = None, None

323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
        # Prepare the tool parser if it's needed
        try:
            if tool_choice_auto and self.tool_parser:
                tool_parsers: List[Optional[ToolParser]] = [
                    self.tool_parser(tokenizer)
                ] * num_choices
            else:
                tool_parsers = [None] * num_choices
        except RuntimeError as e:
            logger.error("Error in tool parser creation: %s", e)
            data = self.create_streaming_error_response(str(e))
            yield f"data: {data}\n\n"
            yield "data: [DONE]\n\n"
            return

338
339
        try:
            async for res in result_generator:
340
341
                if res.prompt_token_ids is not None:
                    num_prompt_tokens = len(res.prompt_token_ids)
342
343
                    if res.encoder_prompt_token_ids is not None:
                        num_prompt_tokens += len(res.encoder_prompt_token_ids)
344

345
346
347
348
                # 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:
349
350
                    # Send first response for each request.n (index) with
                    # the role
351
                    role = self.get_chat_request_role(request)
352
353
354

                    # NOTE num_choices defaults to 1 so this usually executes
                    # once per request
355
                    for i in range(num_choices):
356
                        tool_parser = tool_parsers[i]
357
358
                        choice_data = ChatCompletionResponseStreamChoice(
                            index=i,
359
360
361
362
                            delta=DeltaMessage(
                                role=role,
                                content="",
                            ),
363
364
365
366
367
368
369
370
                            logprobs=None,
                            finish_reason=None)
                        chunk = ChatCompletionStreamResponse(
                            id=request_id,
                            object=chunk_object_type,
                            created=created_time,
                            choices=[choice_data],
                            model=model_name)
371
372

                        # if usage should be included
373
374
                        if (request.stream_options
                                and request.stream_options.include_usage):
375
376
                            # if continuous usage stats are requested, add it
                            if request.stream_options.continuous_usage_stats:
377
378
379
380
                                usage = UsageInfo(
                                    prompt_tokens=num_prompt_tokens,
                                    completion_tokens=0,
                                    total_tokens=num_prompt_tokens)
381
                                chunk.usage = usage
382
                            # otherwise don't
383
384
385
                            else:
                                chunk.usage = None

386
387
388
                        data = chunk.model_dump_json(exclude_unset=True)
                        yield f"data: {data}\n\n"

389
390
                    # Send response to echo the input portion of the
                    # last message
391
                    if request.echo or request.continue_final_message:
392
393
394
395
                        last_msg_content: str = ""
                        if conversation and "content" in conversation[
                                -1] and conversation[-1].get("role") == role:
                            last_msg_content = conversation[-1]["content"] or ""
396
397

                        if last_msg_content:
398
                            for i in range(num_choices):
399
400
401
402
403
                                choice_data = (
                                    ChatCompletionResponseStreamChoice(
                                        index=i,
                                        delta=DeltaMessage(
                                            content=last_msg_content),
404
                                        logprobs=None,
405
                                        finish_reason=None))
406
407
408
409
410
411
                                chunk = ChatCompletionStreamResponse(
                                    id=request_id,
                                    object=chunk_object_type,
                                    created=created_time,
                                    choices=[choice_data],
                                    model=model_name)
412
413
                                if (request.stream_options and
                                        request.stream_options.include_usage):
414
415
416
                                    if (request.stream_options.
                                            continuous_usage_stats):
                                        usage = UsageInfo(
417
                                            prompt_tokens=num_prompt_tokens,
418
                                            completion_tokens=0,
419
                                            total_tokens=num_prompt_tokens)
420
421
422
423
                                        chunk.usage = usage
                                    else:
                                        chunk.usage = None

424
425
426
427
428
429
430
                                data = chunk.model_dump_json(
                                    exclude_unset=True)
                                yield f"data: {data}\n\n"
                    first_iteration = False

                for output in res.outputs:
                    i = output.index
431
                    tool_parser = tool_parsers[i]
432
433
434
435

                    if finish_reason_sent[i]:
                        continue

436
                    if request.logprobs and request.top_logprobs is not None:
437
                        assert output.logprobs is not None, (
438
                            "Did not output logprobs")
439
                        logprobs = self._create_chat_logprobs(
440
441
                            token_ids=output.token_ids,
                            top_logprobs=output.logprobs,
442
                            tokenizer=tokenizer,
443
                            num_output_top_logprobs=request.top_logprobs,
444
445
446
447
                        )
                    else:
                        logprobs = None

448
449
                    delta_text = output.text
                    delta_message: Optional[DeltaMessage]
450

451
                    # handle streaming deltas for tools with named tool_choice
452
                    if tool_choice_function_name:
453
                        delta_message = DeltaMessage(tool_calls=[
454
                            DeltaToolCall(function=DeltaFunctionCall(
455
                                name=tool_choice_function_name,
456
457
                                arguments=delta_text),
                                          index=i)
458
                        ])
459
460

                    # handle streaming deltas for tools with "auto" tool choice
461
462
463
464
465
466
467
468
469
470
471
                    elif tool_choice_auto:
                        assert previous_texts is not None
                        assert all_previous_token_ids is not None
                        assert tool_parser is not None
                        #TODO optimize manipulation of these lists
                        previous_text = previous_texts[i]
                        previous_token_ids = all_previous_token_ids[i]
                        current_text = previous_text + delta_text
                        current_token_ids = previous_token_ids + list(
                            output.token_ids)

472
473
                        delta_message = (
                            tool_parser.extract_tool_calls_streaming(
474
475
                                previous_text=previous_text,
                                current_text=current_text,
476
                                delta_text=delta_text,
477
478
                                previous_token_ids=previous_token_ids,
                                current_token_ids=current_token_ids,
479
480
                                delta_token_ids=output.token_ids,
                                request=request))
481
482
483
484

                        # update the previous values for the next iteration
                        previous_texts[i] = current_text
                        all_previous_token_ids[i] = current_token_ids
485
486

                    # handle streaming just a content delta
487
488
489
                    else:
                        delta_message = DeltaMessage(content=delta_text)

490
                    # set the previous values for the next iteration
491
                    previous_num_tokens[i] += len(output.token_ids)
492
493
494
495
496
497
498
499

                    # 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:
                        continue

500
501
                    if output.finish_reason is None:
                        # Send token-by-token response for each request.n
502

503
504
                        choice_data = ChatCompletionResponseStreamChoice(
                            index=i,
505
                            delta=delta_message,
506
507
508
509
510
511
512
513
                            logprobs=logprobs,
                            finish_reason=None)
                        chunk = ChatCompletionStreamResponse(
                            id=request_id,
                            object=chunk_object_type,
                            created=created_time,
                            choices=[choice_data],
                            model=model_name)
514
515

                        # handle usage stats if requested & if continuous
516
517
                        if (request.stream_options
                                and request.stream_options.include_usage):
518
                            if request.stream_options.continuous_usage_stats:
519
520
                                completion_tokens = len(output.token_ids)
                                usage = UsageInfo(
521
                                    prompt_tokens=num_prompt_tokens,
522
                                    completion_tokens=completion_tokens,
523
                                    total_tokens=num_prompt_tokens +
524
525
526
527
528
529
                                    completion_tokens,
                                )
                                chunk.usage = usage
                            else:
                                chunk.usage = None

530
531
                        data = chunk.model_dump_json(exclude_unset=True)
                        yield f"data: {data}\n\n"
532
533

                    # if the model is finished generating
534
                    else:
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
                        # check to make sure we haven't "forgotten" to stream
                        #   any tokens that were generated but previously
                        #   matched by partial json parsing
                        # only happens if we are NOT using guided decoding
                        if tool_parser:
                            index = len(
                                tool_parser.prev_tool_call_arr) - 1 if len(
                                    tool_parser.prev_tool_call_arr) > 0 else 0
                        else:
                            index = 0

                        if self._should_check_for_unstreamed_tool_arg_tokens(
                                delta_message, output) and tool_parser:
                            # get the expected call based on partial JSON
                            # parsing which "autocompletes" the JSON
                            expected_call = json.dumps(
                                tool_parser.prev_tool_call_arr[index].get(
                                    "arguments", {}))

554
                            # get what we've streamed so far for arguments
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
                            # for the current tool
                            actual_call = tool_parser.streamed_args_for_tool[
                                index]

                            # check to see if there's anything left to stream
                            remaining_call = expected_call.replace(
                                actual_call, "", 1)

                            # set that as a delta message
                            delta_message = DeltaMessage(tool_calls=[
                                DeltaToolCall(index=index,
                                              function=DeltaFunctionCall(
                                                  arguments=remaining_call).
                                              model_dump(exclude_none=True))
                            ])

571
572
573
                        # Send the finish response for each request.n only once
                        choice_data = ChatCompletionResponseStreamChoice(
                            index=i,
574
                            delta=delta_message,
575
                            logprobs=logprobs,
576
577
578
579
                            finish_reason=output.finish_reason
                            if not (tool_parser
                                    and len(tool_parser.prev_tool_call_arr))
                            else "tool_calls",
580
                            stop_reason=output.stop_reason)
581
582
583
584
585
586
                        chunk = ChatCompletionStreamResponse(
                            id=request_id,
                            object=chunk_object_type,
                            created=created_time,
                            choices=[choice_data],
                            model=model_name)
587
588
                        if (request.stream_options
                                and request.stream_options.include_usage):
589
                            if request.stream_options.continuous_usage_stats:
590
591
                                completion_tokens = len(output.token_ids)
                                usage = UsageInfo(
592
                                    prompt_tokens=num_prompt_tokens,
593
                                    completion_tokens=completion_tokens,
594
                                    total_tokens=num_prompt_tokens +
595
596
597
598
599
                                    completion_tokens,
                                )
                                chunk.usage = usage
                            else:
                                chunk.usage = None
600
                        data = chunk.model_dump_json(exclude_unset=True)
601
602
                        yield f"data: {data}\n\n"
                        finish_reason_sent[i] = True
603

604
605
            # once the final token is handled, if stream_options.include_usage
            # is sent, send the usage
606
607
            if (request.stream_options
                    and request.stream_options.include_usage):
608
                completion_tokens = previous_num_tokens[i]
609
                final_usage = UsageInfo(
610
611
612
                    prompt_tokens=num_prompt_tokens,
                    completion_tokens=completion_tokens,
                    total_tokens=num_prompt_tokens + completion_tokens,
613
614
615
616
617
618
619
620
621
622
623
624
                )

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

626
627
628
629
630
631
632
            # 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,
                total_tokens=num_prompt_tokens + num_completion_tokens)

633
634
        except ValueError as e:
            # TODO: Use a vllm-specific Validation Error
635
            logger.error("error in chat completion stream generator: %s", e)
636
637
            data = self.create_streaming_error_response(str(e))
            yield f"data: {data}\n\n"
638
639
640
641
        # Send the final done message after all response.n are finished
        yield "data: [DONE]\n\n"

    async def chat_completion_full_generator(
642
643
644
645
646
        self,
        request: ChatCompletionRequest,
        result_generator: AsyncIterator[RequestOutput],
        request_id: str,
        conversation: List[ConversationMessage],
647
        tokenizer: AnyTokenizer,
648
        request_metadata: RequestResponseMetadata,
649
    ) -> Union[ErrorResponse, ChatCompletionResponse]:
650

651
        model_name = self.base_model_paths[0].name
652
        created_time = int(time.time())
653
        final_res: Optional[RequestOutput] = None
654

655
656
657
658
659
660
        try:
            async for res in result_generator:
                final_res = res
        except asyncio.CancelledError:
            return self.create_error_response("Client disconnected")

661
662
        assert final_res is not None

663
        choices: List[ChatCompletionResponseChoice] = []
664

665
666
        role = self.get_chat_request_role(request)
        for output in final_res.outputs:
667
            token_ids = output.token_ids
668
            out_logprobs = output.logprobs
669

670
671
            if request.logprobs and request.top_logprobs is not None:
                assert out_logprobs is not None, "Did not output logprobs"
672
                logprobs = self._create_chat_logprobs(
673
                    token_ids=token_ids,
674
                    top_logprobs=out_logprobs,
675
                    num_output_top_logprobs=request.top_logprobs,
676
                    tokenizer=tokenizer,
677
678
679
680
                )
            else:
                logprobs = None

681
682
683
684
685
            # by default, tools are not used.
            tools_called = False

            # if auto tools are not enabled, and a named tool choice using
            #   outlines is not being used
686
            if (not self.enable_auto_tools
687
688
689
690
691
692
693
                    or not self.tool_parser) and not isinstance(
                        request.tool_choice,
                        ChatCompletionNamedToolChoiceParam):
                message = ChatMessage(role=role, content=output.text)

            # if the request uses tools and specified a tool choice
            elif request.tool_choice and type(
694
                    request.tool_choice) is ChatCompletionNamedToolChoiceParam:
695

696
697
698
699
700
701
702
703
                message = ChatMessage(
                    role=role,
                    content="",
                    tool_calls=[
                        ToolCall(function=FunctionCall(
                            name=request.tool_choice.function.name,
                            arguments=output.text))
                    ])
704
705
706
707
                tools_called = True

            # if the request doesn't use tool choice
            # OR specifies to not use a tool
708
            elif not request.tool_choice or request.tool_choice == "none":
709
710
711
712
713
714
715
716
717

                message = ChatMessage(role=role, content=output.text)

            # handle when there are tools and tool choice is auto
            elif request.tools and (
                    request.tool_choice == "auto"
                    or request.tool_choice is None) and self.enable_auto_tools \
                    and self.tool_parser:

718
719
720
721
722
723
                try:
                    tool_parser = self.tool_parser(tokenizer)
                except RuntimeError as e:
                    logger.error("Error in tool parser creation: %s", e)
                    return self.create_error_response(str(e))

724
725
                tool_call_info = tool_parser.extract_tool_calls(
                    output.text, request=request)
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
                tools_called = tool_call_info.tools_called
                if tool_call_info.tools_called:
                    message = ChatMessage(role=role,
                                          content=tool_call_info.content,
                                          tool_calls=tool_call_info.tool_calls)

                else:
                    # FOR NOW make it a chat message; we will have to detect
                    # the type to make it later.
                    message = ChatMessage(role=role, content=output.text)

            # 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 "
                    "completion.")
743
744
                message = ChatMessage(role=role, content=output.text)

745
746
            choice_data = ChatCompletionResponseChoice(
                index=output.index,
747
                message=message,
748
                logprobs=logprobs,
749
750
                finish_reason="tool_calls" if tools_called else
                output.finish_reason if output.finish_reason else "stop",
751
                stop_reason=output.stop_reason)
752
753
            choices.append(choice_data)

754
        if request.echo or request.continue_final_message:
755
            last_msg_content = ""
756
757
            if conversation and "content" in conversation[-1] and conversation[
                    -1].get("role") == role:
758
                last_msg_content = conversation[-1]["content"] or ""
759
760

            for choice in choices:
761
762
                full_message = last_msg_content + (choice.message.content
                                                   or "")
763
764
                choice.message.content = full_message

765
        assert final_res.prompt_token_ids is not None
766
        num_prompt_tokens = len(final_res.prompt_token_ids)
767
768
        if final_res.encoder_prompt_token_ids is not None:
            num_prompt_tokens += len(final_res.encoder_prompt_token_ids)
769
770
771
772
773
774
775
        num_generated_tokens = sum(
            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,
        )
776
777
778

        request_metadata.final_usage_info = usage

779
780
781
782
783
784
        response = ChatCompletionResponse(
            id=request_id,
            created=created_time,
            model=model_name,
            choices=choices,
            usage=usage,
785
            prompt_logprobs=final_res.prompt_logprobs,
786
787
        )

788
        return response
789
790

    def _get_top_logprobs(
791
            self, logprobs: Dict[int, Logprob], top_logprobs: Optional[int],
792
            tokenizer: AnyTokenizer) -> List[ChatCompletionLogProb]:
793
        return [
794
795
796
797
798
799
800
801
            ChatCompletionLogProb(token=(token := self._get_decoded_token(
                p[1],
                p[0],
                tokenizer,
                return_as_token_id=self.return_tokens_as_token_ids)),
                                  logprob=max(p[1].logprob, -9999.0),
                                  bytes=list(
                                      token.encode("utf-8", errors="replace")))
802
803
804
805
806
807
808
809
            for i, p in enumerate(logprobs.items())
            if top_logprobs and i < top_logprobs
        ]

    def _create_chat_logprobs(
        self,
        token_ids: GenericSequence[int],
        top_logprobs: GenericSequence[Optional[Dict[int, Logprob]]],
810
        tokenizer: AnyTokenizer,
811
812
813
        num_output_top_logprobs: Optional[int] = None,
    ) -> ChatCompletionLogProbs:
        """Create OpenAI-style logprobs."""
814
        logprobs_content: List[ChatCompletionLogProbsContent] = []
815
816
817
818

        for i, token_id in enumerate(token_ids):
            step_top_logprobs = top_logprobs[i]
            if step_top_logprobs is None:
819
                token = tokenizer.decode(token_id)
820
821
                if self.return_tokens_as_token_ids:
                    token = f"token_id:{token_id}"
822

823
824
                logprobs_content.append(
                    ChatCompletionLogProbsContent(
825
                        token=token,
826
827
                        bytes=list(token.encode("utf-8", errors="replace")),
                    ))
828
            else:
829
830
831
                step_token = step_top_logprobs[token_id]
                step_decoded = step_token.decoded_token

832
833
                logprobs_content.append(
                    ChatCompletionLogProbsContent(
834
                        token=self._get_decoded_token(
835
836
837
838
839
840
841
842
                            step_token,
                            token_id,
                            tokenizer,
                            self.return_tokens_as_token_ids,
                        ),
                        logprob=max(step_token.logprob, -9999.0),
                        bytes=None if step_decoded is None else list(
                            step_decoded.encode("utf-8", errors="replace")),
843
                        top_logprobs=self._get_top_logprobs(
844
845
846
847
848
                            step_top_logprobs,
                            num_output_top_logprobs,
                            tokenizer,
                        ),
                    ))
849
850

        return ChatCompletionLogProbs(content=logprobs_content)
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879

    def _should_stream_with_auto_tool_parsing(self,
                                              request: ChatCompletionRequest):
        """
        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.
        """
        return (request.tools and self.tool_parser and self.enable_auto_tools
                and request.tool_choice in ['auto', None])

    def _should_check_for_unstreamed_tool_arg_tokens(
        self,
        delta_message: Optional[DeltaMessage],
        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.
        """

        # yapf: disable
        return bool(
            # if there is a delta message that includes tool calls which
            # include a function that has arguments
880
881
            output.finish_reason is not None
            and self.enable_auto_tools and self.tool_parser and delta_message
882
883
884
885
            and delta_message.tool_calls and delta_message.tool_calls[0]
            and delta_message.tool_calls[0].function
            and delta_message.tool_calls[0].function.arguments is not None
        )