endpoint_request_func.py 27 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
"""The request function for API endpoints."""

5
import io
6
7
8
9
10
import json
import os
import sys
import time
import traceback
11
from collections.abc import Awaitable
12
from dataclasses import dataclass, field
13
from typing import Any, Literal, Protocol
14
15

import aiohttp
16
import regex as re
17
18
19
20
21
from tqdm.asyncio import tqdm

AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60)


22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
class StreamedResponseHandler:
    """Handles streaming HTTP responses by accumulating chunks until complete
    messages are available."""

    def __init__(self):
        self.buffer = ""

    def add_chunk(self, chunk_bytes: bytes) -> list[str]:
        """Add a chunk of bytes to the buffer and return any complete
        messages."""
        chunk_str = chunk_bytes.decode("utf-8")
        self.buffer += chunk_str

        messages = []

        # Split by double newlines (SSE message separator)
        while "\n\n" in self.buffer:
            message, self.buffer = self.buffer.split("\n\n", 1)
            message = message.strip()
            if message:
                messages.append(message)

        # if self.buffer is not empty, check if it is a complete message
        # by removing data: prefix and check if it is a valid JSON
        if self.buffer.startswith("data: "):
            message_content = self.buffer.removeprefix("data: ").strip()
            if message_content == "[DONE]":
                messages.append(self.buffer.strip())
                self.buffer = ""
            elif message_content:
                try:
                    json.loads(message_content)
                    messages.append(self.buffer.strip())
                    self.buffer = ""
                except json.JSONDecodeError:
                    # Incomplete JSON, wait for more chunks.
                    pass

        return messages


63
64
65
@dataclass
class RequestFuncInput:
    """The input for the request function."""
66

67
    prompt: str | list[str]
68
69
70
71
    api_url: str
    prompt_len: int
    output_len: int
    model: str
72
73
74
75
76
    model_name: str | None = None
    logprobs: int | None = None
    extra_headers: dict | None = None
    extra_body: dict | None = None
    multi_modal_content: dict | list[dict] | None = None
77
    ignore_eos: bool = False
78
79
    language: str | None = None
    request_id: str | None = None
80
81
82
83
84


@dataclass
class RequestFuncOutput:
    """The output of the request function including metrics."""
85

86
87
88
89
90
    generated_text: str = ""
    success: bool = False
    latency: float = 0.0
    output_tokens: int = 0
    ttft: float = 0.0  # Time to first token
91
    itl: list[float] = field(default_factory=list)  # list of inter-token latencies
92
93
94
    tpot: float = 0.0  # avg next-token latencies
    prompt_len: int = 0
    error: str = ""
95
    start_time: float = 0.0
96
    input_audio_duration: float = 0.0  # in seconds
97
98


99
100
101
102
103
class RequestFunc(Protocol):
    def __call__(
        self,
        request_func_input: RequestFuncInput,
        session: aiohttp.ClientSession,
104
        pbar: tqdm | None = None,
105
    ) -> Awaitable[RequestFuncOutput]: ...
106
107


108
109
110
def _validate_api_url(
    api_url: str,
    api_name: str,
111
    expected_suffixes: str | set[str],
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
) -> None:
    if isinstance(expected_suffixes, str):
        expected_suffixes = {expected_suffixes}

    expected_suffixes = {*expected_suffixes, "profile"}

    if not api_url.endswith(tuple(expected_suffixes)):
        raise ValueError(f"{api_name} URL must end with one of: {expected_suffixes}.")


def _update_payload_common(
    payload: dict[str, Any],
    request_func_input: RequestFuncInput,
) -> None:
    if request_func_input.ignore_eos:
        payload["ignore_eos"] = request_func_input.ignore_eos
    if request_func_input.extra_body:
        payload.update(request_func_input.extra_body)


def _update_headers_common(
    headers: dict[str, Any],
    request_func_input: RequestFuncInput,
) -> None:
    if request_func_input.extra_headers:
        headers |= request_func_input.extra_headers
    if request_func_input.request_id:
        headers["x-request-id"] = request_func_input.request_id


142
143
144
145
146
147
148
149
150
151
def _get_headers(content_type: str | None = None) -> dict[str, str]:
    headers = {}
    if content_type:
        headers["Content-Type"] = content_type
    api_key = os.environ.get("OPENAI_API_KEY")
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    return headers


152
153
async def async_request_openai_completions(
    request_func_input: RequestFuncInput,
154
    session: aiohttp.ClientSession,
155
    pbar: tqdm | None = None,
156
157
158
159
160
161
162
163
164
165
166
) -> RequestFuncOutput:
    """The async request function for the OpenAI Completions API.

    Args:
        request_func_input: The input for the request function.
        pbar: The progress bar to display the progress.

    Returns:
        The output of the request function.
    """
    api_url = request_func_input.api_url
167
    _validate_api_url(api_url, "OpenAI Completions API", "completions")
168

169
    payload = {
170
        "model": request_func_input.model_name
171
172
        if request_func_input.model_name
        else request_func_input.model,
173
174
175
176
177
178
179
180
181
        "prompt": request_func_input.prompt,
        "repetition_penalty": 1.0,
        "max_tokens": request_func_input.output_len,
        "logprobs": request_func_input.logprobs,
        "stream": True,
        "stream_options": {
            "include_usage": True,
        },
    }
182
183
    _update_payload_common(payload, request_func_input)

184
    headers = _get_headers()
185
    _update_headers_common(headers, request_func_input)
186
187
188
189
190
191

    output = RequestFuncOutput()
    output.prompt_len = request_func_input.prompt_len

    generated_text = ""
    st = time.perf_counter()
192
    output.start_time = st
193
194
    most_recent_timestamp = st
    try:
195
        async with session.post(url=api_url, json=payload, headers=headers) as response:
196
197
            if response.status == 200:
                first_chunk_received = False
198
199
200
                handler = StreamedResponseHandler()

                async for chunk_bytes in response.content.iter_any():
201
202
203
204
                    chunk_bytes = chunk_bytes.strip()
                    if not chunk_bytes:
                        continue

205
206
207
208
209
210
211
                    messages = handler.add_chunk(chunk_bytes)
                    for message in messages:
                        # NOTE: SSE comments (often used as pings) start with
                        # a colon. These are not JSON data payload and should
                        # be skipped.
                        if message.startswith(":"):
                            continue
212

213
                        chunk = message.removeprefix("data: ")
214

215
216
                        if chunk != "[DONE]":
                            data = json.loads(chunk)
217

218
219
220
221
222
223
224
225
226
227
228
229
230
                            # NOTE: Some completion API might have a last
                            # usage summary response without a token so we
                            # want to check a token was generated
                            if choices := data.get("choices"):
                                # Note that text could be empty here
                                # e.g. for special tokens
                                text = choices[0].get("text")
                                timestamp = time.perf_counter()
                                # First token
                                if not first_chunk_received:
                                    first_chunk_received = True
                                    ttft = time.perf_counter() - st
                                    output.ttft = ttft
231

232
233
                                # Decoding phase
                                else:
234
                                    output.itl.append(timestamp - most_recent_timestamp)
235
236
237
238

                                most_recent_timestamp = timestamp
                                generated_text += text or ""
                            elif usage := data.get("usage"):
239
                                output.output_tokens = usage.get("completion_tokens")
240
241
                                if (pt := usage.get("prompt_tokens")) is not None:
                                    output.prompt_len = pt
242
243
                if first_chunk_received:
                    output.success = True
244
245
                else:
                    output.success = False
246
247
                    output.error = (
                        "Never received a valid chunk to calculate TTFT."
248
249
                        "This response will be marked as failed!"
                    )
250
251
252
253
254
255
256
257
258
                output.generated_text = generated_text
                output.latency = most_recent_timestamp - st
            else:
                output.error = response.reason or ""
                output.success = False
    except Exception:
        output.success = False
        exc_info = sys.exc_info()
        output.error = "".join(traceback.format_exception(*exc_info))
259
260
261
262
263
264

    if pbar:
        pbar.update(1)
    return output


265
def _get_chat_content(
266
    request_func_input: RequestFuncInput,
267
268
269
    mm_position: Literal["first", "last"] = "last",
) -> list[dict[str, Any]]:
    text_contents = [{"type": "text", "text": request_func_input.prompt}]
270

271
    mm_contents = []
272
    if request_func_input.multi_modal_content:
273
274
        mm_content = request_func_input.multi_modal_content
        if isinstance(mm_content, list):
275
            mm_contents.extend(request_func_input.multi_modal_content)
276
        elif isinstance(mm_content, dict):
277
            mm_contents.append(request_func_input.multi_modal_content)
278
279
        else:
            raise TypeError(
280
                "multi_modal_content must be a dict or list[dict] for openai-chat"
281
            )
282
283
284
285
286
287
288
289
290
291

    if mm_position == "first":
        return mm_contents + text_contents

    return text_contents + mm_contents


async def async_request_openai_chat_completions(
    request_func_input: RequestFuncInput,
    session: aiohttp.ClientSession,
292
    pbar: tqdm | None = None,
293
294
295
296
297
298
299
    mm_position: Literal["first", "last"] = "last",
) -> RequestFuncOutput:
    api_url = request_func_input.api_url
    _validate_api_url(api_url, "OpenAI Chat Completions API", "chat/completions")

    content = _get_chat_content(request_func_input, mm_position=mm_position)

300
    payload = {
301
302
303
        "model": request_func_input.model_name
        if request_func_input.model_name
        else request_func_input.model,
304
        "messages": [
305
            {"role": "user", "content": content},
306
        ],
307
308
        "max_completion_tokens": request_func_input.output_len,
        "stream": True,
309
310
311
312
        "stream_options": {
            "include_usage": True,
        },
    }
313
314
    _update_payload_common(payload, request_func_input)

315
    headers = _get_headers("application/json")
316
    _update_headers_common(headers, request_func_input)
317
318
319
320
321
322
323

    output = RequestFuncOutput()
    output.prompt_len = request_func_input.prompt_len

    generated_text = ""
    ttft = 0.0
    st = time.perf_counter()
324
    output.start_time = st
325
326
    most_recent_timestamp = st
    try:
327
        async with session.post(url=api_url, json=payload, headers=headers) as response:
328
            if response.status == 200:
329
330
                handler = StreamedResponseHandler()
                async for chunk_bytes in response.content.iter_any():
331
332
333
334
                    chunk_bytes = chunk_bytes.strip()
                    if not chunk_bytes:
                        continue

335
336
337
338
339
340
341
342
343
                    messages = handler.add_chunk(chunk_bytes)
                    for message in messages:
                        # NOTE: SSE comments (often used as pings) start with
                        # a colon. These are not JSON data payload and should
                        # be skipped.
                        if message.startswith(":"):
                            continue

                        chunk = message.removeprefix("data: ")
344

345
346
347
                        if chunk != "[DONE]":
                            timestamp = time.perf_counter()
                            data = json.loads(chunk)
348

349
350
351
352
353
354
                            if choices := data.get("choices"):
                                content = choices[0]["delta"].get("content")
                                # First token
                                if ttft == 0.0:
                                    ttft = timestamp - st
                                    output.ttft = ttft
355

356
357
                                # Decoding phase
                                else:
358
                                    output.itl.append(timestamp - most_recent_timestamp)
359

360
361
                                generated_text += content or ""
                            elif usage := data.get("usage"):
362
                                output.output_tokens = usage.get("completion_tokens")
363
364
                                if (pt := usage.get("prompt_tokens")) is not None:
                                    output.prompt_len = pt
365

366
                            most_recent_timestamp = timestamp
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386

                output.generated_text = generated_text
                output.success = True
                output.latency = most_recent_timestamp - st
            else:
                output.error = response.reason or ""
                output.success = False
    except Exception:
        output.success = False
        exc_info = sys.exc_info()
        output.error = "".join(traceback.format_exception(*exc_info))

    if pbar:
        pbar.update(1)
    return output


async def async_request_openai_audio(
    request_func_input: RequestFuncInput,
    session: aiohttp.ClientSession,
387
    pbar: tqdm | None = None,
388
389
390
391
392
) -> RequestFuncOutput:
    # Lazy import without PlaceholderModule to avoid vllm dep.
    import soundfile

    api_url = request_func_input.api_url
393
    _validate_api_url(api_url, "OpenAI Audio API", {"transcriptions", "translations"})
394
395
396

    content = [{"type": "text", "text": request_func_input.prompt}]
    payload = {
397
398
399
400
401
402
        "model": request_func_input.model_name
        if request_func_input.model_name
        else request_func_input.model,
        "max_completion_tokens": request_func_input.output_len,
        "stream": True,
        "language": "en",
403
        # Flattened due to multipart/form-data
404
405
        "stream_include_usage": True,
        "stream_continuous_usage_stats": True,
406
    }
407
408
    _update_payload_common(payload, request_func_input)

409
    headers = _get_headers()
