serving_chat.py 38.9 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
        num_prompt_tokens = 0

305
306
307
        tool_parsers: List[Optional[ToolParser]] = [
            self.tool_parser(tokenizer) if self.tool_parser else None
        ] * num_choices
308

309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
        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

327
328
        try:
            async for res in result_generator:
329
330
                if res.prompt_token_ids is not None:
                    num_prompt_tokens = len(res.prompt_token_ids)
331
332
                    if res.encoder_prompt_token_ids is not None:
                        num_prompt_tokens += len(res.encoder_prompt_token_ids)
333

334
335
336
337
                # 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:
338
339
                    # Send first response for each request.n (index) with
                    # the role
340
                    role = self.get_chat_request_role(request)
341
342
343

                    # NOTE num_choices defaults to 1 so this usually executes
                    # once per request
344
                    for i in range(num_choices):
345
                        tool_parser = tool_parsers[i]
346
347
                        choice_data = ChatCompletionResponseStreamChoice(
                            index=i,
348
349
350
351
                            delta=DeltaMessage(
                                role=role,
                                content="",
                            ),
352
353
354
355
356
357
358
359
                            logprobs=None,
                            finish_reason=None)
                        chunk = ChatCompletionStreamResponse(
                            id=request_id,
                            object=chunk_object_type,
                            created=created_time,
                            choices=[choice_data],
                            model=model_name)
360
361

                        # if usage should be included
362
363
                        if (request.stream_options
                                and request.stream_options.include_usage):
364
365
                            # if continuous usage stats are requested, add it
                            if request.stream_options.continuous_usage_stats:
366
367
368
369
                                usage = UsageInfo(
                                    prompt_tokens=num_prompt_tokens,
                                    completion_tokens=0,
                                    total_tokens=num_prompt_tokens)
370
                                chunk.usage = usage
371
                            # otherwise don't
372
373
374
                            else:
                                chunk.usage = None

375
376
377
                        data = chunk.model_dump_json(exclude_unset=True)
                        yield f"data: {data}\n\n"

378
379
                    # Send response to echo the input portion of the
                    # last message
380
                    if request.echo or request.continue_final_message:
381
382
383
384
                        last_msg_content: str = ""
                        if conversation and "content" in conversation[
                                -1] and conversation[-1].get("role") == role:
                            last_msg_content = conversation[-1]["content"] or ""
385
386

                        if last_msg_content:
387
                            for i in range(num_choices):
388
389
390
391
392
                                choice_data = (
                                    ChatCompletionResponseStreamChoice(
                                        index=i,
                                        delta=DeltaMessage(
                                            content=last_msg_content),
393
                                        logprobs=None,
394
                                        finish_reason=None))
395
396
397
398
399
400
                                chunk = ChatCompletionStreamResponse(
                                    id=request_id,
                                    object=chunk_object_type,
                                    created=created_time,
                                    choices=[choice_data],
                                    model=model_name)
401
402
                                if (request.stream_options and
                                        request.stream_options.include_usage):
403
404
405
                                    if (request.stream_options.
                                            continuous_usage_stats):
                                        usage = UsageInfo(
406
                                            prompt_tokens=num_prompt_tokens,
407
                                            completion_tokens=0,
408
                                            total_tokens=num_prompt_tokens)
409
410
411
412
                                        chunk.usage = usage
                                    else:
                                        chunk.usage = None

413
414
415
416
417
418
419
                                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
420
                    tool_parser = tool_parsers[i]
421
422
423
424

                    if finish_reason_sent[i]:
                        continue

425
                    if request.logprobs and request.top_logprobs is not None:
426
                        assert output.logprobs is not None, (
427
                            "Did not output logprobs")
428
                        logprobs = self._create_chat_logprobs(
429
430
                            token_ids=output.token_ids,
                            top_logprobs=output.logprobs,
431
                            tokenizer=tokenizer,
432
                            num_output_top_logprobs=request.top_logprobs,
433
434
435
436
                        )
                    else:
                        logprobs = None

437
438
                    delta_text = output.text
                    delta_message: Optional[DeltaMessage]
439

440
                    # handle streaming deltas for tools with named tool_choice
441
                    if tool_choice_function_name:
442
                        delta_message = DeltaMessage(tool_calls=[
443
                            DeltaToolCall(function=DeltaFunctionCall(
444
                                name=tool_choice_function_name,
445
446
                                arguments=delta_text),
                                          index=i)
447
                        ])
