serving_chat.py 25.1 KB
Newer Older
1
import codecs
2
import time
3
4
5
from dataclasses import dataclass, field
from typing import (AsyncGenerator, AsyncIterator, Awaitable, Dict, Iterable,
                    List, Optional)
6
7
from typing import Sequence as GenericSequence
from typing import TypedDict, Union, cast, final
8

9
from fastapi import Request
10
11
from openai.types.chat import (ChatCompletionContentPartImageParam,
                               ChatCompletionContentPartTextParam)
12

13
from vllm.config import ModelConfig, VisionLanguageConfig
14
15
from vllm.engine.async_llm_engine import AsyncLLMEngine
from vllm.entrypoints.openai.protocol import (
16
17
    ChatCompletionContentPartParam, ChatCompletionLogProb,
    ChatCompletionLogProbs, ChatCompletionLogProbsContent,
18
19
    ChatCompletionMessageParam, ChatCompletionNamedToolChoiceParam,
    ChatCompletionRequest, ChatCompletionResponse,
20
21
    ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice,
    ChatCompletionStreamResponse, ChatMessage, DeltaMessage, ErrorResponse,
22
    FunctionCall, ToolCall, UsageInfo)
23
24
from vllm.entrypoints.openai.serving_engine import (LoRAModulePath,
                                                    OpenAIServing)
25
from vllm.inputs import PromptInputs
26
from vllm.logger import init_logger
27
28
from vllm.model_executor.guided_decoding import (
    get_guided_decoding_logits_processor)
29
30
31
from vllm.multimodal.image import ImagePixelData
from vllm.multimodal.utils import (async_get_and_parse_image,
                                   get_full_image_text_prompt)
32
from vllm.outputs import RequestOutput
33
from vllm.sequence import Logprob
34
from vllm.utils import random_uuid
35
36
37
38

logger = init_logger(__name__)


39
40
41
42
43
44
@final  # So that it should be compatible with Dict[str, str]
class ConversationMessage(TypedDict):
    role: str
    content: str


45
46
47
@dataclass(frozen=True)
class ChatMessageParseResult:
    messages: List[ConversationMessage]
48
49
    image_futures: List[Awaitable[ImagePixelData]] = field(
        default_factory=list)
50
51


52
53
54
55
class OpenAIServingChat(OpenAIServing):

    def __init__(self,
                 engine: AsyncLLMEngine,
56
                 model_config: ModelConfig,
57
                 served_model_names: List[str],
58
                 response_role: str,
59
60
                 lora_modules: Optional[List[LoRAModulePath]] = None,
                 chat_template: Optional[str] = None):
61
        super().__init__(engine=engine,
62
                         model_config=model_config,
63
                         served_model_names=served_model_names,
64
                         lora_modules=lora_modules)
65

66
        self.response_role = response_role
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
        self._load_chat_template(chat_template)

    def _load_chat_template(self, chat_template: Optional[str]):
        tokenizer = self.tokenizer

        if chat_template is not None:
            try:
                with open(chat_template, "r") as f:
                    tokenizer.chat_template = f.read()
            except OSError as e:
                JINJA_CHARS = "{}\n"
                if not any(c in chat_template for c in JINJA_CHARS):
                    msg = (f"The supplied chat template ({chat_template}) "
                           f"looks like a file path, but it failed to be "
                           f"opened. Reason: {e}")
                    raise ValueError(msg) from e

                # If opening a file fails, set chat template to be args to
                # ensure we decode so our escape are interpreted correctly
                tokenizer.chat_template = codecs.decode(
                    chat_template, "unicode_escape")

            logger.info("Using supplied chat template:\n%s",
                        tokenizer.chat_template)
        elif tokenizer.chat_template is not None:
            logger.info("Using default chat template:\n%s",
                        tokenizer.chat_template)
        else:
            logger.warning(
                "No chat template provided. Chat API will not work.")
97

98
    def _parse_chat_message_content_parts(
99
        self,
100
101
102
        role: str,
        parts: Iterable[ChatCompletionContentPartParam],
    ) -> ChatMessageParseResult:
103
        texts: List[str] = []
104
        image_futures: List[Awaitable[ImagePixelData]] = []
105

106
107
108
109
110
        vlm_config: Optional[VisionLanguageConfig] = getattr(
            self.engine.engine, "vision_language_config", None)
        model_config = getattr(self.engine.engine, "model_config", None)

        for part in parts:
