diff --git a/lightllm/server/api_anthropic.py b/lightllm/server/api_anthropic.py index 4460b21df..0afe708ee 100644 --- a/lightllm/server/api_anthropic.py +++ b/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}, + } + ) diff --git a/lightllm/server/api_http.py b/lightllm/server/api_http.py index d985b5e66..8c85fd240 100755 --- a/lightllm/server/api_http.py +++ b/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"]: diff --git a/lightllm/server/api_openai.py b/lightllm/server/api_openai.py index e8570369a..e89ec5c5a 100644 --- a/lightllm/server/api_openai.py +++ b/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 )