scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

Merge ae4f9e72ac into f32ac94ad1

Browse source
This commit is contained in:
quietjules 2026-10-30 06:33:33 +00:00 • committed by GitHub
commit 7eb56f9cf9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 983 additions and 231 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/google_genai/adapters/transformation.py
View file

@ -249,6 +249,10 @@ def _genai_tool_call(part: object) -> ChatCompletionAssistantToolCall | None:
)
def _system_texts(parts: object) -> tuple[str, ...]:
return tuple(part["text"] for part in _genai_message_parts(_genai_content_parts(parts)) if part["type"] == "text")
def _genai_content_messages(content: object) -> tuple[AllMessageValues, ...]:
fields: Final = _validated(_CONFIG_FIELDS, content)
if fields is None:
@ -632,11 +636,9 @@ class GoogleGenAIAdapter:
system_instruction: _GenAISystemInstruction | None = None,
) -> list[AllMessageValues]:
"""Transform Google GenAI contents to OpenAI messages format"""
system_parts: Final = system_instruction.get("parts", ()) if system_instruction else ()
system_texts: Final = _system_texts(system_instruction.get("parts") if system_instruction else None)
system_messages: Final = (
(ChatCompletionSystemMessage(role="system", content=system_parts[0]["text"]),)
if system_parts and "text" in system_parts[0]
else ()
(ChatCompletionSystemMessage(role="system", content="\n".join(system_texts)),) if system_texts else ()
)
return [*system_messages, *chain.from_iterable(_genai_content_messages(content) for content in contents)]

