main
py 101 lines 3.15 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5 from http import HTTPStatus
6
7 from fastapi import APIRouter, Depends, FastAPI, Request
8 from fastapi.responses import JSONResponse
9 from typing_extensions import assert_never
10
11 from vllm.entrypoints.serve.engine.protocol import ErrorResponse
12 from vllm.entrypoints.serve.tokenize.protocol import (
13 DetokenizeRequest,
14 DetokenizeResponse,
15 TokenizeRequest,
16 TokenizeResponse,
17 )
18 from vllm.entrypoints.serve.tokenize.serving import ServingTokenization
19 from vllm.entrypoints.serve.utils.api_utils import (
20 validate_json_request,
21 with_cancellation,
22 )
23 from vllm.logger import init_logger
24
25 logger = init_logger(__name__)
26
27
28 def tokenization(request: Request) -> ServingTokenization:
29 return request.app.state.serving_tokenization
30
31
32 router = APIRouter()
33
34
35 @router.post(
36 "/tokenize",
37 dependencies=[Depends(validate_json_request)],
38 responses={
39 HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
40 HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
41 HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
42 HTTPStatus.NOT_IMPLEMENTED.value: {"model": ErrorResponse},
43 },
44 )
45 @with_cancellation
46 async def tokenize(request: TokenizeRequest, raw_request: Request):
47 handler = tokenization(raw_request)
48
49 generator = await handler.create_tokenize(request, raw_request)
50
51 if isinstance(generator, ErrorResponse):
52 return JSONResponse(
53 content=generator.model_dump(), status_code=generator.error.code
54 )
55 elif isinstance(generator, TokenizeResponse):
56 return JSONResponse(content=generator.model_dump())
57
58 assert_never(generator)
59
60
61 @router.post(
62 "/detokenize",
63 dependencies=[Depends(validate_json_request)],
64 responses={
65 HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
66 HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
67 HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
68 },
69 )
70 @with_cancellation
71 async def detokenize(request: DetokenizeRequest, raw_request: Request):
72 handler = tokenization(raw_request)
73
74 generator = await handler.create_detokenize(request, raw_request)
75
76 if isinstance(generator, ErrorResponse):
77 return JSONResponse(
78 content=generator.model_dump(), status_code=generator.error.code
79 )
80 elif isinstance(generator, DetokenizeResponse):
81 return JSONResponse(content=generator.model_dump())
82
83 assert_never(generator)
84
85
86 def attach_router(app: FastAPI):
87 if getattr(app.state.args, "enable_tokenizer_info_endpoint", False):
88 """Conditionally register the tokenizer info endpoint if enabled."""
89
90 @router.get("/tokenizer_info")
91 async def get_tokenizer_info(raw_request: Request):
92 """Get comprehensive tokenizer information."""
93 result = await tokenization(raw_request).get_tokenizer_info()
94 return JSONResponse(
95 content=result.model_dump(),
96 status_code=result.error.code
97 if isinstance(result, ErrorResponse)
98 else 200,
99 )
100
101 app.include_router(router)