A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 11:30 UTC
[Feat] Supports Anthropic Messages count_tokens API (#35588)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
Chauncey committed
Mar 2, 2026 at 17:48 UTC
9a87b0578fc3bf0a1e80c0fc55a31e8db36df2c9
3 files changed
+332
-133
vllm/entrypoints/anthropic/api_router.py
+55
-11
@@ -8,6 +8,8 @@ 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,
@@ -31,6 +33,18 @@ def messages(request: Request) -> AnthropicServingMessages:
33
return request.app.state.anthropic_serving_messages
34
35
36
+def translate_error_response(response: ErrorResponse) -> JSONResponse:
37
+ anthropic_error = AnthropicErrorResponse(
38
+ error=AnthropicError(
39
+ type=response.error.type,
40
+ message=response.error.message,
41
+ )
42
+ )
43
+ return JSONResponse(
44
+ status_code=response.error.code, content=anthropic_error.model_dump()
45
+ )
46
+
47
+
48
@router.post(
49
"/v1/messages",
50
dependencies=[Depends(validate_json_request)],
@@ -44,17 +58,6 @@ def messages(request: Request) -> AnthropicServingMessages:
58
@with_cancellation
59
@load_aware_call
60
async def create_messages(request: AnthropicMessagesRequest, raw_request: Request):
47
- def translate_error_response(response: ErrorResponse) -> JSONResponse:
48
- anthropic_error = AnthropicErrorResponse(
49
- error=AnthropicError(
50
- type=response.error.type,
51
- message=response.error.message,
52
- )
53
- )
54
- return JSONResponse(
55
- status_code=response.error.code, content=anthropic_error.model_dump()
56
- )
57
-
61
handler = messages(raw_request)
62
if handler is None:
63
base_server = raw_request.app.state.openai_serving_tokenization
@@ -88,5 +91,46 @@ async def create_messages(request: AnthropicMessagesRequest, raw_request: Reques
91
return StreamingResponse(content=generator, media_type="text/event-stream")
92
93
94
+@router.post(
95
+ "/v1/messages/count_tokens",
96
+ dependencies=[Depends(validate_json_request)],
97
+ responses={
98
+ HTTPStatus.OK.value: {"model": AnthropicCountTokensResponse},
99
+ HTTPStatus.BAD_REQUEST.value: {"model": AnthropicErrorResponse},
100
+ HTTPStatus.NOT_FOUND.value: {"model": AnthropicErrorResponse},
101
+ HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": AnthropicErrorResponse},
102
+ },
103
+)
104
+@load_aware_call
105
+@with_cancellation
106
+async def count_tokens(request: AnthropicCountTokensRequest, raw_request: Request):
107
+ handler = messages(raw_request)
108
+ if handler is None:
109
+ base_server = raw_request.app.state.openai_serving_tokenization
110
+ error = base_server.create_error_response(
111
+ message="The model does not support Messages API"
112
+ )
113
+ return translate_error_response(error)
114
+
115
+ try:
116
+ response = await handler.count_tokens(request, raw_request)
117
+ except Exception as e:
118
+ logger.exception("Error in count_tokens: %s", e)
119
+ return JSONResponse(
120
+ status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
121
+ content=AnthropicErrorResponse(
122
+ error=AnthropicError(
123
+ type="internal_error",
124
+ message=str(e),
125
+ )
126
+ ).model_dump(),
127
+ )
128
+
129
+ if isinstance(response, ErrorResponse):
130
+ return translate_error_response(response)
131
+
132
+ return JSONResponse(content=response.model_dump(exclude_none=True))
133
+
134
+
135
def attach_router(app: FastAPI):
136
app.include_router(router)
vllm/entrypoints/anthropic/protocol.py
+30
@@ -175,3 +175,33 @@ class AnthropicMessagesResponse(BaseModel):
175
def model_post_init(self, __context):
176
if not self.id:
177
self.id = f"msg_{int(time.time() * 1000)}"
178
+
179
+
180
+class AnthropicContextManagement(BaseModel):
181
+ """Context management information for token counting."""
182
+
183
+ original_input_tokens: int
184
+
185
+
186
+class AnthropicCountTokensRequest(BaseModel):
187
+ """Anthropic messages.count_tokens request"""
188
+
189
+ model: str
190
+ messages: list[AnthropicMessage]
191
+ system: str | list[AnthropicContentBlock] | None = None
192
+ tool_choice: AnthropicToolChoice | None = None
193
+ tools: list[AnthropicTool] | None = None
194
+
195
+ @field_validator("model")
196
+ @classmethod
197
+ def validate_model(cls, v):
198
+ if not v:
199
+ raise ValueError("Model is required")
200
+ return v
201
+
202
+
203
+class AnthropicCountTokensResponse(BaseModel):
204
+ """Anthropic messages.count_tokens response"""
205
+
206
+ input_tokens: int
207
+ context_management: AnthropicContextManagement | None = None
vllm/entrypoints/anthropic/serving.py
+247
-122
@@ -17,6 +17,9 @@ from fastapi import Request
17
from vllm.engine.protocol import EngineClient
18
from vllm.entrypoints.anthropic.protocol import (
19
AnthropicContentBlock,
20
+ AnthropicContextManagement,
21
+ AnthropicCountTokensRequest,
22
+ AnthropicCountTokensResponse,
23
AnthropicDelta,
24
AnthropicError,
25
AnthropicMessagesRequest,
@@ -109,135 +112,202 @@ class AnthropicServingMessages(OpenAIServingChat):
112
113
@classmethod
114
def _convert_anthropic_to_openai_request(
112
- cls, anthropic_request: AnthropicMessagesRequest
115
+ cls, anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest
116
) -> ChatCompletionRequest:
117
"""Convert Anthropic message format to OpenAI format"""
115
- openai_messages = []
118
+ openai_messages: list[dict[str, Any]] = []
119
+
120
+ cls._convert_system_message(anthropic_request, openai_messages)
121
+ cls._convert_messages(anthropic_request.messages, openai_messages)
122
+ req = cls._build_base_request(anthropic_request, openai_messages)
123
+ cls._handle_streaming_options(req, anthropic_request)
124
+ cls._convert_tool_choice(anthropic_request, req)
125
+ cls._convert_tools(anthropic_request, req)
126
+ return req
127
117
- # Add system message if provided
118
- if anthropic_request.system:
119
- if isinstance(anthropic_request.system, str):
120
- openai_messages.append(
121
- {"role": "system", "content": anthropic_request.system}
122
- )
123
- else:
124
- system_prompt = ""
125
- for block in anthropic_request.system:
126
- if block.type == "text" and block.text:
127
- system_prompt += block.text
128
- openai_messages.append({"role": "system", "content": system_prompt})
128
+ @classmethod
129
+ def _convert_system_message(
130
+ cls,
131
+ anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest,
132
+ openai_messages: list[dict[str, Any]],
133
+ ) -> None:
134
+ """Convert Anthropic system message to OpenAI format"""
135
+ if not anthropic_request.system:
136
+ return
137
+
138
+ if isinstance(anthropic_request.system, str):
139
+ openai_messages.append(
140
+ {"role": "system", "content": anthropic_request.system}
141
+ )
142
+ else:
143
+ system_prompt = ""
144
+ for block in anthropic_request.system:
145
+ if block.type == "text" and block.text:
146
+ system_prompt += block.text
147
+ openai_messages.append({"role": "system", "content": system_prompt})
148
130
- for msg in anthropic_request.messages:
149
+ @classmethod
150
+ def _convert_messages(
151
+ cls, messages: list, openai_messages: list[dict[str, Any]]
152
+ ) -> None:
153
+ """Convert Anthropic messages to OpenAI format"""
154
+ for msg in messages:
155
openai_msg: dict[str, Any] = {"role": msg.role} # type: ignore
156
+
157
if isinstance(msg.content, str):
158
openai_msg["content"] = msg.content
159
else:
135
- # Handle complex content blocks
136
- content_parts: list[dict[str, Any]] = []
137
- tool_calls: list[dict[str, Any]] = []
138
- reasoning_parts: list[str] = []
139
-
140
- for block in msg.content:
141
- if block.type == "text" and block.text:
142
- content_parts.append({"type": "text", "text": block.text})
143
- elif block.type == "image" and block.source:
144
- image_url = cls._convert_image_source_to_url(block.source)
145
- content_parts.append(
146
- {
147
- "type": "image_url",
148
- "image_url": {"url": image_url},
149
- }
150
- )
151
- elif block.type == "thinking" and block.thinking is not None:
152
- reasoning_parts.append(block.thinking)
153
- elif block.type == "tool_use":
154
- # Convert tool use to function call format
155
- tool_call = {
156
- "id": block.id or f"call_{int(time.time())}",
157
- "type": "function",
158
- "function": {
159
- "name": block.name or "",
160
- "arguments": json.dumps(block.input or {}),
161
- },
162
- }
163
- tool_calls.append(tool_call)
164
- elif block.type == "tool_result":
165
- if msg.role == "user":
166
- # Parse tool_result content which can be
167
- # a string or a list of content blocks
168
- # (text, image, etc.)
169
- tool_text = ""
170
- tool_image_urls: list[str] = []
171
- if isinstance(block.content, str):
172
- tool_text = block.content
173
- elif isinstance(block.content, list):
174
- text_parts: list[str] = []
175
- for item in block.content:
176
- if not isinstance(item, dict):
177
- continue
178
- item_type = item.get("type")
179
- if item_type == "text":
180
- text_parts.append(item.get("text", ""))
181
- elif item_type == "image":
182
- source = item.get("source", {})
183
- url = cls._convert_image_source_to_url(source)
184
- if url:
185
- tool_image_urls.append(url)
186
- tool_text = "\n".join(text_parts)
187
- openai_messages.append(
188
- {
189
- "role": "tool",
190
- "tool_call_id": block.tool_use_id or "",
191
- "content": tool_text or "",
192
- }
193
- )
194
- # OpenAI tool messages only support string
195
- # content, so inject images from tool
196
- # results as a follow-up user message
197
- if tool_image_urls:
198
- openai_messages.append(
199
- {
200
- "role": "user",
201
- "content": [ # type: ignore[dict-item]
202
- {
203
- "type": "image_url",
204
- "image_url": {"url": img},
205
- }
206
- for img in tool_image_urls
207
- ],
208
- }
209
- )
210
- else:
211
- # Assistant tool result becomes regular text
212
- tool_result_text = (
213
- str(block.content) if block.content else ""
214
- )
215
- content_parts.append(
216
- {
217
- "type": "text",
218
- "text": f"Tool result: {tool_result_text}",
219
- }
220
- )
160
+ cls._convert_message_content(msg, openai_msg, openai_messages)
161
+
162
+ openai_messages.append(openai_msg)
163
222
- if reasoning_parts:
223
- openai_msg["reasoning"] = "".join(reasoning_parts)
164
+ @classmethod
165
+ def _convert_message_content(
166
+ cls,
167
+ msg,
168
+ openai_msg: dict[str, Any],
169
+ openai_messages: list[dict[str, Any]],
170
+ ) -> None:
171
+ """Convert complex message content blocks"""
172
+ content_parts: list[dict[str, Any]] = []
173
+ tool_calls: list[dict[str, Any]] = []
174
+ reasoning_parts: list[str] = []
175
+
176
+ for block in msg.content:
177
+ cls._convert_block(
178
+ block,
179
+ msg.role,
180
+ content_parts,
181
+ tool_calls,
182
+ reasoning_parts,
183
+ openai_messages,
184
+ )
185
225
- # Add tool calls to the message if any
226
- if tool_calls:
227
- openai_msg["tool_calls"] = tool_calls # type: ignore
186
+ if reasoning_parts:
187
+ openai_msg["reasoning"] = "".join(reasoning_parts)
188
229
- # Add content parts if any
230
- if content_parts:
231
- if len(content_parts) == 1 and content_parts[0]["type"] == "text":
232
- openai_msg["content"] = content_parts[0]["text"]
233
- else:
234
- openai_msg["content"] = content_parts # type: ignore
235
- elif not tool_calls and not reasoning_parts:
189
+ if tool_calls:
190
+ openai_msg["tool_calls"] = tool_calls # type: ignore
191
+
192
+ if content_parts:
193
+ if len(content_parts) == 1 and content_parts[0]["type"] == "text":
194
+ openai_msg["content"] = content_parts[0]["text"]
195
+ else:
196
+ openai_msg["content"] = content_parts # type: ignore
197
+ elif not tool_calls and not reasoning_parts:
198
+ return
199
+
200
+ @classmethod
201
+ def _convert_block(
202
+ cls,
203
+ block,
204
+ role: str,
205
+ content_parts: list[dict[str, Any]],
206
+ tool_calls: list[dict[str, Any]],
207
+ reasoning_parts: list[str],
208
+ openai_messages: list[dict[str, Any]],
209
+ ) -> None:
210
+ """Convert individual content block"""
211
+ if block.type == "text" and block.text:
212
+ content_parts.append({"type": "text", "text": block.text})
213
+ elif block.type == "image" and block.source:
214
+ image_url = cls._convert_image_source_to_url(block.source)
215
+ content_parts.append({"type": "image_url", "image_url": {"url": image_url}})
216
+ elif block.type == "thinking" and block.thinking is not None:
217
+ reasoning_parts.append(block.thinking)
218
+ elif block.type == "tool_use":
219
+ cls._convert_tool_use_block(block, tool_calls)
220
+ elif block.type == "tool_result":
221
+ cls._convert_tool_result_block(block, role, openai_messages, content_parts)
222
+
223
+ @classmethod
224
+ def _convert_tool_use_block(cls, block, tool_calls: list[dict[str, Any]]) -> None:
225
+ """Convert tool_use block to OpenAI function call format"""
226
+ tool_call = {
227
+ "id": block.id or f"call_{int(time.time())}",
228
+ "type": "function",
229
+ "function": {
230
+ "name": block.name or "",
231
+ "arguments": json.dumps(block.input or {}),
232
+ },
233
+ }
234
+ tool_calls.append(tool_call)
235
+
236
+ @classmethod
237
+ def _convert_tool_result_block(
238
+ cls,
239
+ block,
240
+ role: str,
241
+ openai_messages: list[dict[str, Any]],
242
+ content_parts: list[dict[str, Any]],
243
+ ) -> None:
244
+ """Convert tool_result block to OpenAI format"""
245
+ if role == "user":
246
+ cls._convert_user_tool_result(block, openai_messages)
247
+ else:
248
+ tool_result_text = str(block.content) if block.content else ""
249
+ content_parts.append(
250
+ {"type": "text", "text": f"Tool result: {tool_result_text}"}
251
+ )
252
+
253
+ @classmethod
254
+ def _convert_user_tool_result(
255
+ cls, block, openai_messages: list[dict[str, Any]]
256
+ ) -> None:
257
+ """Convert user tool_result with text and image support"""
258
+ tool_text = ""
259
+ tool_image_urls: list[str] = []
260
+
261
+ if isinstance(block.content, str):
262
+ tool_text = block.content
263
+ elif isinstance(block.content, list):
264
+ text_parts: list[str] = []
265
+ for item in block.content:
266
+ if not isinstance(item, dict):
267
continue
268
+ item_type = item.get("type")
269
+ if item_type == "text":
270
+ text_parts.append(item.get("text", ""))
271
+ elif item_type == "image":
272
+ source = item.get("source", {})
273
+ url = cls._convert_image_source_to_url(source)
274
+ if url:
275
+ tool_image_urls.append(url)
276
+ tool_text = "\n".join(text_parts)
277
+
278
+ openai_messages.append(
279
+ {
280
+ "role": "tool",
281
+ "tool_call_id": block.tool_use_id or "",
282
+ "content": tool_text or "",
283
+ }
284
+ )
285
238
- openai_messages.append(openai_msg)
286
+ if tool_image_urls:
287
+ openai_messages.append(
288
+ {
289
+ "role": "user",
290
+ "content": [ # type: ignore[dict-item]
291
+ {"type": "image_url", "image_url": {"url": img}}
292
+ for img in tool_image_urls
293
+ ],
294
+ }
295
+ )
296
240
- req = ChatCompletionRequest(
297
+ @classmethod
298
+ def _build_base_request(
299
+ cls,
300
+ anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest,
301
+ openai_messages: list[dict[str, Any]],
302
+ ) -> ChatCompletionRequest:
303
+ """Build base ChatCompletionRequest"""
304
+ if isinstance(anthropic_request, AnthropicCountTokensRequest):
305
+ return ChatCompletionRequest(
306
+ model=anthropic_request.model,
307
+ messages=openai_messages,
308
+ )
309
+
310
+ return ChatCompletionRequest(
311
model=anthropic_request.model,
312
messages=openai_messages,
313
max_tokens=anthropic_request.max_tokens,
@@ -248,19 +318,38 @@ class AnthropicServingMessages(OpenAIServingChat):
318
top_k=anthropic_request.top_k,
319
)
320
321
+ @classmethod
322
+ def _handle_streaming_options(
323
+ cls,
324
+ req: ChatCompletionRequest,
325
+ anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest,
326
+ ) -> None:
327
+ """Handle streaming configuration"""
328
+ if isinstance(anthropic_request, AnthropicCountTokensRequest):
329
+ return
330
if anthropic_request.stream:
331
req.stream = anthropic_request.stream
253
- req.stream_options = StreamOptions.validate(
332
+ req.stream_options = StreamOptions.model_validate(
333
{"include_usage": True, "continuous_usage_stats": True}
334
)
335
336
+ @classmethod
337
+ def _convert_tool_choice(
338
+ cls,
339
+ anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest,
340
+ req: ChatCompletionRequest,
341
+ ) -> None:
342
+ """Convert Anthropic tool_choice to OpenAI format"""
343
if anthropic_request.tool_choice is None:
344
req.tool_choice = None
259
- elif anthropic_request.tool_choice.type == "auto":
345
+ return
346
+
347
+ tool_choice_type = anthropic_request.tool_choice.type
348
+ if tool_choice_type == "auto":
349
req.tool_choice = "auto"
261
- elif anthropic_request.tool_choice.type == "any":
350
+ elif tool_choice_type == "any":
351
req.tool_choice = "required"
263
- elif anthropic_request.tool_choice.type == "tool":
352
+ elif tool_choice_type == "tool":
353
req.tool_choice = ChatCompletionNamedToolChoiceParam.model_validate(
354
{
355
"type": "function",
@@ -268,9 +357,17 @@ class AnthropicServingMessages(OpenAIServingChat):
357
}
358
)
359
271
- tools = []
360
+ @classmethod
361
+ def _convert_tools(
362
+ cls,
363
+ anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest,
364
+ req: ChatCompletionRequest,
365
+ ) -> None:
366
+ """Convert Anthropic tools to OpenAI format"""
367
if anthropic_request.tools is None:
273
- return req
368
+ return
369
+
370
+ tools = []
371
for tool in anthropic_request.tools:
372
tools.append(
373
ChatCompletionToolsParam.model_validate(
@@ -284,10 +381,10 @@ class AnthropicServingMessages(OpenAIServingChat):
381
}
382
)
383
)
384
+
385
if req.tool_choice is None:
386
req.tool_choice = "auto"
387
req.tools = tools
290
- return req
388
389
async def create_messages(
390
self,
@@ -670,3 +767,31 @@ class AnthropicServingMessages(OpenAIServingChat):
767
data = error_response.model_dump_json(exclude_unset=True)
768
yield wrap_data_with_event(data, "error")
769
yield "data: [DONE]\n\n"
770
+
771
+ async def count_tokens(
772
+ self,
773
+ request: AnthropicCountTokensRequest,
774
+ raw_request: Request | None = None,
775
+ ) -> AnthropicCountTokensResponse | ErrorResponse:
776
+ """Implements Anthropic's messages.count_tokens endpoint."""
777
+ chat_req = self._convert_anthropic_to_openai_request(request)
778
+ result = await self.render_chat_request(chat_req)
779
+ if isinstance(result, ErrorResponse):
780
+ return result
781
+
782
+ _, engine_prompts = result
783
+
784
+ input_tokens = sum( # type: ignore
785
+ len(prompt["prompt_token_ids"]) # type: ignore[typeddict-item, misc]
786
+ for prompt in engine_prompts
787
+ if "prompt_token_ids" in prompt
788
+ )
789
+
790
+ response = AnthropicCountTokensResponse(
791
+ input_tokens=input_tokens,
792
+ context_management=AnthropicContextManagement(
793
+ original_input_tokens=input_tokens
794
+ ),
795
+ )
796
+
797
+ return response