111
112
113
            part_type = part["type"]
            if part_type == "text":
                text = cast(ChatCompletionContentPartTextParam, part)["text"]
114
115

                texts.append(text)
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
            elif part_type == "image_url":
                if vlm_config is None:
                    raise ValueError(
                        "'image_url' input is not supported as the loaded "
                        "model is not multimodal.")

                elif len(image_futures) == 0:
                    assert self.tokenizer is not None
                    image_url = cast(ChatCompletionContentPartImageParam,
                                     part)["image_url"]

                    if image_url.get("detail", "auto") != "auto":
                        logger.warning(
                            "'image_url.detail' is currently not supported and "
                            "will be ignored.")

                    image_future = async_get_and_parse_image(image_url["url"])
                    image_futures.append(image_future)

                else:
                    raise NotImplementedError(
                        "Multiple 'image_url' input is currently not supported."
                    )

140
            else:
141
142
                raise NotImplementedError(f"Unknown part type: {part_type}")

143
144
145
146
147
148
        text_prompt = "\n".join(texts)

        if vlm_config is not None and len(image_futures):

            (image_token_prompt,
             image_token_str) = vlm_config.get_image_token_text(self.tokenizer)
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
            # NOTE: If image token string (e.g, <image>) is already present
            # in the text prompt, we assume it follows the same format required
            # by the engine.
            if image_token_str in text_prompt:
                logger.warning(
                    "Detected image token string in the text prompt. "
                    "Skipping prompt formatting.")
                messages = [
                    ConversationMessage(role=role, content=text_prompt)
                ]

            else:
                full_prompt = get_full_image_text_prompt(
                    image_prompt=image_token_prompt,
                    text_prompt=text_prompt,
                    config=model_config)
                messages = [
                    ConversationMessage(role=role, content=full_prompt)
                ]
        else:
            messages = [ConversationMessage(role=role, content=text_prompt)]

        return ChatMessageParseResult(messages=messages,
                                      image_futures=image_futures)
174
175
176
177
178
179
180
181
182

    def _parse_chat_message_content(
        self,
        message: ChatCompletionMessageParam,
    ) -> ChatMessageParseResult:
        role = message["role"]
        content = message.get("content")

        if content is None:
183
            return ChatMessageParseResult(messages=[], image_futures=[])
184
185
        if isinstance(content, str):
            messages = [ConversationMessage(role=role, content=content)]
186
            return ChatMessageParseResult(messages=messages, image_futures=[])
187

188
        return self._parse_chat_message_content_parts(role, content)
189

