serving_chat.py 24.7 KB
Newer Older
1
import time
2
3
from typing import (AsyncGenerator, AsyncIterator, Awaitable, Dict, List,
                    Optional)
4
from typing import Sequence as GenericSequence
5
from typing import Union
6

7
from fastapi import Request
8
from transformers import PreTrainedTokenizer
9

10
from vllm.config import ModelConfig
11
from vllm.engine.async_llm_engine import AsyncLLMEngine
12
13
14
from vllm.entrypoints.chat_utils import (ConversationMessage,
                                         load_chat_template,
                                         parse_chat_message_content)
15
from vllm.entrypoints.logger import RequestLogger
16
from vllm.entrypoints.openai.protocol import (
17
18
    ChatCompletionLogProb, ChatCompletionLogProbs,
    ChatCompletionLogProbsContent, ChatCompletionNamedToolChoiceParam,
19
    ChatCompletionRequest, ChatCompletionResponse,
20
21
    ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice,
    ChatCompletionStreamResponse, ChatMessage, DeltaMessage, ErrorResponse,
22
    FunctionCall, ToolCall, UsageInfo)
23
from vllm.entrypoints.openai.serving_engine import (LoRAModulePath,
24
25
                                                    OpenAIServing,
                                                    PromptAdapterPath)
26
from vllm.inputs import PromptInputs
27
from vllm.logger import init_logger
28
29
from vllm.model_executor.guided_decoding import (
    get_guided_decoding_logits_processor)
30
from vllm.multimodal import MultiModalDataDict
31
from vllm.outputs import RequestOutput
32
from vllm.sequence import Logprob
33
34
from vllm.tracing import (contains_trace_headers, extract_trace_headers,
                          log_tracing_disabled_warning)
35
from vllm.utils import random_uuid
36
37
38
39
40
41

logger = init_logger(__name__)


class OpenAIServingChat(OpenAIServing):

42
43
44
45
46
47
48
49
50
51
52
53
    def __init__(
        self,
        engine: AsyncLLMEngine,
        model_config: ModelConfig,
        served_model_names: List[str],
        response_role: str,
        *,
        lora_modules: Optional[List[LoRAModulePath]],
        prompt_adapters: Optional[List[PromptAdapterPath]],
        request_logger: Optional[RequestLogger],
        chat_template: Optional[str],
    ):
54
        super().__init__(engine=engine,
55
                         model_config=model_config,
56
                         served_model_names=served_model_names,
57
58
59
                         lora_modules=lora_modules,
                         prompt_adapters=prompt_adapters,
                         request_logger=request_logger)
60

61
        self.response_role = response_role
62
63
64

        # If this is None we use the tokenizer's default chat template
        self.chat_template = load_chat_template(chat_template)
65

66
    async def create_chat_completion(
67
68
69
        self,
        request: ChatCompletionRequest,
        raw_request: Optional[Request] = None
70
71
72
73
    ) -> Union[ErrorResponse, AsyncGenerator[str, None],
               ChatCompletionResponse]:
        """Completion API similar to OpenAI's API.

74
75
76
        See https://platform.openai.com/docs/api-reference/chat/create
        for the API specification. This API mimics the OpenAI
        ChatCompletion API.
77

78
        NOTE: Currently we do not support the following feature:
79
80
81
82
83
84
85
            - function_call (Users should implement this by themselves)
        """
        error_check_ret = await self._check_model(request)
        if error_check_ret is not None:
            return error_check_ret

        try:
86
87
88
89
90
91
            (
                lora_request,
                prompt_adapter_request,
            ) = self._maybe_get_adapters(request)

            model_config = self.model_config
92
93
            tokenizer = await self.engine.get_tokenizer(lora_request)

94
            conversation: List[ConversationMessage] = []
95
            mm_futures: List[Awaitable[MultiModalDataDict]] = []
96

97
            for msg in request.messages:
98
                chat_parsed_result = parse_chat_message_content(
99
                    msg, model_config, tokenizer)
100

101
                conversation.extend(chat_parsed_result.messages)
