serving_chat.py 24.9 KB
Newer Older
1
import asyncio
2
import time
3
from typing import AsyncGenerator, AsyncIterator, Dict, Final, List, Optional
4
from typing import Sequence as GenericSequence
5
from typing import Union
6

7
from fastapi import Request
8

9
from vllm.config import ModelConfig
10
from vllm.engine.protocol import AsyncEngineClient
11
from vllm.entrypoints.chat_utils import (ConversationMessage,
12
                                         apply_chat_template,
13
                                         load_chat_template,
14
                                         parse_chat_messages_futures)
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
                                                    OpenAIServing,
25
26
                                                    PromptAdapterPath,
                                                    TextTokensPrompt)
27
from vllm.inputs import TokensPrompt
28
29
from vllm.logger import init_logger
from vllm.outputs import RequestOutput
30
from vllm.sequence import Logprob
31
32
from vllm.tracing import (contains_trace_headers, extract_trace_headers,
                          log_tracing_disabled_warning)
33
from vllm.transformers_utils.tokenizer import AnyTokenizer
34
from vllm.utils import iterate_with_cancellation, random_uuid
35
36
37
38
39
40

logger = init_logger(__name__)


class OpenAIServingChat(OpenAIServing):

41
42
    def __init__(
        self,
43
        async_engine_client: AsyncEngineClient,
44
45
46
47
48
49
50
51
        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],
52
        return_tokens_as_token_ids: bool = False,
53
    ):
54
        super().__init__(async_engine_client=async_engine_client,
55
                         model_config=model_config,
56
                         served_model_names=served_model_names,
57
58
                         lora_modules=lora_modules,
                         prompt_adapters=prompt_adapters,
59
60
                         request_logger=request_logger,
                         return_tokens_as_token_ids=return_tokens_as_token_ids)
61

62
        self.response_role = response_role
63
64
65

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

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

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

79
        NOTE: Currently we do not support the following feature:
80
81
82
83
84
85
86
            - 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:
87
88
89
90
91
92
            (
                lora_request,
                prompt_adapter_request,
            ) = self._maybe_get_adapters(request)

            model_config = self.model_config
93
94
            tokenizer = await self.async_engine_client.get_tokenizer(
                lora_request)
95

96
            conversation, mm_data_future = parse_chat_messages_futures(
97
                request.messages, model_config, tokenizer)
98

99
100
101
102
            tool_dicts = None if request.tools is None else [
                tool.model_dump() for tool in request.tools
            ]

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

116
        try:
117
            mm_data = await mm_data_future
118
        except Exception as e:
119
            logger.error("Error in loading multi-modal data: %s", e)
120
121
            return self.create_error_response(str(e))

122
        request_id = f"chat-{random_uuid()}"
123
        try:
124
            guided_decode_logits_processor = (
125
                await self._guided_decode_logits_processor(request, tokenizer))
126

127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
            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
143

144
145
146
147
148
149
            sampling_params = request.to_sampling_params(
                tokenizer,
                guided_decode_logits_processor,
                default_max_tokens=self.max_model_len -
                len(prompt_inputs["prompt_token_ids"]))

150
151
152
153
154
155
            self._log_inputs(request_id,
                             prompt_inputs,
                             params=sampling_params,
                             lora_request=lora_request,
                             prompt_adapter_request=prompt_adapter_request)

156
157
            engine_inputs = TokensPrompt(
                prompt_token_ids=prompt_inputs["prompt_token_ids"])
158
159
160
            if mm_data is not None:
                engine_inputs["multi_modal_data"] = mm_data

161
162
            is_tracing_enabled = (
                await self.async_engine_client.is_tracing_enabled())
163
164
165
166
167
168
169
            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()

170
            result_generator = self.async_engine_client.generate(
171
172
173
174
175
176
177
                engine_inputs,
                sampling_params,
                request_id,
                lora_request=lora_request,
                trace_headers=trace_headers,
                prompt_adapter_request=prompt_adapter_request,
            )
178
        except ValueError as e:
179
            # TODO: Use a vllm-specific Validation Error
180
181
            return self.create_error_response(str(e))

182
183
184
185
        if raw_request:
            result_generator = iterate_with_cancellation(
                result_generator, raw_request.is_disconnected)

186
187
188
        # Streaming response
        if request.stream:
            return self.chat_completion_stream_generator(
189
                request, result_generator, request_id, conversation, tokenizer)