25
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -35426,6 +35426,11 @@
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Vertex Discovery Proxy Route",
"tags": [
"llm_passthrough"
@ -35465,6 +35470,11 @@
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Vertex Discovery Proxy Route",
"tags": [
"llm_passthrough"
@ -35504,6 +35514,11 @@
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Vertex Discovery Proxy Route",
"tags": [
"llm_passthrough"
@ -35543,6 +35558,11 @@
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Vertex Discovery Proxy Route",
"tags": [
"llm_passthrough"
@ -35582,6 +35602,11 @@
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Vertex Discovery Proxy Route",
"tags": [
"llm_passthrough"

18
litellm/proxy/pass_through_endpoints/common_utils.py
View file

@ -1,18 +0,0 @@
from typing import Final
from fastapi import Request
def get_litellm_virtual_key(request: Request) -> str:
"""
Extract and format API key from request headers.
Prioritizes x-litellm-api-key over Authorization header.
Vertex JS SDK uses `Authorization` header, we use `x-litellm-api-key` to pass litellm virtual key
"""
litellm_api_key: Final = request.headers.get("x-litellm-api-key")
if litellm_api_key:
return f"Bearer {litellm_api_key}"
return request.headers.get("Authorization", "")

35
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -23,7 +23,7 @@ from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast
import httpx
import openai
from fastapi import APIRouter, Depends, HTTPException, Request, Response, WebSocket
from fastapi import APIRouter, Depends, HTTPException, Request, Response, Security, WebSocket
from fastapi.responses import StreamingResponse
from pydantic import ConfigDict, TypeAdapter, ValidationError
from starlette.websockets import WebSocketState
@ -79,6 +79,8 @@ from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports
_get_bearer_token, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
api_key_header,
custom_litellm_key_header,
get_bearer_token,
is_no_auth_dev_mode,
user_api_key_auth,
@ -113,7 +115,6 @@ from litellm.proxy.common_utils.token_count_log import (
logged_provider_token_count,
rejected_token_count,
)
from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key
from litellm.proxy.pass_through_endpoints.injected_stream_usage import (
stream_usage_relay,
without_injected_usage_frames,
@ -3053,7 +3054,7 @@ async def _base_vertex_proxy_route(
fastapi_response: Response,
get_vertex_pass_through_handler: BaseVertexAIPassThroughHandler,
general_settings: Mapping[str, object],
user_api_key_dict: UserAPIKeyAuth | None = None,
user_api_key_dict: UserAPIKeyAuth,
router_credentials: LiteLLM_ManagedVectorStore | None = None,
):
"""
@ -3069,7 +3070,7 @@ async def _base_vertex_proxy_route(
fastapi_response: FastAPI response object
get_vertex_pass_through_handler: Handler for the specific Vertex AI service
general_settings: Proxy general settings, read per request for the anthropic-beta values to drop
user_api_key_dict: User API key authentication dict
user_api_key_dict: The caller the route authenticated
router_credentials: Optional vector store credentials from registry (LiteLLM_ManagedVectorStore)
"""
from litellm.llms.vertex_ai.common_utils import (
@ -3083,18 +3084,6 @@ async def _base_vertex_proxy_route(
encoded_endpoint = httpx.URL(endpoint).path
verbose_proxy_logger.debug("requested endpoint %s", endpoint)
headers: Mapping[str, str] = {}
api_key_to_use = get_litellm_virtual_key(request=request)
user_api_key_dict = await user_api_key_auth(
request=request,
api_key=api_key_to_use,
)
if user_api_key_dict is None:
api_key_to_use = get_litellm_virtual_key(request=request)
user_api_key_dict = await user_api_key_auth(
request=request,
api_key=api_key_to_use,
)
if router_credentials is not None:
await assert_user_can_access_vector_store(
vector_store=router_credentials,
@ -3206,6 +3195,16 @@ async def _base_vertex_proxy_route(
return received_value
async def _authenticate_vertex_caller(
request: Request,
authorization: Annotated[str | None, Security(api_key_header)],
litellm_api_key: Annotated[str | None, Security(custom_litellm_key_header)],
) -> UserAPIKeyAuth:
return await user_api_key_auth(
request=request, api_key=authorization or "", custom_litellm_key_header=litellm_api_key
)
@router.api_route(
"/vertex_ai/discovery/{endpoint:path}",
methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
@ -3215,6 +3214,7 @@ async def vertex_discovery_proxy_route(
endpoint: str,
request: Request,
fastapi_response: Response,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(_authenticate_vertex_caller)],
):
"""
Call any vertex discovery endpoint using the proxy.
@ -3253,6 +3253,7 @@ async def vertex_discovery_proxy_route(
fastapi_response=fastapi_response,
get_vertex_pass_through_handler=discovery_handler,
general_settings=_proxy_general_settings(),
user_api_key_dict=user_api_key_dict,
router_credentials=vector_store_credentials,
)
@ -3272,7 +3273,7 @@ async def vertex_proxy_route(
endpoint: str,
request: Request,
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(_authenticate_vertex_caller)],
):
"""
Call LiteLLM proxy via Vertex AI SDK.

163
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -4,10 +4,10 @@ import asyncio
import json
import math
import time
from collections.abc import AsyncIterator, Mapping, Sequence
from collections.abc import AsyncIterator, Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from itertools import chain
from itertools import chain, product
from types import MappingProxyType
from typing import Final, NoReturn, SupportsFloat, SupportsIndex, SupportsInt, cast
@ -18,6 +18,7 @@ import litellm
from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import select_tier_for_input, tier_rate
from litellm.proxy._types import (
@ -54,6 +55,7 @@ from litellm.router import Router
from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget
from litellm.types.router import DeploymentTypedDict
from litellm.types.utils import CallTypes
from litellm.utils import ProviderConfigManager
@dataclass
@ -183,7 +185,7 @@ async def _apply_over_budget_reservation_policy(
entry: dict[str, float | str],
applied_entries: list[dict[str, float | str]],
reservation_cost: float,
current_spend: float,
spend_before_reservation: float,
fail_closed_budget_enforcement: bool = False,
) -> float:
"""
@ -199,13 +201,13 @@ async def _apply_over_budget_reservation_policy(
applied_entries.remove(entry)
return reservation_cost
remaining_before_reservation: Final = counter.max_budget - (current_spend - reservation_cost)
remaining_before_reservation: Final = counter.max_budget - spend_before_reservation
if remaining_before_reservation <= 1e-12:
_raise_counter_budget_exceeded(counter=counter, current_cost=current_spend)
if fail_closed_budget_enforcement and current_spend - counter.max_budget > 1e-12:
_raise_counter_budget_exceeded(counter=counter, current_cost=spend_before_reservation + reservation_cost)
if fail_closed_budget_enforcement and reservation_cost - remaining_before_reservation > 1e-12:
_raise_counter_budget_exceeded(
counter=counter,
current_cost=current_spend - reservation_cost,
current_cost=spend_before_reservation,
estimated_cost=reservation_cost,
)
await _resize_applied_reservation(
@ -996,7 +998,7 @@ async def _reserve_reservable_counters(
current_spend_by_counter_key=current_spend_by_counter_key,
reservation_cost=reservation_cost,
):
charged_cost = reservation_cost
charged_cost = min(reservation_cost, max(counter.max_budget for counter in group))
entries = tuple(_counter_to_reservation_entry(counter=counter, reserved_cost=charged_cost) for counter in group)
applied_entries.extend(entries)
reserved_values = await _reserve_counters(counters=group, entries=entries, reservation_cost=charged_cost)
@ -1009,18 +1011,19 @@ async def _reserve_reservable_counters(
for counter, entry, reserved_value in zip(group, entries, reserved_values):
if entry not in applied_entries:
continue
if reserved_value is not None:
current_spend = reserved_value - (charged_cost - reservation_cost)
else:
current_spend = current_spend_by_counter_key[counter.counter_key] + reservation_cost
if current_spend > counter.max_budget:
spend_before_reservation = (
current_spend_by_counter_key[counter.counter_key]
if reserved_value is None
else reserved_value - charged_cost
)
if spend_before_reservation + reservation_cost > counter.max_budget:
reservation_cost = await _apply_over_budget_reservation_policy(
counter=counter,
valid_token=valid_token,
entry=entry,
applied_entries=applied_entries,
reservation_cost=reservation_cost,
current_spend=current_spend,
spend_before_reservation=spend_before_reservation,
fail_closed_budget_enforcement=fail_closed_budget_enforcement,
)
return reservation_cost
@ -1401,12 +1404,16 @@ def _estimate_request_max_cost_for_model(
llm_router: Router | None,
input_tokens: int | None = None,
) -> float | None:
output_multiplier: Final = _get_output_multiplier(
_requested_output_counts(request_body=request_body, route=route, model=model, llm_router=llm_router)
)
estimates: Final = [
_max_cost_for_cost_info(
request_body=request_body,
route=route,
model=model,
model_info=model_info,
output_multiplier=output_multiplier,
input_tokens=input_tokens,
)
for model_info in _get_model_cost_infos(model=model, llm_router=llm_router)
@ -1420,6 +1427,7 @@ def _max_cost_for_cost_info(
route: str,
model: str,
model_info: Mapping[str, object],
output_multiplier: int,
input_tokens: int | None = None,
) -> float | None:
image_cost: Final = _estimate_image_generation_cost(
@ -1444,7 +1452,6 @@ def _max_cost_for_cost_info(
if estimated_input_tokens is None or output_tokens is None:
return None
output_multiplier: Final = _get_output_multiplier(request_body=request_body)
tiered_pricing: Final = model_info.get("tiered_pricing")
if isinstance(tiered_pricing, list) and tiered_pricing:
tier: Final = select_tier_for_input(tiered_pricing=tiered_pricing, input_tokens=estimated_input_tokens)
@ -1453,8 +1460,8 @@ def _max_cost_for_cost_info(
tier_rate(tier, "output_cost_per_token"),
tier_rate(tier, "output_cost_per_reasoning_token"),
)
return (estimated_input_tokens * tier_rate(tier, "input_cost_per_token")) + (
output_tokens * output_multiplier * output_rate
return (estimated_input_tokens * tier_rate(tier, "input_cost_per_token")) + _output_cost(
output_tokens, output_multiplier, output_rate
)
input_cost_per_token: Final = _to_float(model_info.get("input_cost_per_token"))
@ -1471,7 +1478,7 @@ def _max_cost_for_cost_info(
# under-reserving reasoning-heavy requests.
output_rate = max(output_cost_per_token or 0.0, output_cost_per_reasoning_token or 0.0)
if output_cost_per_token is not None or output_cost_per_reasoning_token is not None:
cost += output_tokens * output_multiplier * output_rate
cost += _output_cost(output_tokens, output_multiplier, output_rate)
elif output_tokens > 0:
return None
@ -1617,6 +1624,7 @@ async def count_request_input_tokens(
request_body=request_body,
raw_body=raw_body,
models=tuple(dict.fromkeys(chain.from_iterable(counted_models.values()))),
gemini_route=RouteChecks.is_generate_content_route(route),
)
chat_pdf_files: Final = (
inline_pdf_file_count(request_body.get("messages"))
@ -1649,7 +1657,9 @@ def _estimate_input_tokens(
counted: Final = (
input_tokens
if input_tokens is not None
else count_input_tokens_for_model(request_body=request_body, model=model)
else count_input_tokens_for_model(
request_body=request_body, model=model, gemini_route=RouteChecks.is_generate_content_route(route)
)
)
if counted is not None:
return counted
@ -1672,7 +1682,7 @@ def _estimate_output_tokens(
if _is_input_only_route(route=route):
return 0
requested: Final = _requested_output_tokens(request_body)
requested: Final = _requested_output_tokens(request_body, route)
# Clamp at min(requested-or-default, model_max-or-default). Two purposes:
# (1) Without an explicit cap we still need a finite reservation so the
@ -1687,24 +1697,113 @@ def _estimate_output_tokens(
_OUTPUT_TOKEN_FIELDS: Final = ("max_completion_tokens", "max_tokens", "max_output_tokens")
_GEMINI_CONFIG_FIELDS: Final = ("generationConfig", "generation_config")
def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None:
inference_config: Final = request_body.get("inferenceConfig")
def _nested_field(request_body: Mapping[str, object], config_field: str, field: str) -> object:
config: Final = as_str_mapping(request_body.get(config_field))
return config.get(field) if config is not None else None
def _gemini_candidate_counts(request_body: Mapping[str, object]) -> tuple[object, ...]:
return tuple(
_nested_field(request_body, config_field, field)
for config_field, field in product(_GEMINI_CONFIG_FIELDS, ("candidateCount", "candidate_count"))
)
def _requested_output_tokens(request_body: Mapping[str, object], route: str) -> int | None:
if RouteChecks.is_generate_content_route(route):
gemini_cap: Final = _to_int(_nested_field(request_body, "generationConfig", "maxOutputTokens"))
return gemini_cap if gemini_cap is not None and gemini_cap > 0 else None
candidates: Final = (
*(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS),
inference_config.get("maxTokens") if isinstance(inference_config, Mapping) else None,
_nested_field(request_body, "inferenceConfig", "maxTokens"),
)
return next((tokens for tokens in map(_to_int, candidates) if tokens is not None), None)
return next((tokens for tokens in map(_to_int, candidates) if tokens is not None and tokens >= 0), None)
def _get_output_multiplier(request_body: dict) -> int:
output_multiplier = 1
for key in ("n", "best_of"):
value = _to_int(request_body.get(key))
if value is not None:
output_multiplier = max(output_multiplier, value)
return output_multiplier
def _requested_output_counts(
request_body: Mapping[str, object],
route: str,
model: str,
llm_router: Router | None,
) -> tuple[object, ...]:
if not RouteChecks.is_generate_content_route(route):
n: Final = request_body.get("n")
dropped: Final = n is not None and _every_deployment(_drops_n, _deployment_params(model, llm_router))
return (None if dropped else n, request_body.get("best_of"))
candidate_counts: Final = _gemini_candidate_counts(request_body)
if all(count is None for count in candidate_counts):
return ()
generated: Final = _every_deployment(_serves_generate_content_natively, _deployment_params(model, llm_router))
return candidate_counts if generated else ()
def _deployment_params(model: str, llm_router: Router | None) -> tuple[Mapping[str, object], ...]:
if llm_router is None:
return ({"model": model},)
deployments: Final = llm_router.registry().deployments(model) or ()
return tuple(
params for deployment in deployments if (params := as_str_mapping(_get_value(deployment, "litellm_params")))
)
def _every_deployment(
check: Callable[[Mapping[str, object]], bool], deployment_params: tuple[Mapping[str, object], ...]
) -> bool:
return bool(deployment_params) and all(map(check, deployment_params))
def _resolved_provider(litellm_params: Mapping[str, object]) -> tuple[str, str] | None:
model: Final = litellm_params.get("model")
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
if not isinstance(model, str):
return None
try:
backend_model, provider, _, _ = litellm.get_llm_provider(
model=model,
custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None,
)
except litellm.BadRequestError:
return None
return backend_model, provider
def _serves_generate_content_natively(litellm_params: Mapping[str, object]) -> bool:
resolved: Final = _resolved_provider(litellm_params)
if resolved is None:
return False
backend_model, provider = resolved
try:
native_config: Final = ProviderConfigManager.get_provider_google_genai_generate_content_config(
model=backend_model,
provider=litellm.LlmProviders(provider),
litellm_params={key: litellm_params[key] for key in ("base_model", "api_base") if key in litellm_params},
)
except ValueError:
return False
return native_config is not None
def _drops_n(litellm_params: Mapping[str, object]) -> bool:
resolved: Final = _resolved_provider(litellm_params)
if resolved is None:
return False
backend_model, provider = resolved
supported_params: Final = litellm.get_supported_openai_params(model=backend_model, custom_llm_provider=provider)
return supported_params is not None and "n" not in supported_params
def _get_output_multiplier(requested_counts: tuple[object, ...]) -> int:
return max((1, *(count for count in map(_to_int, requested_counts) if count is not None)))
def _output_cost(output_tokens: int, output_multiplier: int, output_rate: float) -> float:
try:
return output_tokens * output_multiplier * output_rate
except OverflowError:
return output_tokens * output_rate
def _is_input_only_route(route: str) -> bool:
@ -1732,7 +1831,7 @@ def _to_int(value: object) -> int | None:
return None
try:
return int(value)
except (TypeError, ValueError):
except (TypeError, ValueError, OverflowError):
return None

121
litellm/proxy/spend_tracking/input_tokens.py
View file

@ -16,11 +16,15 @@ import json
from collections.abc import Iterator, Mapping, Sequence
from itertools import chain
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
from typing import TYPE_CHECKING, Final, TypeAlias
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import OPENAI_CHAT_PDF_FILE_TOKENS, PDF_DATA_URL_PREFIX
from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter
from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items
from litellm.litellm_core_utils.token_counter import (
deployment_openai_model,
pdf_page_pricing,
@ -49,7 +53,30 @@ TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS: Final = 30_000
_OPENAI_PREFIX: Final = "openai/"
_INPUT_SIZE_FIELDS: Final = ("messages", "prompt", "input", "query", "documents", "tools", "tool_choice")
_GEMINI_SYSTEM_INSTRUCTION_FIELDS: Final = ("systemInstruction", "system_instruction")
_GEMINI_CACHED_CONTENT_FIELDS: Final = ("cachedContent", "cached_content")
_GEMINI_FUNCTION_RESPONSE_FIELDS: Final = ("functionResponse", "function_response")
_GEMINI_TEXT_PART_FIELDS: Final = frozenset({"text", "thought", "thoughtSignature", "thought_signature"})
_GEMINI_COUNTED_PART_FIELDS: Final = _GEMINI_TEXT_PART_FIELDS | {
"functionCall",
"function_call",
*_GEMINI_FUNCTION_RESPONSE_FIELDS,
}
_GEMINI_COUNTED_ROLES: Final = frozenset({"user", "model"})
_GEMINI_COUNTED_TOOL_FIELDS: Final = frozenset({"functionDeclarations"})
_GeminiContents: TypeAlias = list[dict[str, object]] | dict[str, object]
_GEMINI_CONTENTS: Final = TypeAdapter[_GeminiContents](_GeminiContents)
_INPUT_SIZE_FIELDS: Final = (
"messages",
"prompt",
"input",
"query",
"documents",
"contents",
*_GEMINI_SYSTEM_INSTRUCTION_FIELDS,
"tools",
"tool_choice",
)
def _approximate_input_size(request_body: Mapping[str, object]) -> int:
@ -62,9 +89,10 @@ def _approximate_input_size(request_body: Mapping[str, object]) -> int:
async def count_input_tokens(
request_body: dict,
request_body: dict[str, object],
raw_body: bytes | None,
models: Sequence[str],
gemini_route: bool = False,
) -> Mapping[str, int]:
"""Input-token count per model, sharing one native count across models that
select the same tokenizer."""
@ -72,13 +100,14 @@ async def count_input_tokens(
(model, rust_tokenizer(model)) for model in models
)
groups: Final[tuple[RustTokenizer | None, ...]] = tuple(dict.fromkeys(tokenizer for _, tokenizer in tokenizers))
counted_raw_body: Final = None if holds_compaction_summary(request_body.get("input")) else raw_body
counted_raw_body: Final = None if gemini_route or holds_compaction_summary(request_body.get("input")) else raw_body
group_counts: Final = [
await _count_group(
request_body=request_body,
raw_body=counted_raw_body,
tokenizer=tokenizer,
models=tuple(model for model, selected in tokenizers if selected == tokenizer),
gemini_route=gemini_route,
)
for tokenizer in groups
]
@ -88,18 +117,20 @@ async def count_input_tokens(
async def _count_group(
request_body: dict,
request_body: dict[str, object],
raw_body: bytes | None,
tokenizer: RustTokenizer | None,
models: tuple[str, ...],
gemini_route: bool,
) -> Mapping[str, int]:
async def python() -> Mapping[str, int]:
if _approximate_input_size(request_body) < TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS:
return _count_input_tokens_for_models(request_body=request_body, models=models)
return _count_input_tokens_for_models(request_body=request_body, models=models, gemini_route=gemini_route)
return await asyncio.to_thread(
_count_input_tokens_for_models,
request_body=request_body,
models=models,
gemini_route=gemini_route,
)
if tokenizer is None or raw_body is None:
@ -132,14 +163,20 @@ async def _native_counts(
def _count_input_tokens_for_models(
request_body: dict,
request_body: dict[str, object],
models: Sequence[str],
gemini_route: bool,
) -> Mapping[str, int]:
return MappingProxyType(
{
model: tokens
for model in models
if (tokens := count_input_tokens_for_model(request_body=request_body, model=model)) is not None
if (
tokens := count_input_tokens_for_model(
request_body=request_body, model=model, gemini_route=gemini_route
)
)
is not None
}
)
@ -221,8 +258,14 @@ def _served_by_openai(deployment: object) -> bool:
return served_provider == "openai"
def count_input_tokens_for_model(request_body: dict, model: str) -> int | None:
def count_input_tokens_for_model(request_body: dict, model: str, gemini_route: bool = False) -> int | None:
try:
if gemini_route:
return (
_count_gemini_input_tokens(request_body=request_body, model=model)
if "contents" in request_body
else None
)
if "messages" in request_body:
try:
return litellm.token_counter(
@ -251,6 +294,66 @@ def count_input_tokens_for_model(request_body: dict, model: str) -> int | None:
return None
def _count_gemini_input_tokens(request_body: Mapping[str, object], model: str) -> int | None:
system_instruction: Final = next(
(request_body[field] for field in _GEMINI_SYSTEM_INSTRUCTION_FIELDS if field in request_body), None
)
contents: Final = _GEMINI_CONTENTS.validate_python(request_body["contents"])
tools: Final = request_body.get("tools")
if not _gemini_body_is_countable(request_body, contents, system_instruction, tools):
return None
translate: Final = GoogleGenAIAdapter().translate_generate_content_to_completion # pyright: ignore[reportUnknownMemberType] # adapter's **kwargs is untyped
chat_request: Final = translate(model=model, contents=contents, systemInstruction=system_instruction, tools=tools)
return count_input_tokens_for_model(request_body=chat_request, model=model)
def _gemini_body_is_countable(
request_body: Mapping[str, object],
contents: _GeminiContents,
system_instruction: object,
tools: object,
) -> bool:
content_items: Final = contents if isinstance(contents, list) else (contents,)
system: Final = as_str_mapping(system_instruction)
return (
not any(field in request_body for field in _GEMINI_CACHED_CONTENT_FIELDS)
and all(_gemini_content_is_countable(content) for content in content_items)
and (
system_instruction is None
or (system is not None and _gemini_parts_are_countable(system.get("parts", ()), _GEMINI_TEXT_PART_FIELDS))
)
and all(_gemini_tool_is_countable(tool) for tool in iter_items(tools or ()))
)
def _gemini_content_is_countable(content: Mapping[str, object]) -> bool:
return content.get("role", "user") in _GEMINI_COUNTED_ROLES and _gemini_parts_are_countable(
content.get("parts", ()), _GEMINI_COUNTED_PART_FIELDS
)
def _gemini_parts_are_countable(parts: object, fields: frozenset[str]) -> bool:
single_part: Final = as_str_mapping(parts)
items: Final = (single_part,) if single_part is not None else iter_items(parts)
return all(_gemini_part_is_countable(part, fields) for part in items)
def _gemini_part_is_countable(part: object, fields: frozenset[str]) -> bool:
if isinstance(part, str):
return True
mapping: Final = as_str_mapping(part)
return (
mapping is not None
and mapping.keys() <= fields
and not any("parts" in (as_str_mapping(mapping.get(field)) or ()) for field in _GEMINI_FUNCTION_RESPONSE_FIELDS)
)
def _gemini_tool_is_countable(tool: object) -> bool:
mapping: Final = as_str_mapping(tool)
return mapping is not None and mapping.keys() <= _GEMINI_COUNTED_TOOL_FIELDS
def _count_text_tokens(model: str, text: object) -> int:
if text is None:
return 0

56
tests/unit/google_genai/test_google_genai_adapter.py
View file

@ -194,6 +194,62 @@ def test_native_function_call_preserves_arguments_and_matching_response(args: ob
]
@pytest.mark.parametrize(
("system_texts", "expected_content"),
[
(("Answer tersely",), "Answer tersely"),
(("Answer tersely", "Use British spelling"), "Answer tersely\nUse British spelling"),
],
)
def test_every_system_instruction_text_part_reaches_the_system_message(
system_texts: tuple[str, ...], expected_content: str
) -> None:
request: Final = GoogleGenAIAdapter().translate_generate_content_to_completion(
model="gemini/gemini-3.8-flash",
contents=[{"role": "user", "parts": [{"text": "hi"}]}],
systemInstruction={"parts": [{"text": text} for text in system_texts]},
)
assert request["messages"][0] == {"role": "system", "content": expected_content}
@pytest.mark.parametrize(
("system_parts", "expected_messages"),
[
pytest.param(
[
{"text": "Answer tersely"},
{"text": None},
{"text": 3},
{"text": ["x"]},
{"text": "Use British spelling"},
],
[{"role": "system", "content": "Answer tersely\nUse British spelling"}, {"role": "user", "content": "hi"}],
id="malformed_parts_skipped_like_contents",
),
pytest.param(
{"text": "Answer tersely"},
[{"role": "system", "content": "Answer tersely"}, {"role": "user", "content": "hi"}],
id="single_object_parts",
),
pytest.param(
["Answer tersely", {"text": "Use British spelling"}],
[{"role": "system", "content": "Answer tersely\nUse British spelling"}, {"role": "user", "content": "hi"}],
id="bare_string_part_read_like_contents",
),
pytest.param(None, [{"role": "user", "content": "hi"}], id="no_parts"),
],
)
def test_system_instruction_parts_follow_the_contents_part_rules(
system_parts: object, expected_messages: list[dict[str, str]]
) -> None:
request: Final = GoogleGenAIAdapter().translate_generate_content_to_completion(
model="gemini/gemini-3.8-flash",
contents=[{"role": "user", "parts": [{"text": "hi"}]}],
systemInstruction={"parts": system_parts},
)
assert request["messages"] == expected_messages
def test_adapter_import():
"""Test that the adapter can be imported successfully"""
from litellm.google_genai.adapters.handler import GenerateContentToCompletionHandler

317
tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
View file

@ -294,12 +294,6 @@ class TestVertexAIPassThroughHandler:
mock.patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route"
) as mock_create_route,
mock.patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_litellm_virtual_key"
) as mock_get_virtual_key,
mock.patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth"
) as mock_user_auth,
mock.patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_vertex_pass_through_handler"
) as mock_get_handler,
@ -310,8 +304,6 @@ class TestVertexAIPassThroughHandler:
# Setup mocks
mock_load_auth.return_value = (mock_credentials, test_project)
mock_get_virtual_key.return_value = "Bearer test-key"
mock_user_auth.return_value = {"api_key": "test-key"}
# Mock the vertex handler
mock_handler = Mock()
@ -396,12 +388,6 @@ class TestVertexAIPassThroughHandler:
mock.patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route"
) as mock_create_route,
mock.patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_litellm_virtual_key"
) as mock_get_virtual_key,
mock.patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth"
) as mock_user_auth,
mock.patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_vertex_pass_through_handler"
) as mock_get_handler,
@ -412,8 +398,6 @@ class TestVertexAIPassThroughHandler:
# Setup mocks
mock_load_auth.return_value = (mock_credentials, test_project)
mock_get_virtual_key.return_value = "Bearer test-key"
mock_user_auth.return_value = {"api_key": "test-key"}
# Mock the vertex handler for global location
mock_handler = Mock()
@ -498,17 +482,12 @@ class TestVertexAIPassThroughHandler:
mock.patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_vertex_pass_through_handler"
) as mock_get_handler,
mock.patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
new_callable=AsyncMock,
) as mock_auth,
):
# Mock credentials object with necessary attributes
mock_credentials = Mock()
mock_credentials.token = default_credentials
mock_load_auth.return_value = (mock_credentials, default_project)
mock_auth.return_value = MagicMock()
# Mock the vertex handler
mock_handler = Mock()
@ -526,6 +505,7 @@ class TestVertexAIPassThroughHandler:
endpoint=endpoint,
request=mock_request,
fastapi_response=mock_response,
user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
)
except Exception as e:
traceback.print_exc()
@ -600,20 +580,16 @@ class TestVertexAIPassThroughHandler:
mock.patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route"
) as mock_create_route,
mock.patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
new_callable=AsyncMock,
) as mock_auth,
):
mock_ensure_token.return_value = ("test-auth-header", test_project)
mock_get_token.return_value = (test_token, "")
mock_auth.return_value = UserAPIKeyAuth(api_key="sk-test-creds")
with pytest.raises(HTTPException) as exc_info:
await vertex_proxy_route(
endpoint=endpoint,
request=mock_request,
fastapi_response=mock_response,
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test-creds"),
)
assert exc_info.value.status_code == 401
@ -625,8 +601,9 @@ class TestVertexAIPassThroughHandler:
Critical
This is how Vertex AI JS SDK will Auth to Litellm Proxy: the virtual key
arrives in x-litellm-api-key and must reach user_api_key_auth. With no Vertex
credential configured, that virtual key must not be forwarded to Google, so
arrives in x-litellm-api-key and authenticates the caller (covered through the
real auth dependency in TestVertexPassthroughAuthenticatesOncePerRequest). With no
Vertex credential configured, that virtual key must not be forwarded to Google, so
the request fails with a clean 401 (LIT-5997).
"""
# Mock dependencies
@ -635,26 +612,21 @@ class TestVertexAIPassThroughHandler:
mock_request.method = "POST"
mock_response = Mock()
with patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth") as mock_auth:
mock_auth.return_value = {"api_key": "test-key-123"}
with patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route"
) as mock_pass_through:
mock_pass_through.return_value = AsyncMock(return_value={"status": "success"})
with patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route"
) as mock_pass_through:
mock_pass_through.return_value = AsyncMock(return_value={"status": "success"})
with pytest.raises(HTTPException) as exc_info:
await vertex_proxy_route(
endpoint="v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent",
request=mock_request,
fastapi_response=mock_response,
user_api_key_dict=UserAPIKeyAuth(api_key="test-key-123"),
)
with pytest.raises(HTTPException) as exc_info:
await vertex_proxy_route(
endpoint="v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent",
request=mock_request,
fastapi_response=mock_response,
)
assert exc_info.value.status_code == 401
mock_pass_through.assert_not_called()
mock_auth.assert_called_once()
call_args = mock_auth.call_args[1]
assert call_args["api_key"] == "Bearer test-key-123"
assert exc_info.value.status_code == 401
mock_pass_through.assert_not_called()
def test_vertex_passthrough_handler_multimodal_embedding_response(self):
"""
@ -1220,12 +1192,6 @@ class TestVertexAIDiscoveryPassThroughHandler:
mock.patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route"
) as mock_create_route,
mock.patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_litellm_virtual_key"
) as mock_get_virtual_key,
mock.patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth"
) as mock_user_auth,
mock.patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_vertex_pass_through_handler"
) as mock_get_handler,
@ -1236,8 +1202,6 @@ class TestVertexAIDiscoveryPassThroughHandler:
# Setup mocks
mock_load_auth.return_value = (mock_credentials, test_project)
mock_get_virtual_key.return_value = "Bearer test-key"
mock_user_auth.return_value = {"api_key": "test-key"}
# Mock the discovery handler
mock_handler = Mock()
@ -1253,6 +1217,7 @@ class TestVertexAIDiscoveryPassThroughHandler:
endpoint=endpoint,
request=mock_request,
fastapi_response=mock_response,
user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
)
# Verify create_pass_through_route was called with correct arguments
@ -1267,36 +1232,29 @@ class TestVertexAIDiscoveryPassThroughHandler:
@pytest.mark.asyncio
async def test_vertex_discovery_proxy_route_api_key_auth(self):
"""
The virtual key arrives in x-litellm-api-key and must reach user_api_key_auth.
With no Vertex credential configured, that virtual key must not be forwarded to
With no Vertex credential configured, the caller's virtual key must not be forwarded to
Google, so the request fails with a clean 401 (LIT-5997).
"""
# Mock dependencies
mock_request = Mock()
mock_request.headers = {"x-litellm-api-key": "test-key-123"}
mock_request.method = "POST"
mock_response = Mock()
with patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth") as mock_auth:
mock_auth.return_value = {"api_key": "test-key-123"}
with patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route"
) as mock_pass_through:
mock_pass_through.return_value = AsyncMock(return_value={"status": "success"})
with patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route"
) as mock_pass_through:
mock_pass_through.return_value = AsyncMock(return_value={"status": "success"})
with pytest.raises(HTTPException) as exc_info:
await vertex_discovery_proxy_route(
endpoint="v1/projects/test-project/locations/us-central1/dataStores/default/servingConfigs/default:search",
request=mock_request,
fastapi_response=mock_response,
user_api_key_dict=UserAPIKeyAuth(api_key="test-key-123"),
)
with pytest.raises(HTTPException) as exc_info:
await vertex_discovery_proxy_route(
endpoint="v1/projects/test-project/locations/us-central1/dataStores/default/servingConfigs/default:search",
request=mock_request,
fastapi_response=mock_response,
)
assert exc_info.value.status_code == 401
mock_pass_through.assert_not_called()
mock_auth.assert_called_once()
call_args = mock_auth.call_args[1]
assert call_args["api_key"] == "Bearer test-key-123"
assert exc_info.value.status_code == 401
mock_pass_through.assert_not_called()
@pytest.mark.asyncio
@ -4481,8 +4439,6 @@ class TestVertexRawPredictStreamingClassification:
return_value=(mock_credentials, "test-project"),
),
mock.patch(f"{module}.create_pass_through_route", side_effect=fake_create_pass_through_route),
mock.patch(f"{module}.get_litellm_virtual_key", return_value="Bearer test-key"),
mock.patch(f"{module}.user_api_key_auth", new=AsyncMock(return_value=UserAPIKeyAuth(api_key="test-key"))),
mock.patch(f"{module}.get_vertex_pass_through_handler", return_value=mock_handler),
):
await vertex_proxy_route(
@ -8418,6 +8374,213 @@ class TestAzureSpeechRawBodyThroughRealAuth:
assert "Invalid JSON payload" in response.text
class _CountingCustomAuth:
def __init__(self, only_accepts: str | None = None) -> None:
self.only_accepts: Final = only_accepts
self.api_keys_seen: tuple[str, ...] = ()
async def __call__(self, request: Request, api_key: str) -> UserAPIKeyAuth:
self.api_keys_seen = (*self.api_keys_seen, api_key)
if self.only_accepts is not None and api_key != self.only_accepts:
raise ValueError("unknown caller")
return UserAPIKeyAuth(user_id=f"custom-auth-caller-{len(self.api_keys_seen)}")
class _ServedCallerRecorder(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.served_user_ids: tuple[str | None, ...] = ()
async def async_pre_call_hook(
self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral
) -> dict:
self.served_user_ids = (*self.served_user_ids, user_api_key_dict.user_id)
return data
class TestVertexPassthroughAuthenticatesOncePerRequest:
"""Each Vertex pass-through request runs the proxy's auth once and is served as the caller that auth resolved"""
PROJECT: Final = "vertex-proj"
VIRTUAL_KEY: Final = "sk-vertex-caller"
GENERATE_CONTENT: Final = (
"v1/projects/vertex-proj/locations/us-central1/publishers/google/models/gemini-3.8-flash:generateContent"
)
DISCOVERY_SEARCH: Final = (
"v1/projects/vertex-proj/locations/global/collections/default_collection/engines/support-engine"
"/servingConfigs/default_search:search"
)
GEMINI_ANSWER: Final = {
"candidates": [{"content": {"role": "model", "parts": [{"text": "hello"}]}, "finishReason": "STOP"}],
"usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2},
}
def _post(
self,
monkeypatch: pytest.MonkeyPatch,
path: str,
headers: Mapping[str, str],
custom_auth: _CountingCustomAuth,
credential_locations: tuple[str, ...],
) -> tuple[httpx.Response, _ServedCallerRecorder]:
from litellm.proxy.pass_through_endpoints.passthrough_endpoint_router import (
PassthroughEndpointRouter,
)
from litellm.proxy.proxy_server import app
router: Final = PassthroughEndpointRouter()
for location in credential_locations:
router.add_vertex_credentials(
project_id=self.PROJECT, location=location, vertex_credentials="vertex-service-account"
)
recorder: Final = _ServedCallerRecorder()
monkeypatch.setattr(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router", router
)
monkeypatch.delitem(app.dependency_overrides, user_api_key_auth, raising=False)
monkeypatch.setattr(litellm, "callbacks", [recorder])
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
with (
patch.multiple( # test-quality-ok: the real user_api_key_auth reads proxy_server module globals (custom auth, master_key, caches) that have no injection seam
"litellm.proxy.proxy_server", **{**_azure_speech_real_auth_attrs(), "user_custom_auth": custom_auth}
),
patch( # test-quality-ok: the route mints its Google token through its own VertexBase, nothing injects the credential loader
"litellm.llms.vertex_ai.vertex_llm_base.VertexBase.load_auth",
return_value=(Mock(token="google-access-token"), self.PROJECT),
),
):
response: Final = TestClient(app).post(
path, json={"contents": [{"role": "user", "parts": [{"text": "hi"}]}]}, headers=dict(headers)
)
return response, recorder
@pytest.mark.parametrize("route_prefix", ["/vertex_ai", "/vertex-ai"])
def test_generate_content_runs_custom_auth_once_and_is_served_as_its_caller(
self, monkeypatch: pytest.MonkeyPatch, route_prefix: str
) -> None:
custom_auth: Final = _CountingCustomAuth()
vertex_url: Final = f"https://us-central1-aiplatform.googleapis.com/{self.GENERATE_CONTENT}"
with respx.mock(assert_all_called=True) as upstream:
vertex: Final = upstream.post(vertex_url).mock(return_value=httpx.Response(200, json=self.GEMINI_ANSWER))
response, recorder = self._post(
monkeypatch,
f"{route_prefix}/{self.GENERATE_CONTENT}",
{"Authorization": f"Bearer {self.VIRTUAL_KEY}"},
custom_auth,
("us-central1",),
)
assert response.status_code == 200, response.text
assert custom_auth.api_keys_seen == (self.VIRTUAL_KEY,)
assert recorder.served_user_ids == ("custom-auth-caller-1",)
assert vertex.call_count == 1
def test_discovery_search_runs_custom_auth_once_and_is_served_as_its_caller(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
custom_auth: Final = _CountingCustomAuth()
with respx.mock(assert_all_called=True) as upstream:
discovery: Final = upstream.post(f"https://discoveryengine.googleapis.com/{self.DISCOVERY_SEARCH}").mock(
return_value=httpx.Response(200, json={"results": []})
)
response, recorder = self._post(
monkeypatch,
f"/vertex_ai/discovery/{self.DISCOVERY_SEARCH}",
{"Authorization": f"Bearer {self.VIRTUAL_KEY}"},
custom_auth,
("global",),
)
assert (response.status_code, response.json()) == (200, {"results": []})
assert custom_auth.api_keys_seen == (self.VIRTUAL_KEY,)
assert recorder.served_user_ids == ("custom-auth-caller-1",)
assert discovery.call_count == 1
@pytest.mark.parametrize("scheme", ["bearer", "Basic"])
@pytest.mark.parametrize(
"route, upstream_url, credential_location, answer",
[
(
f"/vertex_ai/{GENERATE_CONTENT}",
f"https://us-central1-aiplatform.googleapis.com/{GENERATE_CONTENT}",
"us-central1",
GEMINI_ANSWER,
),
(
f"/vertex_ai/discovery/{DISCOVERY_SEARCH}",
f"https://discoveryengine.googleapis.com/{DISCOVERY_SEARCH}",
"global",
{"results": []},
),
],
ids=["generate_content", "discovery_search"],
)
def test_x_litellm_api_key_with_any_auth_scheme_authenticates_once_on_every_vertex_route(
self,
monkeypatch: pytest.MonkeyPatch,
scheme: str,
route: str,
upstream_url: str,
credential_location: str,
answer: Mapping[str, object],
) -> None:
custom_auth: Final = _CountingCustomAuth(only_accepts=self.VIRTUAL_KEY)
with respx.mock(assert_all_called=True) as upstream:
google: Final = upstream.post(upstream_url).mock(return_value=httpx.Response(200, json=answer))
response, recorder = self._post(
monkeypatch,
route,
{"x-litellm-api-key": f"{scheme} {self.VIRTUAL_KEY}"},
custom_auth,
(credential_location,),
)
assert response.status_code == 200, response.text
assert custom_auth.api_keys_seen == (self.VIRTUAL_KEY,)
assert recorder.served_user_ids == ("custom-auth-caller-1",)
assert google.call_count == 1
def test_x_litellm_api_key_authenticates_once_and_is_never_forwarded_without_a_vertex_credential(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
custom_auth: Final = _CountingCustomAuth()
with respx.mock(assert_all_called=False) as upstream:
google: Final = upstream.route().mock(return_value=httpx.Response(200))
response, _ = self._post(
monkeypatch,
f"/vertex_ai/{self.GENERATE_CONTENT}",
{"x-litellm-api-key": self.VIRTUAL_KEY},
custom_auth,
(),
)
assert response.status_code == 401, response.text
assert custom_auth.api_keys_seen == (self.VIRTUAL_KEY,)
assert not google.called
@pytest.mark.parametrize("header", ["x-goog-api-key", "x-api-key", "API-Key", "Ocp-Apim-Subscription-Key"])
def test_key_outside_litellm_headers_is_rejected_and_never_reaches_google_without_a_vertex_credential(
self, monkeypatch: pytest.MonkeyPatch, header: str
) -> None:
custom_auth: Final = _CountingCustomAuth(only_accepts=self.VIRTUAL_KEY)
with respx.mock(assert_all_called=False) as upstream:
google: Final = upstream.route().mock(return_value=httpx.Response(200))
response, recorder = self._post(
monkeypatch, f"/vertex_ai/{self.GENERATE_CONTENT}", {header: self.VIRTUAL_KEY}, custom_auth, ()
)
assert response.status_code == 401, response.text
assert self.VIRTUAL_KEY not in custom_auth.api_keys_seen
assert recorder.served_user_ids == ()
assert not google.called
class TestTypeSafePassthroughRoute:
@staticmethod
def _request(body: object, query_params: Mapping[str, str] | None = None) -> MagicMock:

40
tests/unit/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py
View file

@ -1,46 +1,6 @@
import json
import traceback
from unittest import mock
from unittest.mock import MagicMock, patch
import httpx
import pytest
from fastapi import Request, Response
from fastapi.testclient import TestClient
from litellm.passthrough.utils import CommonUtils
from unittest.mock import Mock
from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key
@pytest.mark.asyncio
async def test_get_litellm_virtual_key():
"""
Test that the get_litellm_virtual_key function correctly handles the API key authentication
"""
# Test with x-litellm-api-key
mock_request = Mock()
mock_request.headers = {"x-litellm-api-key": "test-key-123"}
result = get_litellm_virtual_key(mock_request)
assert result == "Bearer test-key-123"
# Test with Authorization header
mock_request.headers = {"Authorization": "Bearer auth-key-456"}
result = get_litellm_virtual_key(mock_request)
assert result == "Bearer auth-key-456"
# Test with both headers (x-litellm-api-key should take precedence)
mock_request.headers = {
"x-litellm-api-key": "test-key-123",
"Authorization": "Bearer auth-key-456",
}
result = get_litellm_virtual_key(mock_request)
assert result == "Bearer test-key-123"
def test_encode_bedrock_runtime_modelid_arn():
# Test application-inference-profile ARN
endpoint = "model/arn:aws:bedrock:us-east-1:123456789123:application-inference-profile/r742sbn2zckd/converse"

24
tests/unit/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py
View file

@ -67,10 +67,6 @@ async def test_vertex_passthrough_load_balancing():
patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route"
) as mock_create_route,
patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
new_callable=AsyncMock,
) as mock_auth,
):
# Setup additional mocks to avoid side effects
@ -84,7 +80,6 @@ async def test_vertex_passthrough_load_balancing():
mock_endpoint_func = AsyncMock()
mock_create_route.return_value = mock_endpoint_func
mock_auth.return_value = {}
# Execute
await _base_vertex_proxy_route(
@ -93,6 +88,7 @@ async def test_vertex_passthrough_load_balancing():
fastapi_response=mock_response,
get_vertex_pass_through_handler=mock_handler,
general_settings={},
user_api_key_dict=UserAPIKeyAuth(),
)
# Verify
@ -375,10 +371,6 @@ async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens(
patch( # test-quality-ok: the upstream call is captured here, the route offers no injection point
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route"
) as mock_create_route,
patch( # test-quality-ok: the route calls auth directly rather than through Depends
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
new_callable=AsyncMock,
) as mock_auth,
patch( # test-quality-ok: the route reads the request body for this, a MagicMock request has none
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_streaming_request_fn",
new_callable=AsyncMock,
@ -397,7 +389,6 @@ async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens(
"global",
)
mock_create_route.return_value = AsyncMock()
mock_auth.return_value = UserAPIKeyAuth(api_key="sk-litellm-secret-key")
await _base_vertex_proxy_route(
endpoint=f"{VERTEX_ANTHROPIC_MODELS_PREFIX}{model_segment}",
@ -405,6 +396,7 @@ async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens(
fastapi_response=MagicMock(),
get_vertex_pass_through_handler=VertexAIPassThroughHandler(),
general_settings={},
user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-secret-key"),
)
upstream_headers = mock_create_route.call_args.kwargs["custom_headers"]
@ -723,10 +715,6 @@ async def test_vertex_passthrough_custom_model_name_replaced_in_url():
patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route"
) as mock_create_route,
patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
new_callable=AsyncMock,
) as mock_auth,
):
mock_pt_router.get_vertex_credentials.return_value = MagicMock()
@ -738,7 +726,6 @@ async def test_vertex_passthrough_custom_model_name_replaced_in_url():
)
mock_endpoint_func = AsyncMock()
mock_create_route.return_value = mock_endpoint_func
mock_auth.return_value = {}
mock_handler.get_default_base_target_url.return_value = (
"https://global-aiplatform.googleapis.com"
@ -750,6 +737,7 @@ async def test_vertex_passthrough_custom_model_name_replaced_in_url():
fastapi_response=mock_response,
get_vertex_pass_through_handler=mock_handler,
general_settings={},
user_api_key_dict=UserAPIKeyAuth(),
)
# Verify the router was called with the custom model name (extracted from URL)
@ -825,11 +813,6 @@ async def test_vertex_passthrough_attributes_the_call_to_the_resolved_deployment
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route",
return_value=relay_returning_logging_kwargs,
),
patch( # test-quality-ok: the route calls auth directly rather than through Depends
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
new_callable=AsyncMock,
return_value=UserAPIKeyAuth(api_key="hashed-key"),
),
):
mock_pt_router.get_vertex_credentials.return_value = MagicMock()
@ -839,6 +822,7 @@ async def test_vertex_passthrough_attributes_the_call_to_the_resolved_deployment
fastapi_response=MagicMock(),
get_vertex_pass_through_handler=mock_handler,
general_settings={},
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
)
assert logging_kwargs["litellm_params"]["metadata"]["model_info"]["id"] == "vertex-gemini-38-flash-dep"

350
tests/unit/proxy/spend_tracking/test_budget_reservation.py
View file

@ -1,5 +1,6 @@
from __future__ import annotations
import json
import math
from datetime import datetime, timedelta, timezone
from typing import Final
@ -25,6 +26,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
)
from litellm.proxy.spend_tracking.budget_reservation import (
_get_team_member_budget_counter,
estimate_request_input_cost,
estimate_request_max_cost,
get_budget_window_start,
release_unbound_budget_reservation,
@ -177,6 +179,354 @@ def test_bedrock_converse_body_reserves_the_prompt_not_the_context_window():
assert invoke_cost < converse_cost < 2 * invoke_cost
GEMINI_PROMPT: Final = "Say hello in five words"
GEMINI_SYSTEM_PARTS: Final = ("Answer tersely", "Use British spelling")
@pytest.mark.parametrize(
"route",
(
"/vertex_ai/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent",
"/gemini/v1beta/models/gemini-3.8-flash:generateContent",
"/v1beta/models/gemini-3.8-flash:generateContent",
),
)
@pytest.mark.parametrize(
("native_output_config", "chat_output_limit"),
(
({"generationConfig": {"maxOutputTokens": 5}}, {"max_tokens": 5}),
({"generationConfig": {"maxOutputTokens": 5, "candidateCount": 3}}, {"max_tokens": 5, "n": 3}),
({"generation_config": {"max_output_tokens": 5, "candidate_count": 3}}, {"n": 3}),
({"generationConfig": {"candidate_count": 3}}, {"n": 3}),
({}, {}),
),
)
@pytest.mark.parametrize("system_texts", (GEMINI_SYSTEM_PARTS[:1], GEMINI_SYSTEM_PARTS))
def test_gemini_generate_content_body_reserves_the_prompt_not_the_context_window(
route: str,
native_output_config: dict[str, object],
chat_output_limit: dict[str, object],
system_texts: tuple[str, ...],
) -> None:
native_cost: Final = estimate_request_max_cost(
request_body={
"systemInstruction": {"parts": [{"text": text} for text in system_texts]},
"contents": [{"role": "user", "parts": [{"text": GEMINI_PROMPT}]}],
**native_output_config,
},
route=route,
llm_router=None,
input_token_counts={},
)
chat_cost: Final = estimate_request_max_cost(
request_body={
"model": "gemini-3.8-flash",
"messages": [
{"role": "system", "content": "\n".join(system_texts)},
{"role": "user", "content": GEMINI_PROMPT},
],
**chat_output_limit,
},
route="/chat/completions",
llm_router=None,
input_token_counts={},
)
assert native_cost is not None
assert native_cost == chat_cost
GEMINI_ROUTE: Final = "/v1beta/models/gemini-3.8-flash:generateContent"
GEMINI_TEXT_CONTENT: Final = {"role": "user", "parts": [{"text": GEMINI_PROMPT}]}
GEMINI_AUDIO_PART: Final = {"inlineData": {"mimeType": "audio/wav", "data": "UklGRg" * 40_000}}
def _gemini_hold(body: dict[str, object], input_tokens: int | None = None) -> float | None:
return estimate_request_max_cost(
request_body=body,
route=GEMINI_ROUTE,
llm_router=None,
input_token_counts={} if input_tokens is None else {"gemini-3.8-flash": input_tokens},
)
@pytest.mark.parametrize(
"body",
(
{"contents": [{"role": "user", "parts": [{"text": GEMINI_PROMPT}, GEMINI_AUDIO_PART]}]},
{"contents": [{"role": "user", "parts": [{"fileData": {"mimeType": "video/mp4", "fileUri": "gs://b/v.mp4"}}]}]},
{"contents": [{"role": "user", "parts": {"fileData": {"mimeType": "video/mp4", "fileUri": "gs://b/v.mp4"}}}]},
{"contents": [{"role": "user", "parts": [{"inline_data": {"mime_type": "image/png", "data": "iVBORw0K"}}]}]},
{"contents": [{"role": "model", "parts": [{"executableCode": {"language": "PYTHON", "code": "print(1)"}}]}]},
{"contents": [GEMINI_TEXT_CONTENT, {"role": "function", "parts": [{"text": GEMINI_PROMPT * 500}]}]},
{
"contents": [
{
"role": "user",
"parts": [{"functionResponse": {"name": "listen", "response": {}, "parts": [GEMINI_AUDIO_PART]}}],
}
]
},
{"systemInstruction": {"parts": [GEMINI_AUDIO_PART]}, "contents": [GEMINI_TEXT_CONTENT]},
{"systemInstruction": {"parts": GEMINI_AUDIO_PART}, "contents": [GEMINI_TEXT_CONTENT]},
{"tools": [{"function_declarations": [{"name": "lookup"}]}], "contents": [GEMINI_TEXT_CONTENT]},
{"tools": [{"googleSearch": {}}], "contents": [GEMINI_TEXT_CONTENT]},
{"cachedContent": "cachedContents/abc", "contents": [GEMINI_TEXT_CONTENT]},
),
)
def test_gemini_body_the_prompt_count_cannot_see_whole_keeps_the_context_window_hold(body: dict[str, object]) -> None:
context_window: Final = litellm.get_model_info("gemini-3.8-flash")["max_input_tokens"]
assert context_window is not None
assert _gemini_hold(body) == _gemini_hold({"contents": [GEMINI_TEXT_CONTENT]}, input_tokens=context_window)
def test_gemini_tool_turns_and_thought_signatures_count_as_prompt() -> None:
context_window: Final = litellm.get_model_info("gemini-3.8-flash")["max_input_tokens"]
assert context_window is not None
tool_turns: Final = {
"tools": [{"functionDeclarations": [{"name": "lookup", "parameters": {"type": "object"}}]}],
"contents": [
GEMINI_TEXT_CONTENT,
{"role": "model", "parts": [{"functionCall": {"name": "lookup", "args": {}}, "thoughtSignature": "c2ln"}]},
{"role": "user", "parts": [{"functionResponse": {"name": "lookup", "response": {"answer": "hi"}}}]},
],
}
text_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT]})
tool_turns_hold: Final = _gemini_hold(tool_turns)
context_window_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT]}, input_tokens=context_window)
assert text_hold is not None and tool_turns_hold is not None and context_window_hold is not None
assert text_hold < tool_turns_hold < context_window_hold
GEMINI_CHAT_BODY: Final = {
"model": "gemini-3.8-flash",
"messages": [{"role": "user", "content": GEMINI_PROMPT}],
}
def _chat_hold(body: dict[str, object]) -> float | None:
return estimate_request_max_cost(
request_body=body,
route="/chat/completions",
llm_router=None,
input_token_counts={},
)
@pytest.mark.parametrize(
"gemini_output_fields",
(
{"generationConfig": {"maxOutputTokens": 5}},
{"generationConfig": {"candidateCount": 3}},
{"generation_config": {"candidate_count": 3}},
{"contents": [], "generationConfig": {"maxOutputTokens": 5}},
),
)
def test_gemini_output_fields_on_a_chat_body_leave_its_hold_alone(gemini_output_fields: dict[str, object]) -> None:
plain_hold: Final = _chat_hold(GEMINI_CHAT_BODY)
assert plain_hold is not None
assert _chat_hold({**GEMINI_CHAT_BODY, **gemini_output_fields}) == plain_hold
@pytest.mark.parametrize("cap", (0, -1000))
def test_a_gemini_output_cap_of_zero_or_below_holds_the_default_output(cap: int) -> None:
uncapped_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT]})
assert uncapped_hold is not None
assert (
_gemini_hold({"contents": [GEMINI_TEXT_CONTENT], "generationConfig": {"maxOutputTokens": cap}})
== uncapped_hold
)
UNPRICEABLE_COUNTS: Final = (10**400, "9" * 400, 1e308, math.inf)
@pytest.mark.parametrize("count", UNPRICEABLE_COUNTS)
def test_a_candidate_count_too_large_to_price_holds_one_candidate(count: object) -> None:
one_candidate_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT]})
assert one_candidate_hold is not None
assert (
_gemini_hold({"contents": [GEMINI_TEXT_CONTENT], "generationConfig": {"candidateCount": count}})
== one_candidate_hold
)
@pytest.mark.parametrize("count", UNPRICEABLE_COUNTS)
def test_an_n_too_large_to_price_holds_one_choice(count: object) -> None:
one_choice_hold: Final = _chat_hold(GEMINI_CHAT_BODY)
assert one_choice_hold is not None
assert _chat_hold({**GEMINI_CHAT_BODY, "n": count}) == one_choice_hold
def _group_generate_content_hold(router: Router, generation_config: dict[str, object]) -> float | None:
return estimate_request_max_cost(
request_body={"contents": [GEMINI_TEXT_CONTENT], "generationConfig": generation_config},
route="/v1beta/models/assistant:generateContent",
llm_router=router,
input_token_counts={"assistant": 10},
)
def _assistant_router(*backend_models: str) -> Router:
return Router(
model_list=[
{"model_name": "assistant", "litellm_params": {"model": model, "api_key": "sk-test"}}
for model in backend_models
]
)
@pytest.mark.parametrize(
("backend_models", "candidates_are_generated"),
(
(("vertex_ai/claude-haiku-5-5",), False),
(("anthropic/claude-haiku-5-5",), False),
(("gemini/gemini-3.8-flash",), True),
(("vertex_ai/gemini-3.8-flash",), True),
(("gemini/gemini-3.8-flash", "vertex_ai/gemini-3.8-flash"), True),
(("vertex_ai/claude-haiku-5-5", "gemini/gemini-3.8-flash"), False),
(("gemini/gemini-3.8-flash", "anthropic/claude-haiku-5-5"), False),
),
)
def test_a_candidate_count_is_held_only_when_every_deployment_generates_the_candidates(
backend_models: tuple[str, ...], candidates_are_generated: bool
) -> None:
router: Final = _assistant_router(*backend_models)
one_candidate_hold: Final = _group_generate_content_hold(router, {"maxOutputTokens": 100})
three_candidate_hold: Final = _group_generate_content_hold(router, {"maxOutputTokens": 100, "candidateCount": 3})
assert one_candidate_hold is not None
assert three_candidate_hold is not None
assert (three_candidate_hold > one_candidate_hold) is candidates_are_generated
assert three_candidate_hold >= one_candidate_hold
def test_a_candidate_count_for_a_gemini_name_no_deployment_serves_holds_one_candidate() -> None:
router: Final = _assistant_router("anthropic/claude-haiku-5-5")
def hold(generation_config: dict[str, object]) -> float | None:
return estimate_request_max_cost(
request_body={"contents": [GEMINI_TEXT_CONTENT], "generationConfig": generation_config},
route=GEMINI_ROUTE,
llm_router=router,
input_token_counts={"gemini-3.8-flash": 10},
)
one_candidate_hold: Final = hold({"maxOutputTokens": 100})
assert one_candidate_hold is not None
assert hold({"maxOutputTokens": 100, "candidateCount": 3}) == one_candidate_hold
@pytest.mark.parametrize(
("backend_models", "choices_are_generated"),
(
(("anthropic/claude-haiku-5-5",), False),
(("vertex_ai/claude-haiku-5-5", "anthropic/claude-haiku-5-5"), False),
(("gemini/gemini-3.8-flash",), True),
(("anthropic/claude-haiku-5-5", "gemini/gemini-3.8-flash"), True),
),
)
def test_an_n_is_held_unless_every_deployment_drops_it(
backend_models: tuple[str, ...], choices_are_generated: bool
) -> None:
router: Final = _assistant_router(*backend_models)
def hold(output_fields: dict[str, object]) -> float | None:
return estimate_request_max_cost(
request_body={"model": "assistant", "messages": GEMINI_CHAT_BODY["messages"], **output_fields},
route="/v1/chat/completions",
llm_router=router,
input_token_counts={"assistant": 10},
)
one_choice_hold: Final = hold({"max_tokens": 100})
three_choice_hold: Final = hold({"max_tokens": 100, "n": 3})
assert one_choice_hold is not None
assert three_choice_hold is not None
assert (three_choice_hold > one_choice_hold) is choices_are_generated
assert three_choice_hold >= one_choice_hold
HOLDS_FAR_OVER_A_TINY_BUDGET: Final = (
(GEMINI_ROUTE, {"contents": [GEMINI_TEXT_CONTENT], "generationConfig": {"candidateCount": 1e300}}),
("/v1/chat/completions", {**GEMINI_CHAT_BODY, "n": 1e300}),
)
@pytest.mark.asyncio
@pytest.mark.parametrize(("route", "request_body"), HOLDS_FAR_OVER_A_TINY_BUDGET)
async def test_a_hold_far_over_the_budget_keeps_the_spend_already_counted(
spend_counter_cache: DualCache, route: str, request_body: dict[str, object]
) -> None:
counter_key: Final = f"spend:key:{TINY_BUDGET_KEY_TOKEN}"
counted_spend: Final = 0.004
hold_estimate: Final = estimate_request_max_cost(request_body=request_body, route=route, llm_router=None)
assert hold_estimate is not None and 1e20 < hold_estimate < math.inf
await spend_counter_cache.async_set_cache(key=counter_key, value=counted_spend)
reservation: Final = await _reserve_for_tiny_budget_key(route, request_body)
assert reservation is not None
assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(0.01)
await release_unbound_budget_reservation(reservation)
assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(counted_spend)
def test_a_negative_max_tokens_holds_the_default_output() -> None:
uncapped_hold: Final = _chat_hold(GEMINI_CHAT_BODY)
assert uncapped_hold is not None
assert _chat_hold({**GEMINI_CHAT_BODY, "max_tokens": -1000}) == uncapped_hold
def test_a_text_completion_with_max_tokens_zero_holds_only_its_prompt() -> None:
body: Final = {"model": "gemini-3.8-flash", "prompt": GEMINI_PROMPT * 50, "max_tokens": 0}
prompt_cost: Final = estimate_request_input_cost(request_body=body, route="/v1/completions", llm_router=None)
assert prompt_cost is not None and prompt_cost > 0
assert estimate_request_max_cost(request_body=body, route="/v1/completions", llm_router=None) == prompt_cost
GEMINI_LONG_CONTENTS: Final = {"contents": [{"role": "user", "parts": [{"text": GEMINI_PROMPT * 500}]}]}
@pytest.mark.parametrize(
"stray_fields",
(
{"prompt": ""},
{"messages": [], "n": 5},
{"input": "", "max_tokens": 1},
{"query": "", "documents": [], "max_completion_tokens": 1},
{"inferenceConfig": {"maxTokens": 1}},
),
)
def test_fields_of_other_apis_on_a_gemini_route_leave_its_hold_alone(stray_fields: dict[str, object]) -> None:
body: Final = {**GEMINI_LONG_CONTENTS, "generationConfig": {"maxOutputTokens": 4000, "candidateCount": 2}}
gemini_hold: Final = _gemini_hold(body)
assert gemini_hold is not None
assert _gemini_hold({**body, **stray_fields}) == gemini_hold
@pytest.mark.parametrize("stray_input", ({"prompt": ""}, {"messages": []}, {"input": ""}))
@pytest.mark.parametrize("model", ("gpt-4", "claude-sonnet-5-5"))
@pytest.mark.asyncio
async def test_a_stray_input_field_on_a_gemini_route_never_hides_the_counted_contents(
model: str, stray_input: dict[str, object]
) -> None:
async def counted(body: dict[str, object]) -> dict[str, int]:
return dict(
await count_request_input_tokens(
request_body=body,
route=f"/v1beta/models/{model}:generateContent",
llm_router=None,
team_id=None,
raw_body=json.dumps(body).encode(),
)
)
contents_counts: Final = await counted(GEMINI_LONG_CONTENTS)
assert contents_counts[model] > 1000
assert await counted({**GEMINI_LONG_CONTENTS, **stray_input}) == contents_counts
def _tiered_deployment(input_cost_per_token: float) -> Deployment:
return Deployment(
model_name="tiered-group",

29
tests/unit/proxy/spend_tracking/test_input_tokens.py
View file

@ -280,3 +280,32 @@ def test_a_compaction_item_counts_only_its_summary(model: str) -> None:
assert count_input_tokens_for_model(request_body=padded_body, model=model) == count_input_tokens_for_model(
request_body=plain_body, model=model
)
@pytest.mark.parametrize("input_field", ("prompt", "input", "query"))
def test_a_stray_contents_field_never_replaces_the_request_input_count(input_field: str) -> None:
body: Final = {"model": CL100K_MODEL, input_field: "count every one of these words " * 200}
count: Final = count_input_tokens_for_model(request_body=body, model=CL100K_MODEL)
assert count is not None and count > 1000
assert count_input_tokens_for_model(request_body={**body, "contents": []}, model=CL100K_MODEL) == count
GEMINI_TEXT: Final = "count every one of these words " * 200
GEMINI_GREETING: Final = {"role": "user", "parts": [{"text": "hi"}]}
def _gemini_count(body: dict[str, object], model: str) -> int | None:
return count_input_tokens_for_model(request_body=body, model=model, gemini_route=True)
@pytest.mark.parametrize("parts", ([GEMINI_TEXT], {"text": GEMINI_TEXT}), ids=("bare_string", "single_object"))
@pytest.mark.parametrize("model", (CL100K_MODEL, "claude-sonnet-5-5"))
def test_gemini_parts_count_the_same_in_every_shape_the_adapter_reads(model: str, parts: object) -> None:
listed_contents: Final = _gemini_count({"contents": [{"role": "user", "parts": [{"text": GEMINI_TEXT}]}]}, model)
listed_system: Final = _gemini_count(
{"systemInstruction": {"parts": [{"text": GEMINI_TEXT}]}, "contents": [GEMINI_GREETING]}, model
)
assert listed_contents is not None and listed_contents > 1000
assert listed_system is not None and listed_system > 1000
assert _gemini_count({"contents": [{"role": "user", "parts": parts}]}, model) == listed_contents
assert _gemini_count({"systemInstruction": {"parts": parts}, "contents": [GEMINI_GREETING]}, model) == listed_system

2
tests/unit/proxy/test_budget_reservation.py
View file

@ -207,7 +207,7 @@ async def test_over_budget_window_counter_tags_clean_entity_id():
entry={"counter_key": counter.counter_key},
applied_entries=[],
reservation_cost=0.5,
current_spend=2.0,
spend_before_reservation=1.5,
)
assert exc_info.value.entity_type == "key"
assert exc_info.value.entity_id == "test-token"