102
                mm_futures.extend(chat_parsed_result.mm_futures)
103

104
105
106
107
            tool_dicts = None if request.tools is None else [
                tool.model_dump() for tool in request.tools
            ]

108
            prompt = tokenizer.apply_chat_template(
109
                conversation=conversation,
110
                tokenize=False,
111
                add_generation_prompt=request.add_generation_prompt,
112
113
                tools=tool_dicts,
                documents=request.documents,
114
                chat_template=request.chat_template or self.chat_template,
115
                **(request.chat_template_kwargs or {}),
116
            )
117
        except Exception as e:
118
            logger.error("Error in applying chat template from request: %s", e)
119
120
            return self.create_error_response(str(e))

121
        mm_data: Optional[MultiModalDataDict] = None
122
        try:
123
124
            if len(mm_futures):
                # since we support only single mm data currently
125
126
127
                assert len(
                    mm_futures
                ) == 1, "Multiple 'image_url' input is currently not supported."
128
                mm_data = await mm_futures[0]
129
        except Exception as e:
130
            logger.error("Error in loading multi-modal data: %s", e)
131
132
            return self.create_error_response(str(e))

133
        request_id = f"chat-{random_uuid()}"
134
        try:
135
            sampling_params = request.to_sampling_params()
136
            decoding_config = await self.engine.get_decoding_config()
137
138
            guided_decoding_backend = request.guided_decoding_backend \
                or decoding_config.guided_decoding_backend
139
            guided_decode_logits_processor = (
140
141
142
                await
                get_guided_decoding_logits_processor(guided_decoding_backend,
                                                     request, tokenizer))
143
144
145
146
147
            if guided_decode_logits_processor:
                if sampling_params.logits_processors is None:
                    sampling_params.logits_processors = []
                sampling_params.logits_processors.append(
                    guided_decode_logits_processor)
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184

            prompt_inputs = self._tokenize_prompt_input(
                request,
                tokenizer,
                prompt,
                truncate_prompt_tokens=sampling_params.truncate_prompt_tokens,
                add_special_tokens=request.add_special_tokens,
            )

            self._log_inputs(request_id,
                             prompt_inputs,
                             params=sampling_params,
                             lora_request=lora_request,
                             prompt_adapter_request=prompt_adapter_request)

            engine_inputs: PromptInputs = {
                "prompt_token_ids": prompt_inputs["prompt_token_ids"],
            }
            if mm_data is not None:
                engine_inputs["multi_modal_data"] = mm_data

            is_tracing_enabled = await self.engine.is_tracing_enabled()
            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()

            result_generator = self.engine.generate(
                engine_inputs,
                sampling_params,
                request_id,
                lora_request=lora_request,
                trace_headers=trace_headers,
                prompt_adapter_request=prompt_adapter_request,
            )
185
        except ValueError as e:
186
            # TODO: Use a vllm-specific Validation Error
187
188
189
190
191
            return self.create_error_response(str(e))

        # Streaming response
        if request.stream:
            return self.chat_completion_stream_generator(
192
                request, result_generator, request_id, conversation, tokenizer)
193
        else:
194
195
            try:
                return await self.chat_completion_full_generator(
196
                    request, raw_request, result_generator, request_id,
197
                    conversation, tokenizer)
198
199
200
            except ValueError as e:
                # TODO: Use a vllm-specific Validation Error
                return self.create_error_response(str(e))
201
202
203
204
205

    def get_chat_request_role(self, request: ChatCompletionRequest) -> str:
        if request.add_generation_prompt:
            return self.response_role
        else:
206
            return request.messages[-1]["role"]
207
208

    async def chat_completion_stream_generator(
209
210
211
212
213
214
        self,
        request: ChatCompletionRequest,
        result_generator: AsyncIterator[RequestOutput],
        request_id: str,
        conversation: List[ConversationMessage],
        tokenizer: PreTrainedTokenizer,
215
    ) -> AsyncGenerator[str, None]:
216
        model_name = self.served_model_names[0]
217
        created_time = int(time.time())
218
        chunk_object_type = "chat.completion.chunk"
