diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index c1895376675..40ab377dd5e 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -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)] diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index bc64a4c7ef0..8e3d1a596b2 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -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" diff --git a/litellm/proxy/pass_through_endpoints/common_utils.py b/litellm/proxy/pass_through_endpoints/common_utils.py deleted file mode 100644 index 3ae83cbe906..00000000000 --- a/litellm/proxy/pass_through_endpoints/common_utils.py +++ /dev/null @@ -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", "") diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 21078dfc848..a79b1b2648c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -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. diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index fc6fb266d04..76f0d810e6c 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -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 diff --git a/litellm/proxy/spend_tracking/input_tokens.py b/litellm/proxy/spend_tracking/input_tokens.py index 59c7fe2a41b..0c6705fa431 100644 --- a/litellm/proxy/spend_tracking/input_tokens.py +++ b/litellm/proxy/spend_tracking/input_tokens.py @@ -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 diff --git a/tests/unit/google_genai/test_google_genai_adapter.py b/tests/unit/google_genai/test_google_genai_adapter.py index 4ab642dc99d..d622a1b85b9 100644 --- a/tests/unit/google_genai/test_google_genai_adapter.py +++ b/tests/unit/google_genai/test_google_genai_adapter.py @@ -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 diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 806d88b2467..d8b52e72aa7 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -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: diff --git a/tests/unit/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py b/tests/unit/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py index 078bd4dd402..54161b4883b 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py +++ b/tests/unit/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py @@ -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" diff --git a/tests/unit/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py b/tests/unit/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py index 827693168b7..910f360139c 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py +++ b/tests/unit/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py @@ -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" diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index a87ac7fa7c0..c2627f9faf5 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -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", diff --git a/tests/unit/proxy/spend_tracking/test_input_tokens.py b/tests/unit/proxy/spend_tracking/test_input_tokens.py index 40aee7bb00d..87192a47447 100644 --- a/tests/unit/proxy/spend_tracking/test_input_tokens.py +++ b/tests/unit/proxy/spend_tracking/test_input_tokens.py @@ -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 diff --git a/tests/unit/proxy/test_budget_reservation.py b/tests/unit/proxy/test_budget_reservation.py index 535c3da8c3a..25e0ee6a0ea 100644 --- a/tests/unit/proxy/test_budget_reservation.py +++ b/tests/unit/proxy/test_budget_reservation.py @@ -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" diff --git a/tests/unit/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py index 434ea6f3f97..5615db8714f 100644 --- a/tests/unit/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py +++ b/tests/unit/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py @@ -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}