448
449

                    # handle streaming deltas for tools with "auto" tool choice
450
451
452
453
454
455
456
457
458
459
460
                    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)

461
462
                        delta_message = (
                            tool_parser.extract_tool_calls_streaming(
463
464
                                previous_text=previous_text,
                                current_text=current_text,
465
                                delta_text=delta_text,
466
467
                                previous_token_ids=previous_token_ids,
                                current_token_ids=current_token_ids,
468
469
                                delta_token_ids=output.token_ids,
                                request=request))
470
471
472
473

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

                    # handle streaming just a content delta
476
477
478
                    else:
                        delta_message = DeltaMessage(content=delta_text)

479
                    # set the previous values for the next iteration
480
                    previous_num_tokens[i] += len(output.token_ids)
481
482
483
484
485
486
487
488

                    # 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

489
490
                    if output.finish_reason is None:
                        # Send token-by-token response for each request.n
491

492
493
                        choice_data = ChatCompletionResponseStreamChoice(
                            index=i,
494
                            delta=delta_message,
495
496
497
498
499
500
501
502
                            logprobs=logprobs,
                            finish_reason=None)
                        chunk = ChatCompletionStreamResponse(
                            id=request_id,
                            object=chunk_object_type,
                            created=created_time,
                            choices=[choice_data],
                            model=model_name)
503
504

                        # handle usage stats if requested & if continuous
505
506
                        if (request.stream_options
                                and request.stream_options.include_usage):
507
                            if request.stream_options.continuous_usage_stats:
508
509
                                completion_tokens = len(output.token_ids)
                                usage = UsageInfo(
510
                                    prompt_tokens=num_prompt_tokens,
511
                                    completion_tokens=completion_tokens,
512
                                    total_tokens=num_prompt_tokens +
513
514
515
516
517
518
                                    completion_tokens,
                                )
                                chunk.usage = usage
                            else:
                                chunk.usage = None

519
520
                        data = chunk.model_dump_json(exclude_unset=True)
                        yield f"data: {data}\n\n"
521
522

                    # if the model is finished generating
523
                    else:
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
                        # 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", {}))

543
                            # get what we've streamed so far for arguments
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
                            # 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))
                            ])

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

593
594
            # once the final token is handled, if stream_options.include_usage
            # is sent, send the usage
595
596
            if (request.stream_options
                    and request.stream_options.include_usage):
597
                completion_tokens = previous_num_tokens[i]
598
                final_usage = UsageInfo(
599
600
601
                    prompt_tokens=num_prompt_tokens,
                    completion_tokens=completion_tokens,
                    total_tokens=num_prompt_tokens + completion_tokens,
602
603
604
605
606
607
608
609
610
611
612
613
                )

                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"
614

615
616
617
618
619
620
621
            # 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)

622
623
        except ValueError as e:
            # TODO: Use a vllm-specific Validation Error
624
            logger.error("error in chat completion stream generator: %s", e)
625
626
            data = self.create_streaming_error_response(str(e))
            yield f"data: {data}\n\n"
627
628
629
630
        # Send the final done message after all response.n are finished
        yield "data: [DONE]\n\n"

    async def chat_completion_full_generator(
631
632
633
634
635
        self,
        request: ChatCompletionRequest,
        result_generator: AsyncIterator[RequestOutput],
        request_id: str,
        conversation: List[ConversationMessage],
636
        tokenizer: AnyTokenizer,
637
        request_metadata: RequestResponseMetadata,
638
    ) -> Union[ErrorResponse, ChatCompletionResponse]:
639

640
        model_name = self.base_model_paths[0].name
641
        created_time = int(time.time())
642
        final_res: Optional[RequestOutput] = None
643

644
645
646
647
648
649
        try:
            async for res in result_generator:
                final_res = res
        except asyncio.CancelledError:
            return self.create_error_response("Client disconnected")

650
651
        assert final_res is not None

652
        choices: List[ChatCompletionResponseChoice] = []
653

654
655
        role = self.get_chat_request_role(request)
        for output in final_res.outputs:
656
            token_ids = output.token_ids
657
            out_logprobs = output.logprobs
658

659
660
            if request.logprobs and request.top_logprobs is not None:
                assert out_logprobs is not None, "Did not output logprobs"