219
        first_iteration = True
220
221

        # Send response for each token for each request.n (index)
222
223
224
225
226
        num_choices = 1 if request.n is None else request.n
        previous_texts = [""] * num_choices
        previous_num_tokens = [0] * num_choices
        finish_reason_sent = [False] * num_choices

227
228
229
230
231
232
        try:
            async for res in result_generator:
                # 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:
233
234
                    # Send first response for each request.n (index) with
                    # the role
235
                    role = self.get_chat_request_role(request)
236
                    for i in range(num_choices):
237
238
239
240
241
242
243
244
245
246
247
                        choice_data = ChatCompletionResponseStreamChoice(
                            index=i,
                            delta=DeltaMessage(role=role),
                            logprobs=None,
                            finish_reason=None)
                        chunk = ChatCompletionStreamResponse(
                            id=request_id,
                            object=chunk_object_type,
                            created=created_time,
                            choices=[choice_data],
                            model=model_name)
248
249
                        if (request.stream_options
                                and request.stream_options.include_usage):
250
251
252
253
254
255
256
257
258
                            if (request.stream_options.continuous_usage_stats):
                                prompt_tokens = len(res.prompt_token_ids)
                                usage = UsageInfo(prompt_tokens=prompt_tokens,
                                                  completion_tokens=0,
                                                  total_tokens=prompt_tokens)
                                chunk.usage = usage
                            else:
                                chunk.usage = None

259
260
261
                        data = chunk.model_dump_json(exclude_unset=True)
                        yield f"data: {data}\n\n"

262
263
                    # Send response to echo the input portion of the
                    # last message
264
265
                    if request.echo:
                        last_msg_content = ""
266
267
268
269
                        if conversation and conversation[-1].get(
                                "content") and conversation[-1].get(
                                    "role") == role:
                            last_msg_content = conversation[-1]["content"]
270
271

                        if last_msg_content:
272
                            for i in range(num_choices):
273
274
275
276
277
                                choice_data = (
                                    ChatCompletionResponseStreamChoice(
                                        index=i,
                                        delta=DeltaMessage(
                                            content=last_msg_content),
278
                                        logprobs=None,
279
                                        finish_reason=None))
280
281
282
283
284
285
                                chunk = ChatCompletionStreamResponse(
                                    id=request_id,
                                    object=chunk_object_type,
                                    created=created_time,
                                    choices=[choice_data],
                                    model=model_name)
286
287
                                if (request.stream_options and
                                        request.stream_options.include_usage):
288
289
290
291
292
293
294
295
296
297
298
299
                                    if (request.stream_options.
                                            continuous_usage_stats):
                                        prompt_tokens = len(
                                            res.prompt_token_ids)
                                        usage = UsageInfo(
                                            prompt_tokens=prompt_tokens,
                                            completion_tokens=0,
                                            total_tokens=prompt_tokens)
                                        chunk.usage = usage
                                    else:
                                        chunk.usage = None

300
301
302
303
304
305
306
307
308
309
310
311
                                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

                    if finish_reason_sent[i]:
                        continue

                    delta_token_ids = output.token_ids[previous_num_tokens[i]:]
312
                    out_logprobs = output.logprobs[
313
314
                        previous_num_tokens[i]:] if output.logprobs else None

315
316
317
                    if request.logprobs and request.top_logprobs is not None:
                        assert out_logprobs is not None, (
                            "Did not output logprobs")
318
                        logprobs = self._create_chat_logprobs(
319
                            token_ids=delta_token_ids,
320
                            top_logprobs=out_logprobs,
321
                            tokenizer=tokenizer,
322
                            num_output_top_logprobs=request.top_logprobs,
323
324
325
326
327
328
329
                        )
                    else:
                        logprobs = None

                    delta_text = output.text[len(previous_texts[i]):]
                    previous_texts[i] = output.text
                    previous_num_tokens[i] = len(output.token_ids)