190
191
192
193
194
195
        try:
            return await self.chat_completion_full_generator(
                request, result_generator, request_id, conversation, tokenizer)
        except ValueError as e:
            # TODO: Use a vllm-specific Validation Error
            return self.create_error_response(str(e))
196
197
198
199
200

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

    async def chat_completion_stream_generator(
204
205
206
207
208
        self,
        request: ChatCompletionRequest,
        result_generator: AsyncIterator[RequestOutput],
        request_id: str,
        conversation: List[ConversationMessage],
209
        tokenizer: AnyTokenizer,
210
    ) -> AsyncGenerator[str, None]:
211
        model_name = self.served_model_names[0]
212
        created_time = int(time.time())
213
        chunk_object_type: Final = "chat.completion.chunk"
214
        first_iteration = True
215
216

        # Send response for each token for each request.n (index)
217
218
219
220
221
        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

222
223
224
225
226
227
        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:
228
229
                    # Send first response for each request.n (index) with
                    # the role
230
                    role = self.get_chat_request_role(request)
231
                    for i in range(num_choices):
232
233
234
235
236
237
238
239
240
241
242
                        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)
243
244
                        if (request.stream_options
                                and request.stream_options.include_usage):
245
246
247
248
249
250
251
252
253
                            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

254
255
256
                        data = chunk.model_dump_json(exclude_unset=True)
                        yield f"data: {data}\n\n"

257
258
                    # Send response to echo the input portion of the
                    # last message
259
260
                    if request.echo:
                        last_msg_content = ""
261
262
263
264
                        if conversation and conversation[-1].get(
                                "content") and conversation[-1].get(
                                    "role") == role:
                            last_msg_content = conversation[-1]["content"]
265
266

                        if last_msg_content:
267
                            for i in range(num_choices):
268
269
270
271
272
                                choice_data = (
                                    ChatCompletionResponseStreamChoice(
                                        index=i,
                                        delta=DeltaMessage(
                                            content=last_msg_content),
273
                                        logprobs=None,
274
                                        finish_reason=None))
275
276
277
278
279
280
                                chunk = ChatCompletionStreamResponse(
                                    id=request_id,
                                    object=chunk_object_type,
                                    created=created_time,
                                    choices=[choice_data],
                                    model=model_name)
281
282
                                if (request.stream_options and
                                        request.stream_options.include_usage):
283
284
285
286
287
288
289
290
291
292
293
294
                                    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

295
296
297
298
299
300
301
302
303
304
305
306
                                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]:]
307
                    out_logprobs = output.logprobs[
308
309
                        previous_num_tokens[i]:] if output.logprobs else None

310
311
312
                    if request.logprobs and request.top_logprobs is not None:
                        assert out_logprobs is not None, (
                            "Did not output logprobs")
313
                        logprobs = self._create_chat_logprobs(
314
                            token_ids=delta_token_ids,
315
                            top_logprobs=out_logprobs,
316
                            tokenizer=tokenizer,
317
                            num_output_top_logprobs=request.top_logprobs,
318
319
320
321
322
323
324
                        )
                    else:
                        logprobs = None

                    delta_text = output.text[len(previous_texts[i]):]
                    previous_texts[i] = output.text
                    previous_num_tokens[i] = len(output.token_ids)
325
326
327
328
329
330
331
332
333
334
335
336

                    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)

337
338
                    if output.finish_reason is None:
                        # Send token-by-token response for each request.n
339

340
341
                        choice_data = ChatCompletionResponseStreamChoice(
                            index=i,
342
                            delta=delta_message,
343
344
345
346
347
348
349
350
                            logprobs=logprobs,
                            finish_reason=None)
                        chunk = ChatCompletionStreamResponse(
                            id=request_id,
                            object=chunk_object_type,
                            created=created_time,
                            choices=[choice_data],
                            model=model_name)
351
352
                        if (request.stream_options
                                and request.stream_options.include_usage):
353
354
355
356
357
358
359
360
361
362
363
364
365
                            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

366
367
368
369
370
371
372
                        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,
373
                            delta=delta_message,
374
                            logprobs=logprobs,
375
376
                            finish_reason=output.finish_reason,
                            stop_reason=output.stop_reason)