24
tests/unit/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py
View file

@ -442,19 +442,16 @@ async def test_vertex_discovery_denies_other_team_vector_store_credentials():
"team_id": "team-b",
}
with patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
new=AsyncMock(return_value=UserAPIKeyAuth(team_id="team-a")),
):
with pytest.raises(HTTPException) as exc_info:
await _base_vertex_proxy_route(
endpoint="projects/p/locations/us-central1/dataStores/vs_other_team",
request=request,
fastapi_response=Response(),
get_vertex_pass_through_handler=MagicMock(),
general_settings={},
router_credentials=vector_store_credentials,
)
with pytest.raises(HTTPException) as exc_info:
await _base_vertex_proxy_route(
endpoint="projects/p/locations/us-central1/dataStores/vs_other_team",
request=request,
fastapi_response=Response(),
get_vertex_pass_through_handler=MagicMock(),
general_settings={},
user_api_key_dict=UserAPIKeyAuth(team_id="team-a"),
router_credentials=vector_store_credentials,
)
assert exc_info.value.status_code == 403
@ -542,6 +539,7 @@ async def test_vertex_discovery_allows_unregistered_provider_native_datastore_id
endpoint="projects/p/locations/us-central1/dataStores/vs_unknown",
request=request,
fastapi_response=Response(),
user_api_key_dict=UserAPIKeyAuth(),
)
assert response == {"ok": True}