330
331
332
333
334
335
336
337
338
339
340
341

                    if request.tool_choice and type(
                            request.tool_choice
                    ) is ChatCompletionNamedToolChoiceParam:
                        delta_message = DeltaMessage(tool_calls=[
                            ToolCall(function=FunctionCall(
                                name=request.tool_choice.function.name,
                                arguments=delta_text))
                        ])
                    else:
                        delta_message = DeltaMessage(content=delta_text)

342
343
                    if output.finish_reason is None:
                        # Send token-by-token response for each request.n
344

345
346
                        choice_data = ChatCompletionResponseStreamChoice(
                            index=i,
347
                            delta=delta_message,
348
349
350
351
352
353
354
355
                            logprobs=logprobs,
                            finish_reason=None)
                        chunk = ChatCompletionStreamResponse(
                            id=request_id,
                            object=chunk_object_type,
                            created=created_time,
                            choices=[choice_data],
                            model=model_name)
356
357
                        if (request.stream_options
                                and request.stream_options.include_usage):
358
359
360
361
362
363
364
365
366
367
368
369
370
                            if (request.stream_options.continuous_usage_stats):
                                prompt_tokens = len(res.prompt_token_ids)
                                completion_tokens = len(output.token_ids)
                                usage = UsageInfo(
                                    prompt_tokens=prompt_tokens,
                                    completion_tokens=completion_tokens,
                                    total_tokens=prompt_tokens +
                                    completion_tokens,
                                )
                                chunk.usage = usage
                            else:
                                chunk.usage = None

371
372
373
374
375
376
377
                        data = chunk.model_dump_json(exclude_unset=True)
                        yield f"data: {data}\n\n"
                    else:
                        # Send the finish response for each request.n only once
                        prompt_tokens = len(res.prompt_token_ids)
                        choice_data = ChatCompletionResponseStreamChoice(
                            index=i,
378
                            delta=delta_message,
379
                            logprobs=logprobs,
380
381
                            finish_reason=output.finish_reason,
                            stop_reason=output.stop_reason)
382
383
384
385
386
387
                        chunk = ChatCompletionStreamResponse(
                            id=request_id,
                            object=chunk_object_type,
                            created=created_time,
                            choices=[choice_data],
                            model=model_name)
388
389
                        if (request.stream_options
                                and request.stream_options.include_usage):
390
391
392
393
394
395
396
397
398
399
400
401
                            if (request.stream_options.continuous_usage_stats):
                                prompt_tokens = len(res.prompt_token_ids)
                                completion_tokens = len(output.token_ids)
                                usage = UsageInfo(
                                    prompt_tokens=prompt_tokens,
                                    completion_tokens=completion_tokens,
                                    total_tokens=prompt_tokens +
                                    completion_tokens,
                                )
                                chunk.usage = usage
                            else:
                                chunk.usage = None
402
                        data = chunk.model_dump_json(exclude_unset=True)
403
404
                        yield f"data: {data}\n\n"
                        finish_reason_sent[i] = True
405

406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
            if (request.stream_options
                    and request.stream_options.include_usage):
                final_usage = UsageInfo(
                    prompt_tokens=prompt_tokens,
                    completion_tokens=previous_num_tokens[i],
                    total_tokens=prompt_tokens + previous_num_tokens[i],
                )

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

425
426
427
428
        except ValueError as e:
            # TODO: Use a vllm-specific Validation Error
            data = self.create_streaming_error_response(str(e))
            yield f"data: {data}\n\n"
429
430
431
432
        # Send the final done message after all response.n are finished
        yield "data: [DONE]\n\n"

    async def chat_completion_full_generator(
433
434
435
436
437
438
439
        self,
        request: ChatCompletionRequest,
        raw_request: Optional[Request],
        result_generator: AsyncIterator[RequestOutput],
        request_id: str,
        conversation: List[ConversationMessage],
        tokenizer: PreTrainedTokenizer,
440
    ) -> Union[ErrorResponse, ChatCompletionResponse]:
441

442
        model_name = self.served_model_names[0]
443
        created_time = int(time.time())
444
        final_res: Optional[RequestOutput] = None
445
446

        async for res in result_generator:
447
            if raw_request is not None and await raw_request.is_disconnected():