377
378
379
380
381
382
                        chunk = ChatCompletionStreamResponse(
                            id=request_id,
                            object=chunk_object_type,
                            created=created_time,
                            choices=[choice_data],
                            model=model_name)
383
384
                        if (request.stream_options
                                and request.stream_options.include_usage):
385
386
387
388
389
390
391
392
393
394
395
396
                            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
397
                        data = chunk.model_dump_json(exclude_unset=True)
398
399
                        yield f"data: {data}\n\n"
                        finish_reason_sent[i] = True
400

401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
            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"
419

420
421
422
423
        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"
424
425
426
427
        # Send the final done message after all response.n are finished
        yield "data: [DONE]\n\n"

    async def chat_completion_full_generator(
428
429
430
431
432
        self,
        request: ChatCompletionRequest,
        result_generator: AsyncIterator[RequestOutput],
        request_id: str,
        conversation: List[ConversationMessage],
433
        tokenizer: AnyTokenizer,
434
    ) -> Union[ErrorResponse, ChatCompletionResponse]:
435

436
        model_name = self.served_model_names[0]
437
        created_time = int(time.time())
438
        final_res: Optional[RequestOutput] = None
439

440
441
442
443
444
445
        try:
            async for res in result_generator:
                final_res = res
        except asyncio.CancelledError:
            return self.create_error_response("Client disconnected")

446
447
        assert final_res is not None

448
        choices: List[ChatCompletionResponseChoice] = []
449

450
451
        role = self.get_chat_request_role(request)
        for output in final_res.outputs:
452
            token_ids = output.token_ids
453
            out_logprobs = output.logprobs
454

455
456
            if request.logprobs and request.top_logprobs is not None:
                assert out_logprobs is not None, "Did not output logprobs"
457
                logprobs = self._create_chat_logprobs(
458
                    token_ids=token_ids,
459
                    top_logprobs=out_logprobs,
460
                    num_output_top_logprobs=request.top_logprobs,
461
                    tokenizer=tokenizer,
462
463
464
465
                )
            else:
                logprobs = None

466
467
468
469
470
471
472
473
474
475
476
477
478
            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)

479
480
            choice_data = ChatCompletionResponseChoice(
                index=output.index,
481
                message=message,
482
                logprobs=logprobs,
483
                finish_reason=output.finish_reason,
484
                stop_reason=output.stop_reason)
485
486
487
488
            choices.append(choice_data)

        if request.echo:
            last_msg_content = ""
489
490
491
            if conversation and conversation[-1].get(
                    "content") and conversation[-1].get("role") == role:
                last_msg_content = conversation[-1]["content"]
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510

            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,
511
            prompt_logprobs=final_res.prompt_logprobs,
512
513
        )

514
        return response
515
516

    def _get_top_logprobs(
517
            self, logprobs: Dict[int, Logprob], top_logprobs: Optional[int],
518
            tokenizer: AnyTokenizer) -> List[ChatCompletionLogProb]:
519
        return [
520
521
522
523
524
525
526
527
            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")))
528
529
530
531
532
533
534
535
            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]]],
536
        tokenizer: AnyTokenizer,
537
538
539
        num_output_top_logprobs: Optional[int] = None,
    ) -> ChatCompletionLogProbs:
        """Create OpenAI-style logprobs."""
540
        logprobs_content: List[ChatCompletionLogProbsContent] = []
541
542
543
544

        for i, token_id in enumerate(token_ids):
            step_top_logprobs = top_logprobs[i]
            if step_top_logprobs is None:
545
                token = tokenizer.decode(token_id)
546
547
                if self.return_tokens_as_token_ids:
                    token = f"token_id:{token_id}"
548

549
550
                logprobs_content.append(
                    ChatCompletionLogProbsContent(
551
                        token=token,
552
553
                        bytes=list(token.encode("utf-8", errors="replace")),
                    ))
554
            else:
555
556
557
                step_token = step_top_logprobs[token_id]
                step_decoded = step_token.decoded_token

558
559
                logprobs_content.append(
                    ChatCompletionLogProbsContent(
560
                        token=self._get_decoded_token(
561
562
563
564
565
566
567
568
                            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")),
569
                        top_logprobs=self._get_top_logprobs(
570
571
572
573
574
                            step_top_logprobs,
                            num_output_top_logprobs,
                            tokenizer,
                        ),
                    ))
575
576

        return ChatCompletionLogProbs(content=logprobs_content)