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

9
from fastapi import Request
10
from openai.types.chat import ChatCompletionContentPartTextParam
11

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

logger = init_logger(__name__)


34
35
36
37
38
39
@final  # So that it should be compatible with Dict[str, str]
class ConversationMessage(TypedDict):
    role: str
    content: str


40
41
42
43
44
@dataclass(frozen=True)
class ChatMessageParseResult:
    messages: List[ConversationMessage]


45
46
47
48
class OpenAIServingChat(OpenAIServing):

    def __init__(self,
                 engine: AsyncLLMEngine,
49
                 model_config: ModelConfig,
50
                 served_model_names: List[str],
51
                 response_role: str,
52
53
                 lora_modules: Optional[List[LoRAModulePath]] = None,
                 chat_template: Optional[str] = None):
54
        super().__init__(engine=engine,
55
                         model_config=model_config,
56
                         served_model_names=served_model_names,
57
                         lora_modules=lora_modules)
58

59
        self.response_role = response_role
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
        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.")
90

91
    def _parse_chat_message_content_parts(
92
        self,
93
94
95
        role: str,
        parts: Iterable[ChatCompletionContentPartParam],
    ) -> ChatMessageParseResult:
96
        texts: List[str] = []
97
98
99
100
101

        for _, part in enumerate(parts):
            part_type = part["type"]
            if part_type == "text":
                text = cast(ChatCompletionContentPartTextParam, part)["text"]
102
103
104

                texts.append(text)
            else:
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
                raise NotImplementedError(f"Unknown part type: {part_type}")

        messages = [ConversationMessage(role=role, content="\n".join(texts))]

        return ChatMessageParseResult(messages=messages)

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

        if content is None:
            return ChatMessageParseResult(messages=[])
        if isinstance(content, str):
            messages = [ConversationMessage(role=role, content=content)]
            return ChatMessageParseResult(messages=messages)
123

124
        return self._parse_chat_message_content_parts(role, content)
125