448
449
450
451
452
453
                # Abort the request if the client disconnects.
                await self.engine.abort(request_id)
                return self.create_error_response("Client disconnected")
            final_res = res
        assert final_res is not None

454
        choices: List[ChatCompletionResponseChoice] = []
455

456
457
        role = self.get_chat_request_role(request)
        for output in final_res.outputs:
458
            token_ids = output.token_ids
459
            out_logprobs = output.logprobs
460

461
462
            if request.logprobs and request.top_logprobs is not None:
                assert out_logprobs is not None, "Did not output logprobs"
463
                logprobs = self._create_chat_logprobs(
464
                    token_ids=token_ids,
465
                    top_logprobs=out_logprobs,
466
                    num_output_top_logprobs=request.top_logprobs,
467
                    tokenizer=tokenizer,
468
469
470
471
                )
            else:
                logprobs = None

472
473
474
475
476
477
478
479
480
481
482
483
484
            if request.tool_choice and type(
                    request.tool_choice) is ChatCompletionNamedToolChoiceParam:
                message = ChatMessage(
                    role=role,
                    content="",
                    tool_calls=[
                        ToolCall(function=FunctionCall(
                            name=request.tool_choice.function.name,
                            arguments=output.text))
                    ])
            elif not request.tool_choice or request.tool_choice == "none":
                message = ChatMessage(role=role, content=output.text)

485
486
            choice_data = ChatCompletionResponseChoice(
                index=output.index,
487
                message=message,
488
                logprobs=logprobs,
489
                finish_reason=output.finish_reason,
490
                stop_reason=output.stop_reason)
491
492
493
494
            choices.append(choice_data)

        if request.echo:
            last_msg_content = ""
495
496
497
            if conversation and conversation[-1].get(
                    "content") and conversation[-1].get("role") == role:
                last_msg_content = conversation[-1]["content"]
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518

            for choice in choices:
                full_message = last_msg_content + choice.message.content
                choice.message.content = full_message

        num_prompt_tokens = len(final_res.prompt_token_ids)
        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,
        )
        response = ChatCompletionResponse(
            id=request_id,
            created=created_time,
            model=model_name,
            choices=choices,
            usage=usage,
        )

519
        return response
520
521

    def _get_top_logprobs(
522
523
            self, logprobs: Dict[int, Logprob], top_logprobs: Optional[int],
            tokenizer: PreTrainedTokenizer) -> List[ChatCompletionLogProb]:
524
525
        return [
            ChatCompletionLogProb(
526
527
                token=(token := self._get_decoded_token(p[1], p[0],
                                                        tokenizer)),
528
                logprob=max(p[1].logprob, -9999.0),
529
                bytes=list(token.encode("utf-8", errors="replace")))
530
531
532
533
534
535
536
537
            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]]],
538
        tokenizer: PreTrainedTokenizer,
539
540
541
542
543
544
545
546
547
        num_output_top_logprobs: Optional[int] = None,
    ) -> ChatCompletionLogProbs:
        """Create OpenAI-style logprobs."""

        logprobs_content = []

        for i, token_id in enumerate(token_ids):
            step_top_logprobs = top_logprobs[i]
            if step_top_logprobs is None:
548
                token = tokenizer.decode(token_id)
549
550
                logprobs_content.append(
                    ChatCompletionLogProbsContent(
551
552
                        token=token,
                        bytes=list(token.encode("utf-8", errors="replace"))))
553
554
555
556
557
558
559
560
561
562
            else:
                logprobs_content.append(
                    ChatCompletionLogProbsContent(
                        token=step_top_logprobs[token_id].decoded_token,
                        logprob=max(step_top_logprobs[token_id].logprob,
                                    -9999.0),
                        bytes=list(
                            step_top_logprobs[token_id].decoded_token.encode(
                                "utf-8", errors="replace")),
                        top_logprobs=self._get_top_logprobs(
563
564
                            step_top_logprobs, num_output_top_logprobs,
                            tokenizer)))
565
566

        return ChatCompletionLogProbs(content=logprobs_content)