api/vllm_routes/completion.py

import time
import uuid
from functools import partial
from typing import (
    List,
    Dict,
    Any,
    AsyncIterator,
    Optional,
)

import anyio
from fastapi import APIRouter, Depends
from fastapi import HTTPException, Request
from loguru import logger
from openai.types.completion import Completion
from openai.types.completion_choice import CompletionChoice, Logprobs
from openai.types.completion_usage import CompletionUsage
from sse_starlette import EventSourceResponse
from vllm.outputs import RequestOutput

from api.core.vllm_engine import VllmEngine
from api.models import GENERATE_ENGINE
from api.utils.compat import model_dump
from api.utils.protocol import CompletionCreateParams
from api.utils.request import (
    handle_request,
    get_event_publisher,
    check_api_key
)

completion_router = APIRouter()


def get_engine():
    yield GENERATE_ENGINE


@completion_router.post("/completions", dependencies=[Depends(check_api_key)])
async def create_completion(
    request: CompletionCreateParams,
    raw_request: Request,
    engine: VllmEngine = Depends(get_engine),
):
    """Completion API similar to OpenAI's API.

    See https://platform.openai.com/docs/api-reference/completions/create
    for the API specification. This API mimics the OpenAI Completion API.
    """
    if request.echo:
        # We do not support echo since the vLLM engine does not
        # currently support getting the logprobs of prompt tokens.
        raise HTTPException(status_code=400, detail="echo is not currently supported")

    if request.suffix:
        # The language models we currently support do not support suffix.
        raise HTTPException(status_code=400, detail="suffix is not currently supported")

    request.max_tokens = request.max_tokens or 128
    request = await handle_request(request, engine.prompt_adapter.stop, chat=False)

    if isinstance(request.prompt, list):
        request.prompt = request.prompt[0]

    params = model_dump(request, exclude={"prompt"})
    params.update(dict(prompt_or_messages=request.prompt))
    logger.debug(f"==== request ====\n{params}")

    request_id: str = f"cmpl-{str(uuid.uuid4())}"
    generator = engine.generate(params, request_id)

    if request.stream:
        iterator = create_completion_stream(generator, params, request_id, engine.tokenizer)
        send_chan, recv_chan = anyio.create_memory_object_stream(10)
        return EventSourceResponse(
            recv_chan,
            data_sender_callable=partial(
                get_event_publisher,
                request=raw_request,
                inner_send_chan=send_chan,
                iterator=iterator,
            ),
        )
    else:
        # Non-streaming response
        final_res: RequestOutput = None
        async for res in generator:
            if raw_request is not None and await raw_request.is_disconnected():
                await engine.model.abort(request_id)
                return
            final_res = res

        assert final_res is not None
        choices = []
        for output in final_res.outputs:
            output.text = output.text.replace("�", "")
            logprobs = None
            if params.get("logprobs", None) is not None:
                logprobs = create_logprobs(engine.tokenizer, output.token_ids, output.logprobs)

            choice = CompletionChoice(
                index=output.index,
                text=output.text,
                finish_reason=output.finish_reason,
                logprobs=logprobs,
            )
            choices.append(choice)

        num_prompt_tokens = len(final_res.prompt_token_ids)
        num_generated_tokens = sum(len(output.token_ids) for output in final_res.outputs)
        usage = CompletionUsage(
            prompt_tokens=num_prompt_tokens,
            completion_tokens=num_generated_tokens,
            total_tokens=num_prompt_tokens + num_generated_tokens,
        )

        return Completion(
            id=request_id,
            choices=choices,
            created=int(time.time()),
            model=params.get("model", "llm"),
            object="text_completion",
            usage=usage,
        )


def create_logprobs(
    tokenizer,
    token_ids: List[int],
    top_logprobs: Optional[List[Optional[Dict[int, float]]]] = None,
    num_output_top_logprobs: Optional[int] = None,
    initial_text_offset: int = 0,
) -> Logprobs:
    logprobs = Logprobs(text_offset=[], token_logprobs=[], tokens=[], top_logprobs=None)
    last_token_len = 0
    if num_output_top_logprobs:
        logprobs.top_logprobs = []

    for i, token_id in enumerate(token_ids):
        step_top_logprobs = top_logprobs[i]
        if step_top_logprobs is not None:
            token_logprob = step_top_logprobs[token_id]
        else:
            token_logprob = None

        token = tokenizer.convert_ids_to_tokens(token_id)
        logprobs.tokens.append(token)
        logprobs.token_logprobs.append(token_logprob)
        if len(logprobs.text_offset) == 0:
            logprobs.text_offset.append(initial_text_offset)
        else:
            logprobs.text_offset.append(logprobs.text_offset[-1] + last_token_len)
        last_token_len = len(token)

        if num_output_top_logprobs:
            logprobs.top_logprobs.append(
                {
                    tokenizer.convert_ids_to_tokens(i): p
                    for i, p in step_top_logprobs.items()
                }
                if step_top_logprobs else None
            )
    return logprobs


async def create_completion_stream(
    generator: AsyncIterator, params: Dict[str, Any], request_id: str, tokenizer,
) -> AsyncIterator:
    n = params.get("n", 1)
    previous_texts = [""] * n
    previous_num_tokens = [0] * n
    async for res in generator:
        res: RequestOutput
        for output in res.outputs:
            i = output.index
            output.text = output.text.replace("�", "")
            delta_text = output.text[len(previous_texts[i]):]

            if params.get("logprobs") is not None:
                logprobs = create_logprobs(
                    tokenizer,
                    output.token_ids[previous_num_tokens[i]:],
                    output.logprobs[previous_num_tokens[i]:],
                    len(previous_texts[i])
                )
            else:
                logprobs = None

            previous_texts[i] = output.text
            previous_num_tokens[i] = len(output.token_ids)

            choice = CompletionChoice(
                index=i,
                text=delta_text,
                finish_reason="stop",
                logprobs=logprobs,
            )
            yield Completion(
                id=request_id,
                choices=[choice],
                created=int(time.time()),
                model=params.get("model", "llm"),
                object="text_completion",
            )

            if output.finish_reason is not None:
                if params.get("logprobs") is not None:
                    logprobs = Logprobs(
                        text_offset=[], token_logprobs=[], tokens=[], top_logprobs=[]
                    )
                else:
                    logprobs = None

                choice = CompletionChoice(
                    index=i,
                    text=delta_text,
                    finish_reason="stop",
                    logprobs=logprobs,
                )
                yield Completion(
                    id=request_id,
                    choices=[choice],
                    created=int(time.time()),
                    model=params.get("model", "llm"),
                    object="text_completion",
                )