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/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 0b91e57f01e..cfc865a84d1 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -1654,9 +1654,11 @@ _OUTPUT_TOKEN_FIELDS: Final = ("max_completion_tokens", "max_tokens", "max_outpu def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None: inference_config: Final = request_body.get("inferenceConfig") + generation_config: Final = request_body.get("generationConfig") candidates: Final = ( *(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS), inference_config.get("maxTokens") if isinstance(inference_config, Mapping) else None, + generation_config.get("maxOutputTokens") if isinstance(generation_config, Mapping) else None, ) return next((tokens for tokens in map(_to_int, candidates) if tokens is not None), None) diff --git a/litellm/proxy/spend_tracking/input_tokens.py b/litellm/proxy/spend_tracking/input_tokens.py index 927152296e9..a7e9981ac62 100644 --- a/litellm/proxy/spend_tracking/input_tokens.py +++ b/litellm/proxy/spend_tracking/input_tokens.py @@ -15,10 +15,13 @@ import asyncio import json from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Final +from typing import Final, TypeAlias + +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger +from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter from litellm.responses.litellm_completion_transformation.compaction_items import ( billable_input_since_newest_compaction, holds_compaction_summary, @@ -35,7 +38,20 @@ from litellm.rust_bridge.token_counter import ( TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS: Final = 30_000 -_INPUT_SIZE_FIELDS: Final = ("messages", "prompt", "input", "query", "documents", "tools", "tool_choice") +_GEMINI_SYSTEM_INSTRUCTION_FIELDS: Final = ("systemInstruction", "system_instruction") +_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: @@ -48,7 +64,7 @@ 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], ) -> Mapping[str, int]: @@ -74,7 +90,7 @@ 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, ...], @@ -118,7 +134,7 @@ async def _native_counts( def _count_input_tokens_for_models( - request_body: dict, + request_body: dict[str, object], models: Sequence[str], ) -> Mapping[str, int]: return MappingProxyType( @@ -142,6 +158,8 @@ def count_input_tokens_for_model(request_body: dict, model: str) -> int | None: ) except ValueError: return _count_text_tokens(model=model, text=request_body.get("messages")) + if "contents" in request_body: + return _count_gemini_input_tokens(request_body=request_body, model=model) if "prompt" in request_body: return _count_text_tokens(model=model, text=request_body.get("prompt")) if "input" in request_body: @@ -160,6 +178,20 @@ 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 + ) + translate: Final = GoogleGenAIAdapter().translate_generate_content_to_completion # pyright: ignore[reportUnknownMemberType] # adapter's **kwargs is untyped + chat_request: Final = translate( + model=model, + contents=_GEMINI_CONTENTS.validate_python(request_body["contents"]), + systemInstruction=system_instruction, + tools=request_body.get("tools"), + ) + return count_input_tokens_for_model(request_body=chat_request, model=model) + + def _count_text_tokens(model: str, text: object) -> int: if text is None: return 0 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/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 024fe5229be..f1ae7bcb5cd 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -176,6 +176,52 @@ 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: Final = "Answer tersely" + + +@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}), ({}, {})), +) +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] +) -> None: + native_cost: Final = estimate_request_max_cost( + request_body={ + "systemInstruction": {"parts": [{"text": GEMINI_SYSTEM}]}, + "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": GEMINI_SYSTEM}, + {"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 + + def _tiered_deployment(input_cost_per_token: float) -> Deployment: return Deployment( model_name="tiered-group",