mlc-ai--mlc-llm/python/mlc_llm/serve/entrypoints/openai_entrypoints.py

"""OpenAI API-compatible server entrypoints in MLC LLM"""

import base64
import struct
from collections.abc import AsyncGenerator
from http import HTTPStatus
from typing import List, Optional  # noqa: UP035

import fastapi
import numpy as np

from mlc_llm.protocol import error_protocol
from mlc_llm.protocol.openai_api_protocol import (
    ChatCompletionRequest,
    CompletionLogProbs,
    CompletionRequest,
    EmbeddingObject,
    EmbeddingRequest,
    EmbeddingResponse,
    EmbeddingUsage,
    ListResponse,
    LogProbsContent,
    ModelResponse,
)
from mlc_llm.serve import engine_base, engine_utils
from mlc_llm.serve.server import ServerContext


def verify_api_key(request: fastapi.Request):
    """Function to verify API key"""
    server_context = ServerContext.current()
    # Only perform verification when API key is configured
    if server_context is not None and server_context.api_key is not None:
        provided_key = request.headers.get("Authorization", "").replace("Bearer ", "")
        if provided_key != server_context.api_key:
            raise fastapi.HTTPException(status_code=401, detail="Invalid API Key")


app = fastapi.APIRouter(dependencies=[fastapi.Depends(verify_api_key)])


################ v1/embeddings ################


@app.post("/v1/embeddings")
async def request_embedding(request: EmbeddingRequest):
    """OpenAI-compatible embedding API.
    API reference: https://platform.openai.com/docs/api-reference/embeddings/create
    """
    server_context: ServerContext = ServerContext.current()
    embedding_engine = server_context.get_embedding_engine(request.model)
    if embedding_engine is None:
        return error_protocol.create_error_response(
            HTTPStatus.BAD_REQUEST,
            message=f'The requested model "{request.model}" is not served as an embedding model.',
        )

    # Normalize input to List[str]
    inputs: List[str]  # noqa: UP006
    if isinstance(request.input, str):
        inputs = [request.input]
    elif (
        isinstance(request.input, list)
        and len(request.input) > 0
        and isinstance(request.input[0], str)
    ):
        inputs = list(request.input)
    else:
        # Token ID inputs (List[int] or List[List[int]]) — decode back to strings
        if isinstance(request.input[0], int):
            inputs = [embedding_engine.tokenizer.decode(request.input)]
        else:
            inputs = [embedding_engine.tokenizer.decode(ids) for ids in request.input]

    # Run embedding inference (async — does not block the event loop)
    try:
        embeddings, total_tokens = await embedding_engine.async_embed(inputs)
    except Exception as exc:
        return error_protocol.create_error_response(
            HTTPStatus.INTERNAL_SERVER_ERROR,
            message=f"Embedding inference failed: {exc}",
        )

    # Optional: truncate dimensions (Matryoshka-style).
    # This is API-level renormalization after dimension truncation,
    # independent of model metadata normalize. Always renormalize
    # truncated vectors to maintain unit length per OpenAI API contract.
    if request.dimensions is not None:
        for i, emb in enumerate(embeddings):
            vec = np.array(emb[: request.dimensions], dtype=np.float32)
            norm = np.linalg.norm(vec)
            if norm > 1e-12:
                vec = vec / norm
            embeddings[i] = vec.tolist()

    # Build response data
    resp_data = []
    for i, emb in enumerate(embeddings):
        if request.encoding_format == "base64":
            binary = struct.pack(f"<{len(emb)}f", *emb)
            resp_data.append(
                EmbeddingObject(
                    embedding=base64.b64encode(binary).decode("utf-8"),
                    index=i,
                )
            )
        else:
            resp_data.append(EmbeddingObject(embedding=emb, index=i))

    return EmbeddingResponse(
        data=resp_data,
        model=request.model,
        usage=EmbeddingUsage(prompt_tokens=total_tokens, total_tokens=total_tokens),
    )


################ v1/models ################


@app.get("/v1/models")
async def request_models() -> ListResponse:
    """OpenAI-compatible served model query API.
    API reference: https://platform.openai.com/docs/api-reference/models
    """
    server_context: ServerContext = ServerContext.current()
    return ListResponse(data=[ModelResponse(id=model) for model in server_context.get_model_list()])


################ v1/completions ################


