"""The MLC LLM Serving Engine.""" import asyncio import queue import weakref from collections.abc import AsyncGenerator, Iterator from typing import ( # noqa: UP035 Any, Dict, List, Literal, Optional, Tuple, Union, overload, ) from tvm.runtime import Device from mlc_llm.protocol import debug_protocol, openai_api_protocol from mlc_llm.protocol.generation_config import GenerationConfig from mlc_llm.serve import data, engine_utils from mlc_llm.serve.config import EngineConfig from mlc_llm.support import logging from mlc_llm.tokenizers import TextStreamer from . import engine_base logger = logging.getLogger(__name__) # Note: we define both AsyncChat and Chat for Python type analysis. class AsyncChat: """The proxy class to direct to async chat completions.""" def __init__(self, engine: weakref.ReferenceType) -> None: assert isinstance(engine(), AsyncMLCEngine) self.completions = AsyncChatCompletion(engine) class Chat: """The proxy class to direct to chat completions.""" def __init__(self, engine: weakref.ReferenceType) -> None: assert isinstance(engine(), MLCEngine) self.completions = ChatCompletion(engine) class AsyncChatCompletion: """The proxy class to direct to async chat completions.""" engine: weakref.ReferenceType["AsyncMLCEngine"] def __init__(self, engine: weakref.ReferenceType) -> None: self.engine = engine @overload async def create( self, *, messages: List[Dict[str, Any]], # noqa: UP006 stream: Literal[True], model: Optional[str] = None, frequency_penalty: Optional[float] = None, presence_penalty: Optional[float] = None, logprobs: bool = False, top_logprobs: int = 0, logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 max_tokens: Optional[int] = None, n: int = 1, seed: Optional[int] = None, stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 temperature: Optional[float] = None, top_p: Optional[float] = None, tools: Optional[List[Dict[str, Any]]] = None, # noqa: UP006 tool_choice: Optional[Union[Literal["none", "auto"], Dict]] = None, # noqa: UP006 user: Optional[str] = None, response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 request_id: Optional[str] = None, extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 ) -> AsyncGenerator[openai_api_protocol.ChatCompletionStreamResponse, Any]: """Asynchronous streaming chat completion interface with OpenAI API compatibility. The method is a coroutine that streams ChatCompletionStreamResponse that conforms to OpenAI API one at a time via yield. See https://platform.openai.com/docs/api-reference/chat/create for specification. Parameters ---------- request_id : Optional[str] The optional request id. A random one will be generated if it is not given. extra_body: Optional[Dict[str, Any]] = None, Extra body options to pass to the request. Can be used to pass debug config as extra_body["debug_config"] Yields ------ stream_response : ChatCompletionStreamResponse The stream response conforming to OpenAI API. See mlc_llm/protocol/openai_api_protocol.py or https://platform.openai.com/docs/api-reference/chat/streaming for specification. Raises ------ e : BadRequestError BadRequestError is raised when the request is invalid. """ @overload async def create( self, *, messages: List[Dict[str, Any]], # noqa: UP006 model: Optional[str] = None, frequency_penalty: Optional[float] = None, presence_penalty: Optional[float] = None, logprobs: bool = False, top_logprobs: int = 0, logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 max_tokens: Optional[int] = None, n: int = 1, seed: Optional[int] = None, stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 stream: Literal[False] = False, stream_options: Literal[None] = None, temperature: Optional[float] = None, top_p: Optional[float] = None, tools: Optional[List[Dict[str, Any]]] = None, # noqa: UP006 tool_choice: Optional[Union[Literal["none", "auto"], Dict]] = None, # noqa: UP006 user: Optional[str] = None, response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 request_id: Optional[str] = None, extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 ) -> openai_api_protocol.ChatCompletionResponse: """Asynchronous non-streaming chat completion interface with OpenAI API compatibility. The method is a coroutine that streams ChatCompletionStreamResponse that conforms to OpenAI API one at a time via yield. See https://platform.openai.com/docs/api-reference/chat/create for specification. Parameters ---------- request_id : Optional[str] The optional request id. A random one will be generated if it is not given. extra_body: Optional[Dict[str, Any]] = None, Extra body options to pass to the request. Can be used to pass debug config as extra_body["debug_config"] Returns ------- response : ChatCompletionResponse The chat completion response conforming to OpenAI API. See mlc_llm/protocol/openai_api_protocol.py or https://platform.openai.com/docs/api-reference/chat/object for specification. Raises ------ e : BadRequestError BadRequestError is raised when the request is invalid. """ async def create( self, *, messages: List[Dict[str, Any]], # noqa: UP006 model: Optional[str] = None, frequency_penalty: Optional[float] = None, presence_penalty: Optional[float] = None, logprobs: bool = False, top_logprobs: int = 0, logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 max_tokens: Optional[int] = None, n: int = 1, seed: Optional[int] = None, stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 stream: bool = False, stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 temperature: Optional[float] = None, top_p: Optional[float] = None, tools: Optional[List[Dict[str, Any]]] = None, # noqa: UP006 tool_choice: Optional[Union[Literal["none", "auto"], Dict]] = None, # noqa: UP006 user: Optional[str] = None, response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 request_id: Optional[str] = None, extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 ) -> Union[ AsyncGenerator[openai_api_protocol.ChatCompletionStreamResponse, Any], openai_api_protocol.ChatCompletionResponse, ]: """Asynchronous chat completion interface with OpenAI API compatibility. See https://platform.openai.com/docs/api-reference/chat/create for specification. Parameters ---------- request_id : Optional[str] The optional request id. A random one will be generated if it is not given. extra_body: Optional[Dict[str, Any]] = None, Extra body options to pass to the request. Can be used to pass debug config as extra_body["debug_config"] Raises ------ e : BadRequestError BadRequestError is raised when the request is invalid. """ return await self.engine()._chat_completion( messages=messages, model=model, frequency_penalty=frequency_penalty, presence_penalty=presence_penalty, logprobs=logprobs, top_logprobs=top_logprobs, logit_bias=logit_bias, max_tokens=max_tokens, n=n, seed=seed, stop=stop, stream=stream, stream_options=( openai_api_protocol.StreamOptions.model_validate(stream_options) if stream_options is not None else None ), temperature=temperature, top_p=top_p, tools=tools, tool_choice=tool_choice, user=user, response_format=response_format, request_id=request_id, debug_config=(extra_body.get("debug_config", None) if extra_body is not None else None), ) class ChatCompletion: """The proxy class to direct to chat completions.""" engine: weakref.ReferenceType["MLCEngine"] def __init__(self, engine: weakref.ReferenceType) -> None: self.engine = engine @overload def create( self, *, messages: List[Dict[str, Any]], # noqa: UP006 stream: Literal[True], model: Optional[str] = None, frequency_penalty: Optional[float] = None, presence_penalty: Optional[float] = None, logprobs: bool = False, top_logprobs: int = 0, logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 max_tokens: Optional[int] = None, n: int = 1, seed: Optional[int] = None, stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 temperature: Optional[float] = None, top_p: Optional[float] = None, tools: Optional[List[Dict[str, Any]]] = None, # noqa: UP006 tool_choice: Optional[Union[Literal["none", "auto"], Dict]] = None, # noqa: UP006 user: Optional[str] = None, response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 request_id: Optional[str] = None, extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 ) -> Iterator[openai_api_protocol.ChatCompletionStreamResponse]: """Synchronous streaming chat completion interface with OpenAI API compatibility. The method streams back ChatCompletionStreamResponse that conforms to OpenAI API one at a time via yield. See https://platform.openai.com/docs/api-reference/chat/create for specification. Parameters ---------- request_id : Optional[str] The optional request id. A random one will be generated if it is not given. extra_body: Optional[Dict[str, Any]] = None, Extra body options to pass to the request. Can be used to pass debug config as extra_body["debug_config"] Yields ------ stream_response : ChatCompletionStreamResponse The stream response conforming to OpenAI API. See mlc_llm/protocol/openai_api_protocol.py or https://platform.openai.com/docs/api-reference/chat/streaming for specification. Raises ------ e : BadRequestError BadRequestError is raised when the request is invalid. """ @overload def create( self, *, messages: List[Dict[str, Any]], # noqa: UP006 model: Optional[str] = None, frequency_penalty: Optional[float] = None, presence_penalty: Optional[float] = None, logprobs: bool = False, top_logprobs: int = 0, logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 max_tokens: Optional[int] = None, n: int = 1, seed: Optional[int] = None, stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 stream: Literal[False] = False, stream_options: Literal[None] = None, temperature: Optional[float] = None, top_p: Optional[float] = None, tools: Optional[List[Dict[str, Any]]] = None, # noqa: UP006 tool_choice: Optional[Union[Literal["none", "auto"], Dict]] = None, # noqa: UP006 user: Optional[str] = None, response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 request_id: Optional[str] = None, extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 ) -> openai_api_protocol.ChatCompletionResponse: """Synchronous non-streaming chat completion interface with OpenAI API compatibility. See https://platform.openai.com/docs/api-reference/chat/create for specification. Parameters ---------- request_id : Optional[str] The optional request id. A random one will be generated if it is not given. extra_body: Optional[Dict[str, Any]] = None, Extra body options to pass to the request. Can be used to pass debug config as extra_body["debug_config"] Returns ------ response : ChatCompletionResponse The chat completion response conforming to OpenAI API. See mlc_llm/protocol/openai_api_protocol.py or https://platform.openai.com/docs/api-reference/chat/object for specification. Raises ------ e : BadRequestError BadRequestError is raised when the request is invalid. """ def create( self, *, messages: List[Dict[str, Any]], # noqa: UP006 model: Optional[str] = None, frequency_penalty: Optional[float] = None, presence_penalty: Optional[float] = None, logprobs: bool = False, top_logprobs: int = 0, logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 max_tokens: Optional[int] = None, n: int = 1, seed: Optional[int] = None, stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 stream: bool = False, stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 temperature: Optional[float] = None, top_p: Optional[float] = None, tools: Optional[List[Dict[str, Any]]] = None, # noqa: UP006 tool_choice: Optional[Union[Literal["none", "auto"], Dict]] = None, # noqa: UP006 user: Optional[str] = None, response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 request_id: Optional[str] = None, extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 ) -> Union[ Iterator[openai_api_protocol.ChatCompletionStreamResponse], openai_api_protocol.ChatCompletionResponse, ]: """Synchronous chat completion interface with OpenAI API compatibility. See https://platform.openai.com/docs/api-reference/chat/create for specification. Parameters ---------- request_id : Optional[str] The optional request id. A random one will be generated if it is not given. extra_body: Optional[Dict[str, Any]] = None, Extra body options to pass to the request. Can be used to pass debug config as extra_body["debug_config"] Raises ------ e : BadRequestError BadRequestError is raised when the request is invalid. """ return self.engine()._chat_completion( messages=messages, model=model, frequency_penalty=frequency_penalty, presence_penalty=presence_penalty, logprobs=logprobs, top_logprobs=top_logprobs, logit_bias=logit_bias, max_tokens=max_tokens, n=n, seed=seed, stop=stop, stream=stream, stream_options=( openai_api_protocol.StreamOptions.model_validate(stream_options) if stream_options is not None else None ), temperature=temperature, top_p=top_p, tools=tools, tool_choice=tool_choice, user=user, response_format=response_format, request_id=request_id, debug_config=(extra_body.get("debug_config", None) if extra_body is not None else None), ) class AsyncCompletion: """The proxy class to direct to async completions.""" engine: weakref.ReferenceType["AsyncMLCEngine"] def __init__(self, engine: weakref.ReferenceType) -> None: self.engine = engine @overload async def create( self, *, prompt: Union[str, List[int]], # noqa: UP006 stream: Literal[True], model: Optional[str] = None, best_of: int = 1, echo: bool = False, frequency_penalty: Optional[float] = None, presence_penalty: Optional[float] = None, logprobs: Optional[int] = None, logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 max_tokens: Optional[int] = None, n: int = 1, seed: Optional[int] = None, stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 suffix: Optional[str] = None, temperature: Optional[float] = None, top_p: Optional[float] = None, user: Optional[str] = None, response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 request_id: Optional[str] = None, extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 ) -> AsyncGenerator[openai_api_protocol.CompletionResponse, Any]: """Asynchronous streaming completion interface with OpenAI API compatibility. The method is a coroutine that streams CompletionResponse that conforms to OpenAI API one at a time via yield. See https://platform.openai.com/docs/api-reference/completions/create for specification. Parameters ---------- request_id : Optional[str] The optional request id. A random one will be generated if it is not given. extra_body: Optional[Dict[str, Any]] = None, Extra body options to pass to the request. Can be used to pass debug config as extra_body["debug_config"] Yields ------ stream_response : CompletionResponse The stream response conforming to OpenAI API. See mlc_llm/protocol/openai_api_protocol.py or https://platform.openai.com/docs/api-reference/completions/object for specification. Raises ------ e : BadRequestError BadRequestError is raised when the request is invalid. """ @overload async def create( self, *, prompt: Union[str, List[int]], # noqa: UP006 model: Optional[str] = None, best_of: int = 1, echo: bool = False, frequency_penalty: Optional[float] = None, presence_penalty: Optional[float] = None, logprobs: Optional[int] = None, logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 max_tokens: Optional[int] = None, n: int = 1, seed: Optional[int] = None, stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 stream: Literal[False] = False, stream_options: Literal[None] = None, suffix: Optional[str] = None, temperature: Optional[float] = None, top_p: Optional[float] = None, user: Optional[str] = None, response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 request_id: Optional[str] = None, extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 ) -> openai_api_protocol.CompletionResponse: """Asynchronous non-streaming completion interface with OpenAI API compatibility. See https://platform.openai.com/docs/api-reference/completions/create for specification. Parameters ---------- request_id : Optional[str] The optional request id. A random one will be generated if it is not given. extra_body: Optional[Dict[str, Any]] = None, Extra body options to pass to the request. Can be used to pass debug config as extra_body["debug_config"] Returns ------ response : CompletionResponse The completion response conforming to OpenAI API. See mlc_llm/protocol/openai_api_protocol.py or https://platform.openai.com/docs/api-reference/completions/object for specification. Raises ------ e : BadRequestError BadRequestError is raised when the request is invalid. """ async def create( self, *, prompt: Union[str, List[int]], # noqa: UP006 model: Optional[str] = None, best_of: int = 1, echo: bool = False, frequency_penalty: Optional[float] = None, presence_penalty: Optional[float] = None, logprobs: Optional[int] = None, logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 max_tokens: Optional[int] = None, n: int = 1, seed: Optional[int] = None, stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 stream: bool = False, stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 suffix: Optional[str] = None, temperature: Optional[float] = None, top_p: Optional[float] = None, user: Optional[str] = None, response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 request_id: Optional[str] = None, extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 ) -> Union[ AsyncGenerator[openai_api_protocol.CompletionResponse, Any], openai_api_protocol.CompletionResponse, ]: """Asynchronous completion interface with OpenAI API compatibility. See https://platform.openai.com/docs/api-reference/completions/create for specification. Parameters ---------- request_id : Optional[str] The optional request id. A random one will be generated if it is not given. extra_body: Optional[Dict[str, Any]] = None, Extra body options to pass to the request. Can be used to pass debug config as extra_body["debug_config"] Raises ------ e : BadRequestError BadRequestError is raised when the request is invalid. """ return await self.engine()._completion( model=model, prompt=prompt, best_of=best_of, echo=echo, frequency_penalty=frequency_penalty, presence_penalty=presence_penalty, logprobs=logprobs, logit_bias=logit_bias, max_tokens=max_tokens, n=n, seed=seed, stop=stop, stream=stream, stream_options=( openai_api_protocol.StreamOptions.model_validate(stream_options) if stream_options is not None else None ), suffix=suffix, temperature=temperature, top_p=top_p, user=user, response_format=response_format, request_id=request_id, debug_config=(extra_body.get("debug_config", None) if extra_body is not None else None), ) class Completion: """The proxy class to direct to completions.""" engine: weakref.ReferenceType["MLCEngine"] def __init__(self, engine: weakref.ReferenceType) -> None: self.engine = engine @overload def create( self, *, prompt: Union[str, List[int]], # noqa: UP006 stream: Literal[True], model: Optional[str] = None, best_of: int = 1, echo: bool = False, frequency_penalty: Optional[float] = None, presence_penalty: Optional[float] = None, logprobs: Optional[int] = None, logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 max_tokens: Optional[int] = None, n: int = 1, seed: Optional[int] = None, stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 suffix: Optional[str] = None, temperature: Optional[float] = None, top_p: Optional[float] = None, user: Optional[str] = None, response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 request_id: Optional[str] = None, extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 ) -> Iterator[openai_api_protocol.CompletionResponse]: """Synchronous streaming completion interface with OpenAI API compatibility. The method streams back CompletionResponse that conforms to OpenAI API one at a time via yield. See https://platform.openai.com/docs/api-reference/completions/create for specification. Parameters ---------- request_id : Optional[str] The optional request id. A random one will be generated if it is not given. extra_body: Optional[Dict[str, Any]] = None, Extra body options to pass to the request. Can be used to pass debug config as extra_body["debug_config"] Yields ------ stream_response : CompletionResponse The stream response conforming to OpenAI API. See mlc_llm/protocol/openai_api_protocol.py or https://platform.openai.com/docs/api-reference/completions/object for specification. Raises ------ e : BadRequestError BadRequestError is raised when the request is invalid. """ @overload def create( self, *, prompt: Union[str, List[int]], # noqa: UP006 model: Optional[str] = None, best_of: int = 1, echo: bool = False, frequency_penalty: Optional[float] = None, presence_penalty: Optional[float] = None, logprobs: Optional[int] = None, logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 max_tokens: Optional[int] = None, n: int = 1, seed: Optional[int] = None, stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 stream: Literal[False] = False, stream_options: Literal[None] = None, suffix: Optional[str] = None, temperature: Optional[float] = None, top_p: Optional[float] = None, user: Optional[str] = None, response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 request_id: Optional[str] = None, extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 ) -> openai_api_protocol.CompletionResponse: """Synchronous non-streaming completion interface with OpenAI API compatibility. See https://platform.openai.com/docs/api-reference/completions/create for specification. Parameters ---------- request_id : Optional[str] The optional request id. A random one will be generated if it is not given. extra_body: Optional[Dict[str, Any]] = None, Extra body options to pass to the request. Can be used to pass debug config as extra_body["debug_config"] Returns ------- response : CompletionResponse The completion response conforming to OpenAI API. See mlc_llm/protocol/openai_api_protocol.py or https://platform.openai.com/docs/api-reference/completions/object for specification. Raises ------ e : BadRequestError BadRequestError is raised when the request is invalid. """ def create( self, *, prompt: Union[str, List[int]], # noqa: UP006 model: Optional[str] = None, best_of: int = 1, echo: bool = False, frequency_penalty: Optional[float] = None, presence_penalty: Optional[float] = None, logprobs: Optional[int] = None, logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 max_tokens: Optional[int] = None, n: int = 1, seed: Optional[int] = None, stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 stream: bool = False, stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 suffix: Optional[str] = None, temperature: Optional[float] = None, top_p: Optional[float] = None, user: Optional[str] = None, response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 request_id: Optional[str] = None, extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 ) -> Union[ Iterator[openai_api_protocol.CompletionResponse], openai_api_protocol.CompletionResponse, ]: """Synchronous completion interface with OpenAI API compatibility. See https://platform.openai.com/docs/api-reference/completions/create for specification. Parameters ---------- request_id : Optional[str] The optional request id. A random one will be generated if it is not given. extra_body: Optional[Dict[str, Any]] = None, Extra body options to pass to the request. Can be used to pass debug config as extra_body["debug_config"] Raises ------ e : BadRequestError BadRequestError is raised when the request is invalid. """ return self.engine()._completion( model=model, prompt=prompt, best_of=best_of, echo=echo, frequency_penalty=frequency_penalty, presence_penalty=presence_penalty, logprobs=logprobs, logit_bias=logit_bias, max_tokens=max_tokens, n=n, seed=seed, stop=stop, stream=stream, stream_options=( openai_api_protocol.StreamOptions.model_validate(stream_options) if stream_options is not None else None ), suffix=suffix, temperature=temperature, top_p=top_p, user=user, response_format=response_format, request_id=request_id, debug_config=(extra_body.get("debug_config", None) if extra_body is not None else None), ) class AsyncMLCEngine(engine_base.MLCEngineBase): """The AsyncMLCEngine in MLC LLM that provides the asynchronous interfaces with regard to OpenAI API. Parameters ---------- model : str A path to ``mlc-chat-config.json``, or an MLC model directory that contains `mlc-chat-config.json`. It can also be a link to a HF repository pointing to an MLC compiled model. device: Union[str, Device] The device used to deploy the model such as "cuda" or "cuda:0". Will default to "auto" and detect from local available GPUs if not specified. model_lib : Optional[str] The full path to the model library file to use (e.g. a ``.so`` file). If unspecified, we will use the provided ``model`` to search over possible paths. It the model lib is not found, it will be compiled in a JIT manner. mode : Literal["local", "interactive", "server"] The engine mode in MLC LLM. We provide three preset modes: "local", "interactive" and "server". The default mode is "local". The choice of mode decides the values of "max_num_sequence", "max_total_sequence_length" and "prefill_chunk_size" when they are not explicitly specified. 1. Mode "local" refers to the local server deployment which has low request concurrency. So the max batch size will be set to 4, and max total sequence length and prefill chunk size are set to the context window size (or sliding window size) of the model. 2. Mode "interactive" refers to the interactive use of server, which has at most 1 concurrent request. So the max batch size will be set to 1, and max total sequence length and prefill chunk size are set to the context window size (or sliding window size) of the model. 3. Mode "server" refers to the large server use case which may handle many concurrent request and want to use GPU memory as much as possible. In this mode, we will automatically infer the largest possible max batch size and max total sequence length. You can manually specify arguments "max_num_sequence", "max_total_sequence_length" and "prefill_chunk_size" to override the automatic inferred values. engine_config : Optional[EngineConfig] Additional configurable arguments of MLC engine. See class "EngineConfig" for more detail. enable_tracing : bool A boolean indicating if to enable event logging for requests. """ def __init__( self, model: str, device: Union[str, Device] = "auto", *, model_lib: Optional[str] = None, mode: Literal["local", "interactive", "server"] = "local", engine_config: Optional[EngineConfig] = None, enable_tracing: bool = False, ) -> None: super().__init__( "async", model=model, device=device, model_lib=model_lib, mode=mode, engine_config=engine_config, enable_tracing=enable_tracing, ) self.chat = AsyncChat(weakref.ref(self)) self.completions = AsyncCompletion(weakref.ref(self)) async def abort(self, request_id: str) -> None: """Generation abortion interface. Parameters --------- request_id : str The id of the request to abort. """ self._abort(request_id) async def metrics(self) -> engine_base.EngineMetrics: """Get engine metrics Returns ------- metrics: EngineMetrics The engine metrics """ return await engine_base._async_query_engine_metrics(self) async def _chat_completion( self, *, messages: List[Dict[str, Any]], # noqa: UP006 model: Optional[str] = None, frequency_penalty: Optional[float] = None, presence_penalty: Optional[float] = None, logprobs: bool = False, top_logprobs: int = 0, logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 max_tokens: Optional[int] = None, n: int = 1, seed: Optional[int] = None, stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 stream: bool = False, stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 temperature: Optional[float] = None, top_p: Optional[float] = None, tools: Optional[List[Dict[str, Any]]] = None, # noqa: UP006 tool_choice: Optional[Union[Literal["none", "auto"], Dict]] = None, # noqa: UP006 user: Optional[str] = None, response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 request_id: Optional[str] = None, debug_config: Optional[Dict[str, Any]] = None, # noqa: UP006 ) -> Union[ AsyncGenerator[openai_api_protocol.ChatCompletionStreamResponse, Any], openai_api_protocol.ChatCompletionResponse, ]: """Asynchronous chat completion internal interface with OpenAI API compatibility. See https://platform.openai.com/docs/api-reference/chat/create for specification. Parameters ---------- request_id : Optional[str] The optional request id. A random one will be generated if it is not given. Extra body options to pass to the request. Can be used to pass debug config as extra_body["debug_config"] debug_config: Optional[Dict[str, Any]] = None, Debug config body options to pass to the request. Raises ------ e : BadRequestError BadRequestError is raised when the request is invalid. """ if request_id is None: request_id = f"chatcmpl-{engine_utils.random_uuid()}" chatcmpl_generator = self._handle_chat_completion( openai_api_protocol.ChatCompletionRequest( messages=[ openai_api_protocol.ChatCompletionMessage.model_validate(message) for message in messages ], model=model, frequency_penalty=frequency_penalty, presence_penalty=presence_penalty, logprobs=logprobs, top_logprobs=top_logprobs, logit_bias=logit_bias, max_tokens=max_tokens, n=n, seed=seed, stop=stop, stream=stream, stream_options=( openai_api_protocol.StreamOptions.model_validate(stream_options) if stream_options is not None else None ), temperature=temperature, top_p=top_p, tools=( [openai_api_protocol.ChatTool.model_validate(tool) for tool in tools] if tools is not None else None ), tool_choice=tool_choice, user=user, response_format=( openai_api_protocol.RequestResponseFormat.model_validate(response_format) if response_format is not None else None ), debug_config=( debug_protocol.DebugConfig.model_validate(debug_config) if debug_config is not None else None ), ), request_id=request_id, request_final_usage_include_extra=True, ) if stream: # Stream response. return chatcmpl_generator # Normal response. output_texts = ["" for _ in range(n)] finish_reasons: List[Optional[str]] = [None for _ in range(n)] # noqa: UP006 logprob_results: Optional[List[List[openai_api_protocol.LogProbsContent]]] = ( # noqa: UP006 [[] for _ in range(n)] if logprobs else None ) request_final_usage = None try: async for response in chatcmpl_generator: # when usage is not None this is the last chunk if response.usage is not None: request_final_usage = response.usage continue for choice in response.choices: assert isinstance(choice.delta.content, str) output_texts[choice.index] += choice.delta.content if choice.finish_reason is not None and finish_reasons[choice.index] is None: finish_reasons[choice.index] = choice.finish_reason if choice.logprobs is not None: assert logprob_results is not None logprob_results[choice.index] += choice.logprobs.content except asyncio.CancelledError: # for cancelled error, we can simply pass it through raise except Exception as err: logger.error("Error in chat completion with request ID %s: %s", request_id, err) raise assert all(finish_reason is not None for finish_reason in finish_reasons) use_function_calling, tool_calls_list = engine_base.process_function_call_output( output_texts, finish_reasons ) return engine_base.wrap_chat_completion_response( request_id=request_id, model=model, output_texts=output_texts, finish_reasons=finish_reasons, tool_calls_list=tool_calls_list, logprob_results=logprob_results, use_function_calling=use_function_calling, usage=request_final_usage, ) async def _completion( self, *, prompt: Union[str, List[int]], # noqa: UP006 model: Optional[str] = None, best_of: int = 1, echo: bool = False, frequency_penalty: Optional[float] = None, presence_penalty: Optional[float] = None, logprobs: Optional[int] = None, logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 max_tokens: Optional[int] = None, n: int = 1, seed: Optional[int] = None, stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 stream: bool = False, stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 suffix: Optional[str] = None, temperature: Optional[float] = None, top_p: Optional[float] = None, user: Optional[str] = None, response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 request_id: Optional[str] = None, debug_config: Optional[Dict[str, Any]] = None, # noqa: UP006 ) -> Union[ AsyncGenerator[openai_api_protocol.CompletionResponse, Any], openai_api_protocol.CompletionResponse, ]: """Asynchronous completion internal interface with OpenAI API compatibility. See https://platform.openai.com/docs/api-reference/completions/create for specification. Parameters ---------- request_id : Optional[str] The optional request id. A random one will be generated if it is not given. debug_config: Optional[Dict[str, Any]] = None, Extra debug options to pass to the request. Raises ------ e : BadRequestError BadRequestError is raised when the request is invalid. """ if request_id is None: request_id = f"cmpl-{engine_utils.random_uuid()}" cmpl_generator = self._handle_completion( openai_api_protocol.CompletionRequest( model=model, prompt=prompt, best_of=best_of, echo=echo, frequency_penalty=frequency_penalty, presence_penalty=presence_penalty, logprobs=logprobs, logit_bias=logit_bias, max_tokens=max_tokens, n=n, seed=seed, stop=stop, stream=stream, stream_options=( openai_api_protocol.StreamOptions.model_validate(stream_options) if stream_options is not None else None ), suffix=suffix, temperature=temperature, top_p=top_p, user=user, response_format=( openai_api_protocol.RequestResponseFormat.model_validate(response_format) if response_format is not None else None ), debug_config=( debug_protocol.DebugConfig.model_validate(debug_config) if debug_config is not None else None ), ), request_id=request_id, request_final_usage_include_extra=True, ) if stream: # Stream response. return cmpl_generator # Normal response. request_final_usage = None output_texts = [""] * n finish_reasons: List[Optional[str]] = [None] * n # noqa: UP006 logprob_results: List[Optional[openai_api_protocol.CompletionLogProbs]] = [None] * n # noqa: UP006 async for response in cmpl_generator: # this is the final chunk if response.usage is not None: request_final_usage = response.usage continue for choice in response.choices: output_texts[choice.index] += choice.text if choice.finish_reason is not None and finish_reasons[choice.index] is None: finish_reasons[choice.index] = choice.finish_reason if choice.logprobs is not None: logprob_results[choice.index] = choice.logprobs assert all(finish_reason is not None for finish_reason in finish_reasons) return engine_base.wrap_completion_response( request_id=request_id, model=model, output_texts=output_texts, finish_reasons=finish_reasons, logprob_results=logprob_results, usage=request_final_usage, ) async def _handle_chat_completion( self, request: openai_api_protocol.ChatCompletionRequest, request_id: str, request_final_usage_include_extra: bool, ) -> AsyncGenerator[openai_api_protocol.ChatCompletionStreamResponse, Any]: """The implementation fo asynchronous ChatCompletionRequest handling. Yields ------ stream_response : ChatCompletionStreamResponse The stream response conforming to OpenAI API. See mlc_llm/protocol/openai_api_protocol.py or https://platform.openai.com/docs/api-reference/chat/streaming for specification. Raises ------ e : BadRequestError BadRequestError is raised when the request is invalid. """ ( prompts, generation_cfg, use_function_calling, prompt_length, ) = engine_base.process_chat_completion_request( request, request_id, self.state, self.model_config_dicts[0], self.tokenizer.encode, self.max_input_sequence_length, self.conv_template.model_copy(deep=True), ) # prompt length is not used _ = prompt_length finish_reasons: List[Optional[str]] = [None for _ in range(generation_cfg.n)] # noqa: UP006 self.state.record_event(request_id, event="invoke generate") try: async for delta_outputs in self._generate( prompts, generation_cfg, request_id, ): response = engine_base.process_chat_completion_stream_output( delta_outputs, request, request_id, self.state, use_function_calling, finish_reasons, ) if response is not None: if response.usage is not None: if not request_final_usage_include_extra: response.usage.extra = None yield response self.state.record_event(request_id, event="finish") except asyncio.CancelledError: # for cancelled error, we can simply pass it through raise except Exception as err: logger.error("Error in _handle_chat_completion for request %s: %s", request_id, err) raise async def _handle_completion( self, request: openai_api_protocol.CompletionRequest, request_id: str, request_final_usage_include_extra: bool, ) -> AsyncGenerator[openai_api_protocol.CompletionResponse, Any]: """The implementation fo asynchronous CompletionRequest handling. Yields ------ stream_response : CompletionResponse The stream response conforming to OpenAI API. See mlc_llm/protocol/openai_api_protocol.py or https://platform.openai.com/docs/api-reference/completions/object for specification. Raises ------ e : BadRequestError BadRequestError is raised when the request is invalid. """ ( prompt, generation_cfg, prompt_length, echo_response, ) = engine_base.process_completion_request( request, request_id, self.state, self.tokenizer, self.max_input_sequence_length, self.conv_template.model_copy(deep=True), ) _ = prompt_length if echo_response is not None: yield echo_response finish_reasons: List[Optional[str]] = [None] * generation_cfg.n # noqa: UP006 self.state.record_event(request_id, event="invoke generate") try: async for delta_outputs in self._generate( prompt, generation_cfg, request_id, ): response = engine_base.process_completion_stream_output( delta_outputs, request, request_id, self.state, finish_reasons, ) if response is not None: if response.usage is not None: if not request_final_usage_include_extra: response.usage.extra = None yield response suffix_response = engine_base.create_completion_suffix_response( request, request_id, finish_reasons ) if suffix_response is not None: yield suffix_response self.state.record_event(request_id, event="finish") except asyncio.CancelledError: # for cancelled error, we can simply pass it through raise except Exception as err: logger.error("Error in _handle_completion for request %s: %s", request_id, err) raise async def _generate( self, prompt: Union[str, List[int], List[Union[str, List[int], data.Data]]], # noqa: UP006 generation_config: GenerationConfig, request_id: str, ) -> AsyncGenerator[List[engine_base.CallbackStreamOutput], Any]: # noqa: UP006 """Internal asynchronous text generation interface of AsyncMLCEngine. The method is a coroutine that streams a list of CallbackStreamOutput at a time via yield. The returned list length is the number of parallel generations specified by `generation_config.n`. Parameters ---------- prompt : Union[str, List[int], List[Union[str, List[int], data.Data]]] The input prompt in forms of text strings, lists of token ids or data. generation_config : GenerationConfig The generation config of the request. request_id : str The unique identifier (in string) or this generation request. Yields ------ request_output : List[engine_base.CallbackStreamOutput] The delta generated outputs in a list. The number of list elements equals to `generation_config.n`, and each element corresponds to the delta output of a parallel generation. """ if self._terminated: raise ValueError("The AsyncThreadedEngine has terminated.") self.state.async_lazy_init_event_loop() # Create the request with the given id, input data, generation # config and the created callback. input_data = engine_utils.convert_prompts_to_data(prompt) request = self._ffi["create_request"]( request_id, input_data, generation_config.model_dump_json(by_alias=True) ) # Create the unique async request stream of the request. stream = engine_base.AsyncRequestStream() if request_id in self.state.async_streamers: # Report error in the stream if the request id already exists. stream.push( RuntimeError( f'The request id "{request_id} already exists. ' 'Please make sure the request id is unique."' ) ) else: # Record the stream in the tracker self.state.async_streamers[request_id] = ( stream, [TextStreamer(self.tokenizer) for _ in range(generation_config.n)], ) self._ffi["add_request"](request) def abort_request(): """clean up""" self._abort(request_id) logger.info("request %s cancelled", request_id) with engine_utils.ErrorCleanupScope(abort_request): # Iterate the stream asynchronously and yield the output. try: async for request_output in stream: yield request_output except asyncio.CancelledError: # for cancelled error, we can simply pass it through raise except Exception as exception: logger.error("Exception in _generate for request %s: %s", request_id, exception) raise def _abort(self, request_id: str): """Internal implementation of request abortion.""" self.state.async_streamers.pop(request_id, None) self._ffi["abort_request"](request_id) class MLCEngine(engine_base.MLCEngineBase): """The MLCEngine in MLC LLM that provides the synchronous interfaces with regard to OpenAI API. Parameters ---------- model : str A path to ``mlc-chat-config.json``, or an MLC model directory that contains `mlc-chat-config.json`. It can also be a link to a HF repository pointing to an MLC compiled model. device: Union[str, Device] The device used to deploy the model such as "cuda" or "cuda:0". Will default to "auto" and detect from local available GPUs if not specified. model_lib : Optional[str] The full path to the model library file to use (e.g. a ``.so`` file). If unspecified, we will use the provided ``model`` to search over possible paths. It the model lib is not found, it will be compiled in a JIT manner. mode : Literal["local", "interactive", "server"] The engine mode in MLC LLM. We provide three preset modes: "local", "interactive" and "server". The default mode is "local". The choice of mode decides the values of "max_num_sequence", "max_total_sequence_length" and "prefill_chunk_size" when they are not explicitly specified. 1. Mode "local" refers to the local server deployment which has low request concurrency. So the max batch size will be set to 4, and max total sequence length and prefill chunk size are set to the context window size (or sliding window size) of the model. 2. Mode "interactive" refers to the interactive use of server, which has at most 1 concurrent request. So the max batch size will be set to 1, and max total sequence length and prefill chunk size are set to the context window size (or sliding window size) of the model. 3. Mode "server" refers to the large server use case which may handle many concurrent request and want to use GPU memory as much as possible. In this mode, we will automatically infer the largest possible max batch size and max total sequence length. You can manually specify arguments "max_num_sequence", "max_total_sequence_length" and "prefill_chunk_size" to override the automatic inferred values. engine_config : Optional[EngineConfig] Additional configurable arguments of MLC engine. See class "EngineConfig" for more detail. enable_tracing : bool A boolean indicating if to enable event logging for requests. """ def __init__( self, model: str, device: Union[str, Device] = "auto", *, model_lib: Optional[str] = None, mode: Literal["local", "interactive", "server"] = "local", engine_config: Optional[EngineConfig] = None, enable_tracing: bool = False, ) -> None: super().__init__( "sync", model=model, device=device, model_lib=model_lib, mode=mode, engine_config=engine_config, enable_tracing=enable_tracing, ) self.chat = Chat(weakref.ref(self)) self.completions = Completion(weakref.ref(self)) def abort(self, request_id: str) -> None: """Generation abortion interface. Parameters --------- request_id : str The id of the request to abort. """ self._ffi["abort_request"](request_id) def metrics(self) -> engine_base.EngineMetrics: """Get engine metrics Returns ------- metrics: EngineMetrics The engine metrics """ return engine_base._query_engine_metrics(self) def _chat_completion( self, *, messages: List[Dict[str, Any]], # noqa: UP006 model: Optional[str] = None, frequency_penalty: Optional[float] = None, presence_penalty: Optional[float] = None, logprobs: bool = False, top_logprobs: int = 0, logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 max_tokens: Optional[int] = None, n: int = 1, seed: Optional[int] = None, stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 stream: bool = False, stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 temperature: Optional[float] = None, top_p: Optional[float] = None, tools: Optional[List[Dict[str, Any]]] = None, # noqa: UP006 tool_choice: Optional[Union[Literal["none", "auto"], Dict]] = None, # noqa: UP006 user: Optional[str] = None, response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 request_id: Optional[str] = None, debug_config: Optional[Dict[str, Any]] = None, # noqa: UP006 ) -> Union[ Iterator[openai_api_protocol.ChatCompletionStreamResponse], openai_api_protocol.ChatCompletionResponse, ]: """Synchronous chat completion internal interface with OpenAI API compatibility. See https://platform.openai.com/docs/api-reference/chat/create for specification. Parameters ---------- request_id : Optional[str] The optional request id. A random one will be generated if it is not given. debug_config: Optional[Dict[str, Any]] = None, Extra debug options to pass to the request. Raises ------ e : BadRequestError BadRequestError is raised when the request is invalid. """ if request_id is None: request_id = f"chatcmpl-{engine_utils.random_uuid()}" chatcmpl_generator = self._handle_chat_completion( openai_api_protocol.ChatCompletionRequest( messages=[ openai_api_protocol.ChatCompletionMessage.model_validate(message) for message in messages ], model=model, frequency_penalty=frequency_penalty, presence_penalty=presence_penalty, logprobs=logprobs, top_logprobs=top_logprobs, logit_bias=logit_bias, max_tokens=max_tokens, n=n, seed=seed, stop=stop, stream=stream, stream_options=( openai_api_protocol.StreamOptions.model_validate(stream_options) if stream_options is not None else None ), temperature=temperature, top_p=top_p, tools=( [openai_api_protocol.ChatTool.model_validate(tool) for tool in tools] if tools is not None else None ), tool_choice=tool_choice, user=user, response_format=( openai_api_protocol.RequestResponseFormat.model_validate(response_format) if response_format is not None else None ), debug_config=( debug_protocol.DebugConfig.model_validate(debug_config) if debug_config is not None else None ), ), request_id=request_id, ) if stream: # Stream response. return chatcmpl_generator # Normal response. request_final_usage = None output_texts = ["" for _ in range(n)] finish_reasons: List[Optional[str]] = [None for _ in range(n)] # noqa: UP006 logprob_results: Optional[List[List[openai_api_protocol.LogProbsContent]]] = ( # noqa: UP006 [[] for _ in range(n)] if logprobs else None ) for response in chatcmpl_generator: # if usage is not None, this is the last chunk if response.usage is not None: request_final_usage = response.usage continue for choice in response.choices: assert isinstance(choice.delta.content, str) output_texts[choice.index] += choice.delta.content if choice.finish_reason is not None and finish_reasons[choice.index] is None: finish_reasons[choice.index] = choice.finish_reason if choice.logprobs is not None: assert logprob_results is not None logprob_results[choice.index] += choice.logprobs.content assert all(finish_reason is not None for finish_reason in finish_reasons) use_function_calling, tool_calls_list = engine_base.process_function_call_output( output_texts, finish_reasons ) return engine_base.wrap_chat_completion_response( request_id=request_id, model=model, output_texts=output_texts, finish_reasons=finish_reasons, tool_calls_list=tool_calls_list, logprob_results=logprob_results, use_function_calling=use_function_calling, usage=request_final_usage, ) def _completion( self, *, prompt: Union[str, List[int]], # noqa: UP006 model: Optional[str] = None, best_of: int = 1, echo: bool = False, frequency_penalty: Optional[float] = None, presence_penalty: Optional[float] = None, logprobs: Optional[int] = None, logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 max_tokens: Optional[int] = None, n: int = 1, seed: Optional[int] = None, stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 stream: bool = False, stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 suffix: Optional[str] = None, temperature: Optional[float] = None, top_p: Optional[float] = None, user: Optional[str] = None, response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 request_id: Optional[str] = None, debug_config: Optional[Dict[str, Any]] = None, # noqa: UP006 ) -> Union[ Iterator[openai_api_protocol.CompletionResponse], openai_api_protocol.CompletionResponse, ]: """Synchronous completion internal interface with OpenAI API compatibility. See https://platform.openai.com/docs/api-reference/completions/create for specification. Parameters ---------- request_id : Optional[str] The optional request id. A random one will be generated if it is not given. debug_config: Optional[Dict[str, Any]] = None, Extra debug options to pass to the request. Raises ------ e : BadRequestError BadRequestError is raised when the request is invalid. """ if request_id is None: request_id = f"cmpl-{engine_utils.random_uuid()}" cmpl_generator = self._handle_completion( openai_api_protocol.CompletionRequest( model=model, prompt=prompt, best_of=best_of, echo=echo, frequency_penalty=frequency_penalty, presence_penalty=presence_penalty, logprobs=logprobs, logit_bias=logit_bias, max_tokens=max_tokens, n=n, seed=seed, stop=stop, stream=stream, stream_options=( openai_api_protocol.StreamOptions.model_validate(stream_options) if stream_options is not None else None ), suffix=suffix, temperature=temperature, top_p=top_p, user=user, response_format=( openai_api_protocol.RequestResponseFormat.model_validate(response_format) if response_format is not None else None ), debug_config=( debug_protocol.DebugConfig.model_validate(debug_config) if debug_config is not None else None ), ), request_id=request_id, ) if stream: # Stream response. return cmpl_generator # Normal response. request_final_usage = None output_texts = [""] * n finish_reasons: List[Optional[str]] = [None] * n # noqa: UP006 logprob_results: List[Optional[openai_api_protocol.CompletionLogProbs]] = [None] * n # noqa: UP006 for response in cmpl_generator: # this is the final chunk if response.usage is not None: request_final_usage = response.usage continue for choice in response.choices: output_texts[choice.index] += choice.text if choice.finish_reason is not None and finish_reasons[choice.index] is None: finish_reasons[choice.index] = choice.finish_reason if choice.logprobs is not None: logprob_results[choice.index] = choice.logprobs assert all(finish_reason is not None for finish_reason in finish_reasons) return engine_base.wrap_completion_response( request_id=request_id, model=model, output_texts=output_texts, finish_reasons=finish_reasons, logprob_results=logprob_results, usage=request_final_usage, ) def _handle_chat_completion( self, request: openai_api_protocol.ChatCompletionRequest, request_id: str ) -> Iterator[openai_api_protocol.ChatCompletionStreamResponse]: """The implementation fo synchronous ChatCompletionRequest handling. Yields ------ stream_response : CompletionResponse The stream response conforming to OpenAI API. See mlc_llm/protocol/openai_api_protocol.py or https://platform.openai.com/docs/api-reference/chat/streaming for specification. Raises ------ e : BadRequestError BadRequestError is raised when the request is invalid. """ ( prompts, generation_cfg, use_function_calling, prompt_length, ) = engine_base.process_chat_completion_request( request, request_id, self.state, self.model_config_dicts[0], self.tokenizer.encode, self.max_input_sequence_length, self.conv_template.model_copy(deep=True), ) _ = prompt_length finish_reasons: List[Optional[str]] = [None for _ in range(generation_cfg.n)] # noqa: UP006 self.state.record_event(request_id, event="invoke generate") for delta_outputs in self._generate(prompts, generation_cfg, request_id): response = engine_base.process_chat_completion_stream_output( delta_outputs, request, request_id, self.state, use_function_calling, finish_reasons, ) if response is not None: yield response self.state.record_event(request_id, event="finish") def _handle_completion( self, request: openai_api_protocol.CompletionRequest, request_id: str ) -> Iterator[openai_api_protocol.CompletionResponse]: """The implementation for synchronous CompletionRequest handling. Yields ------ stream_response : CompletionResponse The stream response conforming to OpenAI API. See mlc_llm/protocol/openai_api_protocol.py or https://platform.openai.com/docs/api-reference/completions/object for specification. Raises ------ e : BadRequestError BadRequestError is raised when the request is invalid. """ ( prompt, generation_cfg, prompt_length, echo_response, ) = engine_base.process_completion_request( request, request_id, self.state, self.tokenizer, self.max_input_sequence_length, self.conv_template.model_copy(deep=True), ) _ = prompt_length if echo_response is not None: yield echo_response finish_reasons: List[Optional[str]] = [None for _ in range(generation_cfg.n)] # noqa: UP006 self.state.record_event(request_id, event="invoke generate") for delta_outputs in self._generate(prompt, generation_cfg, request_id): response = engine_base.process_completion_stream_output( delta_outputs, request, request_id, self.state, finish_reasons, ) if response is not None: yield response suffix_response = engine_base.create_completion_suffix_response( request, request_id, finish_reasons ) if suffix_response is not None: yield suffix_response self.state.record_event(request_id, event="finish") def _generate( self, prompt: Union[str, List[int], List[Union[str, List[int], data.Data]]], # noqa: UP006 generation_config: GenerationConfig, request_id: str, ) -> Iterator[List[engine_base.CallbackStreamOutput]]: # noqa: UP006 """Internal synchronous text generation interface of MLCEngine. The method is a coroutine that streams a list of CallbackStreamOutput at a time via yield. The returned list length is the number of parallel generations specified by `generation_config.n` except for the final chunk(which is always an List of size 1 and comes with usage) Parameters ---------- prompt : Union[str, List[int], List[Union[str, List[int], data.Data]]] The input prompt in forms of text strings, lists of token ids or data. generation_config : GenerationConfig The generation config of the request. request_id : str The unique identifier (in string) or this generation request. Yields ------ request_output : List[engine_base.CallbackStreamOutput] The delta generated outputs in a list. Except for the final chunk, the number of list elements equals to `generation_config.n`, and each element corresponds to the delta output of a parallel generation. """ if self._terminated: raise ValueError("The engine has terminated.") # Create the request with the given id, input data, generation # config and the created callback. input_data = engine_utils.convert_prompts_to_data(prompt) request = self._ffi["create_request"]( request_id, input_data, generation_config.model_dump_json(by_alias=True) ) # Record the stream in the tracker self.state.sync_output_queue = queue.Queue() self.state.sync_text_streamers = [ TextStreamer(self.tokenizer) for _ in range(generation_config.n) ] self._ffi["add_request"](request) def abort_request(): """clean up request if exception happens""" self.abort(request_id) # Iterate the stream asynchronously and yield the token. with engine_utils.ErrorCleanupScope(abort_request): while True: delta_outputs = self.state.sync_output_queue.get() request_outputs, request_final_usage_json_str = self._request_stream_callback_impl( delta_outputs ) for request_output in request_outputs: yield request_output if request_final_usage_json_str is not None: # final chunk, we can break output = engine_base.CallbackStreamOutput( delta_text="", delta_logprob_json_strs=None, finish_reason=None, request_final_usage_json_str=request_final_usage_json_str, ) yield [output] break def _request_stream_callback_impl( self, delta_outputs: List[data.RequestStreamOutput], # noqa: UP006 ) -> Tuple[List[List[engine_base.CallbackStreamOutput]], Optional[str]]: # noqa: UP006 """The underlying implementation of request stream callback of MLCEngine.""" batch_outputs: List[List[engine_base.CallbackStreamOutput]] = [] # noqa: UP006 for delta_output in delta_outputs: request_id, stream_outputs = delta_output.unpack() self.state.record_event(request_id, event="start callback") # final chunk is now always indicated by a chunk # where usage json is present # the backend engine always streams back this chunk # regardless of include_usage option is_final_chunk = stream_outputs[0].request_final_usage_json_str is not None if is_final_chunk: return (batch_outputs, stream_outputs[0].request_final_usage_json_str) outputs: List[engine_base.CallbackStreamOutput] = [] # noqa: UP006 for stream_output, text_streamer in zip(stream_outputs, self.state.sync_text_streamers): self.state.record_event(request_id, event="start detokenization") delta_text = stream_output.extra_prefix_string + ( text_streamer.put(stream_output.delta_token_ids) if len(stream_output.delta_token_ids) > 0 else "" ) if stream_output.finish_reason is not None: delta_text += text_streamer.finish() self.state.record_event(request_id, event="finish detokenization") outputs.append( engine_base.CallbackStreamOutput( delta_text=delta_text, delta_logprob_json_strs=stream_output.delta_logprob_json_strs, finish_reason=stream_output.finish_reason, request_final_usage_json_str=None, ) ) batch_outputs.append(outputs) self.state.record_event(request_id, event="finish callback") return (batch_outputs, None)