wehub-resource-sync
43 lines
1.3 KiB
Python
43 lines
1.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
from typing import TYPE_CHECKING
|
|
|
|
from vllm.config import ModelConfig
|
|
from vllm.tasks import SupportedTask
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm.entrypoints.serve.sagemaker.api_router import (
|
|
EndpointFn,
|
|
GetHandlerFn,
|
|
RequestType,
|
|
)
|
|
|
|
|
|
def get_generate_invocation_types(
|
|
supported_tasks: tuple["SupportedTask", ...],
|
|
model_config: ModelConfig | None = None,
|
|
):
|
|
# NOTE: Items defined earlier take higher priority
|
|
invocation_types: list[tuple[RequestType, tuple[GetHandlerFn, EndpointFn]]] = []
|
|
|
|
if "generate" in supported_tasks:
|
|
from vllm.entrypoints.openai.chat_completion.api_router import (
|
|
chat,
|
|
create_chat_completion,
|
|
)
|
|
from vllm.entrypoints.openai.chat_completion.protocol import (
|
|
ChatCompletionRequest,
|
|
)
|
|
from vllm.entrypoints.openai.completion.api_router import (
|
|
completion,
|
|
create_completion,
|
|
)
|
|
from vllm.entrypoints.openai.completion.protocol import CompletionRequest
|
|
|
|
invocation_types += [
|
|
(ChatCompletionRequest, (chat, create_chat_completion)),
|
|
(CompletionRequest, (completion, create_completion)),
|
|
]
|
|
|
|
return invocation_types
|