410
    _update_headers_common(headers, request_func_input)
411
412
413
414
415
416
417
418

    # Send audio file
    def to_bytes(y, sr):
        buffer = io.BytesIO()
        soundfile.write(buffer, y, sr, format="WAV")
        buffer.seek(0)
        return buffer

419
420
421
422
    mm_audio = request_func_input.multi_modal_content
    if not isinstance(mm_audio, dict) or "audio" not in mm_audio:
        raise TypeError("multi_modal_content must be a dict containing 'audio'")
    with to_bytes(*mm_audio["audio"]) as f:
423
424
425
426
        form = aiohttp.FormData()
        form.add_field("file", f, content_type="audio/wav")
        for key, value in payload.items():
            form.add_field(key, str(value))
427
428
429

        output = RequestFuncOutput()
        output.prompt_len = request_func_input.prompt_len
430
431
        output.input_audio_duration = soundfile.info(f).duration
        f.seek(0)
432
433
434
435

        generated_text = ""
        ttft = 0.0
        st = time.perf_counter()
436
        output.start_time = st
437
438
        most_recent_timestamp = st
        try:
439
440
441
            async with session.post(
                url=api_url, data=form, headers=headers
            ) as response:
442
                if response.status == 200:
443
444
445
                    handler = StreamedResponseHandler()

                    async for chunk_bytes in response.content.iter_any():
446
447
448
449
                        chunk_bytes = chunk_bytes.strip()
                        if not chunk_bytes:
                            continue

450
451
                        messages = handler.add_chunk(chunk_bytes)
                        for message in messages:
452
453
454
                            if type(message) is bytes:
                                message = message.decode("utf-8")
                            chunk = message.removeprefix("data: ")
455
456
457
458
459
                            if chunk != "[DONE]":
                                timestamp = time.perf_counter()
                                data = json.loads(chunk)

                                if choices := data.get("choices"):
460
                                    content = choices[0]["delta"].get("content")
461
462
463
464
465
466
467
468
                                    # First token
                                    if ttft == 0.0:
                                        ttft = timestamp - st
                                        output.ttft = ttft

                                    # Decoding phase
                                    else:
                                        output.itl.append(
469
470
                                            timestamp - most_recent_timestamp
                                        )
471
472
473
474

                                    generated_text += content or ""
                                elif usage := data.get("usage"):
                                    output.output_tokens = usage.get(
475
476
                                        "completion_tokens"
                                    )
477
478

                                most_recent_timestamp = timestamp
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495

                    output.generated_text = generated_text
                    output.success = True
                    output.latency = most_recent_timestamp - st
                else:
                    output.error = response.reason or ""
                    output.success = False
        except Exception:
            output.success = False
            exc_info = sys.exc_info()
            output.error = "".join(traceback.format_exception(*exc_info))

    if pbar:
        pbar.update(1)
    return output


496
async def _run_pooling_request(
497
    session: aiohttp.ClientSession,
498
499
500
    api_url: str,
    payload: dict[str, Any],
    headers: dict[str, Any],
501
    pbar: tqdm | None = None,
502
) -> RequestFuncOutput:
503
504
    output = RequestFuncOutput()
    st = time.perf_counter()
505
    output.start_time = st
506
    try:
507
        async with session.post(url=api_url, headers=headers, json=payload) as response:
508
            if response.status == 200:
509
                output.ttft = output.latency = time.perf_counter() - st
510
511
512
513
514
515
516
517

                if payload.get("encoding_format", "float") == "bytes":
                    metadata = json.loads(response.headers["metadata"])
                    usage = metadata.get("usage", {})
                else:
                    data = await response.json()
                    usage = data.get("usage", {})

518
519
                output.success = True
                output.generated_text = ""
520
                output.prompt_len = usage.get("prompt_tokens", 0)
521
522
523
524
525
526
527
528
529
530
531
532
            else:
                output.success = False
                output.error = response.reason or ""
    except Exception as e:
        output.success = False
        output.error = str(e)

    if pbar:
        pbar.update(1)
    return output


533
534
535
async def async_request_openai_embeddings(
    request_func_input: RequestFuncInput,
    session: aiohttp.ClientSession,
536
    pbar: tqdm | None = None,
537
538
539
540
541
542
543
544
545
) -> RequestFuncOutput:
    api_url = request_func_input.api_url
    _validate_api_url(api_url, "OpenAI Embeddings API", "embeddings")

    payload = {
        "model": request_func_input.model_name
        if request_func_input.model_name
        else request_func_input.model,
        "input": request_func_input.prompt,
546
547
548
        # Many embedding models have short context length,
        # this is to avoid dropping some of the requests.
        "truncate_prompt_tokens": -1,
549
550
551
    }
    _update_payload_common(payload, request_func_input)

552
    headers = _get_headers("application/json")
553
554
    _update_headers_common(headers, request_func_input)

555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
    return await _run_pooling_request(
        session,
        api_url,
        payload=payload,
        headers=headers,
        pbar=pbar,
    )


async def async_request_vllm_rerank(
    request_func_input: RequestFuncInput,
    session: aiohttp.ClientSession,
    pbar: tqdm | None = None,
) -> RequestFuncOutput:
    api_url = request_func_input.api_url
    _validate_api_url(api_url, "vLLM score API", "rerank")

    assert (
        isinstance(request_func_input.prompt, list)
        and len(request_func_input.prompt) > 1
    )

    payload = {
        "model": request_func_input.model_name
        if request_func_input.model_name
        else request_func_input.model,
        "query": request_func_input.prompt[0],
        "documents": request_func_input.prompt[1:],
583
584
585
        # Many reranker models have short context length,
        # this is to avoid dropping some of the requests.
        "truncate_prompt_tokens": -1,
586
587
    }

588
    headers = _get_headers("application/json")
589
590
591
    _update_headers_common(headers, request_func_input)

    return await _run_pooling_request(
592
593
594
595
596
597
598
599
600
601
602
        session,
        api_url,
        payload=payload,
        headers=headers,
        pbar=pbar,
    )


async def async_request_openai_embeddings_chat(
    request_func_input: RequestFuncInput,
    session: aiohttp.ClientSession,
603
    pbar: tqdm | None = None,
604
605
606
607
608
609
610
611
612
613
614
615
616
617
    mm_position: Literal["first", "last"] = "last",
) -> RequestFuncOutput:
    api_url = request_func_input.api_url
    _validate_api_url(api_url, "OpenAI Embeddings API", "embeddings")

    content = _get_chat_content(request_func_input, mm_position=mm_position)

    payload = {
        "model": request_func_input.model_name
        if request_func_input.model_name
        else request_func_input.model,
        "messages": [
            {"role": "user", "content": content},
        ],
618
619
620
        # Many embedding models have short context length,
        # this is to avoid dropping some of the requests.
        "truncate_prompt_tokens": -1,
621
622
623
    }
    _update_payload_common(payload, request_func_input)

624
    headers = _get_headers("application/json")
625
626
    _update_headers_common(headers, request_func_input)

627
    return await _run_pooling_request(
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
        session,
        api_url,
        payload=payload,
        headers=headers,
        pbar=pbar,
    )


def _try_extract_request_idx(request_func_input: RequestFuncInput):
    if request_func_input.request_id:
        match = re.search(r"(\d+)$", request_func_input.request_id)
        if match:
            try:
                return int(match.group(1))
            except ValueError:
                pass

    return None


648
649
650
651
652
653
654
def _preprocess_clip(request_func_input: RequestFuncInput):
    if request_func_input.multi_modal_content:
        # Image input
        request_func_input.prompt = ""


def _preprocess_vlm2vec(request_func_input: RequestFuncInput):
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
    if request_func_input.multi_modal_content:
        request_idx = _try_extract_request_idx(request_func_input)

        # Adjust the ratio manually if needed.
        use_image_only_prompt = request_idx is None or request_idx % 2 == 0

        if use_image_only_prompt:
            # Image input
            request_func_input.prompt = "Represent the given image."
        else:
            # Text+Image input
            request_func_input.prompt = (
                f"Represent the given image with the following question: "
                f"{request_func_input.prompt}"
            )

671
672
673
674

async def async_request_openai_embeddings_clip(
    request_func_input: RequestFuncInput,
    session: aiohttp.ClientSession,
675
    pbar: tqdm | None = None,
676
677
678
679
680
681
682
683
684
685
686
687
688
) -> RequestFuncOutput:
    _preprocess_clip(request_func_input)

    return await async_request_openai_embeddings_chat(
        request_func_input,
        session,
        pbar=pbar,
    )


