mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(rate-limit): estimate embedding image and audio data URIs as media in TPM pre-call checks (#49831)
* fix(rate-limit): estimate embedding image and audio data URIs as media in TPM pre-call checks * fix(rate-limit): count embedding media as media only on models that embed it A data URI sent to a text-only embedding model is read as text, so it keeps the text estimate there. On models flagged supports_embedding_image_input, image and audio use the chat media estimates and PDF and video count at most the model's max_input_tokens. The same estimate now feeds project and key limits, the usage-based, latency and cost routers' deployment tpm, and the deployment itpm check. * fix(proxy): count embedding media only where the backend reads it as media Vertex AI, Gemini, Azure AI and Bedrock count a strict base64 data URI as media per item. Cohere counts it as media only when the whole batch is media, and Voyage, OCI and unknown models keep the text estimate. File blocks carrying a data URI count like the bare URI. The router context window pre-call check and the batch limiter now use the same estimate, resolving each deployment or JSONL row model to its backends. The cost map gains supports_embedding_image_input on voyage/voyage-multimodal-3 and oci/cohere.embed-multilingual-image-v3.0. * fix(router): hold each embedding media item to the context window on its own A PDF or video data URI alone already reserves at most the model's max_input_tokens, but the context window pre-call check then ran the blanked input through the Responses transform and added a few framing tokens, so a PDF that Vertex accepts got a 400. When every input item is media, the check now returns the largest item's media estimate, since an embedding request embeds each item separately * fix(proxy): leave embedding media out of the context check and narrow Bedrock media scope The router context window check now counts only the text of an embedding input, so Titan image lists no longer 400. Bedrock Marengo and Nova with a PDF keep the text estimate since they read those inputs as text or reject them, async_invoke Nova counts images as media, and the unused batch limiter change is dropped. * fix(bedrock): bill Cohere embeddings by Bedrock's input token count header * fix(embeddings): check data URI base64 in place and keep Bedrock Cohere text usage The strict base64 check now scans the payload in place instead of decoding and re-encoding a full copy. Bedrock's billed token header is used only for image requests, so text-only Bedrock Cohere usage matches main. The deployment itpm estimator change is reverted, since the router never passes embedding input to that check. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
norkelbyte
GitHub
parent
391c8013c9
commit
3b35ae340d
18 changed files with 837 additions and 66 deletions
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue