main
py 111 lines 3.33 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from http import HTTPStatus
4
5 from fastapi import APIRouter, Depends, Request
6
7 from vllm.entrypoints.serve.engine.protocol import ErrorResponse
8 from vllm.entrypoints.serve.utils.api_utils import (
9 load_aware_call,
10 validate_json_request,
11 with_cancellation,
12 )
13 from vllm.logger import init_logger
14
15 from .protocol import RerankRequest, ScoreRequest
16 from .serving import ServingScores
17
18 router = APIRouter()
19
20 logger = init_logger(__name__)
21
22
23 def score(request: Request) -> ServingScores:
24 handler = getattr(request.app.state, "serving_scores", None)
25 if handler is None:
26 raise NotImplementedError("The model does not support Score API")
27 return handler
28
29
30 def rerank(request: Request) -> ServingScores:
31 handler = getattr(request.app.state, "serving_scores", None)
32 if handler is None:
33 raise NotImplementedError("The model does not support Rerank (Score) API")
34 return handler
35
36
37 @router.post(
38 "/score",
39 dependencies=[Depends(validate_json_request)],
40 responses={
41 HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
42 HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
43 },
44 )
45 @with_cancellation
46 @load_aware_call
47 async def create_score(request: ScoreRequest, raw_request: Request):
48 handler = score(raw_request)
49 return await handler(request, raw_request)
50
51
52 @router.post(
53 "/v1/score",
54 dependencies=[Depends(validate_json_request)],
55 responses={
56 HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
57 HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
58 },
59 )
60 async def create_score_v1(request: ScoreRequest, raw_request: Request):
61 logger.warning(
62 "To indicate that Score API is not part of standard OpenAI API, we "
63 "have moved it to `/score`. Please update your client accordingly."
64 )
65
66 return await create_score(request, raw_request)
67
68
69 @router.post(
70 "/rerank",
71 dependencies=[Depends(validate_json_request)],
72 responses={
73 HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
74 HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
75 },
76 )
77 @with_cancellation
78 @load_aware_call
79 async def do_rerank(request: RerankRequest, raw_request: Request):
80 handler = rerank(raw_request)
81 return await handler(request, raw_request)
82
83
84 @router.post(
85 "/v1/rerank",
86 dependencies=[Depends(validate_json_request)],
87 responses={
88 HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
89 HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
90 },
91 )
92 async def do_rerank_v1(request: RerankRequest, raw_request: Request):
93 logger.warning_once(
94 "To indicate that the rerank API is not part of the standard OpenAI"
95 " API, we have located it at `/rerank`. Please update your client "
96 "accordingly. (Note: Conforms to JinaAI rerank API)"
97 )
98
99 return await do_rerank(request, raw_request)
100
101
102 @router.post(
103 "/v2/rerank",
104 dependencies=[Depends(validate_json_request)],
105 responses={
106 HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
107 HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
108 },
109 )
110 async def do_rerank_v2(request: RerankRequest, raw_request: Request):
111 return await do_rerank(request, raw_request)