async def async_request_openai_embeddings_vlm2vec(
    request_func_input: RequestFuncInput,
    session: aiohttp.ClientSession,
689
    pbar: tqdm | None = None,
690
691
692
) -> RequestFuncOutput:
    _preprocess_vlm2vec(request_func_input)

693
694
695
696
697
698
699
700
    return await async_request_openai_embeddings_chat(
        request_func_input,
        session,
        pbar=pbar,
        mm_position="first",
    )


701
702
703
async def async_request_infinity_embeddings(
    request_func_input: RequestFuncInput,
    session: aiohttp.ClientSession,
704
    pbar: tqdm | None = None,
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
) -> RequestFuncOutput:
    api_url = request_func_input.api_url
    _validate_api_url(api_url, "Infinity Embeddings API", "embeddings")

    payload = {
        "model": request_func_input.model_name
        if request_func_input.model_name
        else request_func_input.model,
    }

    if request_func_input.prompt:
        payload["input"] = request_func_input.prompt
    else:
        mm_content = request_func_input.multi_modal_content
        assert isinstance(mm_content, dict)

        mm_type = mm_content["type"]
        payload["input"] = mm_content[mm_type]["url"]
        payload["modality"] = mm_type.split("_", 1)[0]

    _update_payload_common(payload, request_func_input)

727
    headers = _get_headers("application/json")
728
729
    _update_headers_common(headers, request_func_input)

730
    return await _run_pooling_request(
731
732
733
734
735
736
737
738
739
740
741
        session,
        api_url,
        payload=payload,
        headers=headers,
        pbar=pbar,
    )


async def async_request_infinity_embeddings_clip(
    request_func_input: RequestFuncInput,
    session: aiohttp.ClientSession,
742
    pbar: tqdm | None = None,
743
744
745
746
747
748
749
750
751
752
) -> RequestFuncOutput:
    _preprocess_clip(request_func_input)

    return await async_request_infinity_embeddings(
        request_func_input,
        session,
        pbar=pbar,
    )


753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
async def async_request_vllm_pooling(
    request_func_input: RequestFuncInput,
    session: aiohttp.ClientSession,
    pbar: tqdm | None = None,
) -> RequestFuncOutput:
    api_url = request_func_input.api_url
    _validate_api_url(api_url, "vLLM Pooling API", "pooling")

    payload = {
        "model": request_func_input.model_name
        if request_func_input.model_name
        else request_func_input.model,
        "truncate_prompt_tokens": -1,
    }

    payload = payload | request_func_input.prompt

    _update_payload_common(payload, request_func_input)

    headers = _get_headers("application/json")
    _update_headers_common(headers, request_func_input)

    return await _run_pooling_request(
        session,
        api_url,
        payload=payload,
        headers=headers,
        pbar=pbar,
    )


784
# TODO: Add more request functions for different API protocols.
785
ASYNC_REQUEST_FUNCS: dict[str, RequestFunc] = {
786
787
788
789
    "vllm": async_request_openai_completions,
    "openai": async_request_openai_completions,
    "openai-chat": async_request_openai_chat_completions,
    "openai-audio": async_request_openai_audio,
790
    "openai-embeddings": async_request_openai_embeddings,
791
792
793
    "openai-embeddings-chat": async_request_openai_embeddings_chat,
    "openai-embeddings-clip": async_request_openai_embeddings_clip,
    "openai-embeddings-vlm2vec": async_request_openai_embeddings_vlm2vec,
794
795
796
797
    # Infinity embedding server: https://github.com/michaelfeil/infinity
    "infinity-embeddings": async_request_infinity_embeddings,
    "infinity-embeddings-clip": async_request_infinity_embeddings_clip,
    # (Infinity embedding server does not support vlm2vec)
798
    "vllm-pooling": async_request_vllm_pooling,
799
    "vllm-rerank": async_request_vllm_rerank,
800
}
801

802
803
804
805
806
807
808
809
810
811
812
POOLING_BACKENDS = {
    "openai-embeddings",
    "openai-embeddings-chat",
    "openai-embeddings-clip",
    "openai-embeddings-vlm2vec",
    "infinity-embeddings",
    "infinity-embeddings-clip",
    "vllm-pooling",
    "vllm-rerank",
}

813
OPENAI_COMPATIBLE_BACKENDS = [
814
815
816
    k
    for k, v in ASYNC_REQUEST_FUNCS.items()
    if v in (async_request_openai_completions, async_request_openai_chat_completions)
817
]