serving_chat.py 19.4 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
17
    ChatCompletionContentPartParam, ChatCompletionLogProb,
    ChatCompletionLogProbs, ChatCompletionLogProbsContent,
    ChatCompletionMessageParam, ChatCompletionRequest, ChatCompletionResponse,
18
19
20
    ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice,
    ChatCompletionStreamResponse, ChatMessage, DeltaMessage, ErrorResponse,
    UsageInfo)
21
22
from vllm.entrypoints.openai.serving_engine import (LoRAModulePath,
                                                    OpenAIServing)
23
from vllm.logger import init_logger
24
25
from vllm.model_executor.guided_decoding import (
    get_guided_decoding_logits_processor)
26
from vllm.outputs import RequestOutput
27
from vllm.sequence import Logprob
28
from vllm.utils import random_uuid
29
30
31
32

logger = init_logger(__name__)


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


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


44
45
46
47
class OpenAIServingChat(OpenAIServing):

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

58
        self.response_role = response_role
59
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
        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.")
89

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

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

                texts.append(text)
            else:
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
                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)
122

123
        return self._parse_chat_message_content_parts(role, content)
124

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

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

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

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

150
                conversation.extend(parsed_msg.messages)
151

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

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

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

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

    async def chat_completion_stream_generator(
            self, request: ChatCompletionRequest,
213
214
215
            result_generator: AsyncIterator[RequestOutput], request_id: str,
            conversation: List[ConversationMessage]
    ) -> 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
        assert request.n is not None
223
224
225
        previous_texts = [""] * request.n
        previous_num_tokens = [0] * request.n
        finish_reason_sent = [False] * request.n
226
227
228
229
230
231
        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:
232
233
                    # Send first response for each request.n (index) with
                    # the role
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
                    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"

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

                        if last_msg_content:
                            for i in range(request.n):
261
262
263
264
265
266
                                choice_data = (
                                    ChatCompletionResponseStreamChoice(
                                        index=i,
                                        delta=DeltaMessage(
                                            content=last_msg_content),
                                        finish_reason=None))
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
                                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:
290
                        logprobs = self._create_chat_logprobs(
291
292
                            token_ids=delta_token_ids,
                            top_logprobs=top_logprobs,
293
                            num_output_top_logprobs=request.top_logprobs,
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
                        )
                    else:
                        logprobs = None

                    delta_text = output.text[len(previous_texts[i]):]
                    previous_texts[i] = output.text
                    previous_num_tokens[i] = len(output.token_ids)
                    if output.finish_reason is None:
                        # Send token-by-token response for each request.n
                        choice_data = ChatCompletionResponseStreamChoice(
                            index=i,
                            delta=DeltaMessage(content=delta_text),
                            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,
                            delta=DeltaMessage(content=delta_text),
                            logprobs=logprobs,
329
330
                            finish_reason=output.finish_reason,
                            stop_reason=output.stop_reason)
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
                        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"
347
348
349
350
        # Send the final done message after all response.n are finished
        yield "data: [DONE]\n\n"

    async def chat_completion_full_generator(
351
        self, request: ChatCompletionRequest, raw_request: Optional[Request],
352
353
354
        result_generator: AsyncIterator[RequestOutput], request_id: str,
        conversation: List[ConversationMessage]
    ) -> Union[ErrorResponse, ChatCompletionResponse]:
355

356
        model_name = self.served_model_names[0]
357
        created_time = int(time.time())
358
        final_res: Optional[RequestOutput] = None
359
360

        async for res in result_generator:
361
            if raw_request is not None and await raw_request.is_disconnected():
362
363
364
365
366
367
368
                # 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 = []
369

370
371
        role = self.get_chat_request_role(request)
        for output in final_res.outputs:
372
373
374
375
            token_ids = output.token_ids
            top_logprobs = output.logprobs

            if request.logprobs:
376
                logprobs = self._create_chat_logprobs(
377
378
                    token_ids=token_ids,
                    top_logprobs=top_logprobs,
379
                    num_output_top_logprobs=request.top_logprobs,
380
381
382
383
                )
            else:
                logprobs = None

384
385
386
            choice_data = ChatCompletionResponseChoice(
                index=output.index,
                message=ChatMessage(role=role, content=output.text),
387
                logprobs=logprobs,
388
                finish_reason=output.finish_reason,
389
                stop_reason=output.stop_reason)
390
391
392
393
            choices.append(choice_data)

        if request.echo:
            last_msg_content = ""
394
395
396
            if conversation and conversation[-1].get(
                    "content") and conversation[-1].get("role") == role:
                last_msg_content = conversation[-1]["content"]
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417

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

418
        return response
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466

    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)