Support Anthropic messages count_tokens API

lightllm/server/api_anthropic.py

@@ -1316,3 +1316,41 @@ async def anthropic_messages_impl(raw_request: Request) -> Response:
         logger.error("Failed to translate response to Anthropic format: %s", exc)
         return _anthropic_error_response(HTTPStatus.INTERNAL_SERVER_ERROR, str(exc))
     return JSONResponse(anthropic_dict)
+
+
+async def anthropic_count_tokens_impl(raw_request: Request) -> Response:
+    """Count the fully rendered input for Anthropic's Messages API."""
+    from .api_http import g_objs
+    from .api_models import ChatCompletionRequest
+    from .api_openai import _build_multimodal_params, _select_chat_template_tools
+    from .build_prompt import build_prompt
+    from .core.objs.sampling_params import SamplingParams
+
+    raw_body = await raw_request.json()
+    # LiteLLM validates the body as a Messages creation request, where
+    # max_tokens is required. Anthropic's count_tokens request omits it;
+    # the value has no effect on prompt rendering, so provide a sentinel.
+    raw_body.setdefault("max_tokens", 1)
+    chat_dict, _ = await asyncio.to_thread(_anthropic_to_chat_request, raw_body)
+    chat_request = ChatCompletionRequest(**chat_dict)
+
+    prompt = await build_prompt(chat_request, _select_chat_template_tools(chat_request))
+    multimodal_params = _build_multimodal_params(chat_request)
+    await multimodal_params.verify_and_preload(raw_request)
+
+    sampling_params = SamplingParams()
+    sampling_params.init(tokenizer=g_objs.httpserver_manager.tokenizer, add_special_tokens=False)
+    sampling_params.verify()
+    input_tokens = g_objs.httpserver_manager.tokens(
+        prompt,
+        multimodal_params,
+        sampling_params,
+        {"add_special_tokens": False},
+    )
+
+    return JSONResponse(
+        {
+            "input_tokens": input_tokens,
+            "context_management": {"original_input_tokens": input_tokens},
+        }
+    )

lightllm/server/api_http.py

@@ -420,6 +420,20 @@ async def anthropic_messages(raw_request: Request) -> Response:
         return Response(status_code=499)
 
 
+@app.post("/v1/messages/count_tokens")
+async def anthropic_count_tokens(raw_request: Request) -> Response:
+    from .api_anthropic import _anthropic_error_response, anthropic_count_tokens_impl
+
+    try:
+        return await anthropic_count_tokens_impl(raw_request)
+    except ClientDisconnected as e:
+        logger.warning(str(e))
+        return Response(status_code=499)
+    except Exception as e:
+        logger.error("An error occurred: %s", str(e), exc_info=True)
+        return _anthropic_error_response(HTTPStatus.EXPECTATION_FAILED, f"error: {str(e)}")
+
+
 @app.post("/v1/responses")
 async def openai_responses(raw_request: Request) -> Response:
     if get_env_start_args().run_mode in ["prefill", "decode"]:

lightllm/server/api_openai.py

@@ -217,20 +217,8 @@ def _split_tool_argument_delta(arguments: Optional[str]) -> List[str]:
     return [arguments]
 
 
-async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Request) -> Response:
-    from .api_http import g_objs
-
-    if request.logit_bias is not None:
-        return create_error_response(
-            HTTPStatus.BAD_REQUEST,
-            "The logit_bias parameter is not currently supported",
-        )
-
-    if request.function_call != "none":
-        return create_error_response(HTTPStatus.BAD_REQUEST, "The function call feature is not supported")
-
-    created_time = int(time.time())
-
+def _build_multimodal_params(request: ChatCompletionRequest) -> MultimodalParams:
+    """Build LightLLM multimodal inputs from chat content parts."""
     multimodal_params_dict = {"images": [], "audios": []}
     for message in request.messages:
         if isinstance(message.content, list):
@@ -275,6 +263,11 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req
                     else:
                         raise ValueError("Unrecognized audio input. Supports local path, http url, base64.")
 
+    return MultimodalParams(**multimodal_params_dict)
+
+
+def _select_chat_template_tools(request: ChatCompletionRequest) -> Optional[List[dict]]:
+    """Select the tool schemas exposed to the model's chat template."""
     tools = None
     if request.tools and request.tool_choice != "none":
         # request.skip_special_tokens = False
@@ -293,6 +286,24 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req
         else:
             tools = [item.function.model_dump(exclude_none=True) for item in request.tools]
 
+    return tools
+
+
+async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Request) -> Response:
+    from .api_http import g_objs
+
+    if request.logit_bias is not None:
+        return create_error_response(
+            HTTPStatus.BAD_REQUEST,
+            "The logit_bias parameter is not currently supported",
+        )
+
+    if request.function_call != "none":
+        return create_error_response(HTTPStatus.BAD_REQUEST, "The function call feature is not supported")
+
+    created_time = int(time.time())
+    multimodal_params = _build_multimodal_params(request)
+    tools = _select_chat_template_tools(request)
     prompt = await build_prompt(request, tools)
     sampling_params_dict = {
         "do_sample": request.do_sample,
@@ -340,8 +351,6 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req
     sampling_params.init(tokenizer=g_objs.httpserver_manager.tokenizer, **sampling_params_dict)
 
     sampling_params.verify()
-    multimodal_params = MultimodalParams(**multimodal_params_dict)
-
     results_generator = g_objs.httpserver_manager.generate(
         prompt, sampling_params, multimodal_params, request=raw_request
     )