diff --git a/litellm/litellm_core_utils/deployment_pricing.py b/litellm/litellm_core_utils/deployment_pricing.py index 622cc91ab5e..7aa85921634 100644 --- a/litellm/litellm_core_utils/deployment_pricing.py +++ b/litellm/litellm_core_utils/deployment_pricing.py @@ -1,4 +1,4 @@ -from collections.abc import Iterator, Mapping +from collections.abc import Iterable, Iterator, Mapping from types import MappingProxyType from typing import Final @@ -25,6 +25,14 @@ def deployment_pricing_models(deployment: object) -> tuple[str, ...]: return pricing_models(info.get("base_model") or params.get("base_model"), params.get("model")) +def deployment_backend_model(deployment: object) -> str: + return next(iter(deployment_pricing_models(deployment)), "") + + +def deployment_backend_models(deployments: Iterable[object]) -> tuple[str, ...]: + return tuple(map(deployment_backend_model, deployments)) + + def _named_models(*candidates: object) -> Iterator[str]: return (candidate for candidate in candidates if isinstance(candidate, str) and candidate) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 6ecab166dcd..f48a42963cf 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -399,6 +399,74 @@ def split_media_for_tpm(value: object) -> tuple[object, tuple[Mapping[str, objec return projection.content, projection.media +_TPM_EMBEDDING_FILE_MIME_PREFIXES: Final = ("application/pdf", "video/") + + +def _tpm_embedding_file(data_uri: str) -> tuple[Mapping[str, object], ...]: + from litellm.litellm_core_utils.prompt_templates.factory import parse_mime_type + + mime: Final = (parse_mime_type(data_uri[:256]) or "").strip().lower() + if not mime.startswith(_TPM_EMBEDDING_FILE_MIME_PREFIXES): + return () + return ({"type": "file", "file": {"file_data": data_uri, "format": mime}},) + + +_BASE64_PAYLOAD: Final = re.compile(r"[A-Za-z0-9+/]*={0,2}") + + +def _is_base64_data_uri(value: str) -> bool: + header_end: Final = value.find(",", 0, 256) + if header_end < 0 or not value[:header_end].endswith(";base64"): + return False + payload_start: Final = header_end + 1 + if (len(value) - payload_start) % 4 or _BASE64_PAYLOAD.fullmatch(value, payload_start) is None: + return False + last_quantum: Final = value[max(payload_start, len(value) - 4) :] + return base64.b64encode(base64.b64decode(last_quantum)).decode("ascii") == last_quantum + + +def _tpm_embedding_file_block(part: Mapping[str, object]) -> _TPMMediaProjection: + projection: Final = _tpm_file_media(part) + fields: Final = as_str_mapping(part.get("file")) + data: Final = fields.get("file_data") if fields is not None else None + if projection.media or fields is None or not isinstance(data, str) or not _is_base64_data_uri(data): + return projection + media: Final = _tpm_embedding_file(data) + return _TPMMediaProjection({**part, "file": _tpm_blank_fields(fields, ("file_data",))} if media else part, media) + + +def _tpm_embedding_leaf(value: object) -> _TPMMediaProjection: + part: Final = as_str_mapping(value) + if part is not None and part.get("type") == "file": + return _tpm_embedding_file_block(part) + if not isinstance(value, str) or not value.startswith("data:"): + return _tpm_media_leaf(value) + if not _is_base64_data_uri(value): + return _TPMMediaProjection(value) + media: Final = _tpm_file_media({"file_data": value}).media or _tpm_embedding_file(value) + return _TPMMediaProjection("" if media else value, media) + + +def _is_json_array(value: object) -> bool: + return isinstance(value, (list, tuple)) + + +def _tpm_embedding_item(value: object) -> _TPMMediaProjection: + return _tpm_media_list(value, _tpm_embedding_leaf) if _is_json_array(value) else _tpm_embedding_leaf(value) + + +def split_embedding_media_for_tpm(value: object) -> tuple[object, tuple[Mapping[str, object], ...]]: + projection: Final = ( + _tpm_media_list(value, _tpm_embedding_item) if _is_json_array(value) else _tpm_embedding_leaf(value) + ) + return projection.content, projection.media + + +def embedding_input_is_all_media(value: object) -> bool: + items: Final = cast(Sequence[object], value) if _is_json_array(value) else (value,) # cast-ok: JSON input array + return all(_tpm_embedding_leaf(item).media for item in items) + + def _content_parts_contain_image(parts: Sequence[object]) -> bool: """Depth-bounded frontier walk over nested content lists, iterative because the repo bans recursion; an Anthropic tool_result nests its image parts exactly one level down.""" diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index df8038e756c..52a5026c44f 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -610,6 +610,127 @@ def estimate_image_tokens( return max(_image_policy_tokens(family, dimensions, detail) for family in families) +_MEDIA_FILE_CHARS_PER_TOKEN: Final = 4 +_EmbeddingInputT: Final = TypeVar("_EmbeddingInputT") + + +def _max_input_tokens(model: str) -> int | None: + try: + limit: Final = litellm.get_model_info(model).get("max_input_tokens") + except Exception: # noqa: BLE001 # get_model_info raises a bare Exception for an unmapped model + return None + return limit if isinstance(limit, int) and limit > 0 else None + + +def _media_file_tokens(data: str, models: Sequence[str]) -> int: + text_estimate: Final = len(data) // _MEDIA_FILE_CHARS_PER_TOKEN + ceilings: Final = tuple(map(_max_input_tokens, models)) + known_ceilings: Final = tuple(ceiling for ceiling in ceilings if ceiling is not None) + if not known_ceilings or len(known_ceilings) != len(ceilings): + return text_estimate + return min(text_estimate, max(known_ceilings)) + + +def _block_str_field(block: Mapping[str, object], container: str, field: str) -> str | None: + inner: Final = as_str_mapping(block.get(container)) + value: Final = inner.get(field) if inner is not None else None + return value if isinstance(value, str) else None + + +def estimate_media_block_tokens(block: Mapping[str, object], models: Sequence[str], fallback_model: str | None) -> int: + from litellm.litellm_core_utils.audio_utils.utils import estimate_audio_tokens + + match block.get("type"): + case "input_audio": + return estimate_audio_tokens( + _block_str_field(block, "input_audio", "data"), models, fallback_model=fallback_model + ) + case "file": + return _media_file_tokens(_block_str_field(block, "file", "file_data") or "", models) + case _: + detail: Final = _block_str_field(block, "image_url", "detail") + return estimate_image_tokens( + models, + _block_str_field(block, "image_url", "url") or "", + "auto" if detail is None else detail, + fallback_model=fallback_model, + ) + + +_PER_ITEM_DATA_URI_EMBEDDING_PROVIDERS: Final = frozenset({"vertex_ai", "gemini", "azure_ai"}) + +_DataUriMediaScope: TypeAlias = Literal["per_item", "whole_batch"] + + +def _is_pdf_block(block: Mapping[str, object]) -> bool: + return (_block_str_field(block, "file", "file_data") or "").startswith("data:application/pdf") + + +def _bedrock_data_uri_media_scope( + resolved_model: str, media: Sequence[Mapping[str, object]] +) -> _DataUriMediaScope | None: + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + from litellm.llms.bedrock.common_utils import BedrockModelInfo + from litellm.utils import supports_embedding_image_input + + base_model: Final = BedrockModelInfo.get_base_model(resolved_model.removeprefix("async_invoke/")) + if not supports_embedding_image_input(f"bedrock/{base_model}"): + return None + match BaseAWSLLM.get_bedrock_embedding_provider(base_model): + case "cohere": + return "whole_batch" + case "twelvelabs": + return None + case "nova" if any(map(_is_pdf_block, media)): + return None + case _: + return "per_item" + + +def _data_uri_media_scope(model: str, media: Sequence[Mapping[str, object]]) -> _DataUriMediaScope | None: + from litellm.utils import supports_embedding_image_input + + try: + resolved_model, provider, _, _ = litellm.get_llm_provider(model=model) + except litellm.exceptions.BadRequestError: + return None + if provider == "bedrock": + return _bedrock_data_uri_media_scope(resolved_model, media) + if not supports_embedding_image_input(model): + return None + if provider == "cohere": + return "whole_batch" + return "per_item" if provider in _PER_ITEM_DATA_URI_EMBEDDING_PROVIDERS else None + + +def split_embedding_input_tokens( + value: _EmbeddingInputT, media_models: Callable[[], Sequence[str]], fallback_model: str | None = None +) -> tuple[_EmbeddingInputT, int]: + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + embedding_input_is_all_media, + split_embedding_media_for_tpm, + ) + + text, media = split_embedding_media_for_tpm(value) + if not media: + return value, 0 + models: Final = tuple(media_models()) + scopes: Final = frozenset(_data_uri_media_scope(model, media) for model in models) + if not models or None in scopes or ("whole_batch" in scopes and not embedding_input_is_all_media(value)): + return value, 0 + media_tokens: Final = sum(estimate_media_block_tokens(block, models, fallback_model) for block in media) + return cast(_EmbeddingInputT, text), media_tokens # cast-ok: the projection keeps the input's container shape + + +def embedding_aware_token_count( + messages: Sequence[Mapping[str, str]] | None, + text: str | list[str] | None, # mutable-ok: litellm.token_counter only takes text as list[str] + media_models: Callable[[], Sequence[str]], +) -> int: + embedding_text, media_tokens = split_embedding_input_tokens(text, media_models) + return litellm.token_counter(messages=messages, text=embedding_text) + media_tokens + + TokenCounterFunction = Callable[[str], int] """ Type for a function that counts tokens in a string. diff --git a/litellm/llms/cohere/embed/v1_transformation.py b/litellm/llms/cohere/embed/v1_transformation.py index ad5bfa507e5..056b5bce742 100644 --- a/litellm/llms/cohere/embed/v1_transformation.py +++ b/litellm/llms/cohere/embed/v1_transformation.py @@ -9,6 +9,7 @@ import httpx from litellm import COHERE_DEFAULT_EMBEDDING_INPUT_TYPE from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.custom_httpx.http_handler import header_value from litellm.types.llms.bedrock import ( CohereEmbeddingRequest, CohereEmbeddingRequestWithModel, @@ -23,6 +24,11 @@ class _SupportsEncode(Protocol): def encode(self, text: str, /) -> Sized: ... +def _bedrock_billed_input_tokens(response: httpx.Response) -> int | None: + count: Final = header_value(response.headers, "x-amzn-bedrock-input-token-count") + return int(count) if count is not None and count.isdigit() else None + + class CohereEmbeddingConfig: """ Reference: https://docs.cohere.com/v2/reference/embed @@ -76,32 +82,30 @@ class CohereEmbeddingConfig: ) -> CohereEmbeddingRequestWithModel: return self._transform_request(model, input, inference_params) - def _calculate_usage(self, input: list[str], encoding: _SupportsEncode, meta: dict) -> Usage: - input_tokens = 0 - + def _calculate_usage( + self, input: list[str], encoding: _SupportsEncode, meta: dict, billed_input_tokens: int | None = None + ) -> Usage: text_tokens: Final[int | None] = meta.get("billed_units", {}).get("input_tokens") image_tokens: Final[int | None] = meta.get("billed_units", {}).get("images") - prompt_tokens_details: PromptTokensDetailsWrapper | None = None if image_tokens is None and text_tokens is None: - for text in input: - input_tokens += len(encoding.encode(text)) - else: - prompt_tokens_details = PromptTokensDetailsWrapper( - image_tokens=image_tokens, - text_tokens=text_tokens, + counted_tokens: Final = ( + billed_input_tokens + if billed_input_tokens is not None + else sum(len(encoding.encode(text)) for text in input) ) - if image_tokens: - input_tokens += image_tokens - if text_tokens: - input_tokens += text_tokens + return Usage(prompt_tokens=counted_tokens, completion_tokens=0, total_tokens=counted_tokens) + input_tokens: Final = (image_tokens or 0) + (text_tokens or 0) return Usage( prompt_tokens=input_tokens, completion_tokens=0, total_tokens=input_tokens, - prompt_tokens_details=prompt_tokens_details, + prompt_tokens_details=PromptTokensDetailsWrapper( + image_tokens=image_tokens, + text_tokens=text_tokens, + ), ) def _transform_response( @@ -129,6 +133,7 @@ class CohereEmbeddingConfig: model=model, encoding=encoding, input=input, + billed_input_tokens=_bedrock_billed_input_tokens(response) if "images" in data else None, ) def transform_response( @@ -151,6 +156,7 @@ class CohereEmbeddingConfig: model: str, encoding: _SupportsEncode, input: list, + billed_input_tokens: int | None = None, ) -> EmbeddingResponse: """ Parse a Cohere embed response body into an OpenAI-style EmbeddingResponse. @@ -195,7 +201,7 @@ class CohereEmbeddingConfig: setattr( model_response, "usage", - self._calculate_usage(input, encoding, response_json.get("meta", {})), + self._calculate_usage(input, encoding, response_json.get("meta", {}), billed_input_tokens), ) return model_response diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 85e6ecf204f..77ce188ad57 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -31329,6 +31329,7 @@ "output_cost_per_token": 0, "output_vector_size": 3072, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supports_embedding_image_input": true, "uses_embed_content": true }, "gemini-embedding-2": { @@ -31346,6 +31347,7 @@ "output_cost_per_token": 0, "output_vector_size": 3072, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supports_embedding_image_input": true, "supports_multimodal": true, "uses_embed_content": true }, @@ -31364,6 +31366,7 @@ "output_cost_per_token": 0, "output_vector_size": 3072, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supports_embedding_image_input": true, "supports_multimodal": true, "uses_embed_content": true }, @@ -31382,6 +31385,7 @@ "output_cost_per_token": 0, "output_vector_size": 3072, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supports_embedding_image_input": true, "supports_multimodal": true, "uses_embed_content": true }, @@ -31428,6 +31432,7 @@ "rpm": 10000, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supports_audio_input": true, + "supports_embedding_image_input": true, "supports_multimodal": true, "supports_vision": true, "tpm": 10000000 @@ -31606,6 +31611,7 @@ "rpm": 10000, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supports_audio_input": true, + "supports_embedding_image_input": true, "supports_multimodal": true, "supports_vision": true, "tpm": 10000000 @@ -44433,7 +44439,8 @@ "text", "image", "video" - ] + ], + "supports_embedding_image_input": true }, "multimodalembedding@001": { "deprecation_date": "2027-04-01", @@ -44456,7 +44463,8 @@ "text", "image", "video" - ] + ], + "supports_embedding_image_input": true }, "nscale/Qwen/QwQ-32B": { "input_cost_per_token": 1.8e-07, @@ -46172,6 +46180,7 @@ "mode": "embedding", "output_vector_size": 1024, "source": "https://www.oracle.com/artificial-intelligence/enterprise-ai/cost-estimator/", + "supports_embedding_image_input": true, "supports_vision": true }, "oci/cohere.command-a-reasoning-08-2025": { @@ -54620,7 +54629,8 @@ "litellm_provider": "voyage", "max_input_tokens": 32000, "mode": "embedding", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "supports_embedding_image_input": true }, "wandb/openai/gpt-oss-120b": { "supports_reasoning": true, diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 1cbe5eb39e8..5c2f6bab650 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -55,7 +55,6 @@ from litellm.litellm_core_utils.add_retry_fallback_headers import ( ensure_response_additional_headers, response_has_hidden_params, ) -from litellm.litellm_core_utils.audio_utils.utils import estimate_audio_tokens from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, @@ -64,7 +63,11 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.litellm_core_utils.ptu_pricing import is_ptu_cost_attribution_enabled from litellm.litellm_core_utils.ptu_shares import PTUTeamCeiling from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias -from litellm.litellm_core_utils.token_counter import estimate_image_tokens, offload_token_count +from litellm.litellm_core_utils.token_counter import ( + estimate_media_block_tokens, + offload_token_count, + split_embedding_input_tokens, +) from litellm.llms.azure.ptu_capacity import normalized_tokens from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( @@ -156,6 +159,12 @@ def _resolve_media_models(router: "Router | None", model: str, team_id: str | No return tuple(dict.fromkeys(router.registry().names(model, kind="litellm_base_models", team_id=team_id))) or (model,) +def _media_models( + resolver: Callable[[str, str | None], tuple[str, ...]] | None, model: str, team_id: str | None +) -> tuple[str, ...]: + return resolver(model, team_id) if resolver is not None else (model,) + + def _resolve_media_models_via_proxy_router(model: str, team_id: str | None) -> tuple[str, ...]: from litellm.proxy.proxy_server import llm_router @@ -1378,20 +1387,28 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) media_source: Final = messages if messages is not None else original_input + is_embedding_call: Final = call_type in EMBEDDING_API_CALL_TYPES projected, media = ( split_media_for_tpm(media_source) if include_media and (isinstance(messages, list) or call_type in RESPONSES_API_CALL_TYPES) else (original_input, ()) ) - input_text: Final = projected if call_type in RESPONSES_API_CALL_TYPES else original_input fallback_model: Final = request_model if isinstance(request_model, str) else None resolved_models: Final = ( - media_model_resolver(fallback_model or "", team_id) - if media and media_model_resolver is not None - else (fallback_model or "",) + _media_models(media_model_resolver, fallback_model or "", team_id) if media else (fallback_model or "",) ) - media_tokens: Final = sum( - self._estimate_media_block_tokens(block, resolved_models, fallback_model) for block in media + embedding_text, embedding_media_tokens = ( + split_embedding_input_tokens( + original_input, + lambda: _media_models(media_model_resolver, fallback_model or "", team_id), + fallback_model, + ) + if include_media and is_embedding_call + else (original_input, 0) + ) + input_text: Final = projected if call_type in RESPONSES_API_CALL_TYPES else embedding_text + media_tokens: Final = embedding_media_tokens + sum( + estimate_media_block_tokens(block, resolved_models, fallback_model) for block in media ) total_chars: Final = ( @@ -1436,24 +1453,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return estimated_input_tokens, max_tokens_estimate * self.get_output_candidate_count(data, call_type) - @staticmethod - def _estimate_media_block_tokens( - block: Mapping[str, object], models: tuple[str, ...], fallback_model: str | None - ) -> int: - if block.get("type") == "input_audio": - audio: Final = _callback_kwargs_mapping(block.get("input_audio")) - data: Final = audio.get("data") - return estimate_audio_tokens(data if isinstance(data, str) else None, models, fallback_model=fallback_model) - image: Final = _callback_kwargs_mapping(block.get("image_url")) - url: Final = image.get("url") - detail: Final = image.get("detail") - return estimate_image_tokens( - models, - url if isinstance(url, str) else "", - detail if isinstance(detail, str) else "auto", - fallback_model=fallback_model, - ) - def _is_redis_cluster(self) -> bool: """ Check if the dual cache is using Redis cluster. @@ -4098,12 +4097,20 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): None if is_responses_request else self._translate_google_genai_native_request(data, call_type) ) is_embedding_request: Final = self._is_embedding_request(data, call_type) - embedding_text: Final = data.get("input") if is_embedding_request else None + request: Final = as_str_mapping(data) + raw_embedding_input: Final = request.get("input") if is_embedding_request and request is not None else None pretokenized_input_tokens: Final = ( - self._count_pretokenized_embedding_input(embedding_text) if is_embedding_request else None + self._count_pretokenized_embedding_input(raw_embedding_input) if is_embedding_request else None ) if pretokenized_input_tokens is not None: return pretokenized_input_tokens + embedding_text, embedding_media_tokens = ( + split_embedding_input_tokens( + raw_embedding_input, lambda: self._media_model_resolver(model or "", team_id), model + ) + if is_embedding_request + else (None, 0) + ) prompt: Final = data.get("prompt") fallback_text: Final = prompt if prompt is not None else data.get("input") @@ -4153,12 +4160,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): for tokenizer_model in (self._counted_models_resolver(model, team_id) if model else ("",)) ), ) - return estimate + audio_token_estimate + return estimate + audio_token_estimate + embedding_media_tokens except Exception: # noqa: BLE001 # tokenizer failures degrade to the cheap estimate if call_type in RERANK_API_CALL_TYPES and isinstance(selected_text, str): return max(0, len(selected_text) // DEFAULT_CHARS_PER_TOKEN) estimated_input_tokens, _ = self._estimate_input_and_output_tokens( - data=data, call_type=call_type, include_media=False + data=data, + call_type=call_type, + include_media=is_embedding_request, + team_id=team_id, + media_model_resolver=self._media_model_resolver, ) return estimated_input_tokens + audio_token_estimate diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index 7ccb23ebfa8..b9a90c7ea2a 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -135,7 +135,7 @@ from litellm.litellm_core_utils.core_helpers import ( from litellm.litellm_core_utils.coroutine_checker import coroutine_checker from litellm.litellm_core_utils.credential_accessor import CredentialAccessor, reject_api_key_secret_reference from litellm.litellm_core_utils.dd_tracing import tracer -from litellm.litellm_core_utils.deployment_pricing import deployment_pricing_models +from litellm.litellm_core_utils.deployment_pricing import deployment_backend_models, deployment_pricing_models from litellm.litellm_core_utils.get_litellm_params import SUPPLEMENTAL_KWARGS_KEYS from litellm.litellm_core_utils.get_llm_provider_logic import ( declared_authenticating_provider, @@ -14866,6 +14866,7 @@ class PythonCore(DeprecatedRouterMembers): input: str | list[object] | None, request_kwargs: Mapping[str, object] | None = None, tokenizer_model: str = "", + deployments: Sequence[Mapping[str, object]] = (), ) -> int: """ Count input tokens for context-window pre-call checks. @@ -14900,11 +14901,15 @@ class PythonCore(DeprecatedRouterMembers): if input is not None: from openai.types.responses.response_create_params import ResponseInputParam + from litellm.litellm_core_utils.token_counter import split_embedding_input_tokens from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) - typed_input: Final = cast(str | ResponseInputParam, input) # cast-ok: str | list matches transform input + text_input, _ = split_embedding_input_tokens(input, partial(deployment_backend_models, deployments)) + typed_input: Final = cast( # cast-ok: str | list matches transform input + str | ResponseInputParam, text_input + ) input_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( input=typed_input, responses_api_request={"instructions": instructions} if instructions is not None else {}, @@ -14972,6 +14977,7 @@ class PythonCore(DeprecatedRouterMembers): self.served_openai_models(model) or served_openai_models(healthy_deployments) ) or "", + deployments=healthy_deployments, ) except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request verbose_router_logger.error( @@ -15002,6 +15008,7 @@ class PythonCore(DeprecatedRouterMembers): self.served_openai_models(model) or served_openai_models(deployments) ) or "", + deployments=deployments, ) except Exception as e: verbose_router_logger.error( diff --git a/litellm/router_strategy/lowest_cost.py b/litellm/router_strategy/lowest_cost.py index 3f0a5412b62..6a8757f6e6e 100644 --- a/litellm/router_strategy/lowest_cost.py +++ b/litellm/router_strategy/lowest_cost.py @@ -2,16 +2,22 @@ # picks based on response time (for streaming, this is time to first token) from collections.abc import Mapping from datetime import datetime +from functools import partial from typing import Final import litellm -from litellm import ModelResponse, token_counter, verbose_logger +from litellm import ModelResponse, verbose_logger from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type -from litellm.litellm_core_utils.deployment_pricing import deployment_pricing_models, str_keyed +from litellm.litellm_core_utils.deployment_pricing import ( + deployment_backend_models, + deployment_pricing_models, + str_keyed, +) +from litellm.litellm_core_utils.token_counter import embedding_aware_token_count def _deployment_cost_map_rates(deployment: object) -> tuple[float | None, float | None]: @@ -235,7 +241,9 @@ class LowestCostLoggingHandler(CustomLogger): } try: - input_tokens = token_counter(messages=messages, text=input) + input_tokens = embedding_aware_token_count( + messages, input, partial(deployment_backend_models, healthy_deployments) + ) except Exception: input_tokens = 0 diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index c352934db6c..981fb5341dc 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -3,18 +3,21 @@ import random from collections.abc import Sequence from datetime import datetime, timedelta +from functools import partial from math import ceil, isfinite from typing import TYPE_CHECKING, Any, Final from pydantic import Field import litellm -from litellm import ModelResponse, token_counter, verbose_logger +from litellm import ModelResponse, verbose_logger from litellm._internal_context import with_service_target from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs, safe_divide_seconds +from litellm.litellm_core_utils.deployment_pricing import deployment_backend_models +from litellm.litellm_core_utils.token_counter import embedding_aware_token_count from litellm.responses.utils import ResponseAPILoggingUtils from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import LiteLLMPydanticObjectBase @@ -439,7 +442,9 @@ class LowestLatencyLoggingHandler(CustomLogger): } try: - input_tokens = token_counter(messages=messages, text=input) + input_tokens = embedding_aware_token_count( + messages, input, partial(deployment_backend_models, healthy_deployments) + ) except Exception: input_tokens = 0 diff --git a/litellm/router_strategy/lowest_tpm_rpm.py b/litellm/router_strategy/lowest_tpm_rpm.py index 9636d7f009d..e9769c6f595 100644 --- a/litellm/router_strategy/lowest_tpm_rpm.py +++ b/litellm/router_strategy/lowest_tpm_rpm.py @@ -2,14 +2,16 @@ # identifies lowest tpm deployment import traceback from datetime import datetime +from functools import partial from typing import Final -from litellm import token_counter from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type +from litellm.litellm_core_utils.deployment_pricing import deployment_backend_models +from litellm.litellm_core_utils.token_counter import embedding_aware_token_count from litellm.types.utils import LiteLLMPydanticObjectBase from litellm.utils import print_verbose @@ -174,7 +176,9 @@ class LowestTPMLoggingHandler(CustomLogger): verbose_router_logger.debug("tpm_key=%s, tpm_dict: %s, rpm_dict: %s", tpm_key, tpm_dict, rpm_dict) try: - input_tokens = token_counter(messages=messages, text=input) + input_tokens = embedding_aware_token_count( + messages, input, partial(deployment_backend_models, healthy_deployments) + ) except Exception: input_tokens = 0 verbose_router_logger.debug("input_tokens=%s", input_tokens) diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 0719b30d36d..0574c3d7f69 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -5,18 +5,20 @@ from collections.abc import Iterator, Mapping, Sequence from contextlib import contextmanager from contextvars import ContextVar from dataclasses import dataclass +from functools import partial from typing import TYPE_CHECKING, Any, Final import httpx import litellm -from litellm import token_counter from litellm._internal_context import with_service_target from litellm._logging import verbose_logger, verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs +from litellm.litellm_core_utils.deployment_pricing import deployment_backend_models +from litellm.litellm_core_utils.token_counter import embedding_aware_token_count from litellm.types.router import RouterErrors from litellm.types.utils import LiteLLMPydanticObjectBase, StandardLoggingPayload from litellm.utils import get_utc_datetime, print_verbose @@ -422,7 +424,9 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): rpm_dict[rpm_keys[idx].split(":")[0]] = rpm_values[idx] try: - input_tokens = token_counter(messages=messages, text=input) + input_tokens = embedding_aware_token_count( + messages, input, partial(deployment_backend_models, healthy_deployments) + ) except Exception: input_tokens = 0 verbose_router_logger.debug("input_tokens=%s", input_tokens) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 85e6ecf204f..77ce188ad57 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -31329,6 +31329,7 @@ "output_cost_per_token": 0, "output_vector_size": 3072, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supports_embedding_image_input": true, "uses_embed_content": true }, "gemini-embedding-2": { @@ -31346,6 +31347,7 @@ "output_cost_per_token": 0, "output_vector_size": 3072, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supports_embedding_image_input": true, "supports_multimodal": true, "uses_embed_content": true }, @@ -31364,6 +31366,7 @@ "output_cost_per_token": 0, "output_vector_size": 3072, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supports_embedding_image_input": true, "supports_multimodal": true, "uses_embed_content": true }, @@ -31382,6 +31385,7 @@ "output_cost_per_token": 0, "output_vector_size": 3072, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supports_embedding_image_input": true, "supports_multimodal": true, "uses_embed_content": true }, @@ -31428,6 +31432,7 @@ "rpm": 10000, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supports_audio_input": true, + "supports_embedding_image_input": true, "supports_multimodal": true, "supports_vision": true, "tpm": 10000000 @@ -31606,6 +31611,7 @@ "rpm": 10000, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supports_audio_input": true, + "supports_embedding_image_input": true, "supports_multimodal": true, "supports_vision": true, "tpm": 10000000 @@ -44433,7 +44439,8 @@ "text", "image", "video" - ] + ], + "supports_embedding_image_input": true }, "multimodalembedding@001": { "deprecation_date": "2027-04-01", @@ -44456,7 +44463,8 @@ "text", "image", "video" - ] + ], + "supports_embedding_image_input": true }, "nscale/Qwen/QwQ-32B": { "input_cost_per_token": 1.8e-07, @@ -46172,6 +46180,7 @@ "mode": "embedding", "output_vector_size": 1024, "source": "https://www.oracle.com/artificial-intelligence/enterprise-ai/cost-estimator/", + "supports_embedding_image_input": true, "supports_vision": true }, "oci/cohere.command-a-reasoning-08-2025": { @@ -54620,7 +54629,8 @@ "litellm_provider": "voyage", "max_input_tokens": 32000, "mode": "embedding", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "supports_embedding_image_input": true }, "wandb/openai/gpt-oss-120b": { "supports_reasoning": true, diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 146bb08a4fd..a17deedf7a9 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -39,6 +39,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( schema_has_unresolvable_ref, schema_property_count, split_concatenated_json_objects, + split_embedding_media_for_tpm, split_media_for_tpm, strip_encrypted_reasoning_from_messages, strip_litellm_response_only_fields, @@ -2709,3 +2710,22 @@ def test_content_list_flattening_returns_new_messages_and_leaves_the_callers_unt assert flattened == [{"role": "user", "content": "Hello world"}, {"role": "assistant", "content": "Hi"}] assert messages == messages_before + + +@pytest.mark.parametrize("mime", ["application/pdf", "video/mp4", "audio/wav", "image/png"]) +def test_embedding_media_split_holds_at_most_one_copy_of_a_large_payload(mime: str) -> None: + import gc + import tracemalloc + + uri: Final = f"data:{mime};base64," + base64.b64encode(bytes(3 * 1024 * 1024)).decode("ascii") + split_embedding_media_for_tpm(f"data:{mime};base64,AAAA") + gc.collect() + tracemalloc.start() + try: + _, media = split_embedding_media_for_tpm(uri) + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + + assert media + assert peak < 1.5 * len(uri) diff --git a/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py b/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py index 889faf784bf..904ca3bb8ca 100644 --- a/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py @@ -61,6 +61,7 @@ def test_bedrock_embedding_with_api_key_bearer_token(model, input_type, embed_re with patch.object(client, "post") as mock_post: mock_response = Mock() mock_response.status_code = 200 + mock_response.headers = httpx.Headers() mock_response.text = json.dumps(embed_response) mock_response.json = lambda: json.loads(mock_response.text) mock_post.return_value = mock_response @@ -774,6 +775,7 @@ def test_bedrock_embedding_header_forwarding(model, embed_response): with patch.object(client, "post") as mock_post: mock_response = Mock() mock_response.status_code = 200 + mock_response.headers = httpx.Headers() mock_response.text = json.dumps(embed_response) mock_response.json = lambda: json.loads(mock_response.text) mock_post.return_value = mock_response @@ -909,6 +911,7 @@ def test_bedrock_cohere_v4_embedding_response_parsing(): with patch.object(client, "post") as mock_post: mock_response = Mock() mock_response.status_code = 200 + mock_response.headers = httpx.Headers() mock_response.text = json.dumps(cohere_v4_response) mock_response.json = lambda: json.loads(mock_response.text) mock_post.return_value = mock_response @@ -1317,6 +1320,7 @@ def test_bedrock_cohere_embedding_types_wrapped_as_list( with patch.object(client, "post") as mock_post: mock_response = Mock() mock_response.status_code = 200 + mock_response.headers = httpx.Headers() mock_response.text = json.dumps(cohere_embedding_response) mock_response.json = lambda: json.loads(mock_response.text) mock_post.return_value = mock_response @@ -1848,6 +1852,7 @@ def test_bedrock_embedding_models(model, input_type, embed_response): with patch.object(client, "post") as mock_post: mock_response = Mock() mock_response.status_code = 200 + mock_response.headers = httpx.Headers() mock_response.text = json.dumps(embed_response) mock_response.json = lambda: json.loads(mock_response.text) mock_post.return_value = mock_response diff --git a/tests/unit/llms/cohere/embed/test_v1_transformation.py b/tests/unit/llms/cohere/embed/test_v1_transformation.py index 31b4450ed88..63e61eaecd1 100644 --- a/tests/unit/llms/cohere/embed/test_v1_transformation.py +++ b/tests/unit/llms/cohere/embed/test_v1_transformation.py @@ -1,5 +1,7 @@ from unittest.mock import MagicMock +import httpx +import pytest from litellm.llms.cohere.embed.v1_transformation import CohereEmbeddingConfig from litellm.types.utils import EmbeddingResponse @@ -183,6 +185,7 @@ class TestCohereEmbeddingV1Transform: "meta": {}, # No billed_units } mock_response.json = MagicMock(return_value=response_json) + mock_response.headers = httpx.Headers() input_data = ["test text"] data = {"texts": input_data, "input_type": "search_query"} @@ -208,3 +211,31 @@ class TestCohereEmbeddingV1Transform: # Verify encoding was called self.encoding.encode.assert_called() + + @pytest.mark.parametrize( + ("body_key", "item", "expected_tokens"), + [("images", "data:image/png;base64," + "iVBORw0KGgo" * 400, 2047), ("texts", "a short sentence", 5)], + ) + def test_transform_response_counts_bedrock_billed_tokens_only_for_images( + self, body_key: str, item: str, expected_tokens: int + ): + response = httpx.Response( + 200, + json={"id": "e1", body_key: [{}], "embeddings": [[0.1, 0.2]], "response_type": "embeddings_floats"}, + headers={"x-amzn-bedrock-input-token-count": "2047"}, + ) + + result = self.config.transform_response( + response=response, + api_key=None, + logging_obj=self.logging_obj, + data={body_key: [item], "input_type": "image" if body_key == "images" else "search_document"}, + model_response=EmbeddingResponse(), + model="cohere.embed-v4:0", + encoding=self.encoding, + input=[item], + ) + + assert result.usage is not None + assert result.usage.prompt_tokens == expected_tokens + assert result.usage.total_tokens == expected_tokens diff --git a/tests/unit/proxy/hooks/test_tpm_concurrent.py b/tests/unit/proxy/hooks/test_tpm_concurrent.py index bfb7a9d6199..0d517e87ef6 100644 --- a/tests/unit/proxy/hooks/test_tpm_concurrent.py +++ b/tests/unit/proxy/hooks/test_tpm_concurrent.py @@ -20,6 +20,7 @@ import copy import io import json from datetime import datetime, timedelta +from types import MappingProxyType from typing import Any, Dict, Final import pytest @@ -27,7 +28,7 @@ import soundfile as sf from fastapi import HTTPException from PIL import Image -from litellm import token_counter +from litellm import get_model_info, token_counter from litellm.caching.caching import DualCache from litellm.litellm_core_utils.ptu_shares import PTUTeamCeiling from litellm.litellm_core_utils.token_counter import UnsupportedContentTypeError, estimate_image_tokens @@ -3971,6 +3972,334 @@ async def test_media_pre_call_rejects_tight_quota_and_reserves_sufficient_quota( assert expected[0] > 1000 and expected[1] >= expected[0] +_EMBEDDING_MODEL: Final = "gemini-embedding-2" +_EMBEDDING_CAPTION: Final = "a red bicycle leaning on a wall" + + +def _embedding_media_uri(kind: str, compressed: bool) -> str: + mime: Final = "image/png" if kind == "image" else "audio/wav" + return f"data:{mime};base64,{_quota_media_payload(kind, compressed)}" + + +@pytest.mark.parametrize("kind", ["image", "audio"]) +@pytest.mark.parametrize( + "shape", + ["string", "list", "file_block", "with_caption", "combined"], +) +def test_embedding_media_input_is_estimated_independently_of_encoding( + rate_limiter: tuple[RateLimitHandler, DualCache], kind: str, shape: str +) -> None: + handler, _ = rate_limiter + + def embedding_request(uri: str) -> dict[str, object]: + shapes: Final[dict[str, object]] = { + "string": uri, + "list": [uri], + "file_block": [{"type": "file", "file": {"file_data": uri}}], + "with_caption": [_EMBEDDING_CAPTION, uri], + "combined": [[_EMBEDDING_CAPTION, uri]], + } + return {"model": _EMBEDDING_MODEL, "input": shapes[shape]} + + uncompressed: Final = embedding_request(_embedding_media_uri(kind, False)) + original: Final = copy.deepcopy(uncompressed) + estimate: Final = handler._estimate_input_and_output_tokens(uncompressed, call_type="aembedding") + assert estimate == handler._estimate_input_and_output_tokens( + embedding_request(_embedding_media_uri(kind, True)), call_type="aembedding" + ) + assert 0 < estimate[0] < 2_000 and estimate[1] == 0 + assert uncompressed == original + + +@pytest.mark.parametrize("kind", ["image", "audio"]) +def test_embedding_media_input_costs_the_same_as_chat_media( + rate_limiter: tuple[RateLimitHandler, DualCache], kind: str +) -> None: + handler, _ = rate_limiter + uri: Final = _embedding_media_uri(kind, False) + chat_block: Final = ( + {"type": "image_url", "image_url": {"url": uri}} + if kind == "image" + else {"type": "input_audio", "input_audio": {"data": _quota_media_payload(kind, False), "format": "wav"}} + ) + chat_input_tokens, _ = handler._estimate_input_and_output_tokens( + {"model": _EMBEDDING_MODEL, "messages": [{"role": "user", "content": [chat_block]}]}, + call_type="acompletion", + ) + caption_tokens, _ = handler._estimate_input_and_output_tokens( + {"model": _EMBEDDING_MODEL, "input": [_EMBEDDING_CAPTION]}, call_type="aembedding" + ) + + assert handler._estimate_input_and_output_tokens( + {"model": _EMBEDDING_MODEL, "input": uri}, call_type="aembedding" + ) == (chat_input_tokens, 0) + assert handler._estimate_input_and_output_tokens( + {"model": _EMBEDDING_MODEL, "input": [_EMBEDDING_CAPTION, uri]}, call_type="aembedding" + ) == (caption_tokens + chat_input_tokens, 0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("limit_field", ["model_tpm_limit", "model_itpm_limit"]) +async def test_project_token_limit_admits_image_embedding_and_still_rejects_long_text( + rate_limiter: tuple[RateLimitHandler, DualCache], limit_field: str +) -> None: + handler, cache = rate_limiter + auth: Final = UserAPIKeyAuth( + api_key=hash_token(f"sk-embedding-image-{limit_field}"), + project_id=f"proj-embedding-image-{limit_field}", + project_metadata={limit_field: {_EMBEDDING_MODEL: 75_000}}, + ) + image_uri: Final = _embedding_media_uri("image", False) + long_text: Final = "word " * 80_000 + + await handler.async_pre_call_hook( + user_api_key_dict=auth, + cache=cache, + data={"model": _EMBEDDING_MODEL, "input": image_uri}, + call_type="aembedding", + ) + with pytest.raises(HTTPException) as error: + await handler.async_pre_call_hook( + user_api_key_dict=auth, + cache=cache, + data={"model": _EMBEDDING_MODEL, "input": long_text}, + call_type="aembedding", + ) + assert error.value.status_code == 429 + assert len(image_uri) > len(long_text) + + +_EMBEDDING_FILE_MIMES: Final = MappingProxyType({"pdf": "application/pdf", "video": "video/mp4"}) + + +def _embedding_file_uri(kind: str) -> str: + return f"data:{_EMBEDDING_FILE_MIMES[kind]};base64,{base64.b64encode(bytes(range(256)) * 1200).decode()}" + + +def _embedding_uri(kind: str) -> str: + return _embedding_file_uri(kind) if kind in _EMBEDDING_FILE_MIMES else _embedding_media_uri(kind, False) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("limit_field", ["model_tpm_limit", "model_itpm_limit"]) +@pytest.mark.parametrize("kind", ["pdf", "video"]) +@pytest.mark.parametrize("shape", ["data_uri", "file_block"]) +async def test_project_token_limit_admits_document_embedding_up_to_the_model_input_limit( + rate_limiter: tuple[RateLimitHandler, DualCache], limit_field: str, kind: str, shape: str +) -> None: + handler, cache = rate_limiter + uri: Final = _embedding_file_uri(kind) + model_input_limit: Final = get_model_info(_EMBEDDING_MODEL)["max_input_tokens"] + assert isinstance(model_input_limit, int) and len(uri) // 4 > 75_000 > model_input_limit + + def project_auth(limit: int) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key=hash_token(f"sk-embedding-{kind}-{shape}-{limit_field}-{limit}"), + project_id=f"proj-embedding-{kind}-{shape}-{limit_field}-{limit}", + project_metadata={limit_field: {_EMBEDDING_MODEL: limit}}, + ) + + item: Final = uri if shape == "data_uri" else {"type": "file", "file": {"file_data": uri}} + request: Final = {"model": _EMBEDDING_MODEL, "input": [_EMBEDDING_CAPTION, item]} + await handler.async_pre_call_hook( + user_api_key_dict=project_auth(75_000), cache=cache, data=dict(request), call_type="aembedding" + ) + with pytest.raises(HTTPException) as error: + await handler.async_pre_call_hook( + user_api_key_dict=project_auth(model_input_limit // 2), + cache=cache, + data=dict(request), + call_type="aembedding", + ) + assert error.value.status_code == 429 + + +_TEXT_ONLY_EMBEDDING_MODEL: Final = "text-embedding-3-small" + + +@pytest.mark.parametrize("kind", ["image", "audio", "pdf", "video"]) +def test_text_only_embedding_model_counts_data_uri_as_text( + rate_limiter: tuple[RateLimitHandler, DualCache], kind: str +) -> None: + handler, _ = rate_limiter + uri: Final = _embedding_uri(kind) + + assert handler._estimate_input_and_output_tokens( + {"model": _TEXT_ONLY_EMBEDDING_MODEL, "input": [_EMBEDDING_CAPTION, uri]}, call_type="aembedding" + ) == handler._estimate_input_and_output_tokens( + {"model": _TEXT_ONLY_EMBEDDING_MODEL, "input": [_EMBEDDING_CAPTION, "x" * len(uri)]}, + call_type="aembedding", + ) + + +def _project_limited_embedding_call( + limit_field: str, backends: tuple[str, ...], kind: str = "image" +) -> tuple[RateLimitHandler, DualCache, UserAPIKeyAuth, dict[str, object]]: + cache: Final = DualCache() + handler: Final = RateLimitHandler( + internal_usage_cache=InternalUsageCache(cache), media_model_resolver=lambda _model, _team_id: backends + ) + scope: Final = f"{limit_field}-{kind}-{'-'.join(backends)}" + auth: Final = UserAPIKeyAuth( + api_key=hash_token(f"sk-embedding-backends-{scope}"), + project_id=f"proj-embedding-backends-{scope}", + project_metadata={limit_field: {"team-embedder": 75_000}}, + ) + return handler, cache, auth, {"model": "team-embedder", "input": _embedding_uri(kind)} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("limit_field", ["model_tpm_limit", "model_itpm_limit"]) +@pytest.mark.parametrize( + ("backends", "kind"), + [ + ((_TEXT_ONLY_EMBEDDING_MODEL,), "image"), + ((_EMBEDDING_MODEL, _TEXT_ONLY_EMBEDDING_MODEL), "image"), + (("voyage/voyage-multimodal-3.5",), "image"), + (("oci/cohere.embed-v4.0",), "image"), + (("bedrock/us.twelvelabs.marengo-embed-2-7-v1:0",), "image"), + (("bedrock/twelvelabs.marengo-embed-3-0-v1:0",), "image"), + (("bedrock/amazon.nova-2-multimodal-embeddings-v1:0",), "pdf"), + ], + ids=["text", "mixed", "voyage", "oci", "marengo_2_7", "marengo_3", "nova_pdf"], +) +async def test_project_token_limit_rejects_embedding_media_a_backend_reads_as_text( + limit_field: str, backends: tuple[str, ...], kind: str +) -> None: + handler, cache, auth, request = _project_limited_embedding_call(limit_field, backends, kind) + + with pytest.raises(HTTPException) as error: + await handler.async_pre_call_hook(user_api_key_dict=auth, cache=cache, data=request, call_type="aembedding") + assert error.value.status_code == 429 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("limit_field", ["model_tpm_limit", "model_itpm_limit"]) +@pytest.mark.parametrize( + "backend", + [ + _EMBEDDING_MODEL, + "bedrock/amazon.titan-embed-image-v1", + "bedrock/async_invoke/amazon.nova-2-multimodal-embeddings-v1:0", + ], + ids=["gemini", "titan", "nova_async_invoke"], +) +async def test_project_token_limit_admits_image_embedding_through_an_alias_of_an_image_backend( + limit_field: str, backend: str +) -> None: + handler, cache, auth, request = _project_limited_embedding_call(limit_field, (backend,)) + + await handler.async_pre_call_hook(user_api_key_dict=auth, cache=cache, data=request, call_type="aembedding") + with pytest.raises(HTTPException) as error: + await handler.async_pre_call_hook( + user_api_key_dict=auth, + cache=cache, + data={**request, "input": "word " * 100_000}, + call_type="aembedding", + ) + assert error.value.status_code == 429 + + +def _malformed_embedding_uri(shape: str) -> str: + payload: Final = _quota_media_payload("image", False) + whole_quanta: Final = payload.rstrip("=")[: len(payload.rstrip("=")) // 4 * 4] + words: Final = "word " * 100_000 + return { + "no_comma": f"data:audio/wav;base64{words}", + "text_payload": f"data:image/png;base64,{words}", + "no_base64_marker": f"data:image/png,{payload}", + "header_past_256_chars": f"data:image/png;{'x' * 300};base64,{payload}", + "noncanonical_padding": f"data:image/png;base64,{whole_quanta}AB==", + "padding_mid_payload": f"data:image/png;base64,{whole_quanta}AA==AAAA", + "length_not_multiple_of_4": f"data:image/png;base64,{whole_quanta}A", + }[shape] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "shape", + [ + "no_comma", + "text_payload", + "no_base64_marker", + "header_past_256_chars", + "noncanonical_padding", + "padding_mid_payload", + "length_not_multiple_of_4", + ], +) +async def test_project_token_limit_counts_a_malformed_data_uri_as_text( + rate_limiter: tuple[RateLimitHandler, DualCache], shape: str +) -> None: + handler, cache = rate_limiter + auth: Final = UserAPIKeyAuth( + api_key=hash_token(f"sk-embedding-malformed-{shape}"), + project_id=f"proj-embedding-malformed-{shape}", + project_metadata={"model_tpm_limit": {_EMBEDDING_MODEL: 75_000}}, + ) + + with pytest.raises(HTTPException) as error: + await handler.async_pre_call_hook( + user_api_key_dict=auth, + cache=cache, + data={"model": _EMBEDDING_MODEL, "input": _malformed_embedding_uri(shape)}, + call_type="aembedding", + ) + assert error.value.status_code == 429 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ["cohere/embed-v4.0", "cohere.embed-v4:0"], ids=["cohere", "bedrock_cohere"]) +async def test_project_token_limit_counts_a_mixed_batch_as_text_when_the_backend_types_the_whole_batch( + backend: str, +) -> None: + handler, cache, auth, request = _project_limited_embedding_call("model_tpm_limit", (backend,)) + image_uri: Final = request["input"] + + await handler.async_pre_call_hook( + user_api_key_dict=auth, cache=cache, data={**request, "input": [image_uri, image_uri]}, call_type="aembedding" + ) + with pytest.raises(HTTPException) as error: + await handler.async_pre_call_hook( + user_api_key_dict=auth, + cache=cache, + data={**request, "input": [image_uri, _EMBEDDING_CAPTION]}, + call_type="aembedding", + ) + assert error.value.status_code == 429 + + +@pytest.mark.asyncio +async def test_project_itpm_limit_reserves_media_tokens_for_an_image_embedding( + rate_limiter: tuple[RateLimitHandler, DualCache], +) -> None: + handler, cache = rate_limiter + image_uri: Final = _embedding_media_uri("image", False) + image_tokens, _ = handler._estimate_input_and_output_tokens( + {"model": _EMBEDDING_MODEL, "input": image_uri}, call_type="aembedding" + ) + auth: Final = UserAPIKeyAuth( + api_key=hash_token("sk-embedding-image-itpm-media"), + project_id="proj-embedding-image-itpm-media", + project_metadata={"model_itpm_limit": {_EMBEDDING_MODEL: image_tokens // 2}}, + ) + + await handler.async_pre_call_hook( + user_api_key_dict=auth, + cache=cache, + data={"model": _EMBEDDING_MODEL, "input": _EMBEDDING_CAPTION}, + call_type="aembedding", + ) + with pytest.raises(HTTPException) as error: + await handler.async_pre_call_hook( + user_api_key_dict=auth, + cache=cache, + data={"model": _EMBEDDING_MODEL, "input": image_uri}, + call_type="aembedding", + ) + assert error.value.status_code == 429 + + @pytest.mark.asyncio async def test_ptu_pre_call_weights_media_input_and_output() -> None: cache: Final = DualCache() diff --git a/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py b/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py index 129bb10eccc..f04ade1ce2d 100644 --- a/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py +++ b/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py @@ -2,7 +2,9 @@ # This tests the router's ability to pick deployment with lowest tpm using 'usage-based-routing-v2-v2' import asyncio +import base64 import importlib +import io import os import time import traceback @@ -15,6 +17,7 @@ from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outco from unittest.mock import patch import pytest +from PIL import Image import litellm from litellm import Router @@ -23,7 +26,7 @@ from litellm.caching.caching import DualCache from litellm.router_strategy.lowest_tpm_rpm_v2 import ( LowestTPMLoggingHandler_v2 as LowestTPMLoggingHandler, ) -from litellm.types.router import DeploymentTypedDict +from litellm.types.router import DeploymentTypedDict, RouterRateLimitError from litellm.types.utils import StandardLoggingPayload from litellm.utils import get_utc_datetime @@ -347,6 +350,55 @@ def test_router_skip_rate_limited_deployments(): print(f"An exception occurred! {str(e)}") +def _embedding_image_uri() -> str: + image: Final = io.BytesIO() + Image.new("RGB", (1024, 1024)).save(image, format="PNG", compress_level=0) + return f"data:image/png;base64,{base64.b64encode(image.getvalue()).decode()}" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "routing_strategy", + ["usage-based-routing", "usage-based-routing-v2", "latency-based-routing", "cost-based-routing"], +) +@pytest.mark.parametrize( + ("litellm_params", "model_info", "admitted"), + [ + ({"model": "vertex_ai/gemini-embedding-2"}, {}, True), + ( + {"model": "openai/team-embedder", "api_base": "https://embedder.invalid"}, + {"base_model": "vertex_ai/gemini-embedding-2"}, + True, + ), + ({"model": "text-embedding-3-small"}, {}, False), + ], + ids=["image-model", "image-base-model", "text-model"], +) +async def test_deployment_tpm_counts_embedding_image_as_media_only_on_image_models( + routing_strategy: str, litellm_params: Dict[str, str], model_info: Dict[str, str], admitted: bool +) -> None: + uri: Final = _embedding_image_uri() + assert len(uri) // 4 > 75_000 + router: Final = Router( + model_list=[ + { + "model_name": "embedder", + "litellm_params": {**litellm_params, "tpm": 75_000}, + "model_info": {**model_info, "id": f"embedder-{index}"}, + } + for index in range(2) + ], + routing_strategy=routing_strategy, + ) + + if not admitted: + with pytest.raises((litellm.RateLimitError, RouterRateLimitError)): + await router.select("embedder", {}, input=[uri]) + return + deployment: Final = await router.select("embedder", {}, input=[uri]) + assert deployment["model_info"]["id"] in ("embedder-0", "embedder-1") + + @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_multiple_potential_deployments(sync_mode): diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 4f35efdfd6e..3cc97f4507b 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -5,6 +5,7 @@ import functools import gc import itertools import io +import base64 import json import linecache import logging @@ -8174,6 +8175,77 @@ def test_pre_call_checks_count_a_group_mixing_in_another_vendor_by_its_own_name( ) +_EMBEDDING_CONTEXT_WINDOW: Final = 8_192 +_EMBEDDING_IMAGE_INPUT: Final = f"data:image/png;base64,{base64.b64encode(bytes(range(256)) * 1024).decode()}" + + +def _embedding_context_router(backend: str, window: int = _EMBEDDING_CONTEXT_WINDOW) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "embedder", + "litellm_params": {"model": backend, "api_key": "sk-test"}, + "model_info": {"max_input_tokens": window}, + } + ], + enable_pre_call_checks=True, + ) + + +async def test_context_window_pre_call_check_counts_an_embedding_image_as_media_on_an_image_backend(): + router: Final = _embedding_context_router("vertex_ai/gemini-embedding-2") + deployments: Final = router.registry().deployments("embedder") + + kept: Final = router._pre_call_checks(model="embedder", healthy_deployments=deployments, input=_EMBEDDING_IMAGE_INPUT) + counted: Final = await router._acount_pre_call_check_tokens( + model="embedder", healthy_deployments=deployments, messages=None, input=_EMBEDDING_IMAGE_INPUT, request_kwargs=None + ) + + assert len(kept) == 1 + assert counted is not None and 0 < counted < _EMBEDDING_CONTEXT_WINDOW + with pytest.raises(litellm.ContextWindowExceededError): + router._pre_call_checks(model="embedder", healthy_deployments=deployments, input="word " * 20_000) + + +@pytest.mark.parametrize( + ("backend", "shape"), + [ + ("vertex_ai/gemini-embedding-2", "pdf"), + ("vertex_ai/gemini-embedding-2", "pdf_file_block"), + ("vertex_ai/gemini-embedding-2", "two_pdfs"), + ("bedrock/amazon.titan-embed-image-v1", "image"), + ("bedrock/amazon.titan-embed-image-v1", "two_images"), + ], +) +def test_context_window_pre_call_check_leaves_embedding_media_to_the_backend(backend: str, shape: str): + window: Final = litellm.get_model_info(backend)["max_input_tokens"] + assert window is not None + router: Final = _embedding_context_router(backend, window) + pdf: Final = f"data:application/pdf;base64,{base64.b64encode(bytes(window * 8)).decode()}" + inputs: Final = { + "pdf": pdf, + "pdf_file_block": [{"type": "file", "file": {"file_data": pdf}}], + "two_pdfs": [pdf, pdf], + "image": _EMBEDDING_IMAGE_INPUT, + "two_images": [_EMBEDDING_IMAGE_INPUT, _EMBEDDING_IMAGE_INPUT], + } + + kept: Final = router._pre_call_checks( + model="embedder", healthy_deployments=router.registry().deployments("embedder"), input=inputs[shape] + ) + + assert len(kept) == 1 + + +def test_context_window_pre_call_check_counts_an_embedding_image_as_text_on_a_text_only_backend(): + router: Final = _embedding_context_router("text-embedding-3-small") + + with pytest.raises(litellm.ContextWindowExceededError): + router._pre_call_checks( + model="embedder", healthy_deployments=router.registry().deployments("embedder"), input=_EMBEDDING_IMAGE_INPUT + ) + + async def test_acount_pre_call_check_tokens_counts_a_group_of_one_openai_model_with_that_models_tokenizer(): router: Final = _context_limited_router("openai/gpt-6.1-sol")