adapter.py 61.1 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
# Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
14
"""Conversion between OpenAI APIs and native SRT APIs"""
Liangsheng Yin's avatar
Liangsheng Yin committed
15

16
import asyncio
17
import json
18
import logging
19
import os
20
21
import time
import uuid
22
from http import HTTPStatus
23
from typing import Dict, List
24

25
from fastapi import HTTPException, Request, UploadFile
26
from fastapi.responses import ORJSONResponse, StreamingResponse
27
from pydantic import ValidationError
28

29
30
from sglang.lang.chat_template import get_chat_template_by_model_path

31
32
33
34
35
36
37
try:
    from outlines.fsm.json_schema import convert_json_schema_to_str
except ImportError:
    # Before outlines 0.0.47, convert_json_schema_to_str is under
    # outlines.integrations.utils
    from outlines.integrations.utils import convert_json_schema_to_str

38
39
40
41
42
43
44
from sglang.srt.conversation import (
    Conversation,
    SeparatorStyle,
    chat_template_exists,
    generate_chat_conv,
    register_conv_template,
)
YAMY's avatar
YAMY committed
45
from sglang.srt.function_call_parser import TOOLS_TAG_LIST, FunctionCallParser
Ying Sheng's avatar
Ying Sheng committed
46
from sglang.srt.managers.io_struct import EmbeddingReqInput, GenerateReqInput
Mingyi's avatar
Mingyi committed
47
from sglang.srt.openai_api.protocol import (
48
49
    BatchRequest,
    BatchResponse,
50
51
52
53
54
    ChatCompletionRequest,
    ChatCompletionResponse,
    ChatCompletionResponseChoice,
    ChatCompletionResponseStreamChoice,
    ChatCompletionStreamResponse,
55
    ChatCompletionTokenLogprob,
56
    ChatMessage,
57
    ChoiceLogprobs,
58
59
60
61
62
63
    CompletionRequest,
    CompletionResponse,
    CompletionResponseChoice,
    CompletionResponseStreamChoice,
    CompletionStreamResponse,
    DeltaMessage,
Ying Sheng's avatar
Ying Sheng committed
64
    EmbeddingObject,
65
66
    EmbeddingRequest,
    EmbeddingResponse,
67
    ErrorResponse,
68
    FileDeleteResponse,
69
70
    FileRequest,
    FileResponse,
Tanjiro's avatar
Tanjiro committed
71
    FunctionResponse,
72
    LogProbs,
Tanjiro's avatar
Tanjiro committed
73
    ToolCall,
74
    TopLogprob,
75
76
    UsageInfo,
)
77
from sglang.utils import get_exception_traceback
78

79
80
logger = logging.getLogger(__name__)

81
82
chat_template_name = None

Liangsheng Yin's avatar
Liangsheng Yin committed
83

84
85
86
87
88
89
90
91
92
93
class FileMetadata:
    def __init__(self, filename: str, purpose: str):
        self.filename = filename
        self.purpose = purpose


# In-memory storage for batch jobs and files
batch_storage: Dict[str, BatchResponse] = {}
file_id_request: Dict[str, FileMetadata] = {}
file_id_response: Dict[str, FileResponse] = {}
94
# map file id to file path in SGLang backend
95
96
97
98
99
100
file_id_storage: Dict[str, str] = {}

# backend storage directory
storage_dir = None


101
102
103
def create_error_response(
    message: str,
    err_type: str = "BadRequestError",
104
105
106
    status_code: HTTPStatus = HTTPStatus.BAD_REQUEST,
):
    error = ErrorResponse(message=message, type=err_type, code=status_code.value)
107
    return ORJSONResponse(content=error.model_dump(), status_code=error.code)
108
109
110
111
112


def create_streaming_error_response(
    message: str,
    err_type: str = "BadRequestError",
113
114
115
    status_code: HTTPStatus = HTTPStatus.BAD_REQUEST,
) -> str:
    error = ErrorResponse(message=message, type=err_type, code=status_code.value)
116
117
118
119
    json_str = json.dumps({"error": error.model_dump()})
    return json_str


120
def load_chat_template_for_openai_api(tokenizer_manager, chat_template_arg, model_path):
121
122
    global chat_template_name

123
124
125
    logger.info(
        f"Use chat template for the OpenAI-compatible API server: {chat_template_arg}"
    )
126

127
128
129
130
131
132
    if not chat_template_exists(chat_template_arg):
        if not os.path.exists(chat_template_arg):
            raise RuntimeError(
                f"Chat template {chat_template_arg} is not a built-in template name "
                "or a valid chat template file path."
            )
133
134
135
136
137
        if chat_template_arg.endswith(".jinja"):
            with open(chat_template_arg, "r") as f:
                chat_template = "".join(f.readlines()).strip("\n")
            tokenizer_manager.tokenizer.chat_template = chat_template.replace(
                "\\n", "\n"
138
            )
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
            chat_template_name = None
        else:
            assert chat_template_arg.endswith(
                ".json"
            ), "unrecognized format of chat template file"
            with open(chat_template_arg, "r") as filep:
                template = json.load(filep)
                try:
                    sep_style = SeparatorStyle[template["sep_style"]]
                except KeyError:
                    raise ValueError(
                        f"Unknown separator style: {template['sep_style']}"
                    ) from None
                register_conv_template(
                    Conversation(
                        name=template["name"],
                        system_template=template["system"] + "\n{system_message}",
                        system_message=template.get("system_message", ""),
                        roles=(template["user"], template["assistant"]),
                        sep_style=sep_style,
                        sep=template.get("sep", "\n"),
                        stop_str=template["stop_str"],
                    ),
                    override=True,
                )
            chat_template_name = template["name"]
165
166
167
    else:
        chat_template_name = chat_template_arg

168
169
170
171
172
173
174
175
176
177
178
179
    # check chat-template
    chat_template = get_chat_template_by_model_path(model_path)
    if chat_template is not None:
        official_chat_template = chat_template.name
        used_chat_template = chat_template_name
        if official_chat_template != used_chat_template:
            logger.warning(
                f"Using a chat_template: '{used_chat_template}', "
                f"which is different from official chat template: '{official_chat_template}', "
                f"This discrepancy may lead to performance degradation."
            )

180

181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
async def v1_files_create(file: UploadFile, purpose: str, file_storage_pth: str = None):
    try:
        global storage_dir
        if file_storage_pth:
            storage_dir = file_storage_pth
        # Read the file content
        file_content = await file.read()

        # Create an instance of RequestBody
        request_body = FileRequest(file=file_content, purpose=purpose)

        # Save the file to the sglang_oai_storage directory
        os.makedirs(storage_dir, exist_ok=True)
        file_id = f"backend_input_file-{uuid.uuid4()}"
        filename = f"{file_id}.jsonl"
        file_path = os.path.join(storage_dir, filename)

        with open(file_path, "wb") as f:
            f.write(request_body.file)

        # add info to global file map
        file_id_request[file_id] = FileMetadata(filename=file.filename, purpose=purpose)
        file_id_storage[file_id] = file_path

        # Return the response in the required format
        response = FileResponse(
            id=file_id,
            bytes=len(request_body.file),
            created_at=int(time.time()),
            filename=file.filename,
            purpose=request_body.purpose,
        )
        file_id_response[file_id] = response

        return response
    except ValidationError as e:
        return {"error": "Invalid input", "details": e.errors()}


220
221
222
223
224
225
226
227
228
229
230
231
232
233
async def v1_delete_file(file_id: str):
    # Retrieve the file job from the in-memory storage
    file_response = file_id_response.get(file_id)
    if file_response is None:
        raise HTTPException(status_code=404, detail="File not found")
    file_path = file_id_storage.get(file_id)
    if file_path is None:
        raise HTTPException(status_code=404, detail="File not found")
    os.remove(file_path)
    del file_id_response[file_id]
    del file_id_storage[file_id]
    return FileDeleteResponse(id=file_id, deleted=True)


234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
async def v1_batches(tokenizer_manager, raw_request: Request):
    try:
        body = await raw_request.json()

        batch_request = BatchRequest(**body)

        batch_id = f"batch_{uuid.uuid4()}"

        # Create an instance of BatchResponse
        batch_response = BatchResponse(
            id=batch_id,
            endpoint=batch_request.endpoint,
            input_file_id=batch_request.input_file_id,
            completion_window=batch_request.completion_window,
            created_at=int(time.time()),
            metadata=batch_request.metadata,
        )

        batch_storage[batch_id] = batch_response

        # Start processing the batch asynchronously
        asyncio.create_task(process_batch(tokenizer_manager, batch_id, batch_request))

        # Return the initial batch_response
        return batch_response

    except ValidationError as e:
        return {"error": "Invalid input", "details": e.errors()}
    except Exception as e:
        return {"error": str(e)}


async def process_batch(tokenizer_manager, batch_id: str, batch_request: BatchRequest):
    try:
        # Update the batch status to "in_progress"
        batch_storage[batch_id].status = "in_progress"
        batch_storage[batch_id].in_progress_at = int(time.time())

        # Retrieve the input file content
        input_file_request = file_id_request.get(batch_request.input_file_id)
        if not input_file_request:
            raise ValueError("Input file not found")

        # Parse the JSONL file and process each request
        input_file_path = file_id_storage.get(batch_request.input_file_id)
        with open(input_file_path, "r", encoding="utf-8") as f:
            lines = f.readlines()

        total_requests = len(lines)
        completed_requests = 0
        failed_requests = 0

        all_ret = []
        end_point = batch_storage[batch_id].endpoint
        file_request_list = []
        all_requests = []
290
        request_ids = []
291
292
293
294
        for line in lines:
            request_data = json.loads(line)
            file_request_list.append(request_data)
            body = request_data["body"]
295
            request_ids.append(request_data["custom_id"])
296
297
298
299
300
301

            # Although streaming is supported for standalone completions, it is not supported in
            # batch mode (multiple completions in single request).
            if body.get("stream", False):
                raise ValueError("Streaming requests are not supported in batch mode")

302
303
304
305
            if end_point == "/v1/chat/completions":
                all_requests.append(ChatCompletionRequest(**body))
            elif end_point == "/v1/completions":
                all_requests.append(CompletionRequest(**body))
306

307
308
        if end_point == "/v1/chat/completions":
            adapted_request, request = v1_chat_generate_request(
309
                all_requests, tokenizer_manager, request_ids=request_ids
310
311
            )
        elif end_point == "/v1/completions":
312
313
314
315
            adapted_request, request = v1_generate_request(
                all_requests, request_ids=request_ids
            )

316
317
318
319
320
        try:
            ret = await tokenizer_manager.generate_request(adapted_request).__anext__()
            if not isinstance(ret, list):
                ret = [ret]
            if end_point == "/v1/chat/completions":
321
322
323
324
325
                responses = v1_chat_generate_response(
                    request,
                    ret,
                    to_file=True,
                    cache_report=tokenizer_manager.server_args.enable_cache_report,
YAMY's avatar
YAMY committed
326
                    tool_call_parser=tokenizer_manager.server_args.tool_call_parser,
327
                )
328
            else:
yichuan~'s avatar
yichuan~ committed
329
330
331
                responses = v1_generate_response(
                    request, ret, tokenizer_manager, to_file=True
                )
332
333

        except Exception as e:
334
335
            logger.error(f"error: {get_exception_traceback()}")
            responses = []
336
337
338
339
340
341
342
343
344
345
            error_json = {
                "id": f"batch_req_{uuid.uuid4()}",
                "custom_id": request_data.get("custom_id"),
                "response": None,
                "error": {"message": str(e)},
            }
            all_ret.append(error_json)
            failed_requests += len(file_request_list)

        for idx, response in enumerate(responses):
346
            # the batch_req here can be changed to be named within a batch granularity
347
348
349
350
351
352
353
354
            response_json = {
                "id": f"batch_req_{uuid.uuid4()}",
                "custom_id": file_request_list[idx].get("custom_id"),
                "response": response,
                "error": None,
            }
            all_ret.append(response_json)
            completed_requests += 1
355

356
357
358
359
360
361
362
363
364
365
366
367
        # Write results to a new file
        output_file_id = f"backend_result_file-{uuid.uuid4()}"
        global storage_dir
        output_file_path = os.path.join(storage_dir, f"{output_file_id}.jsonl")
        with open(output_file_path, "w", encoding="utf-8") as f:
            for ret in all_ret:
                f.write(json.dumps(ret) + "\n")

        # Update batch response with output file information
        retrieve_batch = batch_storage[batch_id]
        retrieve_batch.output_file_id = output_file_id
        file_id_storage[output_file_id] = output_file_path
368
369
370
371
372
373
374
        file_id_response[output_file_id] = FileResponse(
            id=output_file_id,
            bytes=os.path.getsize(output_file_path),
            created_at=int(time.time()),
            filename=f"{output_file_id}.jsonl",
            purpose="batch_result",
        )
375
376
377
378
379
380
381
382
383
384
        # Update batch status to "completed"
        retrieve_batch.status = "completed"
        retrieve_batch.completed_at = int(time.time())
        retrieve_batch.request_counts = {
            "total": total_requests,
            "completed": completed_requests,
            "failed": failed_requests,
        }

    except Exception as e:
385
        logger.error(f"error: {e}")
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
        # Update batch status to "failed"
        retrieve_batch = batch_storage[batch_id]
        retrieve_batch.status = "failed"
        retrieve_batch.failed_at = int(time.time())
        retrieve_batch.errors = {"message": str(e)}


async def v1_retrieve_batch(batch_id: str):
    # Retrieve the batch job from the in-memory storage
    batch_response = batch_storage.get(batch_id)
    if batch_response is None:
        raise HTTPException(status_code=404, detail="Batch not found")

    return batch_response


402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
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
467
async def v1_cancel_batch(tokenizer_manager, batch_id: str):
    # Retrieve the batch job from the in-memory storage
    batch_response = batch_storage.get(batch_id)
    if batch_response is None:
        raise HTTPException(status_code=404, detail="Batch not found")

    # Only do cancal when status is "validating" or "in_progress"
    if batch_response.status in ["validating", "in_progress"]:
        # Start cancelling the batch asynchronously
        asyncio.create_task(
            cancel_batch(
                tokenizer_manager=tokenizer_manager,
                batch_id=batch_id,
                input_file_id=batch_response.input_file_id,
            )
        )

        # Update batch status to "cancelling"
        batch_response.status = "cancelling"

        return batch_response
    else:
        raise HTTPException(
            status_code=500,
            detail=f"Current status is {batch_response.status}, no need to cancel",
        )


async def cancel_batch(tokenizer_manager, batch_id: str, input_file_id: str):
    try:
        # Update the batch status to "cancelling"
        batch_storage[batch_id].status = "cancelling"

        # Retrieve the input file content
        input_file_request = file_id_request.get(input_file_id)
        if not input_file_request:
            raise ValueError("Input file not found")

        # Parse the JSONL file and process each request
        input_file_path = file_id_storage.get(input_file_id)
        with open(input_file_path, "r", encoding="utf-8") as f:
            lines = f.readlines()

        file_request_list = []
        request_ids = []
        for line in lines:
            request_data = json.loads(line)
            file_request_list.append(request_data)
            request_ids.append(request_data["custom_id"])

        # Cancel requests by request_ids
        for rid in request_ids:
            tokenizer_manager.abort_request(rid=rid)

        retrieve_batch = batch_storage[batch_id]
        retrieve_batch.status = "cancelled"

    except Exception as e:
        logger.error("error in SGLang:", e)
        # Update batch status to "failed"
        retrieve_batch = batch_storage[batch_id]
        retrieve_batch.status = "failed"
        retrieve_batch.failed_at = int(time.time())
        retrieve_batch.errors = {"message": str(e)}


468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
async def v1_retrieve_file(file_id: str):
    # Retrieve the batch job from the in-memory storage
    file_response = file_id_response.get(file_id)
    if file_response is None:
        raise HTTPException(status_code=404, detail="File not found")
    return file_response


async def v1_retrieve_file_content(file_id: str):
    file_pth = file_id_storage.get(file_id)
    if not file_pth or not os.path.exists(file_pth):
        raise HTTPException(status_code=404, detail="File not found")

    def iter_file():
        with open(file_pth, mode="rb") as file_like:
            yield from file_like

    return StreamingResponse(iter_file(), media_type="application/octet-stream")


488
489
490
def v1_generate_request(
    all_requests: List[CompletionRequest], request_ids: List[str] = None
):
491
492
493
494
495
496
497
498
499
500
501
    if len(all_requests) > 1:
        first_prompt_type = type(all_requests[0].prompt)
        for request in all_requests:
            assert (
                type(request.prompt) is first_prompt_type
            ), "All prompts must be of the same type in file input settings"
            if request.n > 1:
                raise ValueError(
                    "Parallel sampling is not supported for completions from files"
                )

502
503
    prompts = []
    sampling_params_list = []
504
    return_logprobs = []
505
    logprob_start_lens = []
506
    top_logprobs_nums = []
507
    lora_paths = []
yichuan~'s avatar
yichuan~ committed
508

509
    for request in all_requests:
510
        # NOTE: with openai API, the prompt's logprobs are always not computed
511
        if request.echo and request.logprobs:
512
            logger.warning(
513
                "Echo is not compatible with logprobs. "
514
                "To compute logprobs of input prompt, please use the native /generate API."
515
516
517
            )

        prompts.append(request.prompt)
518
        lora_paths.append(request.lora_path)
519
520
521
522
        if request.echo and request.logprobs:
            current_logprob_start_len = 0
        else:
            current_logprob_start_len = -1
523
524
525
526
527
528
529
530
        sampling_params_list.append(
            {
                "temperature": request.temperature,
                "max_new_tokens": request.max_tokens,
                "min_new_tokens": request.min_tokens,
                "stop": request.stop,
                "stop_token_ids": request.stop_token_ids,
                "top_p": request.top_p,
531
532
                "top_k": request.top_k,
                "min_p": request.min_p,
533
534
535
536
537
                "presence_penalty": request.presence_penalty,
                "frequency_penalty": request.frequency_penalty,
                "repetition_penalty": request.repetition_penalty,
                "regex": request.regex,
                "json_schema": request.json_schema,
538
                "ebnf": request.ebnf,
539
540
                "n": request.n,
                "no_stop_trim": request.no_stop_trim,
541
542
                "ignore_eos": request.ignore_eos,
                "skip_special_tokens": request.skip_special_tokens,
543
544
            }
        )
545
        return_logprobs.append(request.logprobs is not None)
546
        logprob_start_lens.append(current_logprob_start_len)
547
548
549
        top_logprobs_nums.append(
            request.logprobs if request.logprobs is not None else 0
        )
550
551

    if len(all_requests) == 1:
552
553
554
555
        if isinstance(prompts[0], str) or isinstance(prompts[0][0], str):
            prompt_kwargs = {"text": prompts[0]}
        else:
            prompt_kwargs = {"input_ids": prompts[0]}
556
        sampling_params_list = sampling_params_list[0]
557
        return_logprobs = return_logprobs[0]
558
        logprob_start_lens = logprob_start_lens[0]
559
        top_logprobs_nums = top_logprobs_nums[0]
560
        lora_paths = lora_paths[0]
561
    else:
562
        if isinstance(prompts[0], str) or isinstance(prompts[0][0], str):
563
564
565
            prompt_kwargs = {"text": prompts}
        else:
            prompt_kwargs = {"input_ids": prompts}
yichuan~'s avatar
yichuan~ committed
566

567
    adapted_request = GenerateReqInput(
568
        **prompt_kwargs,
569
        sampling_params=sampling_params_list,
570
571
        return_logprob=return_logprobs,
        top_logprobs_num=top_logprobs_nums,
572
        logprob_start_len=logprob_start_lens,
573
        return_text_in_logprobs=True,
574
        stream=all_requests[0].stream,
575
        rid=request_ids,
576
        lora_path=lora_paths,
577
    )
yichuan~'s avatar
yichuan~ committed
578

579
    return adapted_request, all_requests if len(all_requests) > 1 else all_requests[0]
580
581


yichuan~'s avatar
yichuan~ committed
582
def v1_generate_response(request, ret, tokenizer_manager, to_file=False):
583
584
585
    choices = []
    echo = False

yichuan~'s avatar
yichuan~ committed
586
    if (not isinstance(request, list)) and request.echo:
587
        # TODO: handle the case propmt is token ids
yichuan~'s avatar
yichuan~ committed
588
589
        if isinstance(request.prompt, list) and isinstance(request.prompt[0], str):
            # for the case of multiple str prompts
590
            prompts = request.prompt
yichuan~'s avatar
yichuan~ committed
591
592
593
594
595
596
597
598
599
600
601
602
603
        elif isinstance(request.prompt, list) and isinstance(request.prompt[0], list):
            # for the case of multiple token ids prompts
            prompts = [
                tokenizer_manager.tokenizer.decode(prompt, skip_special_tokens=True)
                for prompt in request.prompt
            ]
        elif isinstance(request.prompt, list) and isinstance(request.prompt[0], int):
            # for the case of single token ids prompt
            prompts = [
                tokenizer_manager.tokenizer.decode(
                    request.prompt, skip_special_tokens=True
                )
            ]
604
        else:
yichuan~'s avatar
yichuan~ committed
605
            # for the case of single str prompt
606
607
608
609
610
            prompts = [request.prompt]
        echo = True

    for idx, ret_item in enumerate(ret):
        text = ret_item["text"]
yichuan~'s avatar
yichuan~ committed
611
        if isinstance(request, list) and request[idx].echo:
612
613
            echo = True
            text = request[idx].prompt + text
614
        if echo and not isinstance(request, list):
yichuan~'s avatar
yichuan~ committed
615
616
            prompt_index = idx // request.n
            text = prompts[prompt_index] + text
617
618

        logprobs = False
619
        if isinstance(request, list) and request[idx].logprobs is not None:
620
            logprobs = True
621
        elif (not isinstance(request, list)) and request.logprobs is not None:
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
            logprobs = True
        if logprobs:
            if echo:
                input_token_logprobs = ret_item["meta_info"]["input_token_logprobs"]
                input_top_logprobs = ret_item["meta_info"]["input_top_logprobs"]
            else:
                input_token_logprobs = None
                input_top_logprobs = None

            logprobs = to_openai_style_logprobs(
                input_token_logprobs=input_token_logprobs,
                input_top_logprobs=input_top_logprobs,
                output_token_logprobs=ret_item["meta_info"]["output_token_logprobs"],
                output_top_logprobs=ret_item["meta_info"]["output_top_logprobs"],
            )
        else:
            logprobs = None

640
641
        finish_reason = ret_item["meta_info"]["finish_reason"]

642
        if to_file:
643
            # to make the choise data json serializable
644
645
646
647
            choice_data = {
                "index": 0,
                "text": text,
                "logprobs": logprobs,
648
649
650
651
652
                "finish_reason": (finish_reason["type"] if finish_reason else ""),
                "matched_stop": (
                    finish_reason["matched"]
                    if finish_reason and "matched" in finish_reason
                    else None
653
                ),
654
655
656
657
658
659
            }
        else:
            choice_data = CompletionResponseChoice(
                index=idx,
                text=text,
                logprobs=logprobs,
660
661
662
663
664
                finish_reason=(finish_reason["type"] if finish_reason else ""),
                matched_stop=(
                    finish_reason["matched"]
                    if finish_reason and "matched" in finish_reason
                    else None
665
                ),
666
667
668
669
670
671
672
673
674
675
676
            )

        choices.append(choice_data)

    if to_file:
        responses = []
        for i, choice in enumerate(choices):
            response = {
                "status_code": 200,
                "request_id": ret[i]["meta_info"]["id"],
                "body": {
677
                    # remain the same but if needed we can change that
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
                    "id": ret[i]["meta_info"]["id"],
                    "object": "text_completion",
                    "created": int(time.time()),
                    "model": request[i].model,
                    "choices": choice,
                    "usage": {
                        "prompt_tokens": ret[i]["meta_info"]["prompt_tokens"],
                        "completion_tokens": ret[i]["meta_info"]["completion_tokens"],
                        "total_tokens": ret[i]["meta_info"]["prompt_tokens"]
                        + ret[i]["meta_info"]["completion_tokens"],
                    },
                    "system_fingerprint": None,
                },
            }
            responses.append(response)
        return responses
    else:
695
696
697
        prompt_tokens = sum(
            ret[i]["meta_info"]["prompt_tokens"] for i in range(0, len(ret), request.n)
        )
698
699
700
701
702
703
        completion_tokens = sum(item["meta_info"]["completion_tokens"] for item in ret)
        response = CompletionResponse(
            id=ret[0]["meta_info"]["id"],
            model=request.model,
            choices=choices,
            usage=UsageInfo(
yichuan~'s avatar
yichuan~ committed
704
                prompt_tokens=prompt_tokens,
705
                completion_tokens=completion_tokens,
yichuan~'s avatar
yichuan~ committed
706
                total_tokens=prompt_tokens + completion_tokens,
707
708
709
710
711
712
713
714
715
            ),
        )
    return response


async def v1_completions(tokenizer_manager, raw_request: Request):
    request_json = await raw_request.json()
    all_requests = [CompletionRequest(**request_json)]
    adapted_request, request = v1_generate_request(all_requests)
716
717
718
719

    if adapted_request.stream:

        async def generate_stream_resp():
720
721
722
723
            stream_buffers = {}
            n_prev_tokens = {}
            prompt_tokens = {}
            completion_tokens = {}
724
725
            try:
                async for content in tokenizer_manager.generate_request(
726
727
                    adapted_request, raw_request
                ):
728
                    index = content.get("index", 0)
729
730
731
732

                    stream_buffer = stream_buffers.get(index, "")
                    n_prev_token = n_prev_tokens.get(index, 0)

733
                    text = content["text"]
734
735
                    prompt_tokens[index] = content["meta_info"]["prompt_tokens"]
                    completion_tokens[index] = content["meta_info"]["completion_tokens"]
736
737
738

                    if not stream_buffer:  # The first chunk
                        if request.echo:
yichuan~'s avatar
yichuan~ committed
739
740
741
                            if isinstance(request.prompt, str):
                                # for the case of single str prompts
                                prompts = request.prompt
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
                            elif isinstance(request.prompt, list):
                                if isinstance(request.prompt[0], str):
                                    # for the case of multiple str prompts
                                    prompts = request.prompt[index // request.n]
                                elif isinstance(request.prompt[0], int):
                                    # for the case of single token ids prompt
                                    prompts = tokenizer_manager.tokenizer.decode(
                                        request.prompt, skip_special_tokens=True
                                    )
                                elif isinstance(request.prompt[0], list) and isinstance(
                                    request.prompt[0][0], int
                                ):
                                    # for the case of multiple token ids prompts
                                    prompts = tokenizer_manager.tokenizer.decode(
                                        request.prompt[index // request.n],
                                        skip_special_tokens=True,
                                    )
yichuan~'s avatar
yichuan~ committed
759

760
                            # Prepend prompt in response text.
yichuan~'s avatar
yichuan~ committed
761
                            text = prompts + text
762

763
                    if request.logprobs is not None:
764
765
                        # The first chunk and echo is enabled.
                        if not stream_buffer and request.echo:
766
767
                            input_token_logprobs = content["meta_info"][
                                "input_token_logprobs"
768
                            ]
769
770
                            input_top_logprobs = content["meta_info"][
                                "input_top_logprobs"
771
772
                            ]
                        else:
773
774
                            input_token_logprobs = None
                            input_top_logprobs = None
775
776

                        logprobs = to_openai_style_logprobs(
777
778
779
780
                            input_token_logprobs=input_token_logprobs,
                            input_top_logprobs=input_top_logprobs,
                            output_token_logprobs=content["meta_info"][
                                "output_token_logprobs"
781
                            ][n_prev_token:],
782
783
                            output_top_logprobs=content["meta_info"][
                                "output_top_logprobs"
784
                            ][n_prev_token:],
785
                        )
786
                        n_prev_token = len(
787
                            content["meta_info"]["output_token_logprobs"]
788
                        )
789
                    else:
790
                        logprobs = None
791

792
                    delta = text[len(stream_buffer) :]
Liangsheng Yin's avatar
Liangsheng Yin committed
793
                    stream_buffer = stream_buffer + delta
794
                    finish_reason = content["meta_info"]["finish_reason"]
795
                    choice_data = CompletionResponseStreamChoice(
796
                        index=index,
797
798
                        text=delta,
                        logprobs=logprobs,
799
800
801
802
803
                        finish_reason=(finish_reason["type"] if finish_reason else ""),
                        matched_stop=(
                            finish_reason["matched"]
                            if finish_reason and "matched" in finish_reason
                            else None
804
                        ),
805
806
807
808
809
810
811
                    )
                    chunk = CompletionStreamResponse(
                        id=content["meta_info"]["id"],
                        object="text_completion",
                        choices=[choice_data],
                        model=request.model,
                    )
812
813
814
815

                    stream_buffers[index] = stream_buffer
                    n_prev_tokens[index] = n_prev_token

816
                    yield f"data: {chunk.model_dump_json()}\n\n"
817
                if request.stream_options and request.stream_options.include_usage:
818
819
820
821
822
823
824
825
                    total_prompt_tokens = sum(
                        tokens
                        for i, tokens in prompt_tokens.items()
                        if i % request.n == 0
                    )
                    total_completion_tokens = sum(
                        tokens for tokens in completion_tokens.values()
                    )
826
                    usage = UsageInfo(
827
828
829
                        prompt_tokens=total_prompt_tokens,
                        completion_tokens=total_completion_tokens,
                        total_tokens=total_prompt_tokens + total_completion_tokens,
830
831
832
833
834
835
836
837
838
839
840
841
                    )

                    final_usage_chunk = CompletionStreamResponse(
                        id=str(uuid.uuid4().hex),
                        choices=[],
                        model=request.model,
                        usage=usage,
                    )
                    final_usage_data = final_usage_chunk.model_dump_json(
                        exclude_unset=True, exclude_none=True
                    )
                    yield f"data: {final_usage_data}\n\n"
842
843
844
            except ValueError as e:
                error = create_streaming_error_response(str(e))
                yield f"data: {error}\n\n"
845
846
            yield "data: [DONE]\n\n"

847
848
849
850
851
        return StreamingResponse(
            generate_stream_resp(),
            media_type="text/event-stream",
            background=tokenizer_manager.create_abort_task(adapted_request),
        )
852
853

    # Non-streaming response.
854
855
    try:
        ret = await tokenizer_manager.generate_request(
856
857
            adapted_request, raw_request
        ).__anext__()
858
859
    except ValueError as e:
        return create_error_response(str(e))
860

861
862
863
    if not isinstance(ret, list):
        ret = [ret]

yichuan~'s avatar
yichuan~ committed
864
    response = v1_generate_response(request, ret, tokenizer_manager)
865
    return response
866

867

868
def v1_chat_generate_request(
869
870
871
    all_requests: List[ChatCompletionRequest],
    tokenizer_manager,
    request_ids: List[str] = None,
872
):
873
    input_ids = []
874
875
    sampling_params_list = []
    image_data_list = []
876
    return_logprobs = []
877
    logprob_start_lens = []
878
    top_logprobs_nums = []
879
    modalities_list = []
880
    lora_paths = []
881
882
883

    # NOTE: with openai API, the prompt's logprobs are always not computed

884
885
886
887
888
889
890
891
    for request in all_requests:
        # Prep the data needed for the underlying GenerateReqInput:
        #  - prompt: The full prompt string.
        #  - stop: Custom stop tokens.
        #  - image_data: None or a list of image strings (URLs or base64 strings).
        #    None skips any image processing in GenerateReqInput.
        if not isinstance(request.messages, str):
            # Apply chat template and its stop strings.
Tanjiro's avatar
Tanjiro committed
892
893
894
895
896
897
898
899
900
901
902
903
            tools = None
            if request.tools and request.tool_choice != "none":
                request.skip_special_tokens = False
                if not isinstance(request.tool_choice, str):
                    tools = [
                        item.function.model_dump()
                        for item in request.tools
                        if item.function.name == request.tool_choice.function.name
                    ]
                else:
                    tools = [item.function.model_dump() for item in request.tools]

904
            if chat_template_name is None:
905
906
907
908
909
910
911
912
913
914
915
916
917
                openai_compatible_messages = []
                for message in request.messages:
                    if isinstance(message.content, str):
                        openai_compatible_messages.append(
                            {"role": message.role, "content": message.content}
                        )
                    else:
                        content_list = message.dict()["content"]
                        for content in content_list:
                            if content["type"] == "text":
                                openai_compatible_messages.append(
                                    {"role": message.role, "content": content["text"]}
                                )
918
919
920
921
922
                if openai_compatible_messages[-1]["role"] == "assistant":
                    assistant_prefix = openai_compatible_messages[-1]["content"]
                    openai_compatible_messages = openai_compatible_messages[:-1]
                else:
                    assistant_prefix = None
YAMY's avatar
YAMY committed
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942

                try:
                    prompt_ids = tokenizer_manager.tokenizer.apply_chat_template(
                        openai_compatible_messages,
                        tokenize=True,
                        add_generation_prompt=True,
                        tools=tools,
                    )
                except:
                    #  This except branch will be triggered when the chosen model
                    #  has a different tools input format that is not compatiable
                    #  with openAI's apply_chat_template tool_call format, like Mistral.
                    tools = [t if "function" in t else {"function": t} for t in tools]
                    prompt_ids = tokenizer_manager.tokenizer.apply_chat_template(
                        openai_compatible_messages,
                        tokenize=True,
                        add_generation_prompt=True,
                        tools=tools,
                    )

943
                if assistant_prefix:
944
945
946
947
948
949
950
                    encoded = tokenizer_manager.tokenizer.encode(assistant_prefix)
                    if (
                        encoded
                        and encoded[0] == tokenizer_manager.tokenizer.bos_token_id
                    ):
                        encoded = encoded[1:]
                    prompt_ids += encoded
951
952
                stop = request.stop
                image_data = None
953
                modalities = []
954
            else:
955
956
957
                conv = generate_chat_conv(request, chat_template_name)
                prompt = conv.get_prompt()
                image_data = conv.image_data
958
                modalities = conv.modalities
959
960
961
962
963
964
                stop = conv.stop_str or []
                if request.stop:
                    if isinstance(request.stop, str):
                        stop.append(request.stop)
                    else:
                        stop.extend(request.stop)
965
                prompt_ids = tokenizer_manager.tokenizer.encode(prompt)
966
        else:
967
            # Use the raw prompt and stop strings if the messages is already a string.
yichuan~'s avatar
yichuan~ committed
968
            prompt_ids = request.messages
969
970
            stop = request.stop
            image_data = None
971
            modalities = []
972
        input_ids.append(prompt_ids)
973
        return_logprobs.append(request.logprobs)
974
        logprob_start_lens.append(-1)
975
        top_logprobs_nums.append(request.top_logprobs or 0)
976
        lora_paths.append(request.lora_path)
977
978
979
980
981
982
983
984

        sampling_params = {
            "temperature": request.temperature,
            "max_new_tokens": request.max_tokens,
            "min_new_tokens": request.min_tokens,
            "stop": stop,
            "stop_token_ids": request.stop_token_ids,
            "top_p": request.top_p,
985
986
            "top_k": request.top_k,
            "min_p": request.min_p,
987
988
989
990
            "presence_penalty": request.presence_penalty,
            "frequency_penalty": request.frequency_penalty,
            "repetition_penalty": request.repetition_penalty,
            "regex": request.regex,
991
            "ebnf": request.ebnf,
992
            "n": request.n,
993
            "no_stop_trim": request.no_stop_trim,
994
            "ignore_eos": request.ignore_eos,
995
            "skip_special_tokens": request.skip_special_tokens,
996
997
998
999
1000
1001
1002
        }
        if request.response_format and request.response_format.type == "json_schema":
            sampling_params["json_schema"] = convert_json_schema_to_str(
                request.response_format.json_schema.schema_
            )
        sampling_params_list.append(sampling_params)

1003
        image_data_list.append(image_data)
1004
        modalities_list.append(modalities)
1005
    if len(all_requests) == 1:
1006
1007
        if isinstance(input_ids[0], str):
            prompt_kwargs = {"text": input_ids[0]}
yichuan~'s avatar
yichuan~ committed
1008
        else:
1009
            prompt_kwargs = {"input_ids": input_ids[0]}
1010
        sampling_params_list = sampling_params_list[0]
1011
        image_data_list = image_data_list[0]
1012
        return_logprobs = return_logprobs[0]
1013
        logprob_start_lens = logprob_start_lens[0]
1014
        top_logprobs_nums = top_logprobs_nums[0]
1015
        modalities_list = modalities_list[0]
1016
        lora_paths = lora_paths[0]
yichuan~'s avatar
yichuan~ committed
1017
1018
1019
1020
1021
    else:
        if isinstance(input_ids[0], str):
            prompt_kwargs = {"text": input_ids}
        else:
            prompt_kwargs = {"input_ids": input_ids}
1022

1023
    adapted_request = GenerateReqInput(
yichuan~'s avatar
yichuan~ committed
1024
        **prompt_kwargs,
1025
        image_data=image_data_list,
1026
        sampling_params=sampling_params_list,
1027
        return_logprob=return_logprobs,
1028
        logprob_start_len=logprob_start_lens,
1029
1030
1031
        top_logprobs_num=top_logprobs_nums,
        stream=all_requests[0].stream,
        return_text_in_logprobs=True,
1032
        rid=request_ids,
1033
        modalities=modalities_list,
1034
        lora_path=lora_paths,
1035
    )
1036
1037

    return adapted_request, all_requests if len(all_requests) > 1 else all_requests[0]
1038

1039

YAMY's avatar
YAMY committed
1040
1041
1042
def v1_chat_generate_response(
    request, ret, to_file=False, cache_report=False, tool_call_parser=None
):
1043
1044
1045
    choices = []

    for idx, ret_item in enumerate(ret):
1046
        logprobs = False
yichuan~'s avatar
yichuan~ committed
1047
        if isinstance(request, list) and request[idx].logprobs:
1048
            logprobs = True
yichuan~'s avatar
yichuan~ committed
1049
        elif (not isinstance(request, list)) and request.logprobs:
1050
1051
1052
1053
1054
1055
1056
            logprobs = True
        if logprobs:
            logprobs = to_openai_style_logprobs(
                output_token_logprobs=ret_item["meta_info"]["output_token_logprobs"],
                output_top_logprobs=ret_item["meta_info"]["output_top_logprobs"],
            )
            token_logprobs = []
1057
1058
1059
            for token_idx, (token, logprob) in enumerate(
                zip(logprobs.tokens, logprobs.token_logprobs)
            ):
1060
1061
1062
                token_bytes = list(token.encode("utf-8"))
                top_logprobs = []
                if logprobs.top_logprobs:
1063
1064
1065
                    for top_token, top_logprob in logprobs.top_logprobs[
                        token_idx
                    ].items():
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
                        top_token_bytes = list(top_token.encode("utf-8"))
                        top_logprobs.append(
                            TopLogprob(
                                token=top_token,
                                bytes=top_token_bytes,
                                logprob=top_logprob,
                            )
                        )
                token_logprobs.append(
                    ChatCompletionTokenLogprob(
                        token=token,
                        bytes=token_bytes,
                        logprob=logprob,
                        top_logprobs=top_logprobs,
                    )
                )

            choice_logprobs = ChoiceLogprobs(content=token_logprobs)
        else:
            choice_logprobs = None
1086

1087
1088
        finish_reason = ret_item["meta_info"]["finish_reason"]

Tanjiro's avatar
Tanjiro committed
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
        tool_calls = None
        text = ret_item["text"]

        if isinstance(request, list):
            tool_choice = request[idx].tool_choice
            tools = request[idx].tools
        else:
            tool_choice = request.tool_choice
            tools = request.tools

        if tool_choice != "none" and any([i in text for i in TOOLS_TAG_LIST]):
            if finish_reason == "stop":
                finish_reason = "tool_calls"
            try:
YAMY's avatar
YAMY committed
1103
1104
                parser = FunctionCallParser(tools, tool_call_parser)
                full_normal_text, call_info_list = parser.parse_non_stream(text)
Tanjiro's avatar
Tanjiro committed
1105
1106
                tool_calls = [
                    ToolCall(
YAMY's avatar
YAMY committed
1107
                        id=str(call_info.tool_index),
Tanjiro's avatar
Tanjiro committed
1108
                        function=FunctionResponse(
YAMY's avatar
YAMY committed
1109
                            name=call_info.name, arguments=call_info.parameters
Tanjiro's avatar
Tanjiro committed
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
                        ),
                    )
                    for call_info in call_info_list
                ]
            except Exception as e:
                logger.error(f"Exception: {e}")
                return create_error_response(
                    HTTPStatus.BAD_REQUEST,
                    "Failed to parse fc related info to json format!",
                )

1121
        if to_file:
1122
            # to make the choice data json serializable
1123
1124
            choice_data = {
                "index": 0,
Tanjiro's avatar
Tanjiro committed
1125
1126
1127
1128
1129
                "message": {
                    "role": "assistant",
                    "content": ret_item["text"] if tool_calls is None else None,
                    "tool_calls": tool_calls,
                },
1130
                "logprobs": choice_logprobs,
1131
1132
1133
1134
1135
                "finish_reason": (finish_reason["type"] if finish_reason else ""),
                "matched_stop": (
                    finish_reason["matched"]
                    if finish_reason and "matched" in finish_reason
                    else None
1136
                ),
1137
            }
1138
        else:
1139
1140
            choice_data = ChatCompletionResponseChoice(
                index=idx,
Tanjiro's avatar
Tanjiro committed
1141
1142
1143
1144
1145
                message=ChatMessage(
                    role="assistant",
                    content=ret_item["text"] if tool_calls is None else None,
                    tool_calls=tool_calls,
                ),
1146
                logprobs=choice_logprobs,
1147
1148
1149
1150
1151
                finish_reason=(finish_reason["type"] if finish_reason else ""),
                matched_stop=(
                    finish_reason["matched"]
                    if finish_reason and "matched" in finish_reason
                    else None
1152
                ),
1153
1154
1155
            )

        choices.append(choice_data)
1156

1157
1158
1159
1160
1161
1162
1163
1164
    if to_file:
        responses = []

        for i, choice in enumerate(choices):
            response = {
                "status_code": 200,
                "request_id": ret[i]["meta_info"]["id"],
                "body": {
1165
                    # remain the same but if needed we can change that
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
                    "id": ret[i]["meta_info"]["id"],
                    "object": "chat.completion",
                    "created": int(time.time()),
                    "model": request[i].model,
                    "choices": choice,
                    "usage": {
                        "prompt_tokens": ret[i]["meta_info"]["prompt_tokens"],
                        "completion_tokens": ret[i]["meta_info"]["completion_tokens"],
                        "total_tokens": ret[i]["meta_info"]["prompt_tokens"]
                        + ret[i]["meta_info"]["completion_tokens"],
                    },
                    "system_fingerprint": None,
                },
            }
            responses.append(response)
        return responses
1182
    else:
1183
1184
1185
1186
        prompt_tokens = sum(
            ret[i]["meta_info"]["prompt_tokens"] for i in range(0, len(ret), request.n)
        )
        completion_tokens = sum(item["meta_info"]["completion_tokens"] for item in ret)
1187
        cached_tokens = sum(item["meta_info"].get("cached_tokens", 0) for item in ret)
1188
1189
1190
1191
1192
        response = ChatCompletionResponse(
            id=ret[0]["meta_info"]["id"],
            model=request.model,
            choices=choices,
            usage=UsageInfo(
1193
1194
1195
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
                total_tokens=prompt_tokens + completion_tokens,
1196
1197
1198
                prompt_tokens_details=(
                    {"cached_tokens": cached_tokens} if cache_report else None
                ),
1199
1200
1201
            ),
        )
        return response
1202

1203
1204
1205
1206
1207

async def v1_chat_completions(tokenizer_manager, raw_request: Request):
    request_json = await raw_request.json()
    all_requests = [ChatCompletionRequest(**request_json)]
    adapted_request, request = v1_chat_generate_request(all_requests, tokenizer_manager)
1208
1209

    if adapted_request.stream:
YAMY's avatar
YAMY committed
1210
        parser_dict = {}
1211
1212

        async def generate_stream_resp():
1213
1214
1215
1216
1217
            is_firsts = {}
            stream_buffers = {}
            n_prev_tokens = {}
            prompt_tokens = {}
            completion_tokens = {}
1218
            try:
1219
1220
1221
                async for content in tokenizer_manager.generate_request(
                    adapted_request, raw_request
                ):
1222
                    index = content.get("index", 0)
YAMY's avatar
YAMY committed
1223
                    text = content["text"]
1224
1225
1226
1227
1228
1229
1230

                    is_first = is_firsts.get(index, True)
                    stream_buffer = stream_buffers.get(index, "")
                    n_prev_token = n_prev_tokens.get(index, 0)

                    prompt_tokens[index] = content["meta_info"]["prompt_tokens"]
                    completion_tokens[index] = content["meta_info"]["completion_tokens"]
yichuan~'s avatar
yichuan~ committed
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
                    if request.logprobs:
                        logprobs = to_openai_style_logprobs(
                            output_token_logprobs=content["meta_info"][
                                "output_token_logprobs"
                            ][n_prev_token:],
                            output_top_logprobs=content["meta_info"][
                                "output_top_logprobs"
                            ][n_prev_token:],
                        )

                        n_prev_token = len(
                            content["meta_info"]["output_token_logprobs"]
                        )
                        token_logprobs = []
                        for token, logprob in zip(
                            logprobs.tokens, logprobs.token_logprobs
                        ):
                            token_bytes = list(token.encode("utf-8"))
                            top_logprobs = []
                            if logprobs.top_logprobs:
                                for top_token, top_logprob in logprobs.top_logprobs[
                                    0
                                ].items():
                                    top_token_bytes = list(top_token.encode("utf-8"))
                                    top_logprobs.append(
                                        TopLogprob(
                                            token=top_token,
                                            bytes=top_token_bytes,
                                            logprob=top_logprob,
                                        )
                                    )
                            token_logprobs.append(
                                ChatCompletionTokenLogprob(
                                    token=token,
                                    bytes=token_bytes,
                                    logprob=logprob,
                                    top_logprobs=top_logprobs,
                                )
                            )

                        choice_logprobs = ChoiceLogprobs(content=token_logprobs)

                    else:
                        choice_logprobs = None

1276
1277
                    finish_reason = content["meta_info"]["finish_reason"]

1278
1279
1280
1281
                    if is_first:
                        # First chunk with role
                        is_first = False
                        choice_data = ChatCompletionResponseStreamChoice(
1282
                            index=index,
1283
                            delta=DeltaMessage(role="assistant", content=""),
1284
                            finish_reason=(
1285
1286
1287
1288
1289
1290
                                finish_reason["type"] if finish_reason else ""
                            ),
                            matched_stop=(
                                finish_reason["matched"]
                                if finish_reason and "matched" in finish_reason
                                else None
1291
                            ),
yichuan~'s avatar
yichuan~ committed
1292
                            logprobs=choice_logprobs,
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
                        )
                        chunk = ChatCompletionStreamResponse(
                            id=content["meta_info"]["id"],
                            choices=[choice_data],
                            model=request.model,
                        )
                        yield f"data: {chunk.model_dump_json()}\n\n"

                    text = content["text"]
                    delta = text[len(stream_buffer) :]
YAMY's avatar
YAMY committed
1303
                    new_stream_buffer = stream_buffer + delta
1304

YAMY's avatar
YAMY committed
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
                    if request.tool_choice != "none" and request.tools:
                        if index not in parser_dict:
                            parser_dict[index] = FunctionCallParser(
                                tools=request.tools,
                                tool_call_parser=tokenizer_manager.server_args.tool_call_parser,
                            )
                        parser = parser_dict[index]

                        # parse_increment => returns (normal_text, calls)
                        normal_text, calls = parser.parse_stream_chunk(delta)

                        # 1) if there's normal_text, output it as normal content
                        if normal_text:
                            choice_data = ChatCompletionResponseStreamChoice(
                                index=index,
                                delta=DeltaMessage(content=normal_text),
                                finish_reason=(
                                    finish_reason["type"] if finish_reason else ""
                                ),
                            )
                            chunk = ChatCompletionStreamResponse(
                                id=content["meta_info"]["id"],
                                choices=[choice_data],
                                model=request.model,
                            )
                            yield f"data: {chunk.model_dump_json()}\n\n"

                        # 2) if we found calls, we output them as separate chunk(s)
                        for call_item in calls:
                            # transform call_item -> FunctionResponse + ToolCall

                            if (
                                content["meta_info"]["finish_reason"]
                                and content["meta_info"]["finish_reason"]["type"]
                                == "stop"
                            ):
                                latest_delta_len = 0
                                if isinstance(call_item.parameters, str):
                                    latest_delta_len = len(call_item.parameters)

                                expected_call = json.dumps(
                                    parser.multi_format_parser.detectors[0]
                                    .prev_tool_call_arr[index]
                                    .get("arguments", {}),
                                    ensure_ascii=False,
                                )
                                actual_call = parser.multi_format_parser.detectors[
                                    0
                                ].streamed_args_for_tool[index]
                                if latest_delta_len > 0:
                                    actual_call = actual_call[:-latest_delta_len]
                                remaining_call = expected_call.replace(
                                    actual_call, "", 1
                                )
                                call_item.parameters = remaining_call

                            tool_call = ToolCall(
                                id=str(call_item.tool_index),
                                function=FunctionResponse(
                                    name=call_item.name,
                                    arguments=call_item.parameters,
                                ),
                            )
                            choice_data = ChatCompletionResponseStreamChoice(
                                index=index,
                                delta=DeltaMessage(
                                    role="assistant", tool_calls=[tool_call]
                                ),
                                finish_reason="tool_call",
                            )
                            chunk = ChatCompletionStreamResponse(
                                id=content["meta_info"]["id"],
                                choices=[choice_data],
                                model=request.model,
                            )
                            yield f"data: {chunk.model_dump_json()}\n\n"
1381

YAMY's avatar
YAMY committed
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
                        stream_buffers[index] = new_stream_buffer
                        is_firsts[index] = is_first

                    else:
                        # No tool calls => just treat this as normal text
                        choice_data = ChatCompletionResponseStreamChoice(
                            index=index,
                            delta=DeltaMessage(content=delta),
                            finish_reason=(
                                finish_reason["type"] if finish_reason else ""
                            ),
                            matched_stop=(
                                finish_reason["matched"]
                                if finish_reason and "matched" in finish_reason
                                else None
                            ),
                            logprobs=choice_logprobs,
                        )
                        chunk = ChatCompletionStreamResponse(
                            id=content["meta_info"]["id"],
                            choices=[choice_data],
                            model=request.model,
                        )
                        yield f"data: {chunk.model_dump_json()}\n\n"
                        stream_buffers[index] = new_stream_buffer
                        is_firsts[index] = is_first
1408
                if request.stream_options and request.stream_options.include_usage:
1409
1410
1411
1412
1413
1414
1415
1416
                    total_prompt_tokens = sum(
                        tokens
                        for i, tokens in prompt_tokens.items()
                        if i % request.n == 0
                    )
                    total_completion_tokens = sum(
                        tokens for tokens in completion_tokens.values()
                    )
1417
                    usage = UsageInfo(
1418
1419
1420
                        prompt_tokens=total_prompt_tokens,
                        completion_tokens=total_completion_tokens,
                        total_tokens=total_prompt_tokens + total_completion_tokens,
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
                    )

                    final_usage_chunk = ChatCompletionStreamResponse(
                        id=str(uuid.uuid4().hex),
                        choices=[],
                        model=request.model,
                        usage=usage,
                    )
                    final_usage_data = final_usage_chunk.model_dump_json(
                        exclude_unset=True, exclude_none=True
                    )
                    yield f"data: {final_usage_data}\n\n"
1433
1434
1435
            except ValueError as e:
                error = create_streaming_error_response(str(e))
                yield f"data: {error}\n\n"
1436
1437
            yield "data: [DONE]\n\n"

1438
1439
1440
1441
1442
        return StreamingResponse(
            generate_stream_resp(),
            media_type="text/event-stream",
            background=tokenizer_manager.create_abort_task(adapted_request),
        )
1443
1444

    # Non-streaming response.
1445
1446
    try:
        ret = await tokenizer_manager.generate_request(
1447
1448
            adapted_request, raw_request
        ).__anext__()
1449
1450
    except ValueError as e:
        return create_error_response(str(e))
1451
1452
1453
    if not isinstance(ret, list):
        ret = [ret]

1454
    response = v1_chat_generate_response(
YAMY's avatar
YAMY committed
1455
1456
1457
1458
        request,
        ret,
        cache_report=tokenizer_manager.server_args.enable_cache_report,
        tool_call_parser=tokenizer_manager.server_args.tool_call_parser,
1459
    )
1460

1461
1462
1463
    return response


1464
1465
1466
def v1_embedding_request(all_requests, tokenizer_manager):
    prompts = []
    sampling_params_list = []
Ying Sheng's avatar
Ying Sheng committed
1467
    first_prompt_type = type(all_requests[0].input)
1468
1469

    for request in all_requests:
Ying Sheng's avatar
Ying Sheng committed
1470
        prompt = request.input
1471
        assert (
1472
            type(prompt) is first_prompt_type
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
        ), "All prompts must be of the same type in file input settings"
        prompts.append(prompt)

    if len(all_requests) == 1:
        prompt = prompts[0]
        if isinstance(prompt, str) or isinstance(prompt[0], str):
            prompt_kwargs = {"text": prompt}
        else:
            prompt_kwargs = {"input_ids": prompt}
    else:
Baoyuan Qi's avatar
Baoyuan Qi committed
1483
        if isinstance(prompts[0], str) or isinstance(prompts[0][0], str):
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
            prompt_kwargs = {"text": prompts}
        else:
            prompt_kwargs = {"input_ids": prompts}

    adapted_request = EmbeddingReqInput(
        **prompt_kwargs,
    )

    if len(all_requests) == 1:
        return adapted_request, all_requests[0]
    return adapted_request, all_requests


Ying Sheng's avatar
Ying Sheng committed
1497
1498
1499
def v1_embedding_response(ret, model_path, to_file=False):
    embedding_objects = []
    prompt_tokens = 0
1500
    for idx, ret_item in enumerate(ret):
Ying Sheng's avatar
Ying Sheng committed
1501
1502
1503
        embedding_objects.append(
            EmbeddingObject(
                embedding=ret[idx]["embedding"],
1504
1505
1506
                index=idx,
            )
        )
Ying Sheng's avatar
Ying Sheng committed
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
        prompt_tokens += ret[idx]["meta_info"]["prompt_tokens"]

    return EmbeddingResponse(
        data=embedding_objects,
        model=model_path,
        usage=UsageInfo(
            prompt_tokens=prompt_tokens,
            total_tokens=prompt_tokens,
        ),
    )
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533


async def v1_embeddings(tokenizer_manager, raw_request: Request):
    request_json = await raw_request.json()
    all_requests = [EmbeddingRequest(**request_json)]
    adapted_request, request = v1_embedding_request(all_requests, tokenizer_manager)

    try:
        ret = await tokenizer_manager.generate_request(
            adapted_request, raw_request
        ).__anext__()
    except ValueError as e:
        return create_error_response(str(e))

    if not isinstance(ret, list):
        ret = [ret]

Ying Sheng's avatar
Ying Sheng committed
1534
    response = v1_embedding_response(ret, tokenizer_manager.model_path)
1535
1536
1537
1538

    return response


1539
def to_openai_style_logprobs(
1540
1541
1542
1543
    input_token_logprobs=None,
    output_token_logprobs=None,
    input_top_logprobs=None,
    output_top_logprobs=None,
1544
1545
1546
1547
1548
1549
1550
1551
):
    ret_logprobs = LogProbs()

    def append_token_logprobs(token_logprobs):
        for logprob, _, token_text in token_logprobs:
            ret_logprobs.tokens.append(token_text)
            ret_logprobs.token_logprobs.append(logprob)

1552
            # Not supported yet
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
            ret_logprobs.text_offset.append(-1)

    def append_top_logprobs(top_logprobs):
        for tokens in top_logprobs:
            if tokens is not None:
                ret_logprobs.top_logprobs.append(
                    {token[2]: token[0] for token in tokens}
                )
            else:
                ret_logprobs.top_logprobs.append(None)

1564
1565
1566
1567
1568
1569
1570
1571
    if input_token_logprobs is not None:
        append_token_logprobs(input_token_logprobs)
    if output_token_logprobs is not None:
        append_token_logprobs(output_token_logprobs)
    if input_top_logprobs is not None:
        append_top_logprobs(input_top_logprobs)
    if output_top_logprobs is not None:
        append_top_logprobs(output_top_logprobs)
1572

Liangsheng Yin's avatar
Liangsheng Yin committed
1573
    return ret_logprobs