@app.post("/v1/completions")
async def request_completion(request: CompletionRequest, raw_request: fastapi.Request):
    """OpenAI-compatible completion API.
    API reference: https://platform.openai.com/docs/api-reference/completions/create
    """
    # - Check the requested model.
    server_context: ServerContext = ServerContext.current()
    request_final_usage_include_extra = server_context.enable_debug
    request_include_debug_config = server_context.enable_debug

    if not request_include_debug_config:
        request.debug_config = None

    async_engine = server_context.get_engine(request.model)
    if async_engine is None:
        return error_protocol.create_error_response(
            HTTPStatus.BAD_REQUEST,
            message=f'The requested model "{request.model}" is not served.',
        )
    # FIXME: This is a temporary solution to make sure
    # prep_recv, remote_send and start_generation process the same request
    request_id = request.user if request.user is not None else f"cmpl-{engine_utils.random_uuid()}"

    # Streaming response.
    if request.stream:
        # We manually get the first response from generator to
        # capture potential exceptions in this scope, rather then
        # the StreamingResponse scope.
        stream_generator = async_engine._handle_completion(
            request,
            request_id,
            request_final_usage_include_extra=request_final_usage_include_extra,
        )
        first_response = await anext(  # noqa: F821
            stream_generator
        )

        async def completion_stream_generator() -> AsyncGenerator[str, None]:
            if isinstance(first_response, StopAsyncIteration):
                yield "data: [DONE]\n\n"
                return
            yield f"data: {first_response.model_dump_json(by_alias=True)}\n\n"
            async for response in stream_generator:
                yield f"data: {response.model_dump_json(by_alias=True)}\n\n"
            yield "data: [DONE]\n\n"

        return fastapi.responses.StreamingResponse(
            completion_stream_generator(), media_type="text/event-stream"
        )

    # Normal response.
    request_final_usage = None
    output_texts = [""] * request.n
    finish_reasons: List[Optional[str]] = [None] * request.n  # noqa: UP006
    logprob_results: List[Optional[CompletionLogProbs]] = [None] * request.n  # noqa: UP006

    async for response in async_engine._handle_completion(
        request,
        request_id,
        request_final_usage_include_extra=request_final_usage_include_extra,
    ):
        if await raw_request.is_disconnected():
            # In non-streaming cases, the engine will not be notified
            # when the request is disconnected.
            # Therefore, we check if it is disconnected each time,
            # and explicitly return.
            # Note that requesta abort is triggered when the async for and funciton scope ends.
            return error_protocol.create_error_response(
                HTTPStatus.BAD_REQUEST, message="The request has disconnected"
            )
        # this is the final chunk
        if response.usage is not None:
            request_final_usage = response.usage
            # remove extra information if debug is not enabled
            if not server_context.enable_debug:
                request_final_usage.extra = None
            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:
                if logprob_results[choice.index] is None:
                    logprob_results[choice.index] = choice.logprobs
                else:
                    logprob_results[choice.index].token_logprobs.extend(
                        choice.logprobs.token_logprobs
                    )
                    logprob_results[choice.index].tokens.extend(choice.logprobs.tokens)
                    logprob_results[choice.index].top_logprobs.extend(choice.logprobs.top_logprobs)

    return engine_base.wrap_completion_response(
        request_id=request_id,
        model=request.model,
        output_texts=output_texts,
        finish_reasons=finish_reasons,
        logprob_results=logprob_results,
        usage=request_final_usage,
    )


################ v1/chat/completions ################


@app.post("/v1/chat/completions")
async def request_chat_completion(request: ChatCompletionRequest, raw_request: fastapi.Request):
    """OpenAI-compatible chat completion API.
    API reference: https://platform.openai.com/docs/api-reference/chat
    """
    # - Check the requested model.
    server_context: ServerContext = ServerContext.current()
    request_final_usage_include_extra = server_context.enable_debug
    request_include_debug_config = server_context.enable_debug

    if not request_include_debug_config:
        request.debug_config = None

    async_engine = server_context.get_engine(request.model)
    if async_engine is None:
        return error_protocol.create_error_response(
            HTTPStatus.BAD_REQUEST,
            message=f'The requested model "{request.model}" is not served.',
        )
    # FIXME: This is a temporary solution to make sure
    # prep_recv, remote_send and start_generation process the same request
    request_id = (
        request.user if request.user is not None else f"chatcmpl-{engine_utils.random_uuid()}"
    )

    # Streaming response.
    if request.stream:
        # We manually get the first response from generator to
        # capture potential exceptions in this scope, rather then
        # the StreamingResponse scope.
        stream_generator = async_engine._handle_chat_completion(
            request,
            request_id,
            request_final_usage_include_extra=request_final_usage_include_extra,
        )
        first_response = await anext(  # noqa: F821
            stream_generator
        )

        async def completion_stream_generator() -> AsyncGenerator[str, None]:
            if isinstance(first_response, StopAsyncIteration):
                yield "data: [DONE]\n\n"
                return
            yield f"data: {first_response.model_dump_json(by_alias=True)}\n\n"
            async for response in stream_generator:
                yield f"data: {response.model_dump_json(by_alias=True)}\n\n"
            yield "data: [DONE]\n\n"

        return fastapi.responses.StreamingResponse(
            completion_stream_generator(), media_type="text/event-stream"
        )

    # Normal response.
    request_final_usage = None
    output_texts = ["" for _ in range(request.n)]
    finish_reasons: List[Optional[str]] = [None for _ in range(request.n)]  # noqa: UP006
    logprob_results: Optional[List[List[LogProbsContent]]] = (  # noqa: UP006
        [[] for _ in range(request.n)] if request.logprobs else None
    )

    async for response in async_engine._handle_chat_completion(
        request,
        request_id,
        request_final_usage_include_extra=request_final_usage_include_extra,
    ):
        if await raw_request.is_disconnected():
            # In non-streaming cases, the engine will not be notified
            # when the request is disconnected.
            # Therefore, we check if it is disconnected each time,
            # no need to explicitly abort, as the chat completion
            # return will trigger abort call
            return error_protocol.create_error_response(
                HTTPStatus.BAD_REQUEST, message="The request has disconnected"
            )
        # usage is always the last chunk
        if response.usage is not None:
            request_final_usage = response.usage
            # remove extra information if debug is not enabled
            if not server_context.enable_debug:
                request_final_usage.extra = None

        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=request.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,
    )