From 748e03088d0415076d4342061002c227113afbec Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:54:57 -0700 Subject: [PATCH] fix(count_tokens)!: answer 400 for Vertex partner models without a count API and OpenAI-shaped provider errors on /v1/responses/input_tokens (#47163) Backport of #47163 to stable/1.104.x. Cherry-picked from ef1d3959a06289eae04b89c94426e7c93ebbb9a5 (main). --- litellm/llms/vertex_ai/common_utils.py | 18 ++ .../count_tokens/handler.py | 36 ++- .../common_utils/openai_error_payload.py | 4 + .../proxy/response_api_endpoints/endpoints.py | 97 +++++++- .../test_openai_count_tokens_tools_wire.py | 25 ++ .../test_vertex_partner_count_tokens_wire.py | 177 ++++++++++++++- .../response_api_endpoints/test_endpoints.py | 214 ++++++++++++++++++ .../vertex_ai/test_vertex_ai_common_utils.py | 29 +++ tests/unit/proxy/test_custom_tokenizer_bug.py | 60 ++++- 9 files changed, 628 insertions(+), 32 deletions(-) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 2c7e07616b5..eb57ef4e6c6 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1242,6 +1242,7 @@ class VertexAITokenCounter(BaseTokenCounter): ) -> TokenCountResponse | None: import copy + from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import count_tokens_publisher from litellm.llms.vertex_ai.vertex_ai_partner_models.main import ( VertexAIError as PartnerVertexAIError, ) @@ -1252,6 +1253,23 @@ class VertexAITokenCounter(BaseTokenCounter): deployment = deployment or {} count_tokens_params_request: Final = copy.deepcopy(deployment.get("litellm_params", {})) + without_count_route: Final = ( + VertexAIPartnerModels.is_vertex_partner_model(model_to_use) and count_tokens_publisher(model_to_use) is None + ) + model_info: Final[Mapping[str, object]] = deployment.get("model_info") or {} + if without_count_route and messages and model_info.get("custom_tokenizer") is not None: + return None + if without_count_route: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="vertex_ai_partner_models", + error=True, + error_message=f"LiteLLM does not support counting tokens for Vertex AI model {model_to_use}.", + status_code=400, + ) + # Check if this is a partner model (Claude, Mistral, etc.) if VertexAIPartnerModels.is_vertex_partner_model(model_to_use): # Use partner models token counter diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py index 285380f5fc7..7100c9efde3 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py @@ -1,5 +1,5 @@ """ -Token counter for Vertex AI Partner Models (Anthropic Claude, Mistral, etc.) +Token counter for Vertex AI Partner Models (Anthropic Claude, Meta Llama) This handler provides token counting for partner models hosted on Vertex AI. Unlike Gemini models which use Google's token counting API, partner models use @@ -12,36 +12,30 @@ from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.vertex_ai.common_utils import get_vertex_base_url from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +_COUNT_TOKENS_PUBLISHERS: Final = ( + ("claude", "anthropic"), + ("llama", "meta"), + ("meta/", "meta"), +) + + +def count_tokens_publisher(model: str) -> str | None: + return next((publisher for marker, publisher in _COUNT_TOKENS_PUBLISHERS if marker in model), None) + class VertexAIPartnerModelsTokenCounter(VertexBase): """ Token counter for Vertex AI Partner Models. - Handles token counting for models like Claude (Anthropic), Mistral, etc. + Handles token counting for models like Claude (Anthropic) and Llama (Meta) that are available through Vertex AI's partner model program. """ def _get_publisher_for_model(self, model: str) -> str: - """ - Determine the publisher name for the given model. - - Args: - model: The model name (e.g., "claude-3-5-sonnet-20241022") - - Returns: - Publisher name to use in the Vertex AI endpoint URL - - Raises: - ValueError: If the model is not a recognized partner model - """ - if "claude" in model: - return "anthropic" - elif "mistral" in model or "codestral" in model: - return "mistralai" - elif "llama" in model or "meta/" in model: - return "meta" - else: + publisher: Final = count_tokens_publisher(model) + if publisher is None: raise ValueError(f"Unknown partner model: {model}") + return publisher def _build_count_tokens_endpoint( self, diff --git a/litellm/proxy/common_utils/openai_error_payload.py b/litellm/proxy/common_utils/openai_error_payload.py index 202c61b620e..b24761ac039 100644 --- a/litellm/proxy/common_utils/openai_error_payload.py +++ b/litellm/proxy/common_utils/openai_error_payload.py @@ -42,6 +42,10 @@ def openai_error_type(exc: object, status_code: int) -> str: carried: Final = attribute_of(exc, "type") if isinstance(carried, str) and carried != STRINGIFIED_NONE: return carried + return openai_error_type_for_status(status_code) + + +def openai_error_type_for_status(status_code: int) -> str: mapped: Final = _OPENAI_ERROR_TYPE_BY_STATUS.get(status_code) if mapped is not None: return mapped diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 75eefb2e73b..021c741b2bc 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -14,11 +14,12 @@ from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse from openai.types.responses import ResponseItemList from openai.types.responses.response_create_params import ResponseInputParam -from pydantic import BaseModel, ConfigDict, ValidationError +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, with_config from starlette.websockets import WebSocket, WebSocketDisconnect -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping from litellm.constants import EMPTY_MAPPING from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.llms.base_llm.guardrail_translation.utils import ( @@ -31,10 +32,16 @@ from litellm.proxy.auth.user_api_key_auth import ( user_api_key_auth_websocket, ) from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_utils.error_body_call_id import error_body_call_id, with_call_id from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_set_request_parsed_body, ) +from litellm.proxy.common_utils.openai_error_payload import ( + LITELLM_CALL_ID_HEADER, + error_status_code, + openai_error_type_for_status, +) from litellm.proxy.route_llm_request import raise_if_required_body_param_missing from litellm.types.llms.openai import ( REASONING_EFFORT, @@ -1103,6 +1110,85 @@ def _responses_invalid_request_response(message: str, param: str | None, code: s return JSONResponse(status_code=400, content=body) +class _ProviderErrorMessage(TypedDict): + message: ReadOnly[str] + + +class _NestedProviderError(TypedDict): + error: ReadOnly[_ProviderErrorMessage] + + +@with_config(ConfigDict(extra="ignore")) +class _ProviderOpenAIErrorDetail(TypedDict): + message: ReadOnly[str] + type: ReadOnly[str] + param: NotRequired[ReadOnly[str | None]] + code: NotRequired[ReadOnly[str | None]] + + +@with_config(ConfigDict(extra="forbid")) +class _ProviderOpenAIError(TypedDict): + error: ReadOnly[_ProviderOpenAIErrorDetail] + + +_PROVIDER_COUNT_ERROR_TYPE: Final = "token_counting_error" +_PROVIDER_OPENAI_ERROR: Final = TypeAdapter(_ProviderOpenAIError) +_NESTED_PROVIDER_ERROR: Final = TypeAdapter(_NestedProviderError) + + +def _provider_error_json(raw_message: str) -> str: + start: Final = raw_message.find("{") + return raw_message if start == -1 else raw_message[start:] + + +def _provider_error_message(raw_message: str) -> str: + error_json: Final = _provider_error_json(raw_message) + try: + return _NESTED_PROVIDER_ERROR.validate_json(error_json)["error"]["message"] + except ValidationError: + flat_message: Final = AnthropicExceptionMapping.extract_error_message(error_json) + return raw_message if flat_message == error_json else flat_message + + +def _provider_count_error_body(raw_message: str, status_code: int) -> _ResponsesApiErrorBody: + try: + provider_error: Final = _PROVIDER_OPENAI_ERROR.validate_json(_provider_error_json(raw_message))["error"] + except ValidationError: + return { + "error": { + "message": _provider_error_message(raw_message), + "type": openai_error_type_for_status(status_code), + "param": None, + "code": None, + } + } + return { + "error": { + "message": provider_error["message"], + "type": provider_error["type"], + "param": provider_error.get("param"), + "code": provider_error.get("code"), + } + } + + +def _provider_count_error_response(exc: ProxyException, request: Request) -> JSONResponse: + from litellm.proxy.proxy_server import ( + _close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; a route that answers an error itself must stamp the OTel server span like the global ProxyException handler does + general_settings_view, + ) + + status_code: Final = error_status_code(exc, 500) + _close_dangling_otel_server_span(request, status_code, exc=exc) + body: Final = _provider_count_error_body(exc.message, status_code) + call_id: Final = error_body_call_id(general_settings_view(), exc.headers.get(LITELLM_CALL_ID_HEADER)) + return JSONResponse( + status_code=status_code, + content={"error": with_call_id({**body["error"]}, call_id)}, + headers=exc.headers, + ) + + def _missing_responses_param_response(param: str) -> JSONResponse: return _responses_invalid_request_response( message=f"Missing required parameter: '{param}'.", @@ -1194,7 +1280,12 @@ async def responses_input_tokens( message=f"Invalid request for token counting: {e}", param=None, code=None ) - token_response: Final = await token_counter(request=token_request, call_endpoint=True) + try: + token_response: Final = await token_counter(request=token_request, call_endpoint=True) + except ProxyException as e: + if e.type != _PROVIDER_COUNT_ERROR_TYPE: + raise + return _provider_count_error_response(e, request) result: Final[_ResponsesInputTokensResult] = { "object": "response.input_tokens", "input_tokens": token_response.total_tokens, diff --git a/tests/integration/providers/test_openai_count_tokens_tools_wire.py b/tests/integration/providers/test_openai_count_tokens_tools_wire.py index 2b23ccdf4f2..16a5b8685ad 100644 --- a/tests/integration/providers/test_openai_count_tokens_tools_wire.py +++ b/tests/integration/providers/test_openai_count_tokens_tools_wire.py @@ -391,3 +391,28 @@ def test_google_count_tokens_sends_the_system_instruction_text(counting_proxy: O ) assert response.status_code == 200, response.text assert _count_body(wire)["instructions"] == "Be brief.\nAnswer in English." + + +@pytest.mark.parametrize( + "error", + [ + pytest.param( + {"message": "Invalid input", "type": "invalid_request_error", "param": "input", "code": "invalid_value"}, + id="full", + ), + pytest.param({"message": "Invalid input", "type": "invalid_request_error"}, id="bare"), + ], +) +def test_responses_input_tokens_passes_an_openai_count_error_through( + counting_proxy: OwnedProxy, error: dict[str, JsonValue] +) -> None: + def rejecting(_request: Request) -> Reply: + return Reply(status=400, body=json.dumps({"error": error}).encode()) + + gateway: Final = counting_proxy.gateway + with wire_server(rejecting) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request("POST", "/v1/responses/input_tokens", {"model": model, "input": "hello"}) + assert response.status_code == 400, response.text + assert response.json() == {"error": {"param": None, "code": None, **error}}, response.text + assert _count_body(wire)["input"] == [{"role": "user", "content": "hello"}] diff --git a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py index ae9b8434d11..1c4b1d1a6e5 100644 --- a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py @@ -7,8 +7,10 @@ import uuid from collections.abc import Callable, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack +from dataclasses import dataclass from pathlib import Path from queue import SimpleQueue +from types import MappingProxyType from typing import Final import httpx @@ -16,10 +18,11 @@ import psutil import pytest import yaml from integration._support.client import Gateway, Scenario, eventually -from integration._support.process import owned_proxy_process +from integration._support.process import graceful_stop_seconds, owned_proxy_process from integration._support.vertex import service_account_json from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter +from tokenizers import Regex, Tokenizer, models, pre_tokenizers _BACKEND: Final = "claude-sonnet-4-6" _PROJECT: Final = "scripted-project" @@ -35,9 +38,19 @@ _REPLY_TEXT: Final = "scripted reply" _OWNED_MODEL: Final = "partner-claude" _OWNED_UNREACHABLE_MODEL: Final = "partner-claude-unreachable" _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_GLM_BACKEND: Final = "zai-org/glm-5.2-maas" +_MISTRAL_BACKEND: Final = "mistral-medium-3" +_LLAMA_BACKEND: Final = "meta/llama-4-maverick-17b-128e-instruct-maas" +_LLAMA_COUNT_TARGET: Final = ( + f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/meta/models/count-tokens:rawPredict" +) +_COUNTED_TEXT: Final = "Count this message" +_TOKENIZER_REPO: Final = "scripted/char-tokenizer" +_TOKENIZER_TARGET: Final = f"/{_TOKENIZER_REPO}/resolve/main/tokenizer.json" +_TOKENIZER_HEADERS: Final = MappingProxyType({"X-Repo-Commit": "0" * 40, "ETag": '"scripted-char-tokenizer"'}) _STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") -_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this message"}] +_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": _COUNTED_TEXT}] _SYSTEM: Final = "You are a terse assistant that answers in one sentence" _SYSTEM_BLOCKS: Final[list[JsonValue]] = [ {"type": "text", "text": "You are a terse assistant"}, @@ -115,16 +128,16 @@ def _counted(_request: Request) -> Reply: return Reply(body=json.dumps({"input_tokens": _PEER_COUNT}).encode()) -def _rejected(status: int) -> Reply: +def _rejected(status: int, error_type: str = "invalid_request_error") -> Reply: return Reply( status=status, - body=json.dumps({"type": "error", "error": {"type": "invalid_request_error", "message": _REJECTION}}).encode(), + body=json.dumps({"type": "error", "error": {"type": error_type, "message": _REJECTION}}).encode(), ) -def _rejecting(status: int) -> Callable[[Request], Reply]: +def _rejecting(status: int, error_type: str = "invalid_request_error") -> Callable[[Request], Reply]: def count(_request: Request) -> Reply: - return _rejected(status) + return _rejected(status, error_type) return count @@ -224,6 +237,78 @@ def _local_count(gateway: Gateway, body: Mapping[str, JsonValue]) -> int: return total +@dataclass(frozen=True, slots=True) +class _CountRoute: + path: Callable[[str], str] + body: Callable[[str], dict[str, JsonValue]] + rejection: Callable[[str], dict[str, JsonValue]] + count_field: str + params: Mapping[str, str] = MappingProxyType({}) + + def send(self, gateway: Gateway, model: str) -> httpx.Response: + return gateway.request("POST", self.path(model), self.body(model), params=self.params) + + +_MESSAGES_ROUTE: Final = _CountRoute( + path=lambda _model: "/v1/messages/count_tokens", + body=_bare, + rejection=lambda message: {"type": "error", "error": {"type": "invalid_request_error", "message": message}}, + count_field="input_tokens", +) +_RESPONSES_ROUTE: Final = _CountRoute( + path=lambda _model: "/v1/responses/input_tokens", + body=lambda model: {"model": model, "input": _COUNTED_TEXT}, + rejection=lambda message: { + "error": {"message": message, "type": "invalid_request_error", "param": None, "code": None} + }, + count_field="input_tokens", +) +_GOOGLE_ROUTE: Final = _CountRoute( + path=lambda model: f"/v1beta/models/{model}:countTokens", + body=lambda _model: {"contents": [{"role": "user", "parts": [{"text": _COUNTED_TEXT}]}]}, + rejection=lambda message: {"error": {"code": 400, "message": message, "status": "INVALID_ARGUMENT"}}, + count_field="totalTokens", +) +_UTILS_ROUTE: Final = _CountRoute( + path=lambda _model: "/utils/token_counter", + body=_bare, + rejection=lambda message: { + "error": {"message": message, "type": "token_counting_error", "param": "model", "code": "400"} + }, + count_field="total_tokens", + params=MappingProxyType({"call_endpoint": "true"}), +) +_COUNT_ROUTES: Final = (_MESSAGES_ROUTE, _RESPONSES_ROUTE, _GOOGLE_ROUTE, _UTILS_ROUTE) + + +def _error_message(payload: Mapping[str, JsonValue]) -> str: + error: Final = payload.get("error") + assert isinstance(error, dict), payload + message: Final = error.get("message") + assert isinstance(message, str), payload + return message + + +def _route_count(gateway: Gateway, route: _CountRoute, model: str) -> tuple[int, JsonValue]: + response: Final = route.send(gateway, model) + return response.status_code, _JSON_OBJECT.validate_json(response.content).get(route.count_field) + + +def _char_tokenizer() -> bytes: + tokenizer: Final = Tokenizer(models.WordLevel({"[UNK]": 0}, unk_token="[UNK]")) + tokenizer.pre_tokenizer = pre_tokenizers.Split(Regex("."), behavior="isolated") + return tokenizer.to_str().encode() + + +def _tokenizer_hub(tokenizer: bytes) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target == _TOKENIZER_TARGET: + return Reply(body=tokenizer, headers=_TOKENIZER_HEADERS) + return Reply(status=404, body=json.dumps({"error": f"unscripted target {request.target}"}).encode()) + + return respond + + def _closed_port_url() -> str: with socket.socket() as reserve: reserve.bind(("127.0.0.1", 0)) @@ -374,11 +459,89 @@ def test_responses_input_tokens_surfaces_the_peer_error(gateway: Gateway) -> Non {"model": model, "input": "Count this message", "instructions": "Be terse", "tools": _RESPONSES_TOOLS}, ) assert response.status_code == 400, response.text - assert "input_tokens" not in response.text, response.text + assert response.json() == _RESPONSES_ROUTE.rejection(_REJECTION), response.text (sent,) = _counted_bodies(wire) assert sent["tools"] == _RESPONSES_TOOLS, sent +def test_responses_input_tokens_types_a_peer_error_by_its_status(gateway: Gateway) -> None: + with wire_server(_peer(_rejecting(404, "not_found_error"))) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = _RESPONSES_ROUTE.send(gateway, model) + assert response.status_code == 404, response.text + assert response.json() == _RESPONSES_ROUTE.rejection(_REJECTION), response.text + assert _counted_bodies(wire) == (_PEER_BARE,) + + +@pytest.mark.parametrize( + ("backend", "route"), + [ + pytest.param(_GLM_BACKEND, _MESSAGES_ROUTE, id="glm-messages"), + pytest.param(_GLM_BACKEND, _RESPONSES_ROUTE, id="glm-responses"), + pytest.param(_GLM_BACKEND, _GOOGLE_ROUTE, id="glm-google"), + pytest.param(_GLM_BACKEND, _UTILS_ROUTE, id="glm-utils"), + pytest.param(_MISTRAL_BACKEND, _MESSAGES_ROUTE, id="mistral-messages"), + ], +) +def test_partner_without_a_count_route_answers_400_without_calling_vertex( + gateway: Gateway, backend: str, route: _CountRoute +) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url, model=f"vertex_ai/{backend}") + response: Final = route.send(gateway, model) + assert response.status_code == 400, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + message: Final = _error_message(payload) + assert backend in message, message + assert payload == route.rejection(message), response.text + assert wire.drain() == () + + +def test_llama_count_still_reaches_the_meta_count_route(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url, model=f"vertex_ai/{_LLAMA_BACKEND}") + response: Final = _count(gateway, _bare(model)) + assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + (sent,) = _count_bodies(wire.drain(), _LLAMA_COUNT_TARGET) + assert sent["messages"] == _MESSAGES, sent + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_custom_tokenizer_partner_counts_locally_on_every_count_route(gateway: Gateway, tmp_path: Path) -> None: + environment: Final = {"HF_HOME": str(tmp_path / "hf-home")} + with ( + wire_server(_tokenizer_hub(_char_tokenizer())) as hub, + wire_server(_peer()) as vertex, + owned_proxy_process(gateway, tmp_path, {**environment, "HF_ENDPOINT": hub.url}) as owned, + owned.gateway.scenario() as scenario, + ): + model: Final = _deployment( + owned.gateway, + scenario, + vertex.url, + model=f"vertex_ai/{_GLM_BACKEND}", + model_info={"custom_tokenizer": {"identifier": _TOKENIZER_REPO, "revision": "main", "auth_token": None}}, + ) + local: Final = _payload( + owned.gateway.request("POST", "/utils/token_counter", _bare(model), params={"call_endpoint": "false"}) + ) + assert local["tokenizer_type"] == "huggingface_tokenizer", local + total: Final = local["total_tokens"] + assert isinstance(total, int) and total >= len(_COUNTED_TEXT), local + counts: Final = tuple(_route_count(owned.gateway, route, model) for route in _COUNT_ROUTES) + assert counts == ((200, total),) * len(_COUNT_ROUTES), counts + contents_only: Final = owned.gateway.request( + "POST", + "/utils/token_counter", + {"model": model, "contents": [{"role": "user", "parts": [{"text": _COUNTED_TEXT}]}]}, + params={"call_endpoint": "true"}, + ) + assert contents_only.status_code == 400, contents_only.text + assert _GLM_BACKEND in _error_message(_JSON_OBJECT.validate_json(contents_only.content)), contents_only.text + assert vertex.drain() == () + assert {request.target for request in hub.drain()} >= {_TOKENIZER_TARGET} + + def test_gemini_count_tokens_route_reaches_the_partner_peer(gateway: Gateway) -> None: with wire_server(_peer()) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index e684aa55b33..d7648a08b6a 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -3,6 +3,7 @@ Test for response_api_endpoints/endpoints.py """ import unittest +from collections.abc import Mapping from typing import Any, Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -2428,6 +2429,219 @@ class TestResponsesInputTokens: assert response.status_code == 429, response.text assert response.json()["error"]["message"] == "rate limited" + @pytest.mark.parametrize( + "provider_body, status_code, expected_error", + [ + pytest.param( + { + "error": { + "message": "Missing required parameter: 'tools[0].name'.", + "type": "invalid_request_error", + "param": "tools[0].name", + "code": "missing_required_parameter", + } + }, + 400, + { + "message": "Missing required parameter: 'tools[0].name'.", + "type": "invalid_request_error", + "param": "tools[0].name", + "code": "missing_required_parameter", + }, + id="openai_error_passes_through", + ), + pytest.param( + { + "error": { + "message": "Missing required parameter: 'tools[0].name'.", + "type": "invalid_request_error", + "param": "tools[0].name", + } + }, + 400, + { + "message": "Missing required parameter: 'tools[0].name'.", + "type": "invalid_request_error", + "param": "tools[0].name", + "code": None, + }, + id="openai_error_without_code_keeps_type_and_param", + ), + pytest.param( + { + "type": "error", + "error": {"type": "invalid_request_error", "message": "tools.0: Input tag 'function' is invalid"}, + "request_id": "req_vrtx_1", + }, + 400, + { + "message": "tools.0: Input tag 'function' is invalid", + "type": "invalid_request_error", + "param": None, + "code": None, + }, + id="anthropic_error_message_unwrapped", + ), + pytest.param( + { + "type": "error", + "error": {"type": "not_found_error", "message": "model: claude-typo"}, + "request_id": "req_vrtx_2", + }, + 404, + {"message": "model: claude-typo", "type": "invalid_request_error", "param": None, "code": None}, + id="anthropic_error_type_follows_the_status", + ), + pytest.param( + {"error": {"code": 404, "message": "Publisher model not found.", "status": "NOT_FOUND"}}, + 404, + {"message": "Publisher model not found.", "type": "invalid_request_error", "param": None, "code": None}, + id="google_error_message_unwrapped", + ), + pytest.param( + "litellm.APIError: Google Gen AI Studio API error: 404 - " + + json.dumps( + {"error": {"code": 404, "message": "models/gemini-typo is not found.", "status": "NOT_FOUND"}} + ), + 404, + { + "message": "models/gemini-typo is not found.", + "type": "invalid_request_error", + "param": None, + "code": None, + }, + id="prefixed_gemini_error_message_unwrapped", + ), + pytest.param( + "litellm.BadRequestError: AzureException - " + + json.dumps( + { + "error": { + "message": "Unrecognized request argument supplied: foo", + "type": "invalid_request_error", + "param": "foo", + "code": "unknown_parameter", + } + } + ), + 400, + { + "message": "Unrecognized request argument supplied: foo", + "type": "invalid_request_error", + "param": "foo", + "code": "unknown_parameter", + }, + id="prefixed_openai_error_passes_through", + ), + pytest.param( + {"message": "The provided model doesn't support counting tokens."}, + 400, + { + "message": "The provided model doesn't support counting tokens.", + "type": "invalid_request_error", + "param": None, + "code": None, + }, + id="bedrock_error_message_unwrapped", + ), + pytest.param( + "BedrockException - " + json.dumps({"message": "The provided model doesn't support counting tokens."}), + 400, + { + "message": "The provided model doesn't support counting tokens.", + "type": "invalid_request_error", + "param": None, + "code": None, + }, + id="prefixed_bedrock_error_message_unwrapped", + ), + pytest.param( + "tool schema {unclosed is invalid", + 400, + { + "message": "tool schema {unclosed is invalid", + "type": "invalid_request_error", + "param": None, + "code": None, + }, + id="plain_text_with_a_brace_kept_whole", + ), + pytest.param( + "upstream count service unavailable", + 503, + { + "message": "upstream count service unavailable", + "type": "internal_server_error", + "param": None, + "code": None, + }, + id="plain_text_error", + ), + ], + ) + def test_provider_count_error_answers_in_openai_error_shape( + self, + provider_body: str | Mapping[str, object], + status_code: int, + expected_error: Mapping[str, str | None], + ) -> None: + from litellm.proxy._types import ProxyException + + failing_counter: Final = AsyncMock( + side_effect=ProxyException( + message=provider_body if isinstance(provider_body, str) else json.dumps(provider_body), + type="token_counting_error", + param="model", + code=str(status_code), + ) + ) + response, _ = self._post_input_tokens({"model": "gpt-4o", "input": "hi"}, counter=failing_counter) + + assert response.status_code == status_code, response.text + assert response.json() == {"error": expected_error} + + def test_non_provider_count_error_keeps_its_own_type(self) -> None: + from litellm.proxy._types import ProxyException + + failing_counter: Final = AsyncMock( + side_effect=ProxyException( + message="Token counting is disabled and no provider API result available", + type="token_counting_disabled", + param="model", + code=503, + ) + ) + response, _ = self._post_input_tokens({"model": "gpt-4o", "input": "hi"}, counter=failing_counter) + + assert response.status_code == 503, response.text + assert response.json()["error"]["type"] == "token_counting_disabled" + + def test_provider_count_error_names_the_call_id_when_opted_in(self) -> None: + from litellm.proxy._types import ProxyException + + failing_counter: Final = AsyncMock( + side_effect=ProxyException( + message="LiteLLM does not support counting tokens for Vertex AI model zai-org/glm-5.2-maas.", + type="token_counting_error", + param="model", + code="400", + headers={"x-litellm-call-id": "call-9564"}, + ) + ) + with patch("litellm.proxy.proxy_server.general_settings", {"include_call_id_in_error_body": True}): + response, _ = self._post_input_tokens({"model": "gpt-4o", "input": "hi"}, counter=failing_counter) + + assert response.status_code == 400, response.text + assert response.json() == { + "error": { + "message": "LiteLLM does not support counting tokens for Vertex AI model zai-org/glm-5.2-maas.", + "type": "invalid_request_error", + "param": None, + "code": None, + "litellm_call_id": "call-9564", + } + } + def test_responses_routes_document_response_models_in_openapi_schema(): from typing import cast diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py index f9561a0d710..71e6ecc1f81 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -47,6 +47,35 @@ async def test_vertex_gemini_rejects_non_text_system_blocks_before_request( assert "text blocks" in (result.error_message or "") +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model", + ( + "zai-org/glm-5.2-maas", + "deepseek-ai/deepseek-v3.2-maas", + "qwen/qwen3-coder-480b-a35b-instruct-maas", + "mistral-medium-3", + "codestral-2501", + "openai/gpt-oss-120b-maas", + "google/gemma-4-26b-a4b-it-maas", + ), +) +async def test_vertex_partner_model_without_an_inference_matching_count_answers_400_naming_the_model(model: str) -> None: + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use=model, + messages=[{"role": "user", "content": "Count this"}], + contents=None, + deployment={"litellm_params": {"vertex_project": "test-project", "vertex_location": "global"}}, + request_model=f"vertex_ai/{model}", + messages_format="anthropic", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert result.error_message == f"LiteLLM does not support counting tokens for Vertex AI model {model}." + + class _VertexCountTokensPeer: def __init__(self, status_code: int) -> None: self.status_code: Final = status_code diff --git a/tests/unit/proxy/test_custom_tokenizer_bug.py b/tests/unit/proxy/test_custom_tokenizer_bug.py index 23b6b302202..29625ea1a32 100644 --- a/tests/unit/proxy/test_custom_tokenizer_bug.py +++ b/tests/unit/proxy/test_custom_tokenizer_bug.py @@ -9,6 +9,7 @@ The HuggingFace download boundary (Tokenizer.from_pretrained) is mocked so these stay hermetic unit tests; the proxy's extraction-and-selection path runs for real. """ +from typing import Final from unittest.mock import MagicMock, patch import pytest @@ -17,7 +18,7 @@ import litellm import litellm.proxy.proxy_server import litellm.utils from litellm import Router -from litellm.proxy._types import TokenCountRequest +from litellm.proxy._types import ProxyException, TokenCountRequest from litellm.proxy.proxy_server import token_counter @@ -102,3 +103,60 @@ async def test_model_without_custom_tokenizer_uses_default(monkeypatch): assert response.tokenizer_type == "openai_tokenizer" assert response.model_used == "gpt-4" assert response.total_tokens > 0 + + + +def _vertex_glm_router_with_its_own_tokenizer() -> Router: + return Router( + model_list=[ + { + "model_name": "vertex-glm", + "litellm_params": { + "model": "vertex_ai/zai-org/glm-5.2-maas", + "vertex_project": "test-project", + "vertex_location": "global", + }, + "model_info": { + "custom_tokenizer": {"identifier": "zai-org/GLM-5.2", "revision": "main", "auth_token": None}, + }, + } + ] + ) + + +@pytest.mark.asyncio +async def test_custom_tokenizer_counts_a_vertex_partner_model_without_a_count_route( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _vertex_glm_router_with_its_own_tokenizer()) + + with patch.object(litellm.utils, "tokenizer_dispatch") as mock_tokenizer_cls: + mock_tokenizer_cls.from_pretrained.return_value = _fake_hf_tokenizer(9) + + response: Final = await token_counter( + request=TokenCountRequest(model="vertex-glm", messages=[{"role": "user", "content": "Return exactly OK."}]), + call_endpoint=True, + ) + + mock_tokenizer_cls.from_pretrained.assert_called_once_with("zai-org/GLM-5.2", revision="main", token=None) + assert response.error is not True + assert response.tokenizer_type == "huggingface_tokenizer" + assert response.total_tokens >= 9 + + +@pytest.mark.asyncio +async def test_contents_only_count_of_a_vertex_partner_model_without_a_count_route_answers_400( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _vertex_glm_router_with_its_own_tokenizer()) + + with pytest.raises(ProxyException) as raised: + await token_counter( + request=TokenCountRequest( + model="vertex-glm", contents=[{"role": "user", "parts": [{"text": "Return exactly OK."}]}] + ), + call_endpoint=True, + ) + + assert raised.value.code == "400" + assert raised.value.message == "LiteLLM does not support counting tokens for Vertex AI model zai-org/glm-5.2-maas."