661
                logprobs = self._create_chat_logprobs(
662
                    token_ids=token_ids,
663
                    top_logprobs=out_logprobs,
664
                    num_output_top_logprobs=request.top_logprobs,
665
                    tokenizer=tokenizer,
666
667
668
669
                )
            else:
                logprobs = None

670
671
672
673
674
            # 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
675
            if (not self.enable_auto_tools
676
677
678
679
680
681
682
                    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(
683
                    request.tool_choice) is ChatCompletionNamedToolChoiceParam:
684

685
686
687
688
689
690
691
692
                message = ChatMessage(
                    role=role,
                    content="",
                    tool_calls=[
                        ToolCall(function=FunctionCall(
                            name=request.tool_choice.function.name,
                            arguments=output.text))
                    ])
693
694
695
696
                tools_called = True

            # if the request doesn't use tool choice
            # OR specifies to not use a tool
697
            elif not request.tool_choice or request.tool_choice == "none":
698
699
700
701
702
703
704
705
706
707

                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:

                tool_parser = self.tool_parser(tokenizer)
708
709
                tool_call_info = tool_parser.extract_tool_calls(
                    output.text, request=request)
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
                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.")
727
728
                message = ChatMessage(role=role, content=output.text)

729
730
            choice_data = ChatCompletionResponseChoice(
                index=output.index,
731
                message=message,
732
                logprobs=logprobs,
733
734
                finish_reason="tool_calls" if tools_called else
                output.finish_reason if output.finish_reason else "stop",
735
                stop_reason=output.stop_reason)
736
737
            choices.append(choice_data)

738
        if request.echo or request.continue_final_message:
739
            last_msg_content = ""
740
741
            if conversation and "content" in conversation[-1] and conversation[
                    -1].get("role") == role:
742
                last_msg_content = conversation[-1]["content"] or ""
743
744

            for choice in choices:
745
746
                full_message = last_msg_content + (choice.message.content
                                                   or "")
747
748
                choice.message.content = full_message

749
        assert final_res.prompt_token_ids is not None
750
        num_prompt_tokens = len(final_res.prompt_token_ids)
751
752
        if final_res.encoder_prompt_token_ids is not None:
            num_prompt_tokens += len(final_res.encoder_prompt_token_ids)
753
754
755
756
757
758
759
        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,
        )
760
761
762

        request_metadata.final_usage_info = usage

763
764
765
766
767
768
        response = ChatCompletionResponse(
            id=request_id,
            created=created_time,
            model=model_name,
            choices=choices,
            usage=usage,
769
            prompt_logprobs=final_res.prompt_logprobs,
770
771
        )

772
        return response
773
774

    def _get_top_logprobs(
775
            self, logprobs: Dict[int, Logprob], top_logprobs: Optional[int],
776
            tokenizer: AnyTokenizer) -> List[ChatCompletionLogProb]:
777
        return [
778
779
780
781
782
783
784
785
            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")))
786
787
788
789
790
791
792
793
            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]]],
794
        tokenizer: AnyTokenizer,
795
796
797
        num_output_top_logprobs: Optional[int] = None,
    ) -> ChatCompletionLogProbs:
        """Create OpenAI-style logprobs."""
798
        logprobs_content: List[ChatCompletionLogProbsContent] = []
799
800
801
802

        for i, token_id in enumerate(token_ids):
            step_top_logprobs = top_logprobs[i]
            if step_top_logprobs is None:
803
                token = tokenizer.decode(token_id)
804
805
                if self.return_tokens_as_token_ids:
                    token = f"token_id:{token_id}"
806

807
808
                logprobs_content.append(
                    ChatCompletionLogProbsContent(
809
                        token=token,
810
811
                        bytes=list(token.encode("utf-8", errors="replace")),
                    ))
812
            else:
813
814
815
                step_token = step_top_logprobs[token_id]
                step_decoded = step_token.decoded_token

816
817
                logprobs_content.append(
                    ChatCompletionLogProbsContent(
818
                        token=self._get_decoded_token(
819
820
821
822
823
824
825
826
                            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")),
827
                        top_logprobs=self._get_top_logprobs(
828
829
830
831
832
                            step_top_logprobs,
                            num_output_top_logprobs,
                            tokenizer,
                        ),
                    ))
833
834

        return ChatCompletionLogProbs(content=logprobs_content)
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863

    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
864
865
            output.finish_reason is not None
            and self.enable_auto_tools and self.tool_parser and delta_message
866
867
868
869
            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
        )