190
    async def create_chat_completion(
191
192
193
        self,
        request: ChatCompletionRequest,
        raw_request: Optional[Request] = None
194
195
196
197
    ) -> Union[ErrorResponse, AsyncGenerator[str, None],
               ChatCompletionResponse]:
        """Completion API similar to OpenAI's API.

198
199
200
        See https://platform.openai.com/docs/api-reference/chat/create
        for the API specification. This API mimics the OpenAI
        ChatCompletion API.
201

202
        NOTE: Currently we do not support the following feature:
203
204
205
206
207
208
209
            - 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:
210
            conversation: List[ConversationMessage] = []
211
            image_futures: List[Awaitable[ImagePixelData]] = []
212

213
            for msg in request.messages:
214
                chat_parsed_result = self._parse_chat_message_content(msg)
215

216
217
                conversation.extend(chat_parsed_result.messages)
                image_futures.extend(chat_parsed_result.image_futures)
218

219
            prompt = self.tokenizer.apply_chat_template(
220
                conversation=conversation,
221
                tokenize=False,
222
223
                add_generation_prompt=request.add_generation_prompt,
            )
224
        except Exception as e:
225
            logger.error("Error in applying chat template from request: %s", e)
226
227
            return self.create_error_response(str(e))

228
229
230
231
232
233
234
235
236
237
238
        # Fetch image data
        image_data: Optional[ImagePixelData] = None
        try:
            if len(image_futures):
                # since we support only single image currently
                assert len(image_futures) == 1
                image_data = await image_futures[0]
        except Exception as e:
            logger.error("Error in loading image data: %s", e)
            return self.create_error_response(str(e))

239
240
        request_id = f"cmpl-{random_uuid()}"
        try:
241
242
            # Tokenize/detokenize depending on prompt format (string/token list)
            prompt_ids, prompt_text = self._validate_prompt_and_tokenize(
243
244
245
                request,
                prompt=prompt,
                add_special_tokens=request.add_special_tokens)
246
            sampling_params = request.to_sampling_params()
247
            lora_request = self._maybe_get_lora(request)
248
            decoding_config = await self.engine.get_decoding_config()
249
250
            guided_decoding_backend = request.guided_decoding_backend \
                or decoding_config.guided_decoding_backend
251
252
            guided_decode_logits_processor = (
                await get_guided_decoding_logits_processor(
253
254
                    guided_decoding_backend, request, await
                    self.engine.get_tokenizer()))
255
256
257
258
259
            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)
260
261
262
        except ValueError as e:
            return self.create_error_response(str(e))

263
264
265
266
267
268
269
        inputs: PromptInputs = {
            "prompt": prompt_text,
            "prompt_token_ids": prompt_ids,
        }
        if image_data is not None:
            inputs["multi_modal_data"] = image_data

270
        result_generator = self.engine.generate(
271
            inputs,
272
273
274
275
            sampling_params,
            request_id,
            lora_request,
        )
276
277
278
        # Streaming response
        if request.stream:
            return self.chat_completion_stream_generator(
279
                request, result_generator, request_id, conversation)
280
        else:
281
282
            try:
                return await self.chat_completion_full_generator(
283
284
                    request, raw_request, result_generator, request_id,
                    conversation)
285
286
287
            except ValueError as e:
                # TODO: Use a vllm-specific Validation Error
                return self.create_error_response(str(e))
288
289
290
291
292

    def get_chat_request_role(self, request: ChatCompletionRequest) -> str:
        if request.add_generation_prompt:
            return self.response_role
        else:
293
            return request.messages[-1]["role"]
294
295
296

    async def chat_completion_stream_generator(
            self, request: ChatCompletionRequest,
297
298
299
            result_generator: AsyncIterator[RequestOutput], request_id: str,
            conversation: List[ConversationMessage]
    ) -> AsyncGenerator[str, None]:
300
        model_name = self.served_model_names[0]
301
        created_time = int(time.time())
302
        chunk_object_type = "chat.completion.chunk"
303
        first_iteration = True
304
305

        # Send response for each token for each request.n (index)
306
        assert request.n is not None
307
308
309
        previous_texts = [""] * request.n
        previous_num_tokens = [0] * request.n
        finish_reason_sent = [False] * request.n
310
311
312
313
314
315
        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:
316
317
                    # Send first response for each request.n (index) with
                    # the role
318
319
320
321
322
323
324
325
326
327
328
329
330
                    role = self.get_chat_request_role(request)
                    for i in range(request.n):
                        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)
331
332
333
                        if (request.stream_options
                                and request.stream_options.include_usage):
                            chunk.usage = None
334
335
336
                        data = chunk.model_dump_json(exclude_unset=True)
                        yield f"data: {data}\n\n"

337
338
                    # Send response to echo the input portion of the
                    # last message
339
340
                    if request.echo:
                        last_msg_content = ""
341
342
343
344
                        if conversation and conversation[-1].get(
                                "content") and conversation[-1].get(
                                    "role") == role:
                            last_msg_content = conversation[-1]["content"]
345
346
347

                        if last_msg_content:
                            for i in range(request.n):
348
349
350
351
352
353
                                choice_data = (
                                    ChatCompletionResponseStreamChoice(
                                        index=i,
                                        delta=DeltaMessage(
                                            content=last_msg_content),
                                        finish_reason=None))
354
355
356
357
358
359
360
                                chunk = ChatCompletionStreamResponse(
                                    id=request_id,
                                    object=chunk_object_type,
                                    created=created_time,
                                    choices=[choice_data],
                                    logprobs=None,
                                    model=model_name)
361
362
363
                                if (request.stream_options and
                                        request.stream_options.include_usage):
                                    chunk.usage = None
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
                                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]:]
                    top_logprobs = output.logprobs[
                        previous_num_tokens[i]:] if output.logprobs else None

                    if request.logprobs:
380
                        logprobs = self._create_chat_logprobs(
381
382
                            token_ids=delta_token_ids,
                            top_logprobs=top_logprobs,
383
                            num_output_top_logprobs=request.top_logprobs,
384
385
386
387
388
389
390
                        )
                    else:
                        logprobs = None

                    delta_text = output.text[len(previous_texts[i]):]
                    previous_texts[i] = output.text
                    previous_num_tokens[i] = len(output.token_ids)
391
392
393
394
395
396
397
398
399
400
401
402

                    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)

403
404
                    if output.finish_reason is None:
                        # Send token-by-token response for each request.n
405

406
407
                        choice_data = ChatCompletionResponseStreamChoice(
                            index=i,
408
                            delta=delta_message,
409
410
411
412
413
414
415
416
                            logprobs=logprobs,
                            finish_reason=None)
                        chunk = ChatCompletionStreamResponse(
                            id=request_id,
                            object=chunk_object_type,
                            created=created_time,
                            choices=[choice_data],
                            model=model_name)
417
418
419
                        if (request.stream_options
                                and request.stream_options.include_usage):
                            chunk.usage = None
420
421
422
423
424
425
426
                        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,
427
                            delta=delta_message,
428
                            logprobs=logprobs,
429
430
                            finish_reason=output.finish_reason,
                            stop_reason=output.stop_reason)
431
432
433
434
435
436
                        chunk = ChatCompletionStreamResponse(
                            id=request_id,
                            object=chunk_object_type,
                            created=created_time,
                            choices=[choice_data],
                            model=model_name)
437
438
439
440
                        if (request.stream_options
                                and request.stream_options.include_usage):
                            chunk.usage = None
                        data = chunk.model_dump_json(exclude_unset=True)
441
442
                        yield f"data: {data}\n\n"
                        finish_reason_sent[i] = True
443

444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
            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"
462

463
464
465
466
        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"
467
468
469
470
        # Send the final done message after all response.n are finished
        yield "data: [DONE]\n\n"

    async def chat_completion_full_generator(
471
        self, request: ChatCompletionRequest, raw_request: Optional[Request],
472
473
474
        result_generator: AsyncIterator[RequestOutput], request_id: str,
        conversation: List[ConversationMessage]
    ) -> Union[ErrorResponse, ChatCompletionResponse]:
475

476
        model_name = self.served_model_names[0]
477
        created_time = int(time.time())
478
        final_res: Optional[RequestOutput] = None
479
480

        async for res in result_generator:
481
            if raw_request is not None and await raw_request.is_disconnected():
482
483
484
485
486
487
488
                # 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

        choices = []
489

490
491
        role = self.get_chat_request_role(request)
        for output in final_res.outputs:
492
493
494
495
            token_ids = output.token_ids
            top_logprobs = output.logprobs

            if request.logprobs:
496
                logprobs = self._create_chat_logprobs(
497
498
                    token_ids=token_ids,
                    top_logprobs=top_logprobs,
499
                    num_output_top_logprobs=request.top_logprobs,
500
501
502
503
                )
            else:
                logprobs = None

504
505
506
507
508
509
510
511
512
513
514
515
516
            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)

517
518
            choice_data = ChatCompletionResponseChoice(
                index=output.index,
519
                message=message,
520
                logprobs=logprobs,
521
                finish_reason=output.finish_reason,
522
                stop_reason=output.stop_reason)
523
524
525
526
            choices.append(choice_data)

        if request.echo:
            last_msg_content = ""
527
528
529
            if conversation and conversation[-1].get(
                    "content") and conversation[-1].get("role") == role:
                last_msg_content = conversation[-1]["content"]
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550

            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,
        )

551
        return response
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599

    def _get_top_logprobs(
            self, logprobs: Dict[int, Logprob],
            top_logprobs: Optional[int]) -> List[ChatCompletionLogProb]:
        return [
            ChatCompletionLogProb(
                token=self._get_decoded_token(p[1], p[0]),
                logprob=max(p[1].logprob, -9999.0),
                bytes=list(
                    self._get_decoded_token(p[1],
                                            p[0]).encode("utf-8",
                                                         errors="replace")))
            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]]],
        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:
                logprobs_content.append(
                    ChatCompletionLogProbsContent(
                        token=self.tokenizer.decode(token_id),
                        bytes=list(
                            self.tokenizer.decode(token_id).encode(
                                "utf-8", errors="replace"))))
            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(
                            step_top_logprobs, num_output_top_logprobs)))

        return ChatCompletionLogProbs(content=logprobs_content)