126
    async def create_chat_completion(
127
128
129
        self,
        request: ChatCompletionRequest,
        raw_request: Optional[Request] = None
130
131
132
133
    ) -> Union[ErrorResponse, AsyncGenerator[str, None],
               ChatCompletionResponse]:
        """Completion API similar to OpenAI's API.

134
135
136
        See https://platform.openai.com/docs/api-reference/chat/create
        for the API specification. This API mimics the OpenAI
        ChatCompletion API.
137

138
        NOTE: Currently we do not support the following feature:
139
140
141
142
143
144
145
            - 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:
146
147
            conversation: List[ConversationMessage] = []

148
149
            for msg in request.messages:
                parsed_msg = self._parse_chat_message_content(msg)
150

151
                conversation.extend(parsed_msg.messages)
152

153
            prompt = self.tokenizer.apply_chat_template(
154
                conversation=conversation,
155
                tokenize=False,
156
157
                add_generation_prompt=request.add_generation_prompt,
            )
158
        except Exception as e:
159
            logger.error("Error in applying chat template from request: %s", e)
160
161
162
163
            return self.create_error_response(str(e))

        request_id = f"cmpl-{random_uuid()}"
        try:
164
165
            # Tokenize/detokenize depending on prompt format (string/token list)
            prompt_ids, prompt_text = self._validate_prompt_and_tokenize(
166
167
168
                request,
                prompt=prompt,
                add_special_tokens=request.add_special_tokens)
169
            sampling_params = request.to_sampling_params()
170
            lora_request = self._maybe_get_lora(request)
171
            decoding_config = await self.engine.get_decoding_config()
172
173
            guided_decoding_backend = request.guided_decoding_backend \
                or decoding_config.guided_decoding_backend
174
175
            guided_decode_logits_processor = (
                await get_guided_decoding_logits_processor(
176
177
                    guided_decoding_backend, request, await
                    self.engine.get_tokenizer()))
178
179
180
181
182
            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)
183
184
185
        except ValueError as e:
            return self.create_error_response(str(e))

186
187
188
189
190
191
192
193
194
        result_generator = self.engine.generate(
            {
                "prompt": prompt_text,
                "prompt_token_ids": prompt_ids
            },
            sampling_params,
            request_id,
            lora_request,
        )
195
196
197
        # Streaming response
        if request.stream:
            return self.chat_completion_stream_generator(
198
                request, result_generator, request_id, conversation)
199
        else:
200
201
            try:
                return await self.chat_completion_full_generator(
202
203
                    request, raw_request, result_generator, request_id,
                    conversation)
204
205
206
            except ValueError as e:
                # TODO: Use a vllm-specific Validation Error
                return self.create_error_response(str(e))
207
208
209
210
211

    def get_chat_request_role(self, request: ChatCompletionRequest) -> str:
        if request.add_generation_prompt:
            return self.response_role
        else:
212
            return request.messages[-1]["role"]
213
214
215

    async def chat_completion_stream_generator(
            self, request: ChatCompletionRequest,
216
217
218
            result_generator: AsyncIterator[RequestOutput], request_id: str,
            conversation: List[ConversationMessage]
    ) -> AsyncGenerator[str, None]:
219
        model_name = self.served_model_names[0]
220
        created_time = int(time.time())
221
        chunk_object_type = "chat.completion.chunk"
222
        first_iteration = True
223
224

        # Send response for each token for each request.n (index)
225
        assert request.n is not None
226
227
228
        previous_texts = [""] * request.n
        previous_num_tokens = [0] * request.n
        finish_reason_sent = [False] * request.n
229
230
231
232
233
234
        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:
235
236
                    # Send first response for each request.n (index) with
                    # the role
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
                    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)
                        data = chunk.model_dump_json(exclude_unset=True)
                        yield f"data: {data}\n\n"

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

                        if last_msg_content:
                            for i in range(request.n):
264
265
266
267
268
269
                                choice_data = (
                                    ChatCompletionResponseStreamChoice(
                                        index=i,
                                        delta=DeltaMessage(
                                            content=last_msg_content),
                                        finish_reason=None))
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
                                chunk = ChatCompletionStreamResponse(
                                    id=request_id,
                                    object=chunk_object_type,
                                    created=created_time,
                                    choices=[choice_data],
                                    logprobs=None,
                                    model=model_name)
                                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:
293
                        logprobs = self._create_chat_logprobs(
294
295
                            token_ids=delta_token_ids,
                            top_logprobs=top_logprobs,
296
                            num_output_top_logprobs=request.top_logprobs,
297
298
299
300
301
302
303
                        )
                    else:
                        logprobs = None

                    delta_text = output.text[len(previous_texts[i]):]
                    previous_texts[i] = output.text
                    previous_num_tokens[i] = len(output.token_ids)
304
305
306
307
308
309
310
311
312
313
314
315

                    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)

316
317
                    if output.finish_reason is None:
                        # Send token-by-token response for each request.n
318

319
320
                        choice_data = ChatCompletionResponseStreamChoice(
                            index=i,
321
                            delta=delta_message,
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
                            logprobs=logprobs,
                            finish_reason=None)
                        chunk = ChatCompletionStreamResponse(
                            id=request_id,
                            object=chunk_object_type,
                            created=created_time,
                            choices=[choice_data],
                            model=model_name)
                        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)
                        final_usage = UsageInfo(
                            prompt_tokens=prompt_tokens,
                            completion_tokens=previous_num_tokens[i],
                            total_tokens=prompt_tokens +
                            previous_num_tokens[i],
                        )
                        choice_data = ChatCompletionResponseStreamChoice(
                            index=i,
343
                            delta=delta_message,
344
                            logprobs=logprobs,
345
346
                            finish_reason=output.finish_reason,
                            stop_reason=output.stop_reason)
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
                        chunk = ChatCompletionStreamResponse(
                            id=request_id,
                            object=chunk_object_type,
                            created=created_time,
                            choices=[choice_data],
                            model=model_name)
                        if final_usage is not None:
                            chunk.usage = final_usage
                        data = chunk.model_dump_json(exclude_unset=True,
                                                     exclude_none=True)
                        yield f"data: {data}\n\n"
                        finish_reason_sent[i] = True
        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"
363
364
365
366
        # Send the final done message after all response.n are finished
        yield "data: [DONE]\n\n"

    async def chat_completion_full_generator(
367
        self, request: ChatCompletionRequest, raw_request: Optional[Request],
368
369
370
        result_generator: AsyncIterator[RequestOutput], request_id: str,
        conversation: List[ConversationMessage]
    ) -> Union[ErrorResponse, ChatCompletionResponse]:
371

372
        model_name = self.served_model_names[0]
373
        created_time = int(time.time())
374
        final_res: Optional[RequestOutput] = None
375
376

        async for res in result_generator:
377
            if raw_request is not None and await raw_request.is_disconnected():
378
379
380
381
382
383
384
                # 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 = []
385

386
387
        role = self.get_chat_request_role(request)
        for output in final_res.outputs:
388
389
390
391
            token_ids = output.token_ids
            top_logprobs = output.logprobs

            if request.logprobs:
392
                logprobs = self._create_chat_logprobs(
393
394
                    token_ids=token_ids,
                    top_logprobs=top_logprobs,
395
                    num_output_top_logprobs=request.top_logprobs,
396
397
398
399
                )
            else:
                logprobs = None

400
401
402
403
404
405
406
407
408
409
410
411
412
            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)

413
414
            choice_data = ChatCompletionResponseChoice(
                index=output.index,
415
                message=message,
416
                logprobs=logprobs,
417
                finish_reason=output.finish_reason,
418
                stop_reason=output.stop_reason)
419
420
421
422
            choices.append(choice_data)

        if request.echo:
            last_msg_content = ""
423
424
425
            if conversation and conversation[-1].get(
                    "content") and conversation[-1].get("role") == role:
                last_msg_content = conversation[-1]["content"]
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446

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

447
        return response
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495

    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)