mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-26 06:03:19 +00:00
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 ef1d3959a0 (main).
This commit is contained in:
Mateo Wang
mateo-berri
parent
da5a9ffd2e
commit
748e03088d
9 changed files with 628 additions and 32 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"}]
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue