main
py 123 lines 4.09 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, StreamingResponse
9
10 from vllm.entrypoints.anthropic.protocol import (
11 AnthropicCountTokensRequest,
12 AnthropicCountTokensResponse,
13 AnthropicError,
14 AnthropicErrorResponse,
15 AnthropicMessagesRequest,
16 AnthropicMessagesResponse,
17 )
18 from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
19 from vllm.entrypoints.serve.engine.protocol import ErrorResponse
20 from vllm.entrypoints.serve.exception_handling.error_response import (
21 create_error_response,
22 )
23 from vllm.entrypoints.serve.utils.api_utils import (
24 load_aware_call,
25 validate_json_request,
26 with_cancellation,
27 )
28 from vllm.logger import init_logger
29
30 logger = init_logger(__name__)
31
32 router = APIRouter()
33
34
35 def messages(request: Request) -> AnthropicServingMessages:
36 return request.app.state.anthropic_serving_messages
37
38
39 def translate_error_response(response: ErrorResponse) -> JSONResponse:
40 anthropic_error = AnthropicErrorResponse(
41 error=AnthropicError(
42 type=response.error.type,
43 message=response.error.message,
44 )
45 )
46 return JSONResponse(
47 status_code=response.error.code, content=anthropic_error.model_dump()
48 )
49
50
51 @router.post(
52 "/v1/messages",
53 dependencies=[Depends(validate_json_request)],
54 responses={
55 HTTPStatus.OK.value: {"content": {"text/event-stream": {}}},
56 HTTPStatus.BAD_REQUEST.value: {"model": AnthropicErrorResponse},
57 HTTPStatus.NOT_FOUND.value: {"model": AnthropicErrorResponse},
58 HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": AnthropicErrorResponse},
59 },
60 )
61 @with_cancellation
62 @load_aware_call
63 async def create_messages(request: AnthropicMessagesRequest, raw_request: Request):
64 handler = messages(raw_request)
65 if handler is None:
66 base_server = raw_request.app.state.serving_tokenization
67 error = base_server.create_error_response(
68 NotImplementedError("The model does not support Messages API")
69 )
70 return translate_error_response(error)
71
72 try:
73 generator = await handler.create_messages(request, raw_request)
74 except Exception as e:
75 logger.exception("Error in create_messages: %s", e)
76 return translate_error_response(create_error_response(e))
77
78 if isinstance(generator, ErrorResponse):
79 return translate_error_response(generator)
80
81 elif isinstance(generator, AnthropicMessagesResponse):
82 resp = generator.model_dump(exclude_none=True)
83 logger.debug("Anthropic Messages Response: %s", resp)
84 return JSONResponse(content=resp)
85
86 return StreamingResponse(content=generator, media_type="text/event-stream")
87
88
89 @router.post(
90 "/v1/messages/count_tokens",
91 dependencies=[Depends(validate_json_request)],
92 responses={
93 HTTPStatus.OK.value: {"model": AnthropicCountTokensResponse},
94 HTTPStatus.BAD_REQUEST.value: {"model": AnthropicErrorResponse},
95 HTTPStatus.NOT_FOUND.value: {"model": AnthropicErrorResponse},
96 HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": AnthropicErrorResponse},
97 },
98 )
99 @with_cancellation
100 @load_aware_call
101 async def count_tokens(request: AnthropicCountTokensRequest, raw_request: Request):
102 handler = messages(raw_request)
103 if handler is None:
104 base_server = raw_request.app.state.serving_tokenization
105 error = base_server.create_error_response(
106 NotImplementedError("The model does not support Messages API")
107 )
108 return translate_error_response(error)
109
110 try:
111 response = await handler.count_tokens(request, raw_request)
112 except Exception as e:
113 logger.exception("Error in count_tokens: %s", e)
114 return translate_error_response(create_error_response(e))
115
116 if isinstance(response, ErrorResponse):
117 return translate_error_response(response)
118
119 return JSONResponse(content=response.model_dump(exclude_none=True))
120
121
122 def attach_router(app: FastAPI):
123 app.include_router(router)