mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(proxy): account for media in combined TPM reservations (#48865)
This commit is contained in:
tin-berri
GitHub
parent
9675c13cb0
commit
ca6c830221
10 changed files with 1169 additions and 31 deletions
|
|
@ -2,8 +2,12 @@
|
|||
Utils used for litellm.transcription() and litellm.atranscription()
|
||||
"""
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
import hashlib
|
||||
import math
|
||||
import os
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
|
|
@ -331,6 +335,45 @@ def calculate_request_duration(file: FileTypes) -> float | None:
|
|||
return None
|
||||
|
||||
|
||||
DEFAULT_AUDIO_TOKEN_ESTIMATE: Final = 300
|
||||
AUDIO_TOKENS_PER_SECOND: Final = 10
|
||||
AUDIO_BYTES_PER_TOKEN: Final = 1600
|
||||
_MAX_AUDIO_ESTIMATE_BYTES: Final = 16 * 1024 * 1024
|
||||
|
||||
|
||||
def _audio_token_rate(model: str) -> int | None:
|
||||
if model.startswith(
|
||||
("gemini-", "gemini/gemini-", "vertex_ai/gemini-", "vertex_ai_beta/gemini-", "openrouter/google/gemini-")
|
||||
):
|
||||
family: Final = model.rsplit("/", 1)[-1]
|
||||
return 25 if family.startswith(("gemini-3-", "gemini-3.")) else 32
|
||||
if model.startswith(("gpt-", "openai/gpt-", "azure/gpt-", "openrouter/openai/gpt-")):
|
||||
return AUDIO_TOKENS_PER_SECOND
|
||||
return None
|
||||
|
||||
|
||||
def estimate_audio_tokens(data: str | None, models: Sequence[str] = (), fallback_model: str | None = None) -> int:
|
||||
"""Estimate audio tokens from container duration within a bounded decoding budget."""
|
||||
if not data:
|
||||
return DEFAULT_AUDIO_TOKEN_ESTIMATE
|
||||
fallback_estimate: Final = max(len(data) * 3 // 4 // AUDIO_BYTES_PER_TOKEN, DEFAULT_AUDIO_TOKEN_ESTIMATE)
|
||||
padding: Final = 2 if data.endswith("==") else 1 if data.endswith("=") else 0
|
||||
if len(data) * 3 // 4 - padding > _MAX_AUDIO_ESTIMATE_BYTES:
|
||||
return fallback_estimate
|
||||
try:
|
||||
duration: Final = calculate_request_duration(base64.b64decode(data, validate=True))
|
||||
except (ValueError, binascii.Error):
|
||||
return fallback_estimate
|
||||
if duration is None or not math.isfinite(duration) or duration <= 0:
|
||||
return fallback_estimate
|
||||
fallback_rate: Final = _audio_token_rate(fallback_model) if fallback_model is not None else None
|
||||
rate: Final = max(
|
||||
(_audio_token_rate(model) or fallback_rate or AUDIO_TOKENS_PER_SECOND for model in models),
|
||||
default=fallback_rate or AUDIO_TOKENS_PER_SECOND,
|
||||
)
|
||||
return max(1, math.ceil(duration * rate))
|
||||
|
||||
|
||||
DEFAULT_SPEECH_MEDIA_TYPE: Final = "audio/mpeg"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ import json
|
|||
import mimetypes
|
||||
import re
|
||||
from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from graphlib import CycleError, TopologicalSorter
|
||||
from itertools import chain, groupby, islice
|
||||
from os import PathLike
|
||||
|
|
@ -236,6 +237,160 @@ _IMAGE_CONTENT_PART_TYPES: Final = frozenset({"image_url", "input_image", "image
|
|||
_IMAGE_SCAN_MAX_DEPTH: Final = 4
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _TPMMediaProjection:
|
||||
content: object
|
||||
media: tuple[Mapping[str, object], ...] = ()
|
||||
|
||||
|
||||
def _tpm_image_part(url: str, detail: object = "auto") -> Mapping[str, object]:
|
||||
return {"type": "image_url", "image_url": {"url": url, "detail": detail}}
|
||||
|
||||
|
||||
def _tpm_blank_fields(fields: Mapping[str, object], names: tuple[str, ...]) -> Mapping[str, object]:
|
||||
return {key: "" if key in names and isinstance(value, str) else value for key, value in fields.items()}
|
||||
|
||||
|
||||
def _tpm_file_media(part: Mapping[str, object]) -> _TPMMediaProjection:
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import parse_mime_type
|
||||
|
||||
nested: Final = part.get("type") == "file"
|
||||
fields: Final = as_str_mapping(part.get("file")) if nested else part
|
||||
if fields is None:
|
||||
return _TPMMediaProjection(part)
|
||||
raw_data: Final = fields.get("file_data")
|
||||
data: Final = raw_data if isinstance(raw_data, str) else ""
|
||||
declared: Final = next(
|
||||
(
|
||||
value
|
||||
for key in ("format", "mime_type", "content_type")
|
||||
if isinstance(value := fields.get(key), str) and value
|
||||
),
|
||||
"",
|
||||
)
|
||||
mime: Final = (declared or parse_mime_type(data[:256]) or "").split(";", 1)[0].strip().lower()
|
||||
if not mime.startswith(("image/", "audio/")):
|
||||
return _TPMMediaProjection(part)
|
||||
raw_url: Final = fields.get("file_url") or fields.get("file_id")
|
||||
url: Final = data or (raw_url if isinstance(raw_url, str) else "")
|
||||
header_end: Final = data.find(",", 0, 256)
|
||||
encoded: Final = (
|
||||
data[header_end + 1 :]
|
||||
if header_end >= 0 and data[:header_end].endswith(";base64")
|
||||
else ""
|
||||
if data.startswith(("data:", "https://", "http://", "gs://"))
|
||||
else data
|
||||
)
|
||||
media: Final = (
|
||||
_tpm_image_part(
|
||||
url
|
||||
if not data or url.startswith(("data:", "https://", "http://", "gs://"))
|
||||
else f"data:{mime};base64,{url}",
|
||||
fields.get("detail", "auto"),
|
||||
)
|
||||
if mime.startswith("image/")
|
||||
else {"type": "input_audio", "input_audio": {"data": encoded, "format": mime.removeprefix("audio/")}}
|
||||
)
|
||||
projected: Final = _tpm_blank_fields(fields, ("file_data", "file_url"))
|
||||
return _TPMMediaProjection({**part, "file": projected} if nested else projected, (media,))
|
||||
|
||||
|
||||
def _tpm_media_leaf(value: object) -> _TPMMediaProjection:
|
||||
part: Final = as_str_mapping(value)
|
||||
if part is None:
|
||||
return _TPMMediaProjection(value)
|
||||
kind: Final = part.get("type")
|
||||
if kind in ("file", "input_file"):
|
||||
return _tpm_file_media(part)
|
||||
if kind == "input_audio":
|
||||
audio: Final = as_str_mapping(part.get("input_audio"))
|
||||
return _TPMMediaProjection(
|
||||
{**part, "input_audio": _tpm_blank_fields(audio, ("data",))} if audio is not None else value,
|
||||
({"type": "input_audio", "input_audio": part.get("input_audio")},),
|
||||
)
|
||||
if kind in ("image_url", "input_image", "image", "computer_screenshot"):
|
||||
image: Final = as_str_mapping(part.get("image_url"))
|
||||
source: Final = as_str_mapping(part.get("source"))
|
||||
canonical: Final = _tpm_image_part(
|
||||
_image_part_url({**part, "type": "input_image"} if kind == "computer_screenshot" else part) or "",
|
||||
(image or part).get("detail", part.get("detail", "auto")),
|
||||
)
|
||||
projected: Final = (
|
||||
{**part, "image_url": _tpm_blank_fields(image, ("url",))}
|
||||
if image is not None and kind in ("image_url", "input_image", "computer_screenshot")
|
||||
else {**part, "source": _tpm_blank_fields(source, ("data", "url"))}
|
||||
if source is not None and kind == "image" and source.get("type") in ("base64", "url")
|
||||
else _tpm_blank_fields(part, ("image_url", "url"))
|
||||
)
|
||||
return _TPMMediaProjection(projected, (canonical,))
|
||||
return _TPMMediaProjection(value)
|
||||
|
||||
|
||||
def _tpm_media_list(value: object, split: Callable[[object], _TPMMediaProjection]) -> _TPMMediaProjection:
|
||||
if not isinstance(value, (list, tuple)):
|
||||
return _TPMMediaProjection(value)
|
||||
parts: Final = cast(Sequence[object], value) # cast-ok: finite JSON content arrays
|
||||
projected: Final = tuple(map(split, parts))
|
||||
media: Final = tuple(chain.from_iterable(part.media for part in projected))
|
||||
content: Final = tuple(part.content for part in projected)
|
||||
return _TPMMediaProjection(parts if not media else content if isinstance(value, tuple) else list(content), media)
|
||||
|
||||
|
||||
def _tpm_replace_field(part: Mapping[str, object], key: str, projection: _TPMMediaProjection) -> _TPMMediaProjection:
|
||||
return _TPMMediaProjection({**part, key: projection.content} if projection.media else part, projection.media)
|
||||
|
||||
|
||||
def _tpm_document_media(value: object) -> _TPMMediaProjection:
|
||||
part: Final = as_str_mapping(value)
|
||||
source: Final = as_str_mapping(part.get("source")) if part is not None else None
|
||||
if part is None or part.get("type") != "document" or source is None or source.get("type") != "content":
|
||||
return _tpm_media_leaf(value)
|
||||
content: Final = _tpm_media_list(source.get("content"), _tpm_media_leaf)
|
||||
return _tpm_replace_field(part, "source", _tpm_replace_field(source, "content", content))
|
||||
|
||||
|
||||
def _tpm_tool_result_media(value: object) -> _TPMMediaProjection:
|
||||
part: Final = as_str_mapping(value)
|
||||
if part is None or part.get("type") != "tool_result":
|
||||
return _tpm_document_media(value)
|
||||
return _tpm_replace_field(part, "content", _tpm_media_list(part.get("content"), _tpm_document_media))
|
||||
|
||||
|
||||
def _tpm_code_output_media(value: object) -> _TPMMediaProjection:
|
||||
part: Final = as_str_mapping(value)
|
||||
url: Final = part.get("url") if part is not None else None
|
||||
if part is None or part.get("type") != "image" or not isinstance(url, str) or not url:
|
||||
return _TPMMediaProjection(value)
|
||||
return _TPMMediaProjection(_tpm_blank_fields(part, ("url",)), (_tpm_image_part(url),))
|
||||
|
||||
|
||||
def _tpm_input_media(value: object) -> _TPMMediaProjection:
|
||||
part: Final = as_str_mapping(value)
|
||||
if part is None:
|
||||
return _TPMMediaProjection(value)
|
||||
kind: Final = part.get("type")
|
||||
if kind in ("function_call_output", "custom_tool_call_output"):
|
||||
return _tpm_replace_field(part, "output", _tpm_media_list(part.get("output"), _tpm_tool_result_media))
|
||||
if kind == "computer_call_output":
|
||||
return _tpm_replace_field(part, "output", _tpm_media_leaf(part.get("output")))
|
||||
if kind == "code_interpreter_call":
|
||||
return _tpm_replace_field(part, "outputs", _tpm_media_list(part.get("outputs"), _tpm_code_output_media))
|
||||
result: Final = part.get("result")
|
||||
if kind == "image_generation_call" and isinstance(result, str) and result:
|
||||
return _TPMMediaProjection(
|
||||
_tpm_blank_fields(part, ("result",)),
|
||||
(_tpm_image_part("data:application/octet-stream;base64," + result),),
|
||||
)
|
||||
if kind == "message" or (kind is None and isinstance(part.get("role"), str)):
|
||||
return _tpm_replace_field(part, "content", _tpm_media_list(part.get("content"), _tpm_tool_result_media))
|
||||
return _tpm_tool_result_media(value)
|
||||
|
||||
|
||||
def split_media_for_tpm(value: object) -> tuple[object, tuple[Mapping[str, object], ...]]:
|
||||
projection: Final = _tpm_media_list(value, _tpm_input_media)
|
||||
return projection.content, projection.media
|
||||
|
||||
|
||||
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."""
|
||||
|
|
|
|||
|
|
@ -3,9 +3,12 @@
|
|||
import base64
|
||||
import io
|
||||
import math
|
||||
import re
|
||||
import struct
|
||||
from bisect import bisect_right
|
||||
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
|
||||
from itertools import accumulate
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal, cast
|
||||
|
||||
import anyio
|
||||
|
|
@ -35,6 +38,7 @@ from litellm.constants import (
|
|||
TOKEN_COUNTER_MAX_EXACT_CHARS,
|
||||
)
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.prompt_templates.chat_template_tools import (
|
||||
ChatTemplateToolSection,
|
||||
chat_template_tool_section,
|
||||
|
|
@ -188,22 +192,23 @@ def resize_image_high_res(
|
|||
|
||||
# Test the function with the given example
|
||||
def calculate_tiles_needed(
|
||||
resized_width,
|
||||
resized_height,
|
||||
tile_width=MAX_TILE_WIDTH,
|
||||
tile_height=MAX_TILE_HEIGHT,
|
||||
):
|
||||
resized_width: int,
|
||||
resized_height: int,
|
||||
tile_width: int = MAX_TILE_WIDTH,
|
||||
tile_height: int = MAX_TILE_HEIGHT,
|
||||
) -> int:
|
||||
tiles_across: Final = (resized_width + tile_width - 1) // tile_width
|
||||
tiles_down: Final = (resized_height + tile_height - 1) // tile_height
|
||||
total_tiles: Final = tiles_across * tiles_down
|
||||
return total_tiles
|
||||
|
||||
|
||||
def high_detail_image_token_upper_bound(base_tokens: int = 85) -> int:
|
||||
def high_detail_image_token_upper_bound(base_tokens: int = 85, tile_tokens: int | None = None) -> int:
|
||||
largest_tile_count: Final = calculate_tiles_needed(
|
||||
MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES, MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES
|
||||
)
|
||||
return base_tokens + (base_tokens * 2) * largest_tile_count
|
||||
per_tile: Final = base_tokens * 2 if tile_tokens is None else tile_tokens
|
||||
return base_tokens + per_tile * largest_tile_count
|
||||
|
||||
|
||||
def _unpack_ints(fmt: str, buffer: bytes) -> tuple[int, ...]:
|
||||
|
|
@ -381,6 +386,219 @@ def calculate_img_tokens(
|
|||
return total_tokens
|
||||
|
||||
|
||||
_IMAGE_HEADER_BASE64_CHARS: Final = 65_536
|
||||
_MAX_IMAGE_PATCHES: Final = 30_000
|
||||
_ANTHROPIC_IMAGE_LIMITS: Final = MappingProxyType(
|
||||
{"claude-standard": (ANTHROPIC_IMAGE_MAX_LONG_EDGE_PX, 1568), "claude-high": (2576, 4784)}
|
||||
)
|
||||
_TILE_IMAGE_TOKEN_COSTS: Final = MappingProxyType(
|
||||
{
|
||||
"gpt-4o": (85, 170),
|
||||
"gpt-4.1": (85, 170),
|
||||
"gpt-4o-mini": (2833, 5667),
|
||||
"gpt-5": (70, 140),
|
||||
"gpt-5.1": (70, 140),
|
||||
"o1": (75, 150),
|
||||
"o1-pro": (75, 150),
|
||||
"o3": (75, 150),
|
||||
}
|
||||
)
|
||||
_PATCH_IMAGE_MULTIPLIERS: Final = MappingProxyType(
|
||||
{
|
||||
"gpt-6-astra": 1.2,
|
||||
"gpt-5.6-sol": 1.2,
|
||||
"gpt-5.6-terra": 1.2,
|
||||
"gpt-5.6-luna": 1.2,
|
||||
"gpt-5.5": 1.2,
|
||||
"gpt-5.4": 1.2,
|
||||
"gpt-5.4-mini": 1.2,
|
||||
"gpt-5.4-nano": 1.2,
|
||||
"gpt-5.2": 1.2,
|
||||
"gpt-4.1-mini": 1.62,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _inline_image_dimensions(data: str) -> tuple[int, int] | None:
|
||||
separator: Final = data.find(",", 0, 256)
|
||||
if not data.startswith("data:") or separator < 0 or not data[:separator].endswith(";base64"):
|
||||
return None
|
||||
encoded_header: Final = data[separator + 1 : separator + 1 + _IMAGE_HEADER_BASE64_CHARS]
|
||||
try:
|
||||
header: Final = base64.b64decode(encoded_header, validate=True)
|
||||
image_type: Final = get_image_type(header)
|
||||
if image_type not in ("png", "gif", "jpeg", "webp"):
|
||||
return None
|
||||
if image_type == "webp" and header[12:16] not in (b"VP8X", b"VP8 ", b"VP8L"):
|
||||
return None
|
||||
dimensions: Final = _image_dimensions_from_bytes(header)
|
||||
except (ValueError, TypeError, struct.error):
|
||||
return None
|
||||
width, height = dimensions
|
||||
return (width, height) if width > 0 and height > 0 else None
|
||||
|
||||
|
||||
def _patch_image_limits(model: str, detail: str) -> tuple[int, int | None] | None:
|
||||
if model in ("gpt-5.2", "gpt-4.1-mini"):
|
||||
return None if detail == "original" else (2048, 6144)
|
||||
original_default: Final = model in ("gpt-6-astra", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna", "gpt-5.5")
|
||||
effective_detail: Final = ("original" if original_default else "high") if detail == "auto" else detail
|
||||
if effective_detail == "low":
|
||||
return (2048, 6144) if model in ("gpt-5.4", "gpt-5.4-mini", "gpt-5.4-nano") else (512, None)
|
||||
if effective_detail == "high":
|
||||
return (65_535 if model == "gpt-6-astra" else 2048, 2500)
|
||||
return (
|
||||
(65_535, None) if model in ("gpt-6-astra", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna") else (6000, 10_000)
|
||||
)
|
||||
|
||||
|
||||
def _image_patch_count(width: int, height: int, max_dimension: int, budget: int | None) -> int:
|
||||
dimension_scale: Final = min(1.0, max_dimension / max(width, height))
|
||||
fitted_width: Final = max(1, round(width * dimension_scale))
|
||||
fitted_height: Final = max(1, round(height * dimension_scale))
|
||||
patches: Final = math.ceil(fitted_width / 32) * math.ceil(fitted_height / 32)
|
||||
if budget is None or patches <= budget:
|
||||
return patches
|
||||
shrink: Final = math.sqrt(32**2 * budget / (fitted_width * fitted_height))
|
||||
width_patches: Final = fitted_width * shrink / 32
|
||||
height_patches: Final = fitted_height * shrink / 32
|
||||
aligned_shrink: Final = shrink * min(
|
||||
math.floor(width_patches) / width_patches, math.floor(height_patches) / height_patches
|
||||
)
|
||||
resized_width: Final = max(1, math.floor(fitted_width * aligned_shrink))
|
||||
resized_height: Final = max(1, math.floor(fitted_height * aligned_shrink))
|
||||
return math.ceil(resized_width / 32) * math.ceil(resized_height / 32)
|
||||
|
||||
|
||||
def _image_policy_model_name(model: str) -> str | None:
|
||||
provider, separator, provider_model = model.partition("/")
|
||||
if not separator:
|
||||
return model
|
||||
if provider == "openai":
|
||||
return provider_model
|
||||
if provider == "openrouter":
|
||||
return provider_model.removeprefix("openai/") if provider_model.startswith("openai/") else None
|
||||
if provider not in ("azure", "azure_ai"):
|
||||
return None
|
||||
catalog_entry: Final = as_str_mapping(
|
||||
cast( # cast-ok: read-only view of the legacy string-keyed registry; as_str_mapping checks the entry
|
||||
Mapping[str, object], litellm.model_cost
|
||||
).get(model)
|
||||
)
|
||||
if catalog_entry is not None and catalog_entry.get("litellm_provider") in ("azure", "azure_ai"):
|
||||
return provider_model.rsplit("/", 1)[-1]
|
||||
return provider_model.removeprefix("responses/").removeprefix("o_series/").removeprefix("openai/")
|
||||
|
||||
|
||||
def _anthropic_image_family(model: str) -> str | None:
|
||||
from litellm.utils import get_potential_model_names
|
||||
|
||||
provider, separator, _ = model.partition("/")
|
||||
if separator and provider not in (
|
||||
"anthropic",
|
||||
"bedrock",
|
||||
"bedrock_converse",
|
||||
"bedrock_mantle",
|
||||
"vertex_ai",
|
||||
"vertex_ai_beta",
|
||||
"azure_ai",
|
||||
"openrouter",
|
||||
):
|
||||
return None
|
||||
candidates: Final = get_potential_model_names(model, provider if separator else "anthropic")
|
||||
catalog: Final = cast(Mapping[str, object], litellm.model_cost) # cast-ok: read-only view of the legacy registry
|
||||
for candidate in dict.fromkeys(candidates.values()):
|
||||
if not isinstance(candidate, str):
|
||||
continue
|
||||
if (entry := as_str_mapping(catalog.get(candidate))) is None or entry.get("supports_vision") is not True:
|
||||
continue
|
||||
if (
|
||||
version := re.search(
|
||||
r"(?:^|[./])claude-(?:(?:haiku|sonnet|opus)-)?(\d+)(?:[-.](\d{1,2})(?!\d))?", candidate
|
||||
)
|
||||
) is None:
|
||||
continue
|
||||
major, minor = int(version[1]), int(version[2] or 0)
|
||||
return "claude-high" if (major, minor) >= (4, 7) else "claude-standard"
|
||||
return None
|
||||
|
||||
|
||||
def _anthropic_image_tokens(dimensions: tuple[int, int] | None, limits: tuple[int, int]) -> int:
|
||||
max_edge, token_budget = limits
|
||||
if dimensions is None:
|
||||
return token_budget
|
||||
long_edge, short_edge = max(dimensions), min(dimensions)
|
||||
aspect_ratio: Final = long_edge / short_edge
|
||||
|
||||
def tokens_at_edge(edge: int) -> int:
|
||||
return math.ceil(edge / 28) * math.ceil(max(round(edge / aspect_ratio), 1) / 28)
|
||||
|
||||
fitted_edge: Final = bisect_right(range(1, min(long_edge, max_edge) + 1), token_budget, key=tokens_at_edge)
|
||||
return tokens_at_edge(max(fitted_edge, 1))
|
||||
|
||||
|
||||
def _image_model_family(model: str) -> str | None:
|
||||
provider_model: Final = _image_policy_model_name(model)
|
||||
if provider_model is None:
|
||||
return _anthropic_image_family(model)
|
||||
family: Final = re.sub(r"-\d{4}-\d{2}-\d{2}$", "", provider_model)
|
||||
canonical: Final = litellm.azure_llms.get(f"azure/{family}", family) if model.startswith("azure/") else family
|
||||
return (
|
||||
canonical
|
||||
if canonical in _TILE_IMAGE_TOKEN_COSTS or canonical in _PATCH_IMAGE_MULTIPLIERS
|
||||
else _anthropic_image_family(model)
|
||||
)
|
||||
|
||||
|
||||
def _image_policy_tokens(family: str | None, dimensions: tuple[int, int] | None, detail: str) -> int:
|
||||
if family is None:
|
||||
return DEFAULT_IMAGE_TOKEN_COUNT
|
||||
anthropic_limits: Final = _ANTHROPIC_IMAGE_LIMITS.get(family)
|
||||
if anthropic_limits is not None:
|
||||
return _anthropic_image_tokens(dimensions, anthropic_limits)
|
||||
tile_costs: Final = _TILE_IMAGE_TOKEN_COSTS.get(family)
|
||||
multiplier: Final = _PATCH_IMAGE_MULTIPLIERS.get(family)
|
||||
if tile_costs is not None:
|
||||
base_tokens, tile_tokens = tile_costs
|
||||
if detail == "original":
|
||||
return DEFAULT_IMAGE_TOKEN_COUNT
|
||||
if detail == "low":
|
||||
return base_tokens
|
||||
if dimensions is None:
|
||||
return high_detail_image_token_upper_bound(base_tokens, tile_tokens)
|
||||
width, height = dimensions
|
||||
scale: Final = min(1.0, 2048 / max(width, height), 768 / min(width, height))
|
||||
tiles: Final = calculate_tiles_needed(max(1, math.floor(width * scale)), max(1, math.floor(height * scale)))
|
||||
return base_tokens + tiles * tile_tokens
|
||||
limits: Final = _patch_image_limits(family, detail)
|
||||
if limits is None or multiplier is None:
|
||||
return DEFAULT_IMAGE_TOKEN_COUNT
|
||||
max_dimension, patch_budget = limits
|
||||
patches: Final = (
|
||||
min(math.ceil(max_dimension / 32) ** 2, patch_budget if patch_budget is not None else _MAX_IMAGE_PATCHES)
|
||||
if dimensions is None
|
||||
else _image_patch_count(*dimensions, max_dimension, patch_budget)
|
||||
)
|
||||
return math.ceil(patches * multiplier)
|
||||
|
||||
|
||||
def estimate_image_tokens(
|
||||
model: str | Sequence[str], data: str, detail: str = "auto", *, fallback_model: str | None = None
|
||||
) -> int:
|
||||
"""Estimate image policies from at most 48 KiB of inline headers, without fetching media."""
|
||||
if detail not in ("low", "high", "auto", "original"):
|
||||
return DEFAULT_IMAGE_TOKEN_COUNT
|
||||
models: Final = (model,) if isinstance(model, str) else model
|
||||
fallback_family: Final = _image_model_family(fallback_model) if fallback_model is not None else None
|
||||
families: Final = tuple(
|
||||
dict.fromkeys(_image_model_family(candidate) or fallback_family for candidate in models or ("",))
|
||||
)
|
||||
if all(family is None for family in families):
|
||||
return DEFAULT_IMAGE_TOKEN_COUNT
|
||||
dimensions: Final = _inline_image_dimensions(data)
|
||||
return max(_image_policy_tokens(family, dimensions, detail) for family in families)
|
||||
|
||||
|
||||
TokenCounterFunction = Callable[[str], int]
|
||||
"""
|
||||
Type for a function that counts tokens in a string.
|
||||
|
|
|
|||
|
|
@ -45,11 +45,14 @@ from litellm.caching.redis_cache import log_redis_failure
|
|||
from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY
|
||||
from litellm.cost_calculator import get_usage_object
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
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,
|
||||
split_media_for_tpm,
|
||||
)
|
||||
from litellm.litellm_core_utils.ptu_pricing import is_ptu_cost_attribution_enabled
|
||||
from litellm.litellm_core_utils.token_counter import offload_token_count
|
||||
from litellm.litellm_core_utils.token_counter import estimate_image_tokens, offload_token_count
|
||||
from litellm.llms.azure.ptu_capacity import normalized_tokens
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_utils import (
|
||||
|
|
@ -98,6 +101,7 @@ if TYPE_CHECKING:
|
|||
from opentelemetry.trace import Span as _Span
|
||||
|
||||
from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache
|
||||
from litellm.router import Router
|
||||
from litellm.types.agents import AgentResponse
|
||||
from litellm.types.caching import RedisPipelineIncrementOperation
|
||||
|
||||
|
|
@ -135,6 +139,18 @@ class RateLimitedModel:
|
|||
return requested_limit if requested_limit is not None else limits.get(self.group)
|
||||
|
||||
|
||||
def _resolve_media_models(router: "Router | None", model: str, team_id: str | None) -> tuple[str, ...]:
|
||||
if router is None:
|
||||
return (model,)
|
||||
return tuple(dict.fromkeys(router.resolved_litellm_models(model, team_id, use_base_model=True))) or (model,)
|
||||
|
||||
|
||||
def _resolve_media_models_via_proxy_router(model: str, team_id: str | None) -> tuple[str, ...]:
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
return _resolve_media_models(llm_router, model, team_id)
|
||||
|
||||
|
||||
def _resolve_model_group_alias_via_proxy_router(model: str) -> str | None:
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
|
|
@ -862,6 +878,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
time_provider: Callable[[], datetime] | None = None,
|
||||
tag_rate_limit_resolver: TagRateLimitResolver = resolve_tag_rate_limits_from_db,
|
||||
model_group_resolver: Callable[[str], str | None] = _resolve_model_group_alias_via_proxy_router,
|
||||
media_model_resolver: Callable[[str, str | None], tuple[str, ...]] = _resolve_media_models_via_proxy_router,
|
||||
ptu_team_ceiling_resolver: Callable[
|
||||
[str, str], PTUTeamCeiling | None
|
||||
] = _resolve_ptu_team_ceiling_via_proxy_router,
|
||||
|
|
@ -872,6 +889,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
self._time_provider = time_provider or datetime.now
|
||||
self._tag_rate_limit_resolver = tag_rate_limit_resolver
|
||||
self._model_group_resolver = model_group_resolver
|
||||
self._media_model_resolver = media_model_resolver
|
||||
self._ptu_team_ceiling_resolver = ptu_team_ceiling_resolver
|
||||
self._fail_closed_resolver = fail_closed_resolver
|
||||
self._force_hash_tag_grouping_resolver = force_hash_tag_grouping_resolver
|
||||
|
|
@ -1133,11 +1151,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
|
||||
def _estimate_tokens_for_request(
|
||||
self,
|
||||
data: dict,
|
||||
data: object,
|
||||
model: str | None = None,
|
||||
min_configured_tpm_limit: int | None = None,
|
||||
call_type: str | None = None,
|
||||
configured_output_tokens: int | None = None,
|
||||
team_id: str | None = None,
|
||||
) -> int:
|
||||
"""
|
||||
Estimate total tokens this request will consume so we can reserve them
|
||||
|
|
@ -1164,6 +1183,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
min_configured_tpm_limit=min_configured_tpm_limit,
|
||||
call_type=call_type,
|
||||
configured_output_tokens=configured_output_tokens,
|
||||
team_id=team_id,
|
||||
media_model_resolver=self._media_model_resolver,
|
||||
)
|
||||
total_estimated: Final = estimated_input_tokens + max_tokens_estimate
|
||||
|
||||
|
|
@ -1184,6 +1205,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
call_type: str | None,
|
||||
configured_output_tokens: int | None,
|
||||
raw_estimate: int,
|
||||
team_id: str | None = None,
|
||||
) -> int:
|
||||
"""The team PTU ceiling counts Azure normalized tokens, so its reservation weighs the
|
||||
output budget the way the ceiling does instead of the raw sum the other scopes reserve."""
|
||||
|
|
@ -1194,6 +1216,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
min_configured_tpm_limit=min_configured_tpm_limit,
|
||||
call_type=call_type,
|
||||
configured_output_tokens=configured_output_tokens,
|
||||
team_id=team_id,
|
||||
media_model_resolver=self._media_model_resolver,
|
||||
)
|
||||
normalized: Final = normalized_tokens(
|
||||
ceiling, prompt_tokens=estimated_input_tokens, completion_tokens=max_tokens_estimate
|
||||
|
|
@ -1206,6 +1230,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
min_configured_tpm_limit: int | None = None,
|
||||
call_type: str | None = None,
|
||||
configured_output_tokens: int | None = None,
|
||||
include_media: bool = True,
|
||||
team_id: str | None = None,
|
||||
media_model_resolver: Callable[[str, str | None], tuple[str, ...]] | None = None,
|
||||
) -> tuple[int, int]:
|
||||
"""
|
||||
Estimate input tokens and output (max_tokens) budget separately, so
|
||||
|
|
@ -1229,10 +1256,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
floor entirely, so the reservation reflects what this tenant's model
|
||||
actually emits rather than one constant shared by every tenant.
|
||||
"""
|
||||
if not isinstance(data, dict):
|
||||
request_data: Final = as_str_mapping(data)
|
||||
if not isinstance(data, dict) or request_data is None:
|
||||
return 0, 0
|
||||
request_model: Final = request_data.get("model")
|
||||
translated_data: Final = self._translate_google_genai_native_request(data, call_type)
|
||||
estimable_data: Final = translated_data if translated_data is not None else data
|
||||
estimable_data: Final = translated_data if translated_data is not None else request_data
|
||||
selected_fields: Final[tuple[object | None, object | None, object | None]] = (
|
||||
(None, None, estimable_data.get("input"))
|
||||
if call_type in RESPONSES_API_CALL_TYPES or call_type in EMBEDDING_API_CALL_TYPES
|
||||
|
|
@ -1249,10 +1278,27 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
)
|
||||
)
|
||||
messages, prompt, raw_input = selected_fields
|
||||
input_text: Final = (
|
||||
original_input: Final = (
|
||||
billable_input_since_newest_compaction(raw_input) if call_type in RESPONSES_API_CALL_TYPES else raw_input
|
||||
)
|
||||
|
||||
media_source: Final = messages if messages is not None else original_input
|
||||
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_tokens: Final = sum(
|
||||
self._estimate_media_block_tokens(block, resolved_models, fallback_model) for block in media
|
||||
)
|
||||
|
||||
total_chars: Final = (
|
||||
len(get_str_from_messages(messages))
|
||||
if isinstance(messages, list) and messages
|
||||
|
|
@ -1267,7 +1313,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
else 0
|
||||
)
|
||||
|
||||
estimated_input_tokens: Final = max(1, total_chars // DEFAULT_CHARS_PER_TOKEN) if total_chars > 0 else 0
|
||||
estimated_input_tokens: Final = (
|
||||
max(1, total_chars // DEFAULT_CHARS_PER_TOKEN) if total_chars > 0 else 0
|
||||
) + media_tokens
|
||||
|
||||
explicit_max_tokens: Final = self._get_explicit_output_cap(data, call_type)
|
||||
generates_no_output: Final = (
|
||||
|
|
@ -1283,7 +1331,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
max_tokens_estimate: Final = (
|
||||
0
|
||||
if generates_no_output
|
||||
or (explicit_max_tokens is None and total_chars == 0 and configured_output_tokens is None)
|
||||
or (explicit_max_tokens is None and estimated_input_tokens == 0 and configured_output_tokens is None)
|
||||
else explicit_max_tokens
|
||||
if explicit_max_tokens is not None
|
||||
else configured_output_tokens
|
||||
|
|
@ -1293,6 +1341,24 @@ 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.
|
||||
|
|
@ -3958,7 +4024,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
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)
|
||||
estimated_input_tokens, _ = self._estimate_input_and_output_tokens(
|
||||
data=data, call_type=call_type, include_media=False
|
||||
)
|
||||
return estimated_input_tokens + audio_token_estimate
|
||||
|
||||
async def _reserve_project_io_tokens_or_raise(
|
||||
|
|
@ -4003,6 +4071,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
data=data,
|
||||
min_configured_tpm_limit=min_configured_otpm_limit,
|
||||
call_type=call_type,
|
||||
include_media=False,
|
||||
)
|
||||
raw_estimated_input_tokens: Final = await offload_token_count(self._estimate_precise_input_tokens)(
|
||||
data=data, model=requested_model, call_type=call_type
|
||||
|
|
@ -4360,12 +4429,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
# Post-call reconciliation refunds the over-reservation
|
||||
# delta when actual usage comes in below the floor.
|
||||
estimated_tokens: Final = max(
|
||||
self._estimate_tokens_for_request(
|
||||
await offload_token_count(self._estimate_tokens_for_request)(
|
||||
data=data,
|
||||
model=requested_model,
|
||||
min_configured_tpm_limit=min_configured_tpm_limit,
|
||||
call_type=call_type,
|
||||
configured_output_tokens=configured_output_tokens,
|
||||
team_id=user_api_key_dict.team_id,
|
||||
),
|
||||
1,
|
||||
)
|
||||
|
|
@ -4383,13 +4453,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
)
|
||||
|
||||
capped_request: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(data)
|
||||
ptu_estimated_tokens: Final = self._estimate_ptu_tokens_for_request(
|
||||
ptu_estimated_tokens: Final = await offload_token_count(self._estimate_ptu_tokens_for_request)(
|
||||
ceiling=stash.ptu_ceiling,
|
||||
data=capped_request,
|
||||
min_configured_tpm_limit=min_configured_tpm_limit,
|
||||
call_type=call_type,
|
||||
configured_output_tokens=configured_output_tokens,
|
||||
raw_estimate=estimated_tokens,
|
||||
team_id=user_api_key_dict.team_id,
|
||||
)
|
||||
if stash.ptu_ceiling is not None and ptu_estimated_tokens > stash.ptu_ceiling.tpm_limit:
|
||||
if not reuses_request_slot:
|
||||
|
|
|
|||
|
|
@ -10364,28 +10364,38 @@ class Router:
|
|||
model_info[field] = backend_value
|
||||
|
||||
@staticmethod
|
||||
def _cost_map_backend_model(deployment: Deployment) -> str:
|
||||
model_info_base_model: Final = deployment.model_info.base_model
|
||||
def _cost_map_backend_model(deployment: Deployment | DeploymentTypedDict) -> str:
|
||||
model_info_base_model: Final = (
|
||||
deployment.model_info.base_model
|
||||
if isinstance(deployment, Deployment)
|
||||
else _configured_model_info(deployment).get("base_model")
|
||||
)
|
||||
if isinstance(model_info_base_model, str) and model_info_base_model:
|
||||
return model_info_base_model
|
||||
params_base_model: Final = deployment.litellm_params.get("base_model")
|
||||
params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(
|
||||
deployment.litellm_params.model_dump(exclude_none=True)
|
||||
if isinstance(deployment, Deployment)
|
||||
else deployment["litellm_params"]
|
||||
)
|
||||
params_base_model: Final = params.get("base_model")
|
||||
if isinstance(params_base_model, str) and params_base_model:
|
||||
return params_base_model
|
||||
params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment.litellm_params.model_dump(exclude_none=True))
|
||||
model: Final = (
|
||||
deployment.litellm_params.model
|
||||
if isinstance(deployment, Deployment)
|
||||
else deployment["litellm_params"].get("model", "")
|
||||
)
|
||||
invoked_model: Final = params.get("model_id")
|
||||
if (
|
||||
isinstance(invoked_model, str)
|
||||
and invoked_model
|
||||
and (
|
||||
deployment.litellm_params.custom_llm_provider == "bedrock"
|
||||
or deployment.litellm_params.model.startswith("bedrock/")
|
||||
)
|
||||
and (params.get("custom_llm_provider") == "bedrock" or model.startswith("bedrock/"))
|
||||
):
|
||||
profile_model: Final = invoked_model if invoked_model.startswith("bedrock/") else f"bedrock/{invoked_model}"
|
||||
profile_info: Final = lookup_model_info_for_registration(model=profile_model, custom_llm_provider="bedrock")
|
||||
if not isinstance(profile_info, ModelInfoLookupMiss):
|
||||
return profile_model
|
||||
return deployment.litellm_params.model
|
||||
return model
|
||||
|
||||
@staticmethod
|
||||
def _inherit_builtin_service_tier_pricing(
|
||||
|
|
@ -13945,7 +13955,9 @@ class Router:
|
|||
]
|
||||
return named or self._get_wildcard_deployments(model_name=model_group)
|
||||
|
||||
def resolved_litellm_models(self, model_name: str, team_id: str | None = None) -> tuple[str, ...]:
|
||||
def resolved_litellm_models(
|
||||
self, model_name: str, team_id: str | None = None, *, use_base_model: bool = False
|
||||
) -> tuple[str, ...]:
|
||||
"""The provider model strings `model_name` can actually be served by on this proxy.
|
||||
|
||||
`get_model_list` composes every channel the request path itself uses (exact name,
|
||||
|
|
@ -13958,6 +13970,12 @@ class Router:
|
|||
deployment behind it - so the fallback for an empty result is the caller's policy,
|
||||
never this function's.
|
||||
"""
|
||||
if use_base_model:
|
||||
return tuple(
|
||||
backend_model
|
||||
for deployment in self.get_model_list(model_name=model_name, team_id=team_id) or ()
|
||||
if (backend_model := self._cost_map_backend_model(deployment))
|
||||
)
|
||||
return tuple(
|
||||
litellm_model
|
||||
for deployment in self.get_model_list(model_name=model_name, team_id=team_id) or ()
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
schema_has_unresolvable_ref,
|
||||
schema_property_count,
|
||||
split_concatenated_json_objects,
|
||||
split_media_for_tpm,
|
||||
strip_encrypted_reasoning_from_messages,
|
||||
strip_litellm_response_only_fields,
|
||||
system_messages_first,
|
||||
|
|
@ -2512,3 +2513,118 @@ def test_shared_image_url_agrees_with_responses_normalization(
|
|||
expected: Final = {"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE"}}
|
||||
assert as_openai_image_part(part) == expected
|
||||
assert as_openai_image_part(normalized) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"part,projected,media",
|
||||
[
|
||||
(
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "high"}},
|
||||
{"type": "image_url", "image_url": {"url": "", "detail": "high"}},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "high"}},
|
||||
),
|
||||
(
|
||||
{"type": "input_image", "image_url": "https://example.com/image.png", "detail": "low"},
|
||||
{"type": "input_image", "image_url": "", "detail": "low"},
|
||||
{"type": "image_url", "image_url": {"url": "https://example.com/image.png", "detail": "low"}},
|
||||
),
|
||||
(
|
||||
{"type": "computer_screenshot", "image_url": "data:image/png;base64,IMAGE"},
|
||||
{"type": "computer_screenshot", "image_url": ""},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "auto"}},
|
||||
),
|
||||
(
|
||||
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "IMAGE"}},
|
||||
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": ""}},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "auto"}},
|
||||
),
|
||||
(
|
||||
{"type": "input_audio", "input_audio": {"data": "AUDIO", "format": "wav"}},
|
||||
{"type": "input_audio", "input_audio": {"data": "", "format": "wav"}},
|
||||
{"type": "input_audio", "input_audio": {"data": "AUDIO", "format": "wav"}},
|
||||
),
|
||||
(
|
||||
{"type": "file", "file": {"file_data": "IMAGE", "format": "image/png", "filename": "image.png"}},
|
||||
{"type": "file", "file": {"file_data": "", "format": "image/png", "filename": "image.png"}},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "auto"}},
|
||||
),
|
||||
(
|
||||
{"type": "input_file", "file_data": "data:audio/wav;base64,AUDIO", "filename": "audio.wav"},
|
||||
{"type": "input_file", "file_data": "", "filename": "audio.wav"},
|
||||
{"type": "input_audio", "input_audio": {"data": "AUDIO", "format": "wav"}},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_tpm_projection_blanks_only_declared_media_transport(
|
||||
part: Mapping[str, object], projected: Mapping[str, object], media: Mapping[str, object],
|
||||
) -> None:
|
||||
text: Final = {"type": "text", "text": "data:image/png;base64,IMAGE"}
|
||||
request: Final = [{"role": "user", "content": [text, {**part, "cache_control": {"type": "ephemeral"}}]}]
|
||||
original: Final = copy.deepcopy(request)
|
||||
assert split_media_for_tpm(request) == (
|
||||
[{"role": "user", "content": [text, {**projected, "cache_control": {"type": "ephemeral"}}]}],
|
||||
(media,),
|
||||
)
|
||||
assert request == original
|
||||
|
||||
|
||||
def test_tpm_projection_retains_nested_and_native_history_around_media() -> None:
|
||||
image: Final = {"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "high"}}
|
||||
empty_image: Final = {"type": "image_url", "image_url": {"url": "", "detail": "high"}}
|
||||
document: Final = {"type": "document", "title": "keep", "source": {"type": "content", "content": [image]}}
|
||||
empty_document: Final = {**document, "source": {"type": "content", "content": [empty_image]}}
|
||||
tool: Final = {"type": "tool_result", "tool_use_id": "call", "content": [document, {"type": "text", "text": "keep"}]}
|
||||
empty_tool: Final = {**tool, "content": [empty_document, {"type": "text", "text": "keep"}]}
|
||||
request: Final = [
|
||||
{"role": "user", "content": [tool]},
|
||||
{"type": "function_call_output", "call_id": "call", "output": [image]},
|
||||
{"type": "custom_tool_call_output", "call_id": "custom", "output": [image]},
|
||||
{"type": "computer_call_output", "call_id": "computer", "output": {
|
||||
"type": "computer_screenshot", "image_url": "data:image/png;base64,IMAGE",
|
||||
}},
|
||||
{"type": "image_generation_call", "id": "ig", "result": "IMAGE", "status": "completed"},
|
||||
{"type": "code_interpreter_call", "id": "ci", "code": "keep()", "outputs": [
|
||||
{"type": "logs", "logs": "keep"}, {"type": "image", "url": "https://example.com/image.png"},
|
||||
]},
|
||||
]
|
||||
original: Final = copy.deepcopy(request)
|
||||
projected, media = split_media_for_tpm(request)
|
||||
assert projected == [
|
||||
{"role": "user", "content": [empty_tool]},
|
||||
{"type": "function_call_output", "call_id": "call", "output": [empty_image]},
|
||||
{"type": "custom_tool_call_output", "call_id": "custom", "output": [empty_image]},
|
||||
{"type": "computer_call_output", "call_id": "computer", "output": {
|
||||
"type": "computer_screenshot", "image_url": "",
|
||||
}},
|
||||
{"type": "image_generation_call", "id": "ig", "result": "", "status": "completed"},
|
||||
{"type": "code_interpreter_call", "id": "ci", "code": "keep()", "outputs": [
|
||||
{"type": "logs", "logs": "keep"}, {"type": "image", "url": ""},
|
||||
]},
|
||||
]
|
||||
assert media == (
|
||||
image, image, image,
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "auto"}},
|
||||
{"type": "image_url", "image_url": {"url": "data:application/octet-stream;base64,IMAGE", "detail": "auto"}},
|
||||
{"type": "image_url", "image_url": {"url": "https://example.com/image.png", "detail": "auto"}},
|
||||
)
|
||||
assert request == original
|
||||
|
||||
|
||||
def test_tpm_projection_preserves_non_media_identity_and_opaque_values() -> None:
|
||||
image: Final = {"type": "input_image", "image_url": "data:image/png;base64,OPAQUE"}
|
||||
request: Final = [
|
||||
{"role": "user", "content": json.dumps(image)},
|
||||
{"type": "custom_tool_call_output", "call_id": "call", "output": json.dumps(image)},
|
||||
{"type": "function_call", "call_id": "call", "name": "keep", "arguments": {"content": [image]}},
|
||||
{"type": "file_search_call", "results": [{"text": "keep", "content": [image]}]},
|
||||
{"type": "image_generation_call", "id": "ig", "status": "in_progress", "result": None},
|
||||
{"role": "user", "content": [
|
||||
{"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "PDF"}},
|
||||
{"type": "document", "source": {"type": "text", "data": json.dumps(image)}},
|
||||
{"type": "file", "file": {"format": "application/pdf", "file_data": "PDF"}},
|
||||
{"type": "text", "text": "keep", "metadata": {"content": [image]}},
|
||||
]},
|
||||
]
|
||||
projected, media = split_media_for_tpm(request)
|
||||
assert projected is request
|
||||
assert media == ()
|
||||
|
|
|
|||
|
|
@ -2,16 +2,23 @@
|
|||
Test the audio utils functionality in litellm_core_utils/audio_utils/utils.py
|
||||
"""
|
||||
|
||||
import base64
|
||||
import io
|
||||
import os
|
||||
import tempfile
|
||||
from typing import Final
|
||||
from unittest.mock import mock_open, patch
|
||||
|
||||
import pytest
|
||||
import soundfile as sf
|
||||
|
||||
from litellm.litellm_core_utils.audio_utils.utils import (
|
||||
AUDIO_BYTES_PER_TOKEN,
|
||||
AUDIO_TOKENS_PER_SECOND,
|
||||
DEFAULT_AUDIO_TOKEN_ESTIMATE,
|
||||
ProcessedAudioFile,
|
||||
calculate_request_duration,
|
||||
estimate_audio_tokens,
|
||||
get_audio_file_content_hash,
|
||||
get_audio_file_for_health_check,
|
||||
get_audio_file_name,
|
||||
|
|
@ -19,6 +26,63 @@ from litellm.litellm_core_utils.audio_utils.utils import (
|
|||
)
|
||||
|
||||
|
||||
def _encoded_audio(seconds: int, sample_rate: int, channels: int, container: str) -> str:
|
||||
buffer: Final = io.BytesIO()
|
||||
with sf.SoundFile(buffer, "w", samplerate=sample_rate, channels=channels, format=container) as audio:
|
||||
audio.buffer_write(b"\0" * (seconds * sample_rate * channels * 2), dtype="int16")
|
||||
return base64.b64encode(buffer.getvalue()).decode()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("sample_rate", "channels", "container"), [(8000, 1, "WAV"), (48000, 2, "WAV"), (24000, 1, "FLAC")]
|
||||
)
|
||||
def test_audio_estimate_uses_duration_independent_of_sample_rate_channels_and_compression(
|
||||
sample_rate: int, channels: int, container: str
|
||||
) -> None:
|
||||
assert estimate_audio_tokens(_encoded_audio(1, sample_rate, channels, container)) == AUDIO_TOKENS_PER_SECOND
|
||||
assert estimate_audio_tokens(_encoded_audio(31, sample_rate, channels, container)) == 31 * AUDIO_TOKENS_PER_SECOND
|
||||
|
||||
|
||||
# Audio rates verified 2026-10-28: https://ai.google.dev/gemini-api/docs/generate-content/media-resolution and
|
||||
# https://developers.openai.com/api/docs/guides/voice-latency-cost
|
||||
@pytest.mark.parametrize(
|
||||
("models", "fallback", "expected"),
|
||||
[
|
||||
(("openai/gpt-audio-1.5",), "gemini-3.8-flash", 10),
|
||||
(("gpt-audio-1.5", "gemini/gemini-3.8-flash"), None, 25),
|
||||
(("vertex_ai/gemini-3.8-flash",), None, 25),
|
||||
(("openrouter/google/gemini-3.8-flash",), None, 25),
|
||||
(("azure/audio-deployment",), "gpt-audio-1.5", 10),
|
||||
(("opaque-deployment",), "gemini-3.8-flash", 25),
|
||||
((), "gemini-3.8-flash", 25),
|
||||
(("gemini/gemini-2.5-flash", "gemini/gemini-3-pro-preview"), None, 32),
|
||||
(("other/gemini-3.8-flash",), None, 10),
|
||||
],
|
||||
)
|
||||
def test_audio_estimate_prefers_each_backend_rate_over_the_requested_model(
|
||||
models: tuple[str, ...], fallback: str | None, expected: int
|
||||
) -> None:
|
||||
assert estimate_audio_tokens(_encoded_audio(1, 8000, 1, "WAV"), models, fallback) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("data", [None, "", "invalid-base64!", base64.b64encode(b"unreadable audio").decode()])
|
||||
def test_audio_estimate_retains_nonzero_fallback_for_unreadable_audio(data: str | None) -> None:
|
||||
assert estimate_audio_tokens(data) == DEFAULT_AUDIO_TOKEN_ESTIMATE
|
||||
|
||||
|
||||
def test_audio_estimate_retains_size_based_fallback_for_large_unreadable_audio() -> None:
|
||||
data: Final = base64.b64encode(b"\0" * (AUDIO_BYTES_PER_TOKEN * 400)).decode()
|
||||
assert estimate_audio_tokens(data) == 400
|
||||
|
||||
|
||||
def test_audio_estimate_bounds_base64_allocation_before_decoding() -> None:
|
||||
data: Final = base64.b64encode(b"\0" * (16 * 1024 * 1024 + 1)).decode()
|
||||
with patch("base64.b64decode", side_effect=AssertionError("oversized audio must not be decoded")) as decode:
|
||||
estimate: Final = estimate_audio_tokens(data)
|
||||
assert estimate == len(data) * 3 // 4 // AUDIO_BYTES_PER_TOKEN
|
||||
decode.assert_not_called()
|
||||
|
||||
|
||||
class TestProcessAudioFile:
|
||||
"""Test the process_audio_file function with various input types"""
|
||||
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ from litellm.litellm_core_utils.token_counter import (
|
|||
_get_extrapolating_count_function,
|
||||
_get_tiktoken_count_function,
|
||||
calculate_img_tokens,
|
||||
estimate_image_tokens,
|
||||
high_detail_image_token_upper_bound,
|
||||
offload_token_count,
|
||||
)
|
||||
|
|
@ -2152,6 +2153,136 @@ def test_high_detail_image_token_upper_bound_is_reached_by_the_largest_high_res_
|
|||
assert calculate_img_tokens(_png_data_url(1, 1), mode="high") < high_detail_image_token_upper_bound()
|
||||
|
||||
|
||||
# Image sizing verified 2026-10-28: https://developers.openai.com/api/docs/guides/images-vision
|
||||
# and https://platform.claude.com/docs/en/build-with-claude/vision-coordinates
|
||||
@pytest.mark.parametrize(
|
||||
("model", "width", "height", "detail", "expected"),
|
||||
[
|
||||
("gpt-4o-mini", 1024, 1024, "low", 2833),
|
||||
("gpt-4o-mini", 1024, 1024, "high", 25501),
|
||||
("gpt-5.1", 512, 512, "auto", 210),
|
||||
("gpt-6-astra", 2048, 2048, "high", 3000),
|
||||
("gpt-6-astra", 4096, 512, "high", 2458),
|
||||
("gpt-6-astra", 2048, 2048, "auto", 4916),
|
||||
("gpt-5.6-sol", 2048, 2048, "low", 308),
|
||||
("gpt-5.5", 6400, 6400, "original", 12000),
|
||||
("gpt-5.4-mini", 2048, 2048, "low", 4916),
|
||||
("gpt-5.4-nano", 2048, 2048, "auto", 3000),
|
||||
("gpt-5.2", 4096, 512, "auto", 615),
|
||||
("gpt-4.1-mini-2025-04-14", 1024, 1024, "low", 1659),
|
||||
("claude-haiku-5-5", 200, 200, "auto", 64),
|
||||
("anthropic/claude-haiku-5-5", 1024, 1024, "auto", 1369),
|
||||
("claude-haiku-5-5", 1920, 1080, "low", 2691),
|
||||
("claude-haiku-4-5", 1920, 1080, "auto", 1560),
|
||||
("claude-haiku-4-5", 1075, 1520, "auto", 1551),
|
||||
("claude-haiku-5-5", 3840, 2160, "auto", 4784),
|
||||
("claude-haiku-5-5", 8000, 100, "auto", 184),
|
||||
],
|
||||
)
|
||||
def test_image_estimate_uses_model_resolution_and_detail(
|
||||
model: str, width: int, height: int, detail: str, expected: int
|
||||
) -> None:
|
||||
assert estimate_image_tokens(model, _png_data_url(width, height), detail) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("models", "fallback", "expected_model"),
|
||||
[
|
||||
(("azure/deployment-prod",), "gpt-4o", "gpt-4o"),
|
||||
(("openai/gpt-4o",), "gpt-4o-mini", "gpt-4o"),
|
||||
(("azure/deployment-prod", "openai/gpt-4o"), "gpt-4o-mini", "gpt-4o-mini"),
|
||||
(("openai/gpt-4o", "openrouter/openai/gpt-4o-mini"), None, "gpt-4o-mini"),
|
||||
(("azure_ai/openai/gpt-6-astra",), None, "gpt-6-astra"),
|
||||
(("azure/gpt-41",), None, "gpt-4.1"),
|
||||
(("azure/eu/gpt-4o-mini",), None, "gpt-4o-mini"),
|
||||
(("unknown", "gpt-4o-mini"), None, "gpt-4o-mini"),
|
||||
((), "gpt-6-astra", "gpt-6-astra"),
|
||||
(("other/gpt-4o-mini",), None, "unregistered-model"),
|
||||
(("azure/opaque/gpt-4o-mini",), None, "unregistered-model"),
|
||||
(("bedrock/anthropic.claude-haiku-5-5",), None, "claude-haiku-5-5"),
|
||||
(("bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0",), None, "claude-haiku-4-5"),
|
||||
(("vertex_ai/claude-haiku-5-5@default",), None, "claude-haiku-5-5"),
|
||||
(("azure_ai/claude-haiku-5-5",), None, "claude-haiku-5-5"),
|
||||
(("anthropic/claude-haiku-5-5", "openai/gpt-4o"), "gpt-4o-mini", "claude-haiku-5-5"),
|
||||
(("other/claude-haiku-5-5",), None, "unregistered-model"),
|
||||
],
|
||||
)
|
||||
def test_image_estimate_prefers_each_backend_policy_over_the_requested_model(
|
||||
models: tuple[str, ...], fallback: str | None, expected_model: str
|
||||
) -> None:
|
||||
data: Final = _png_data_url(1024, 1024)
|
||||
assert estimate_image_tokens(models, data, "high", fallback_model=fallback) == estimate_image_tokens(
|
||||
expected_model, data, "high"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "detail"),
|
||||
[
|
||||
("unknown", "high"),
|
||||
("claude-unregistered", "auto"),
|
||||
("gpt-4o", "original"),
|
||||
("gpt-5.2", "original"),
|
||||
("gpt-6-astra", "invalid"),
|
||||
],
|
||||
)
|
||||
def test_image_estimate_preserves_default_when_policy_is_unavailable(model: str, detail: str) -> None:
|
||||
assert (
|
||||
estimate_image_tokens(model, _png_data_url(1024, 1024), detail) == litellm.constants.DEFAULT_IMAGE_TOKEN_COUNT
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data",
|
||||
[
|
||||
"https://image.invalid/photo.png",
|
||||
"data:image/png;base64,invalid!",
|
||||
"data:application/octet-stream;base64,bm90LWFuLWltYWdl",
|
||||
"data:image/jpeg;base64,/9j/4AAQ",
|
||||
"data:image/webp;base64,UklGRgAAAABXRUJQ",
|
||||
_png_data_url(0, 1024),
|
||||
],
|
||||
)
|
||||
# Accepted patch limits verified 2026-10-28: https://developers.openai.com/api/docs/guides/images-vision
|
||||
# and https://platform.claude.com/docs/en/build-with-claude/vision
|
||||
@pytest.mark.parametrize(
|
||||
("model", "detail", "bound"),
|
||||
[
|
||||
("gpt-4o", "high", high_detail_image_token_upper_bound()),
|
||||
("gpt-6-astra", "high", 3000),
|
||||
("gpt-6-astra", "auto", 36000),
|
||||
("claude-haiku-5-5", "auto", 4784),
|
||||
("claude-haiku-4-5", "auto", 1568),
|
||||
],
|
||||
)
|
||||
def test_image_estimate_uses_bounded_fallback_without_fetching_media(
|
||||
data: str, model: str, detail: str, bound: int
|
||||
) -> None:
|
||||
with patch("httpx.Client.send") as send:
|
||||
estimate: Final = estimate_image_tokens(model, data, detail)
|
||||
send.assert_not_called()
|
||||
assert estimate == bound
|
||||
assert estimate > estimate_image_tokens(model, _png_data_url(1024, 1024), detail)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("container", ["PNG", "JPEG", "GIF", "WEBP"])
|
||||
def test_image_estimate_depends_on_dimensions_and_not_compressed_payload_size(container: str) -> None:
|
||||
output: Final = io.BytesIO()
|
||||
Image.new("RGB", (1024, 1024)).save(output, format=container)
|
||||
data: Final = "data:application/octet-stream;base64," + base64.b64encode(output.getvalue()).decode()
|
||||
assert estimate_image_tokens("gpt-6-astra", data) == estimate_image_tokens("gpt-6-astra", _png_data_url(1024, 1024))
|
||||
|
||||
|
||||
def test_image_estimate_bounds_header_decoding_for_large_inline_payloads() -> None:
|
||||
header: Final = base64.b64decode(_png_data_url(1024, 1024).split(",", 1)[1])
|
||||
data: Final = "data:image/png;base64," + base64.b64encode(header + b"\0" * 4_000_000).decode()
|
||||
with patch("base64.b64decode", wraps=base64.b64decode) as decode_header:
|
||||
estimate: Final = estimate_image_tokens("gpt-6-astra", data)
|
||||
assert estimate == estimate_image_tokens("gpt-6-astra", _png_data_url(1024, 1024))
|
||||
assert decode_header.call_count == 1
|
||||
assert len(decode_header.call_args.args[0]) <= 65_536
|
||||
|
||||
|
||||
HUB_SETUP_SCRIPT: Final = """
|
||||
import json
|
||||
import sys
|
||||
|
|
|
|||
|
|
@ -15,28 +15,38 @@ Redis.
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import copy
|
||||
import io
|
||||
import json
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Any, Dict, Final
|
||||
|
||||
import pytest
|
||||
import soundfile as sf
|
||||
from fastapi import HTTPException
|
||||
from PIL import Image
|
||||
|
||||
from litellm import token_counter
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.litellm_core_utils.token_counter import UnsupportedContentTypeError, estimate_image_tokens
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
|
||||
_AUDIO_BYTES_PER_TOKEN,
|
||||
DEFAULT_AUDIO_TOKEN_ESTIMATE,
|
||||
PROJECT_ITPM_DESCRIPTOR_KEY,
|
||||
PROJECT_OTPM_DESCRIPTOR_KEY,
|
||||
RateLimitedModel,
|
||||
_AUDIO_BYTES_PER_TOKEN,
|
||||
PROXY_MaxParallelRequestsHandler_v3 as RateLimitHandler,
|
||||
)
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
|
||||
_call_id_from_callback_kwargs,
|
||||
_request_stash,
|
||||
get_or_create_request_stash,
|
||||
get_request_stash,
|
||||
)
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
|
||||
PROXY_MaxParallelRequestsHandler_v3 as RateLimitHandler,
|
||||
)
|
||||
from litellm.proxy.utils import InternalUsageCache, hash_token
|
||||
from litellm.router_utils.ptu_shares import PTUTeamCeiling
|
||||
from litellm.types.llms.openai import (
|
||||
InputTokensDetails,
|
||||
ResponseAPIUsage,
|
||||
|
|
@ -3775,3 +3785,284 @@ def test_reservation_counts_a_compaction_item_by_its_summary(rate_limiter):
|
|||
) == handler._estimate_input_and_output_tokens(
|
||||
data={"model": "gpt-4o", "input": _COMPACTED_INPUT[2:], "max_output_tokens": 1}, call_type="aresponses"
|
||||
)
|
||||
|
||||
|
||||
_QUOTA_TEXT: Final = "a" * 400
|
||||
_QUOTA_JSON_TEXT: Final = '{"type":"input_image","image_url":"data:image/png;base64,AAAA"}'
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("call_type", "request_data", "expected"),
|
||||
[
|
||||
("acompletion", {"messages": [{"role": "user", "content": _QUOTA_TEXT}], "max_tokens": 16}, (100, 16)),
|
||||
("acompletion", {"messages": [{"role": "user", "content": _QUOTA_JSON_TEXT}], "max_tokens": 16}, (15, 16)),
|
||||
(
|
||||
"anthropic_messages",
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "t",
|
||||
"content": [{"type": "text", "text": _QUOTA_TEXT}],
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
"max_tokens": 16,
|
||||
},
|
||||
(0, 16),
|
||||
),
|
||||
(
|
||||
"acompletion",
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{"id": "t", "type": "function", "function": {"name": "f", "arguments": _QUOTA_TEXT}}
|
||||
],
|
||||
}
|
||||
],
|
||||
"max_tokens": 16,
|
||||
},
|
||||
(0, 16),
|
||||
),
|
||||
("aresponses", {"input": _QUOTA_TEXT, "max_output_tokens": 16}, (100, 16)),
|
||||
("aresponses", {"input": _QUOTA_TEXT, "instructions": "b" * 400, "max_output_tokens": 16}, (100, 16)),
|
||||
(
|
||||
"aresponses",
|
||||
{"input": [{"role": "user", "content": _QUOTA_TEXT}], "instructions": "b" * 400, "max_output_tokens": 16},
|
||||
(107, 16),
|
||||
),
|
||||
(
|
||||
"aresponses",
|
||||
{
|
||||
"input": [
|
||||
{
|
||||
"type": "file_search_call",
|
||||
"id": "f",
|
||||
"status": "completed",
|
||||
"queries": ["q"],
|
||||
"results": [
|
||||
{
|
||||
"file_id": "file_1",
|
||||
"filename": "result.txt",
|
||||
"attributes": {},
|
||||
"score": 1,
|
||||
"text": _QUOTA_TEXT,
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
"max_output_tokens": 16,
|
||||
},
|
||||
(146, 16),
|
||||
),
|
||||
(
|
||||
"aresponses",
|
||||
{
|
||||
"input": [{"type": "custom_tool_call_output", "call_id": "t", "output": _QUOTA_JSON_TEXT}],
|
||||
"max_output_tokens": 16,
|
||||
},
|
||||
(32, 16),
|
||||
),
|
||||
("aresponses", {"input": "", "instructions": "b" * 400}, (0, 0)),
|
||||
("aembedding", {"input": [_QUOTA_TEXT]}, (100, 0)),
|
||||
("atext_completion", {"prompt": [_QUOTA_TEXT], "max_tokens": 16}, (100, 16)),
|
||||
],
|
||||
)
|
||||
def test_media_correction_preserves_exact_non_media_reservations(
|
||||
rate_limiter: tuple[RateLimitHandler, DualCache],
|
||||
call_type: str,
|
||||
request_data: dict[str, object],
|
||||
expected: tuple[int, int],
|
||||
) -> None:
|
||||
handler, _ = rate_limiter
|
||||
assert handler._estimate_input_and_output_tokens(request_data, call_type=call_type) == expected
|
||||
|
||||
|
||||
def _quota_media_payload(kind: str, compressed: bool) -> str:
|
||||
output: Final = io.BytesIO()
|
||||
if kind == "image":
|
||||
Image.new("RGB", (1024, 1024)).save(output, format="PNG", compress_level=9 if compressed else 0)
|
||||
else:
|
||||
sample_rate: Final = 8000 if compressed else 48000
|
||||
with sf.SoundFile(output, "w", samplerate=sample_rate, channels=1, format="WAV") as audio:
|
||||
audio.buffer_write(b"\0" * (sample_rate * 2), dtype="int16")
|
||||
return base64.b64encode(output.getvalue()).decode()
|
||||
|
||||
|
||||
def _quota_media_request(kind: str, payload: str, call_type: str) -> dict[str, object]:
|
||||
block: Final = (
|
||||
{"type": "input_audio", "input_audio": {"data": payload, "format": "wav"}}
|
||||
if kind == "audio"
|
||||
else {"type": "input_image", "image_url": "data:image/png;base64," + payload}
|
||||
if call_type == "aresponses"
|
||||
else {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": payload}}
|
||||
if call_type == "anthropic_messages"
|
||||
else {"type": "image_url", "image_url": {"url": "data:image/png;base64," + payload}}
|
||||
)
|
||||
content: Final = (
|
||||
[{"type": "tool_result", "tool_use_id": "t", "content": [block]}]
|
||||
if call_type == "anthropic_messages"
|
||||
else [block]
|
||||
)
|
||||
field: Final = "input" if call_type == "aresponses" else "messages"
|
||||
return {
|
||||
"model": "gpt-6-astra",
|
||||
field: [{"role": "user", "content": content}],
|
||||
"max_tokens": 16,
|
||||
"max_output_tokens": 16,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("kind", "call_type"),
|
||||
[
|
||||
("image", "acompletion"),
|
||||
("image", "aresponses"),
|
||||
("image", "anthropic_messages"),
|
||||
("audio", "acompletion"),
|
||||
("audio", "aresponses"),
|
||||
],
|
||||
)
|
||||
def test_media_reservation_is_independent_of_encoding_and_preserves_request(
|
||||
rate_limiter: tuple[RateLimitHandler, DualCache], kind: str, call_type: str
|
||||
) -> None:
|
||||
handler, _ = rate_limiter
|
||||
compressed: Final = _quota_media_request(kind, _quota_media_payload(kind, True), call_type)
|
||||
uncompressed: Final = _quota_media_request(kind, _quota_media_payload(kind, False), call_type)
|
||||
original: Final = copy.deepcopy(uncompressed)
|
||||
estimate: Final = handler._estimate_input_and_output_tokens(uncompressed, call_type=call_type)
|
||||
assert estimate == handler._estimate_input_and_output_tokens(compressed, call_type=call_type)
|
||||
assert estimate[0] > 0 and estimate[1] == 16
|
||||
assert uncompressed == original
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("call_type", ["acompletion", "aresponses"])
|
||||
@pytest.mark.parametrize("limit", [1000, 20000])
|
||||
async def test_media_pre_call_rejects_tight_quota_and_reserves_sufficient_quota(
|
||||
rate_limiter: tuple[RateLimitHandler, DualCache], call_type: str, limit: int
|
||||
) -> None:
|
||||
handler, cache = rate_limiter
|
||||
request: Final = {
|
||||
key: value
|
||||
for key, value in _quota_media_request("image", _quota_media_payload("image", True), call_type).items()
|
||||
if key not in ("max_tokens", "max_output_tokens")
|
||||
}
|
||||
auth: Final = UserAPIKeyAuth(api_key=hash_token("sk-media-quota"), tpm_limit=limit)
|
||||
if limit == 1000:
|
||||
with pytest.raises(HTTPException) as error:
|
||||
await handler.async_pre_call_hook(user_api_key_dict=auth, cache=cache, data=request, call_type=call_type)
|
||||
assert error.value.status_code == 429
|
||||
return
|
||||
await handler.async_pre_call_hook(user_api_key_dict=auth, cache=cache, data=request, call_type=call_type)
|
||||
stash: Final = get_request_stash()
|
||||
assert stash is not None
|
||||
expected: Final = handler._estimate_input_and_output_tokens(
|
||||
request, call_type=call_type, min_configured_tpm_limit=limit
|
||||
)
|
||||
assert stash.reserved_tokens == sum(expected)
|
||||
assert expected[0] > 1000 and expected[1] >= expected[0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ptu_pre_call_weights_media_input_and_output() -> None:
|
||||
cache: Final = DualCache()
|
||||
ceiling: Final = PTUTeamCeiling(
|
||||
model_group="gpt-6-astra", tpm_limit=10000, output_to_input_ratio=4, cached_input_ratio=0
|
||||
)
|
||||
|
||||
def resolve_ceiling(team: str, model: str) -> PTUTeamCeiling | None:
|
||||
return ceiling if (team, model) == ("test-team", "gpt-6-astra") else None
|
||||
|
||||
handler: Final = RateLimitHandler(
|
||||
internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve_ceiling
|
||||
)
|
||||
request: Final = _quota_media_request("image", _quota_media_payload("image", True), "acompletion")
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-media-ptu"), team_id="test-team"),
|
||||
cache=cache,
|
||||
data=request,
|
||||
call_type="acompletion",
|
||||
)
|
||||
stash: Final = get_request_stash()
|
||||
assert stash is not None
|
||||
assert (
|
||||
stash.ptu_reserved_tokens
|
||||
== estimate_image_tokens("gpt-6-astra", "data:image/png;base64," + _quota_media_payload("image", True)) + 4 * 16
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_project_fallback_keeps_legacy_audio_estimate_and_output_policy(
|
||||
rate_limiter: tuple[RateLimitHandler, DualCache],
|
||||
) -> None:
|
||||
handler, cache = rate_limiter
|
||||
unsupported: Final = {"type": "unsupported_content"}
|
||||
with pytest.raises(UnsupportedContentTypeError):
|
||||
token_counter(model="gpt-6-astra", messages=[{"role": "user", "content": [unsupported]}])
|
||||
request: Final = {
|
||||
"model": "gpt-6-astra",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
unsupported,
|
||||
{
|
||||
"type": "input_audio",
|
||||
"input_audio": {"data": _quota_media_payload("audio", True), "format": "wav"},
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
auth: Final = UserAPIKeyAuth(
|
||||
api_key=hash_token("sk-media-project"),
|
||||
project_id="media-project",
|
||||
project_metadata={"model_itpm_limit": {"gpt-6-astra": 1000}, "model_otpm_limit": {"gpt-6-astra": 1000}},
|
||||
)
|
||||
await handler.async_pre_call_hook(user_api_key_dict=auth, cache=cache, data=request, call_type="acompletion")
|
||||
stash: Final = get_request_stash()
|
||||
assert stash is not None
|
||||
assert stash.itpm_reserved_tokens == DEFAULT_AUDIO_TOKEN_ESTIMATE
|
||||
assert stash.otpm_reserved_tokens == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_single_anthropic_image_exceeds_tight_quota_with_explicit_output_cap(
|
||||
rate_limiter: tuple[RateLimitHandler, DualCache],
|
||||
) -> None:
|
||||
handler, cache = rate_limiter
|
||||
request: Final = {
|
||||
"model": "anthropic/claude-haiku-5-5",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image",
|
||||
"source": {
|
||||
"type": "base64",
|
||||
"media_type": "image/png",
|
||||
"data": _quota_media_payload("image", True),
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
"max_tokens": 64,
|
||||
}
|
||||
with pytest.raises(HTTPException) as error:
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-claude-media-quota"), tpm_limit=500),
|
||||
cache=cache,
|
||||
data=request,
|
||||
call_type="anthropic_messages",
|
||||
)
|
||||
assert error.value.status_code == 429
|
||||
|
|
|
|||
|
|
@ -20630,6 +20630,37 @@ def test_resolved_litellm_models_answers_through_every_channel_a_request_uses(
|
|||
assert set(_resolution_router().resolved_litellm_models(model_name)) == set(expected)
|
||||
|
||||
|
||||
def test_resolved_litellm_models_base_model_precedence_is_opt_in() -> None:
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "pooled",
|
||||
"litellm_params": {"model": "openai/deployment-a", "api_key": "sk-test", "base_model": "gpt-5.5"},
|
||||
"model_info": {"base_model": "gpt-6-astra"},
|
||||
},
|
||||
{
|
||||
"model_name": "pooled",
|
||||
"litellm_params": {"model": "openai/deployment-b", "api_key": "sk-test", "base_model": "gpt-5.6-sol"},
|
||||
},
|
||||
{
|
||||
"model_name": "team-deployment",
|
||||
"litellm_params": {"model": "openai/deployment-c", "api_key": "sk-test"},
|
||||
"model_info": {
|
||||
"base_model": "gpt-6-astra",
|
||||
"team_id": "test-team",
|
||||
"team_public_model_name": "team-model",
|
||||
},
|
||||
},
|
||||
],
|
||||
model_group_alias={"nickname": "pooled"},
|
||||
)
|
||||
assert router.resolved_litellm_models("nickname") == ("openai/deployment-a", "openai/deployment-b")
|
||||
assert router.resolved_litellm_models("nickname", use_base_model=True) == ("gpt-6-astra", "gpt-5.6-sol")
|
||||
assert router.resolved_litellm_models("team-model", "test-team") == ("openai/deployment-c",)
|
||||
assert router.resolved_litellm_models("team-model", "test-team", use_base_model=True) == ("gpt-6-astra",)
|
||||
assert router.resolved_litellm_models("never-configured", use_base_model=True) == ()
|
||||
|
||||
|
||||
class TestTierParamsTheTargetAccepts:
|
||||
"""A tier's litellm_params are applied to every request that tier routes, so one the target
|
||||
cannot take raised UnsupportedParamsError before the request left the proxy, turning the whole
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue