From d25b58f744bc151b8dec89a11bd6158a4f1a67e9 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 01:58:46 +0000 Subject: [PATCH 1/8] fix(count_tokens): count Anthropic-format requests for OpenAI the way /v1/messages inference converts them (#46759) * fix(count_tokens): convert Anthropic-format messages for the OpenAI counter the way inference does * test(count_tokens): check /v1/messages/count_tokens sends OpenAI the input /v1/messages inference sends * fix(count_tokens): build the Anthropic count request without a cast and sync dashboard API types * test(count_tokens): compare tools by value in the public api test * fix(count_tokens): count Anthropic bodies the way /v1/messages inference sends them to OpenAI Unwrap affinity-encoded reasoning, count the structured-output schema, flatten replayed web search results, join a list system prompt, and answer a message without a role with a 400 before calling OpenAI. messages_format is now internal to the Anthropic count route instead of a public /utils/token_counter field. * test(count_tokens): expect the image detail the adapter now sends after merging main * fix(count_tokens): type normalize_native_messages without the unimported Any On Python 3.10 to 3.13 the annotation is evaluated at import, so the proxy failed to start with NameError * test(count_tokens): expect the converted Responses input in the budgeted-key count integration test --------- Co-authored-by: Mrinal Chanshetty --- .../anthropic/count_tokens/token_counter.py | 4 +- .../anthropic/pass_through/messages/utils.py | 15 +- .../responses_adapters/transformation.py | 12 +- .../anthropic/count_tokens/token_counter.py | 4 +- litellm/llms/base_llm/base_utils.py | 4 +- .../count_tokens/bedrock_token_counter.py | 4 +- litellm/llms/gemini/common_utils.py | 6 +- .../openai/responses/count_tokens/handler.py | 6 +- .../responses/count_tokens/token_counter.py | 78 ++++++- .../responses/count_tokens/transformation.py | 88 +++++++- litellm/llms/vertex_ai/common_utils.py | 6 +- .../proxy/anthropic_endpoints/endpoints.py | 17 +- litellm/proxy/proxy_server.py | 16 +- litellm/types/utils.py | 4 + .../test_openai_count_tokens_tools_wire.py | 192 +++++++++++++++- .../integration/spend/test_cache_and_quota.py | 7 +- .../count_tokens/test_token_counter.py | 207 +++++++++++++++++- ...test_openai_count_tokens_transformation.py | 108 +++++++++ .../anthropic_endpoints/test_endpoints.py | 36 ++- tests/unit/proxy/test_proxy_token_counter.py | 159 +++++--------- tests/unit/test_count_tokens_public_api.py | 2 +- 21 files changed, 814 insertions(+), 161 deletions(-) diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index a45d6c1184f..d78de0037a0 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -10,7 +10,7 @@ from litellm._logging import verbose_logger from litellm.exceptions import AuthenticationError from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.types.utils import LlmProviders, TokenCountResponse +from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests anthropic_count_tokens_handler: Final = AnthropicCountTokensHandler() @@ -34,6 +34,8 @@ class AnthropicTokenCounter(BaseTokenCounter): request_model: str = "", tools: list[dict[str, JsonValue]] | None = None, system: JsonValue = None, + messages_format: TokenCountMessagesFormat = "anthropic", + output_format: JsonValue = None, ) -> TokenCountResponse | None: """ Count tokens using Anthropic's CountTokens API. diff --git a/litellm/llms/anthropic/pass_through/messages/utils.py b/litellm/llms/anthropic/pass_through/messages/utils.py index ae5ab68242b..d8b96c1b7a0 100644 --- a/litellm/llms/anthropic/pass_through/messages/utils.py +++ b/litellm/llms/anthropic/pass_through/messages/utils.py @@ -129,6 +129,13 @@ def anthropic_system_to_openai_message(system: object) -> ChatCompletionSystemMe return ChatCompletionSystemMessage(role="system", content=system) +def normalize_native_messages(messages: Sequence[object]) -> list[dict[str, object]]: + """Apply the block rewrites every /v1/messages request gets before it reaches a provider.""" + return flatten_unencrypted_web_search_results_in_anthropic_messages( + sanitize_tool_use_ids_in_anthropic_messages(strip_empty_content_blocks_from_anthropic_messages(messages)) + ) + + def prepare_native_messages( messages: list[dict[str, JsonValue]], system: str | list[dict[str, JsonValue]] | None, @@ -140,13 +147,7 @@ def prepare_native_messages( api_base: str | None = None, presanitized: bool = False, ) -> tuple[list[dict[str, JsonValue]], str | list[dict[str, JsonValue]] | None]: - normalized: Final = ( - messages - if presanitized - else flatten_unencrypted_web_search_results_in_anthropic_messages( - sanitize_tool_use_ids_in_anthropic_messages(strip_empty_content_blocks_from_anthropic_messages(messages)) - ) - ) + normalized: Final = messages if presanitized else normalize_native_messages(messages) return cast( # cast-ok: legacy normalizers and injection preserve the JSON message and system shapes tuple[list[dict[str, JsonValue]], str | list[dict[str, JsonValue]] | None], AnthropicCacheControlHook.maybe_inject_cache_control( diff --git a/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py index 156c8d89d81..97f2db36a85 100644 --- a/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py @@ -579,6 +579,14 @@ class LiteLLMAnthropicToResponsesAPIAdapter: result["summary"] = "detailed" return result + @classmethod + def system_blocks_to_instructions(cls, system: object) -> str: + return "\n".join( + text + for block in cls._content_blocks(system) + if block.get("type") == "text" and isinstance(text := block.get("text"), str) and text + ) + def translate_request( self, anthropic_request: AnthropicMessagesRequest, @@ -628,9 +636,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: if isinstance(system, str): responses_kwargs["instructions"] = system elif isinstance(system, list): - responses_kwargs["instructions"] = "\n".join( - filter(None, (b.get("text", "") for b in system if isinstance(b, dict) and b.get("type") == "text")) - ) + responses_kwargs["instructions"] = self.system_blocks_to_instructions(system) # max_tokens -> max_output_tokens max_tokens: Final = anthropic_request.get("max_tokens") diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index 1d6d59c0592..2707d24ee07 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -12,7 +12,7 @@ from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( AzureAIAnthropicCountTokensHandler, ) from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.types.utils import LlmProviders, TokenCountResponse +from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests azure_ai_anthropic_count_tokens_handler: Final = AzureAIAnthropicCountTokensHandler() @@ -36,6 +36,8 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): request_model: str = "", tools: list[dict[str, JsonValue]] | None = None, system: object | None = None, + messages_format: TokenCountMessagesFormat = "anthropic", + output_format: JsonValue = None, ) -> TokenCountResponse | None: """ Count tokens using Azure AI Anthropic's CountTokens API. diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index fee12e34fa3..70ec6572675 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -16,7 +16,7 @@ from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk -from litellm.types.utils import Message, ProviderSpecificModelInfo, TokenCountResponse +from litellm.types.utils import Message, ProviderSpecificModelInfo, TokenCountMessagesFormat, TokenCountResponse class BaseTokenCounter(ABC): @@ -30,6 +30,8 @@ class BaseTokenCounter(ABC): request_model: str = "", tools: list[dict[str, Any]] | None = None, system: Any | None = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> TokenCountResponse | None: pass diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 3084b0fde07..7ff2fddbb57 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -16,7 +16,7 @@ from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountResponse +from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountMessagesFormat, TokenCountResponse RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api" MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api" @@ -107,6 +107,8 @@ class BedrockTokenCounter(BaseTokenCounter): request_model: str = "", tools: Sequence[Mapping[str, object]] | None = None, system: object | None = None, + messages_format: TokenCountMessagesFormat = "anthropic", + output_format: JsonValue = None, ) -> TokenCountResponse | None: if not messages: return None diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index b18c322b3a4..57ee0f33879 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -6,7 +6,7 @@ from collections.abc import Mapping, Sequence from typing import Any, Final import httpx -from pydantic import TypeAdapter +from pydantic import JsonValue, TypeAdapter import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH @@ -14,7 +14,7 @@ from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import TokenCountResponse +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse GEMINI_IMAGE_ASPECT_RATIOS: Final[dict[str, float]] = { "1:1": 1 / 1, @@ -495,6 +495,8 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): request_model: str = "", tools: list[dict[str, object]] | None = None, system: object | None = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> TokenCountResponse | None: import copy diff --git a/litellm/llms/openai/responses/count_tokens/handler.py b/litellm/llms/openai/responses/count_tokens/handler.py index 02a05b060b8..6f911e9c40f 100644 --- a/litellm/llms/openai/responses/count_tokens/handler.py +++ b/litellm/llms/openai/responses/count_tokens/handler.py @@ -5,7 +5,7 @@ Uses httpx for HTTP requests to OpenAI's /v1/responses/input_tokens endpoint. """ import json -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Final import httpx @@ -31,8 +31,9 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): api_key: str, api_base: str | None = None, timeout: float | httpx.Timeout | None = None, - tools: list[dict[str, object]] | None = None, + tools: Sequence[Mapping[str, object]] | None = None, instructions: str | None = None, + text: Mapping[str, object] | None = None, ) -> dict[str, object]: """ Handle a token counting request to OpenAI's Responses API. @@ -53,6 +54,7 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): input=input, tools=tools, instructions=instructions, + text=text, ) endpoint_url: Final = self.get_openai_count_tokens_endpoint(api_base) diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index ab3ddbd179d..b97e1329675 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -3,9 +3,16 @@ OpenAI Token Counter implementation using the Responses API /input_tokens endpoi """ import os +from collections.abc import Mapping, Sequence from typing import Any, Final +from pydantic import JsonValue +from typing_extensions import assert_never + from litellm._logging import verbose_logger +from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.llms.openai.common_utils import OpenAIError from litellm.llms.openai.responses.count_tokens.handler import ( @@ -13,8 +20,9 @@ from litellm.llms.openai.responses.count_tokens.handler import ( ) from litellm.llms.openai.responses.count_tokens.transformation import ( OpenAICountTokensConfig, + ResponsesCountRequestParts, ) -from litellm.types.utils import LlmProviders, TokenCountResponse +from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests openai_count_tokens_handler: Final = OpenAICountTokensHandler() @@ -29,6 +37,37 @@ class OpenAITokenCounter(BaseTokenCounter): ) -> bool: return custom_llm_provider == LlmProviders.OPENAI.value + @staticmethod + def _system_instructions(system: object) -> str | None: + if system is None or isinstance(system, str): + return system + return LiteLLMAnthropicToResponsesAPIAdapter.system_blocks_to_instructions(system) or None + + @staticmethod + def _responses_request_parts( + model_to_use: str, + messages: Sequence[Mapping[str, object]], + tools: Sequence[Mapping[str, object]] | None, + system: object, + messages_format: TokenCountMessagesFormat, + output_format: JsonValue, + ) -> ResponsesCountRequestParts: + match messages_format: + case "anthropic": + return OpenAICountTokensConfig.anthropic_messages_to_responses_input( + model=model_to_use, messages=messages, system=system, tools=tools, output_format=output_format + ) + case "openai": + input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input( + [dict(message) for message in messages] + ) + fallback_instructions: Final = OpenAITokenCounter._system_instructions(system) + return OpenAICountTokensConfig.count_request_parts( + input_items, instructions if instructions is not None else fallback_instructions, tools + ) + case _: + assert_never(messages_format) + async def count_tokens( self, model_to_use: str, @@ -38,6 +77,8 @@ class OpenAITokenCounter(BaseTokenCounter): request_model: str = "", tools: list[dict[str, object]] | None = None, system: object = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> TokenCountResponse | None: """ Count tokens using OpenAI's Responses API /input_tokens endpoint. @@ -45,6 +86,22 @@ class OpenAITokenCounter(BaseTokenCounter): if not messages: return None + malformed_error: Final = ( + OpenAICountTokensConfig.malformed_anthropic_messages_error(messages) + if messages_format == "anthropic" + else None + ) + if malformed_error is not None: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="openai_api", + error=True, + error_message=malformed_error, + status_code=400, + ) + deployment = deployment or {} litellm_params: Final = deployment.get("litellm_params", {}) @@ -59,12 +116,14 @@ class OpenAITokenCounter(BaseTokenCounter): api_base: Final = litellm_params.get("api_base") - # Convert chat messages to Responses API input format - input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages) - - # Use system param if instructions not extracted from messages - if instructions is None and system is not None: - instructions = system if isinstance(system, str) else str(system) + input_items, instructions, counted_tools, text = self._responses_request_parts( + model_to_use=model_to_use, + messages=messages, + tools=tools, + system=system, + messages_format=messages_format, + output_format=output_format, + ) # If no input items were produced (e.g., system-only messages), fall back to local counting if not input_items: @@ -73,11 +132,12 @@ class OpenAITokenCounter(BaseTokenCounter): try: result: Final = await openai_count_tokens_handler.handle_count_tokens_request( model=model_to_use, - input=input_items if input_items is not None else [], + input=input_items, api_key=api_key, api_base=api_base, - tools=tools, + tools=counted_tools, instructions=instructions, + text=text, ) if result is not None: diff --git a/litellm/llms/openai/responses/count_tokens/transformation.py b/litellm/llms/openai/responses/count_tokens/transformation.py index 75684b06d05..90bfa4a8655 100644 --- a/litellm/llms/openai/responses/count_tokens/transformation.py +++ b/litellm/llms/openai/responses/count_tokens/transformation.py @@ -5,10 +5,18 @@ This module handles the transformation of requests to OpenAI's /v1/responses/inp """ from collections.abc import Mapping, Sequence -from typing import Any, Final, Literal +from typing import Any, Final, Literal, TypeAlias +from pydantic import JsonValue, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict +from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages +from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.llms.anthropic import AnthropicMessagesRequest, AnthropicOutputSchema + class ResponsesInputTextPart(TypedDict): type: ReadOnly[Literal["input_text"]] @@ -31,6 +39,30 @@ ResponsesInputPart = ResponsesInputTextPart | ResponsesInputImagePart | Response ResponsesContentRole = Literal["user", "assistant"] +ResponsesCountRequestParts: TypeAlias = tuple[ + tuple[Mapping[str, object], ...], + str | None, + tuple[Mapping[str, object], ...] | None, + Mapping[str, object] | None, +] +_RESPONSES_COUNT_REQUEST_PARTS: Final = TypeAdapter(ResponsesCountRequestParts) +_TRANSLATED_REQUEST: Final = TypeAdapter(Mapping[str, object]) + + +class _AnthropicCountMessage(TypedDict): + role: ReadOnly[str] + + +_ANTHROPIC_COUNT_MESSAGES: Final = TypeAdapter(tuple[_AnthropicCountMessage, ...]) +_ANTHROPIC_OUTPUT_SCHEMA: Final = TypeAdapter(AnthropicOutputSchema) + + +def _anthropic_output_schema(output_format: JsonValue) -> AnthropicOutputSchema | None: + try: + return _ANTHROPIC_OUTPUT_SCHEMA.validate_python(output_format) + except ValidationError: + return None + def _chat_image_block_to_responses_part(image_url: object) -> ResponsesInputImagePart | None: url: Final = image_url.get("url") if isinstance(image_url, Mapping) else image_url @@ -118,8 +150,9 @@ class OpenAICountTokensConfig: self, model: str, input: str | Sequence[object], - tools: list[dict[str, object]] | None = None, + tools: Sequence[Mapping[str, object]] | None = None, instructions: str | None = None, + text: Mapping[str, object] | None = None, ) -> dict[str, object]: """ Transform request to OpenAI Responses API token counting format. @@ -135,7 +168,10 @@ class OpenAICountTokensConfig: request["instructions"] = instructions if tools is not None: - request["tools"] = self._transform_tools_for_responses_api(tools) + request["tools"] = self._transform_tools_for_responses_api([dict(tool) for tool in tools]) + + if text is not None: + request["text"] = text return request @@ -194,7 +230,7 @@ class OpenAICountTokensConfig: @staticmethod def messages_to_responses_input( messages: list[dict[str, Any]], - ) -> tuple: + ) -> tuple[list[dict[str, object]], str | None]: """ Convert standard chat messages format to OpenAI Responses API input format. @@ -257,3 +293,47 @@ class OpenAICountTokensConfig: instructions: Final = "\n".join(instructions_parts) if instructions_parts else None return input_items, instructions + + @staticmethod + def malformed_anthropic_messages_error(messages: Sequence[Mapping[str, object]]) -> str | None: + try: + _ANTHROPIC_COUNT_MESSAGES.validate_python(messages) + except ValidationError as e: + return "; ".join( + f"messages.{'.'.join(str(part) for part in error['loc'])}: {error['msg']}" for error in e.errors() + ) + return None + + @staticmethod + def anthropic_messages_to_responses_input( + model: str, + messages: Sequence[Mapping[str, object]], + system: object, + tools: Sequence[Mapping[str, object]] | None, + output_format: JsonValue = None, + ) -> ResponsesCountRequestParts: + """Convert with the adapter /v1/messages inference uses, so the count sees the same input it bills.""" + anthropic_request: Final = AnthropicMessagesRequest( + model=model, + messages=normalize_native_messages(messages), + system=system if isinstance(system, (str, list)) else None, + tools=[dict(tool) for tool in tools] if tools else None, + output_format=_anthropic_output_schema(output_format), + ) + translated: Final = _TRANSLATED_REQUEST.validate_python( + LiteLLMAnthropicToResponsesAPIAdapter().translate_request( + anthropic_request, include_encrypted_reasoning=False + ) + ) + return OpenAICountTokensConfig.count_request_parts( + ResponsesAPIRequestUtils.restore_encrypted_content_item_ids_in_input(translated.get("input")), + translated.get("instructions"), + translated.get("tools"), + translated.get("text"), + ) + + @staticmethod + def count_request_parts( + input_items: object, instructions: object, tools: object, text: object = None + ) -> ResponsesCountRequestParts: + return _RESPONSES_COUNT_REQUEST_PARTS.validate_python((input_items, instructions, tools, text)) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 381986877c0..0e4cb856632 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -8,7 +8,7 @@ from types import MappingProxyType from typing import Any, Final, Literal, cast, get_type_hints import httpx -from pydantic import TypeAdapter, ValidationError +from pydantic import JsonValue, TypeAdapter, ValidationError from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm @@ -23,7 +23,7 @@ from litellm.types.llms.vertex_ai import ( PartType, Schema, ) -from litellm.types.utils import TokenCountResponse +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse from litellm.utils import supports_response_schema, supports_system_messages GEMINI_VIDEO_METADATA_KEYS: Final = MappingProxyType( @@ -1278,6 +1278,8 @@ class VertexAITokenCounter(BaseTokenCounter): request_model: str = "", tools: list[dict[str, object]] | None = None, system: object | None = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> TokenCountResponse | None: import copy diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index d4727cc8d8d..1f17f0a42d2 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -2,10 +2,12 @@ Unified /v1/messages endpoint - (Anthropic Spec) """ +from collections.abc import Mapping from typing import Final from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse +from pydantic import JsonValue import litellm from litellm.anthropic_interface.exceptions import ( @@ -289,6 +291,14 @@ async def anthropic_response( ) +def _requested_output_format(data: Mapping[str, JsonValue]) -> JsonValue: + output_format: Final = data.get("output_format") + output_config: Final = data.get("output_config") + if isinstance(output_format, dict) or not isinstance(output_config, dict): + return output_format + return output_config.get("format") + + @router.post( "/v1/messages/count_tokens", tags=["[beta] Anthropic Messages Token Counting"], @@ -318,7 +328,7 @@ async def count_tokens( Returns: {"input_tokens": } """ - from litellm.proxy.proxy_server import token_counter as internal_token_counter + from litellm.proxy.proxy_server import count_request_tokens litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: @@ -345,10 +355,11 @@ async def count_tokens( system=data.get("system"), ) - # Call the internal token counter function with direct request flag set to False - token_response: Final = await internal_token_counter( + token_response: Final = await count_request_tokens( request=token_request, call_endpoint=True, + messages_format="anthropic", + output_format=_requested_output_format(data), ) _token_response_dict: dict = {} if isinstance(token_response, TokenCountResponse): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index fb6a3a8a88d..621dd0d8c04 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -173,6 +173,7 @@ from litellm.types.utils import ( ModelResponseStream, StreamingChoices, TextCompletionResponse, + TokenCountMessagesFormat, TokenCountResponse, echoed_cost_map_pricing_fields, is_server_derived_pricing_key, @@ -14415,6 +14416,8 @@ async def _try_provider_token_count( request_model: str, tools: list | None = None, system: str | None = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> Optional["TokenCountResponse"]: """Attempt provider-specific token counting. Returns result on success, None to fall through to local counting.""" if not provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider): @@ -14428,6 +14431,8 @@ async def _try_provider_token_count( request_model=request_model, tools=tools, system=system, + messages_format=messages_format, + output_format=output_format, ) except httpx.HTTPStatusError as e: error_message: Final = getattr(e, "message", None) or str(e) @@ -14470,8 +14475,15 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) Returns: TokenCountResponse """ - global llm_router + return await count_request_tokens(request=request, call_endpoint=call_endpoint) + +async def count_request_tokens( + request: TokenCountRequest, + call_endpoint: bool = False, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, +) -> TokenCountResponse: prompt: Final = request.prompt messages: Final = request.messages contents: Final = request.contents @@ -14530,6 +14542,8 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) request_model=request.model, tools=tools, system=system, + messages_format=messages_format, + output_format=output_format, ) if result is not None: return result diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d3c884b4a1f..a957e6957ee 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -9,6 +9,7 @@ from typing import ( Any, Final, Literal, + TypeAlias, get_args, ) @@ -4471,6 +4472,9 @@ class LiteLLMLoggingBaseClass: pass +TokenCountMessagesFormat: TypeAlias = Literal["openai", "anthropic"] + + class TokenCountResponse(LiteLLMPydanticObjectBase): total_tokens: int request_model: str diff --git a/tests/integration/providers/test_openai_count_tokens_tools_wire.py b/tests/integration/providers/test_openai_count_tokens_tools_wire.py index bfae84e333d..3ab5826006f 100644 --- a/tests/integration/providers/test_openai_count_tokens_tools_wire.py +++ b/tests/integration/providers/test_openai_count_tokens_tools_wire.py @@ -1,5 +1,5 @@ import json -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from typing import Final import pytest @@ -36,6 +36,146 @@ def _count_response(_request: Request) -> Reply: return Reply(body=json.dumps({"input_tokens": _COUNT}).encode()) +_PDF_BASE64: Final = "JVBERi0xLjQK" +_PNG_BASE64: Final = "aW1hZ2U=" +_LOOKUP_TOOL: Final[dict[str, JsonValue]] = { + "name": "lookup", + "description": "Find a record", + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, +} +_RESPONSES_REPLY: Final[Mapping[str, JsonValue]] = { + "id": "resp_count_parity", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": "msg_count_parity", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "OK", "annotations": []}], + } + ], + "usage": {"input_tokens": 10, "output_tokens": 1, "total_tokens": 11}, +} +_WRAPPED_REASONING: Final = "litellm_enc:bW9kZWxfaWQ6ZGVwbG95bWVudC0x;gAAAAoriginal" +_ANSWER_SCHEMA: Final[dict[str, JsonValue]] = {"type": "object", "properties": {"answer": {"type": "string"}}} +_ANTHROPIC_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = { + "structured-output": { + "output_config": {"format": {"type": "json_schema", "schema": _ANSWER_SCHEMA}}, + "messages": [{"role": "user", "content": "Answer in JSON."}], + }, + "affinity-wrapped-reasoning": { + "messages": [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{_WRAPPED_REASONING}"}, + {"type": "text", "text": "OK"}, + ], + }, + {"role": "user", "content": "again"}, + ], + }, + "base64-pdf": { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the attached contract."}, + { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": _PDF_BASE64}, + }, + ], + } + ], + }, + "base64-image": { + "system": [{"type": "text", "text": "Describe images briefly."}], + "messages": [ + { + "role": "user", + "content": [ + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": _PNG_BASE64}}, + {"type": "text", "text": "What is this?"}, + ], + } + ], + }, + "replayed-web-search": { + "messages": [ + {"role": "user", "content": "What is the newest LiteLLM release?"}, + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_01", + "name": "web_search", + "input": {"query": "newest LiteLLM release"}, + }, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": "LiteLLM ships a weekly release candidate and a stable line.", + } + ], + }, + {"type": "text", "text": "The newest release is in the release notes."}, + ], + }, + {"role": "user", "content": "Thanks."}, + ], + }, + "tool-turns": { + "system": "Use the lookup tool.", + "tools": [_LOOKUP_TOOL], + "messages": [ + {"role": "user", "content": "Find record 14."}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Looking it up."}, + {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "14"}}, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "14. Indemnification."}], + }, + ], + }, +} + + +def _responses_peer(request: Request) -> Reply: + if request.target == "/v1/responses/input_tokens": + return _count_response(request) + if request.target == "/v1/responses": + return Reply(body=json.dumps(_RESPONSES_REPLY).encode()) + return Reply(status=404, body=json.dumps({"error": request.target}).encode()) + + +def _body_for(wire_requests: tuple[Request, ...], target: str) -> dict[str, JsonValue]: + matching: Final = tuple(request for request in wire_requests if request.target == target) + assert len(matching) == 1, [request.target for request in wire_requests] + return _JSON_OBJECT.validate_json(matching[0].body) + + +def _counted_fields(body: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: + return {key: body[key] for key in ("input", "instructions", "tools", "text") if body.get(key) is not None} + + def _count_body(wire: Wire) -> dict[str, JsonValue]: requests: Final = wire.drain() assert len(requests) == 1, requests @@ -159,3 +299,53 @@ def test_utils_token_counter_keeps_chat_tool_conversion(counting_proxy: OwnedPro "parameters": _TOOL_SCHEMA, } ] + + +@pytest.mark.parametrize("case", sorted(_ANTHROPIC_BODIES)) +def test_messages_count_tokens_sends_the_input_v1_messages_inference_sends( + counting_proxy: OwnedProxy, case: str +) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_responses_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + body: Final = {"model": model, **_ANTHROPIC_BODIES[case]} + generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16}) + counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body) + assert generated.status_code == 200, generated.text + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT + requests: Final = wire.drain() + inference_body: Final = _body_for(requests, "/v1/responses") + count_body: Final = _body_for(requests, "/v1/responses/input_tokens") + assert _counted_fields(count_body) == _counted_fields(inference_body) + assert count_body["model"] == inference_body["model"] + + +def test_messages_count_tokens_without_a_role_answers_400_before_calling_openai(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": model, "messages": [{"role": "user", "content": "hello"}, {"content": "no role"}]}, + ) + assert response.status_code == 400, response.text + assert "messages.1.role: Field required" in response.text + assert wire.drain() == () + + +def test_google_count_tokens_sends_the_system_instruction_text(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + f"/v1beta/models/{model}:countTokens", + { + "systemInstruction": {"parts": [{"text": "Be brief."}, {"text": "Answer in English."}]}, + "contents": [{"role": "user", "parts": [{"text": "hello"}]}], + }, + ) + assert response.status_code == 200, response.text + assert _count_body(wire)["instructions"] == "Be brief.\nAnswer in English." diff --git a/tests/integration/spend/test_cache_and_quota.py b/tests/integration/spend/test_cache_and_quota.py index d02ebccf710..801a14c87d3 100644 --- a/tests/integration/spend/test_cache_and_quota.py +++ b/tests/integration/spend/test_cache_and_quota.py @@ -369,7 +369,12 @@ def test_repeated_count_tokens_on_budgeted_key_does_not_reserve_budget_or_block_ { "path": f"/{counting_upstream.scenario_id}/v1/responses/input_tokens", "authorization": "Bearer integration-provider-key", - "body": {"model": "gpt-4o-mini", "input": [{"role": "user", "content": "hello!!!"}]}, + "body": { + "model": "gpt-4o-mini", + "input": [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hello!!!"}]} + ], + }, "method": "POST", "api_key": "", } diff --git a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py index 2e0fc55cda4..976ecff01ea 100644 --- a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py +++ b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py @@ -5,11 +5,12 @@ from typing import Final import httpx import pytest import respx -from pydantic import ValidationError +from pydantic import JsonValue, ValidationError import litellm from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter -from litellm.types.utils import TokenCountResponse +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse COUNT_URL: Final = "https://api.openai.com/v1/responses/input_tokens" @@ -163,6 +164,208 @@ async def test_upstream_error_status_becomes_an_error_response_with_that_status_ ) +async def _sent_count_body( + monkeypatch: pytest.MonkeyPatch, + messages: list[dict[str, object]], + system: object, + messages_format: TokenCountMessagesFormat, + output_format: JsonValue = None, +) -> object: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 1})) + await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=messages, + contents=None, + deployment={"litellm_params": {"api_key": "sk-test"}}, + request_model="openai/gpt-4o", + system=system, + messages_format=messages_format, + output_format=output_format, + ) + return json.loads(route.calls.last.request.content) + + +@pytest.mark.asyncio +async def test_anthropic_format_sends_the_document_image_and_tool_turns_to_the_count_endpoint( + monkeypatch: pytest.MonkeyPatch, +) -> None: + messages: Final[list[dict[str, object]]] = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize."}, + {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBE"}}, + {"type": "image", "source": {"type": "url", "url": "https://example.com/page.png"}}, + ], + }, + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "found"}]}, + ] + + sent: Final = await _sent_count_body( + monkeypatch, messages, [{"type": "text", "text": "Be brief."}], messages_format="anthropic" + ) + + assert sent == { + "model": "gpt-4o", + "input": [ + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "Summarize."}, + {"type": "input_file", "filename": "document.pdf", "file_data": "data:application/pdf;base64,JVBE"}, + {"type": "input_image", "image_url": "https://example.com/page.png", "detail": "auto"}, + ], + }, + {"type": "function_call", "call_id": "toolu_01", "name": "lookup", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "toolu_01", "output": "found"}, + ], + "instructions": "Be brief.", + } + + +@pytest.mark.asyncio +async def test_openai_format_keeps_the_chat_conversion(monkeypatch: pytest.MonkeyPatch) -> None: + messages: Final[list[dict[str, object]]] = [ + {"role": "system", "content": "Be brief."}, + {"role": "user", "content": [{"type": "text", "text": "Summarize."}]}, + {"role": "tool", "tool_call_id": "call_1", "content": "found"}, + ] + + sent: Final = await _sent_count_body(monkeypatch, messages, system=None, messages_format="openai") + + assert sent == { + "model": "gpt-4o", + "input": [ + {"role": "user", "content": "Summarize."}, + {"type": "function_call_output", "call_id": "call_1", "output": "found"}, + ], + "instructions": "Be brief.", + } + + +@pytest.mark.asyncio +async def test_anthropic_format_sends_affinity_wrapped_reasoning_unwrapped_as_inference_does( + monkeypatch: pytest.MonkeyPatch, +) -> None: + wrapped: Final = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("gAAAAoriginal", "deployment-1") + messages: Final[list[dict[str, object]]] = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{wrapped}"}, + {"type": "text", "text": "OK"}, + ], + }, + {"role": "user", "content": "again"}, + ] + + sent: Final = await _sent_count_body(monkeypatch, messages, system=None, messages_format="anthropic") + + assert sent == { + "model": "gpt-4o", + "input": [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]}, + {"type": "reasoning", "summary": [], "encrypted_content": "gAAAAoriginal"}, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "OK"}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "again"}]}, + ], + } + + +@pytest.mark.asyncio +async def test_anthropic_format_counts_the_structured_output_schema(monkeypatch: pytest.MonkeyPatch) -> None: + schema: Final[dict[str, JsonValue]] = {"type": "object", "properties": {"answer": {"type": "string"}}} + + sent: Final = await _sent_count_body( + monkeypatch, + [{"role": "user", "content": "Answer in JSON."}], + system=None, + messages_format="anthropic", + output_format={"type": "json_schema", "schema": schema}, + ) + + assert sent == { + "model": "gpt-4o", + "input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Answer in JSON."}]}], + "text": {"format": {"type": "json_schema", "name": "structured_output", "schema": schema, "strict": False}}, + } + + +@pytest.mark.asyncio +async def test_anthropic_format_sends_a_replayed_web_search_result_as_its_text(monkeypatch: pytest.MonkeyPatch) -> None: + snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." + messages: Final[list[dict[str, object]]] = [ + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": snippet, + } + ], + }, + ], + }, + {"role": "user", "content": "Thanks."}, + ] + + sent: Final = json.dumps(await _sent_count_body(monkeypatch, messages, system=None, messages_format="anthropic")) + + assert snippet in sent + assert "https://docs.litellm.ai/release_notes" in sent + + +@pytest.mark.asyncio +async def test_anthropic_format_message_without_a_role_answers_400_without_calling_openai( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 1})) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}, {"content": "no role"}], + contents=None, + deployment={"litellm_params": {"api_key": "sk-test"}}, + request_model="openai/gpt-4o", + messages_format="anthropic", + ) + + assert route.called is False + assert result is not None + assert (result.error, result.status_code, result.error_message) == (True, 400, "messages.1.role: Field required") + + +@pytest.mark.asyncio +async def test_openai_format_sends_a_list_system_as_its_text(monkeypatch: pytest.MonkeyPatch) -> None: + sent: Final = await _sent_count_body( + monkeypatch, + [{"role": "user", "content": "hi"}], + system=[{"type": "text", "text": "Be brief."}, {"type": "text", "text": "Answer in English."}], + messages_format="openai", + ) + + assert sent == { + "model": "gpt-4o", + "input": [{"role": "user", "content": "hi"}], + "instructions": "Be brief.\nAnswer in English.", + } + + @pytest.mark.asyncio async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_decoder_message( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py index 98ddfa08018..0a5039791ee 100644 --- a/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py +++ b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py @@ -2,6 +2,9 @@ from typing import Final import pytest +from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) from litellm.llms.openai.responses.count_tokens.transformation import ( OpenAICountTokensConfig, ) @@ -484,3 +487,108 @@ def test_get_required_headers(): assert headers["Authorization"] == "Bearer sk-test-key" assert headers["Content-Type"] == "application/json" + + +PDF_BASE64: Final = "JVBERi0xLjcK" +PNG_BASE64: Final = "iVBORw0KGgo=" + +ANTHROPIC_AGENT_TURN: Final[list[dict[str, object]]] = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the indemnity terms."}, + { + "type": "document", + "title": "contract.pdf", + "source": {"type": "base64", "media_type": "application/pdf", "data": PDF_BASE64}, + }, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": PNG_BASE64}}, + ], + }, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Let me pull clause 14."}, + {"type": "tool_use", "id": "toolu_01", "name": "lookup_clause", "input": {"id": "14"}}, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "14. Indemnification."}], + }, +] + +ANTHROPIC_TOOLS: Final[list[dict[str, object]]] = [ + { + "name": "lookup_clause", + "description": "Fetch a clause by id", + "input_schema": {"type": "object", "properties": {"id": {"type": "string"}}}, + } +] + + +@pytest.mark.parametrize( + "system", + [None, "You review contracts.", [{"type": "text", "text": "You review contracts."}]], +) +def test_anthropic_messages_convert_exactly_as_v1_messages_inference_converts_them(system: object) -> None: + inference_request: Final[dict[str, object]] = { + "model": "gpt-4o", + "messages": ANTHROPIC_AGENT_TURN, + "tools": ANTHROPIC_TOOLS, + **({"system": system} if system is not None else {}), + } + inference: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_request( + inference_request, include_encrypted_reasoning=False + ) + + counted: Final = OpenAICountTokensConfig.anthropic_messages_to_responses_input( + model="gpt-4o", messages=ANTHROPIC_AGENT_TURN, system=system, tools=ANTHROPIC_TOOLS + ) + + assert counted == ( + tuple(inference["input"]), + inference.get("instructions"), + tuple(inference["tools"]), + inference.get("text"), + ) + + +def test_anthropic_messages_keep_every_block_the_chat_converter_drops() -> None: + input_items, _, _, _ = OpenAICountTokensConfig.anthropic_messages_to_responses_input( + model="gpt-4o", messages=ANTHROPIC_AGENT_TURN, system=None, tools=None + ) + + assert input_items == ( + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "Summarize the indemnity terms."}, + { + "type": "input_file", + "filename": "contract.pdf", + "file_data": f"data:application/pdf;base64,{PDF_BASE64}", + }, + {"type": "input_image", "image_url": f"data:image/png;base64,{PNG_BASE64}", "detail": "auto"}, + ], + }, + {"type": "function_call", "call_id": "toolu_01", "name": "lookup_clause", "arguments": '{"id": "14"}'}, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Let me pull clause 14."}], + }, + {"type": "function_call_output", "call_id": "toolu_01", "output": "14. Indemnification."}, + ) + + +def test_anthropic_list_system_becomes_joined_instruction_text() -> None: + _, instructions, _, _ = OpenAICountTokensConfig.anthropic_messages_to_responses_input( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + system=[{"type": "text", "text": "First rule."}, {"type": "text", "text": "Second rule."}], + tools=None, + ) + + assert instructions == "First rule.\nSecond rule." diff --git a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py index c0e7a4fdb47..69d182667ec 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py @@ -278,13 +278,17 @@ class TestHttpExceptionDictDetail: request.headers = {} with ( - patch.object(ep, "read_request_body", new=AsyncMock(return_value={})), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object( + ep, "read_request_body", new=AsyncMock(return_value={}) + ), # test-quality-ok: endpoint reads the body via a module function; no injection seam patch.object( # test-quality-ok: the guardrail raise happens deep inside this call; the test targets the endpoint's except block ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=AsyncMock(side_effect=exc), ), - patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam + patch.object( + proxy_server, "proxy_logging_obj" + ) as mock_logging, # test-quality-ok: module global imported at call time; no injection seam ): mock_logging.post_call_failure_hook = AsyncMock() response = await ep.anthropic_response( @@ -375,9 +379,15 @@ class TestErrorLogCarriesCallId: request.headers = {} with ( - patch.object(ep, "read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"})), # test-quality-ok: endpoint reads the body via a module function; no injection seam - patch.object(ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process), # test-quality-ok: the provider failure happens inside this call; the test targets the endpoint's except block - patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam + patch.object( + ep, "read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"}) + ), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object( + ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process + ), # test-quality-ok: the provider failure happens inside this call; the test targets the endpoint's except block + patch.object( + proxy_server, "proxy_logging_obj" + ) as mock_logging, # test-quality-ok: module global imported at call time; no injection seam caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), ): mock_logging.post_call_failure_hook = AsyncMock() @@ -408,9 +418,15 @@ class TestErrorLogCarriesCallId: request.headers = {} with ( - patch.object(ep, "read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"})), # test-quality-ok: endpoint reads the body via a module function; no injection seam - patch.object(ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process), # test-quality-ok: the proxy shaped failure happens inside this call; the test targets the endpoint's except block - patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam + patch.object( + ep, "read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"}) + ), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object( + ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process + ), # test-quality-ok: the proxy shaped failure happens inside this call; the test targets the endpoint's except block + patch.object( + proxy_server, "proxy_logging_obj" + ) as mock_logging, # test-quality-ok: module global imported at call time; no injection seam ): mock_logging.post_call_failure_hook = AsyncMock() response = await ep.anthropic_response( @@ -440,7 +456,9 @@ class TestErrorLogCarriesCallId: "read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet", "messages": [{"role": "user", "content": "hi"}]}), ), - patch.object(proxy_server, "token_counter", new=AsyncMock(side_effect=RuntimeError("tokenizer down"))), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block + patch.object( + proxy_server, "count_request_tokens", new=AsyncMock(side_effect=RuntimeError("tokenizer down")) + ), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), pytest.raises(HTTPException) as raised, ): diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index c93c199cfa2..8cd6887d891 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -30,7 +30,6 @@ from litellm.proxy.proxy_server import token_counter from litellm.types.utils import TokenCountResponse - @pytest.mark.asyncio async def test_vLLM_token_counting(): """ @@ -63,9 +62,7 @@ async def test_vLLM_token_counting(): print("response: ", response) - assert ( - response.tokenizer_type == "openai_tokenizer" - ) # SHOULD use the default tokenizer + assert response.tokenizer_type == "openai_tokenizer" # SHOULD use the default tokenizer assert response.model_used == "wolfram/miquliz-120b-v2.0" @@ -98,9 +95,7 @@ async def test_token_counting_model_not_in_model_list(): print("response: ", response) - assert ( - response.tokenizer_type == "openai_tokenizer" - ) # SHOULD use the OpenAI tokenizer + assert response.tokenizer_type == "openai_tokenizer" # SHOULD use the OpenAI tokenizer assert response.model_used == "special-alias" @@ -133,9 +128,7 @@ async def test_gpt_token_counting(): print("response: ", response) - assert ( - response.tokenizer_type == "openai_tokenizer" - ) # SHOULD use the OpenAI tokenizer + assert response.tokenizer_type == "openai_tokenizer" # SHOULD use the OpenAI tokenizer assert response.request_model == "gpt-4" @@ -173,10 +166,8 @@ async def test_anthropic_messages_count_tokens_endpoint(): anthropic_endpoints.read_request_body = mock_read_request_body # Mock the internal token_counter function to return a controlled response - async def mock_token_counter(request, call_endpoint=False): - assert ( - call_endpoint == True - ), "Should be called with call_endpoint=True for Anthropic endpoint" + async def mock_token_counter(request, call_endpoint=False, messages_format="openai", output_format=None): + assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "claude-3-sonnet-20240229" assert request.messages == [{"role": "user", "content": "Hello Claude!"}] @@ -192,8 +183,8 @@ async def test_anthropic_messages_count_tokens_endpoint(): # Patch the imported token_counter function from proxy_server import litellm.proxy.proxy_server as proxy_server - original_token_counter = proxy_server.token_counter - proxy_server.token_counter = mock_token_counter + original_token_counter = proxy_server.count_request_tokens + proxy_server.count_request_tokens = mock_token_counter try: # Call the endpoint @@ -210,7 +201,7 @@ async def test_anthropic_messages_count_tokens_endpoint(): finally: # Restore original functions anthropic_endpoints.read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter @pytest.mark.asyncio @@ -247,10 +238,8 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): anthropic_endpoints.read_request_body = mock_read_request_body # Mock the internal token_counter function to return a controlled response - async def mock_token_counter(request, call_endpoint=True): - assert ( - call_endpoint == True - ), "Should be called with call_endpoint=True for Anthropic endpoint" + async def mock_token_counter(request, call_endpoint=True, messages_format="openai", output_format=None): + assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "gpt-4" assert request.messages == [{"role": "user", "content": "Hello GPT!"}] @@ -266,8 +255,8 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): # Patch the imported token_counter function from proxy_server import litellm.proxy.proxy_server as proxy_server - original_token_counter = proxy_server.token_counter - proxy_server.token_counter = mock_token_counter + original_token_counter = proxy_server.count_request_tokens + proxy_server.count_request_tokens = mock_token_counter try: # Call the endpoint @@ -284,7 +273,7 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): finally: # Restore original functions anthropic_endpoints.read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter @pytest.mark.asyncio @@ -407,9 +396,7 @@ async def test_factory_anthropic_endpoint_calls_anthropic_counter(): # Mock the global handler instance in token_counter module mock_handler = MagicMock() - mock_handler.handle_count_tokens_request = AsyncMock( - return_value={"input_tokens": 42} - ) + mock_handler.handle_count_tokens_request = AsyncMock(return_value={"input_tokens": 42}) with patch( "litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler", @@ -466,9 +453,7 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): # Mock the global handler instance in token_counter module mock_handler = MagicMock() - mock_handler.handle_count_tokens_request = AsyncMock( - return_value={"input_tokens": 42} - ) + mock_handler.handle_count_tokens_request = AsyncMock(return_value={"input_tokens": 42}) with patch( "litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler", @@ -527,9 +512,7 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): # Mock the global handler instance in token_counter module mock_handler = MagicMock() - mock_handler.handle_count_tokens_request = AsyncMock( - return_value={"input_tokens": 42} - ) + mock_handler.handle_count_tokens_request = AsyncMock(return_value={"input_tokens": 42}) with patch( "litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler", @@ -544,9 +527,7 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): mock_router.model_list = [ { "model_name": "claude-3-5-sonnet", - "litellm_params": { - "model": "anthropic/claude-3-5-sonnet-20241022" - }, + "litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"}, "model_info": {}, } ] @@ -555,9 +536,7 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): mock_router.async_get_available_deployment = AsyncMock( return_value={ "model_name": "claude-3-5-sonnet", - "litellm_params": { - "model": "anthropic/claude-3-5-sonnet-20241022" - }, + "litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"}, "model_info": {}, } ) @@ -592,9 +571,7 @@ async def test_factory_registration(): assert counter is not None # Create test deployments - anthropic_deployment = { - "litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"} - } + anthropic_deployment = {"litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"}} non_anthropic_deployment = {"litellm_params": {"model": "openai/gpt-4"}} @@ -630,9 +607,7 @@ async def test_bedrock_count_tokens_endpoint(): model_list=[ { "model_name": "claude-bedrock", - "litellm_params": { - "model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0" - }, + "litellm_params": {"model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0"}, } ] ) @@ -646,9 +621,7 @@ async def test_bedrock_count_tokens_endpoint(): } # Test the mock handler directly to verify correct parameter extraction - await mock_count_tokens_handler( - request_data, {}, "anthropic.claude-3-sonnet-20240229-v1:0" - ) + await mock_count_tokens_handler(request_data, {}, "anthropic.claude-3-sonnet-20240229-v1:0") @pytest.mark.asyncio @@ -710,10 +683,7 @@ async def test_vertex_ai_anthropic_token_counting(): assert call_args is not None assert call_args.kwargs["model"] == "claude-3-5-sonnet-20241022" assert "messages" in call_args.kwargs["request_data"] - assert ( - call_args.kwargs["request_data"]["messages"][0]["content"] - == "Hello Claude on Vertex AI! How are you?" - ) + assert call_args.kwargs["request_data"]["messages"][0]["content"] == "Hello Claude on Vertex AI! How are you?" # Validate response structure assert response.model_used == "claude-3-5-sonnet-20241022" @@ -746,9 +716,7 @@ def test_vertex_ai_partner_models_token_counting_endpoint(vertex_location): if vertex_location == "global": assert endpoint.startswith("https://aiplatform.googleapis.com") else: - assert endpoint.startswith( - f"https://{vertex_location}-aiplatform.googleapis.com" - ) + assert endpoint.startswith(f"https://{vertex_location}-aiplatform.googleapis.com") class _FailingRuntimeHandler(BedrockCountTokensHandler): @@ -834,20 +802,14 @@ async def test_bedrock_handler_httpx_error_status_code_propagation(): with patch.object(handler, "validate_count_tokens_request"): with patch.object(handler, "_get_aws_region_name", return_value="us-west-2"): - with patch.object( - handler, "transform_anthropic_to_bedrock_count_tokens", return_value={} - ): + with patch.object(handler, "transform_anthropic_to_bedrock_count_tokens", return_value={}): with patch.object( handler, "get_bedrock_count_tokens_endpoint", return_value="https://example.com", ): - with patch.object( - handler, "_sign_request", return_value=({}, "{}") - ): - with patch( - "litellm.llms.bedrock.count_tokens.handler.get_async_httpx_client" - ) as mock_client: + with patch.object(handler, "_sign_request", return_value=({}, "{}")): + with patch("litellm.llms.bedrock.count_tokens.handler.get_async_httpx_client") as mock_client: mock_async_client = AsyncMock() mock_async_client.post = AsyncMock(side_effect=http_error) mock_client.return_value = mock_async_client @@ -856,9 +818,7 @@ async def test_bedrock_handler_httpx_error_status_code_propagation(): await handler.handle_count_tokens_request( request_data={ "model": "test", - "messages": [ - {"role": "user", "content": "hello"} - ], + "messages": [{"role": "user", "content": "hello"}], }, litellm_params={}, resolved_model="anthropic.claude-3-sonnet", @@ -866,10 +826,7 @@ async def test_bedrock_handler_httpx_error_status_code_propagation(): assert exc_info.value.status_code == 403 # Message should be the raw response text - assert ( - exc_info.value.message - == "Forbidden - Invalid credentials" - ) + assert exc_info.value.message == "Forbidden - Invalid credentials" @pytest.mark.asyncio @@ -897,9 +854,7 @@ async def test_token_counter_httpx_status_error_raises_proxy_exception(): mock_counter.count_tokens = AsyncMock(side_effect=http_error) # Save originals - original_get_provider_token_counter = ( - litellm.proxy.proxy_server._get_provider_token_counter - ) + original_get_provider_token_counter = litellm.proxy.proxy_server._get_provider_token_counter original_router = litellm.proxy.proxy_server.llm_router try: @@ -907,9 +862,7 @@ async def test_token_counter_httpx_status_error_raises_proxy_exception(): def mock_get_provider_token_counter(deployment, model_to_use): return (mock_counter, "claude-4-6-sonnet", "vertex_ai") - litellm.proxy.proxy_server._get_provider_token_counter = ( - mock_get_provider_token_counter - ) + litellm.proxy.proxy_server._get_provider_token_counter = mock_get_provider_token_counter mock_router = MagicMock() mock_router.async_get_available_deployment = AsyncMock( @@ -937,9 +890,7 @@ async def test_token_counter_httpx_status_error_raises_proxy_exception(): assert exc_info.value.type == "token_counting_error" assert exc_info.value.param == "model" finally: - litellm.proxy.proxy_server._get_provider_token_counter = ( - original_get_provider_token_counter - ) + litellm.proxy.proxy_server._get_provider_token_counter = original_get_provider_token_counter litellm.proxy.proxy_server.llm_router = original_router @@ -975,9 +926,7 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled(): # Save original value and function original_disable = litellm.disable_token_counter - original_get_provider_token_counter = ( - litellm.proxy.proxy_server._get_provider_token_counter - ) + original_get_provider_token_counter = litellm.proxy.proxy_server._get_provider_token_counter try: litellm.disable_token_counter = True @@ -991,9 +940,7 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled(): def mock_get_provider_token_counter(deployment, model_to_use): return (mock_counter, "anthropic.claude-3-sonnet", "bedrock") - litellm.proxy.proxy_server._get_provider_token_counter = ( - mock_get_provider_token_counter - ) + litellm.proxy.proxy_server._get_provider_token_counter = mock_get_provider_token_counter with pytest.raises(ProxyException) as exc_info: await token_counter( @@ -1008,9 +955,7 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled(): assert "Rate limit exceeded" in exc_info.value.message finally: litellm.disable_token_counter = original_disable - litellm.proxy.proxy_server._get_provider_token_counter = ( - original_get_provider_token_counter - ) + litellm.proxy.proxy_server._get_provider_token_counter = original_get_provider_token_counter @pytest.mark.asyncio @@ -1045,9 +990,7 @@ async def test_proxy_token_counter_error_raises_when_enabled(): # Save original value and function original_disable = litellm.disable_token_counter - original_get_provider_token_counter = ( - litellm.proxy.proxy_server._get_provider_token_counter - ) + original_get_provider_token_counter = litellm.proxy.proxy_server._get_provider_token_counter try: litellm.disable_token_counter = False @@ -1061,9 +1004,7 @@ async def test_proxy_token_counter_error_raises_when_enabled(): def mock_get_provider_token_counter(deployment, model_to_use): return (mock_counter, "anthropic.claude-3-sonnet", "bedrock") - litellm.proxy.proxy_server._get_provider_token_counter = ( - mock_get_provider_token_counter - ) + litellm.proxy.proxy_server._get_provider_token_counter = mock_get_provider_token_counter with patch("litellm.token_counter") as mock_local_counter: with pytest.raises(ProxyException) as exc_info: @@ -1081,9 +1022,7 @@ async def test_proxy_token_counter_error_raises_when_enabled(): mock_local_counter.assert_not_called() finally: litellm.disable_token_counter = original_disable - litellm.proxy.proxy_server._get_provider_token_counter = ( - original_get_provider_token_counter - ) + litellm.proxy.proxy_server._get_provider_token_counter = original_get_provider_token_counter @pytest.mark.asyncio @@ -1109,10 +1048,10 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): original_read_request_body = anthropic_endpoints.read_request_body anthropic_endpoints.read_request_body = mock_read_request_body - original_token_counter = proxy_server.token_counter + original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with Bedrock-style error - async def mock_token_counter_error(request, call_endpoint=False): + async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): raise ProxyException( message='{"detail":{"message":"Input is too long for requested model."}}', type="token_counting_error", @@ -1120,7 +1059,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): code=400, ) - proxy_server.token_counter = mock_token_counter_error + proxy_server.count_request_tokens = mock_token_counter_error try: with pytest.raises(HTTPException) as exc_info: @@ -1136,7 +1075,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): assert detail["error"]["message"] == "Input is too long for requested model." finally: anthropic_endpoints.read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter @pytest.mark.asyncio @@ -1161,10 +1100,10 @@ async def test_anthropic_endpoint_403_permission_error_format(): original_read_request_body = anthropic_endpoints.read_request_body anthropic_endpoints.read_request_body = mock_read_request_body - original_token_counter = proxy_server.token_counter + original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with 403 error - async def mock_token_counter_error(request, call_endpoint=False): + async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): raise ProxyException( message='{"Message":"Bearer Token has expired"}', type="token_counting_error", @@ -1172,7 +1111,7 @@ async def test_anthropic_endpoint_403_permission_error_format(): code=403, ) - proxy_server.token_counter = mock_token_counter_error + proxy_server.count_request_tokens = mock_token_counter_error try: with pytest.raises(HTTPException) as exc_info: @@ -1186,7 +1125,7 @@ async def test_anthropic_endpoint_403_permission_error_format(): assert detail["error"]["message"] == "Bearer Token has expired" finally: anthropic_endpoints.read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter @pytest.mark.asyncio @@ -1211,10 +1150,10 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): original_read_request_body = anthropic_endpoints.read_request_body anthropic_endpoints.read_request_body = mock_read_request_body - original_token_counter = proxy_server.token_counter + original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with 429 error - async def mock_token_counter_error(request, call_endpoint=False): + async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): raise ProxyException( message="Rate limit exceeded", type="token_counting_error", @@ -1222,7 +1161,7 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): code=429, ) - proxy_server.token_counter = mock_token_counter_error + proxy_server.count_request_tokens = mock_token_counter_error try: with pytest.raises(HTTPException) as exc_info: @@ -1236,4 +1175,4 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): assert detail["error"]["message"] == "Rate limit exceeded" finally: anthropic_endpoints.read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter diff --git a/tests/unit/test_count_tokens_public_api.py b/tests/unit/test_count_tokens_public_api.py index 84121dd63e4..e33bc1e5907 100644 --- a/tests/unit/test_count_tokens_public_api.py +++ b/tests/unit/test_count_tokens_public_api.py @@ -97,7 +97,7 @@ def test_acount_tokens_with_tools(): assert result.total_tokens == 30 mock_handler.assert_called_once() call_kwargs = mock_handler.call_args - assert call_kwargs.kwargs.get("tools") == tools + assert list(call_kwargs.kwargs.get("tools") or ()) == tools def test_acount_tokens_with_system(): From cf1e6b19ab626349f32960c6f721b7975a9b6a59 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 15 Oct 2026 19:01:04 -0700 Subject: [PATCH 2/8] fix(sso): share SAML sign-in state over Redis without enable_redis_auth_cache (#46807) * fix(sso): share SAML sign-in state over Redis without enable_redis_auth_cache SAML login tickets, the single-use assertion mark and cached IdP metadata lived in user_api_key_cache, which only gets Redis when enable_redis_auth_cache is on. With several replicas or workers the IdP callback often landed on a process that never saw the ticket and got a 401, and IdP-initiated sign-in was refused even with Redis configured. SAML state now has its own DualCache that _attach_redis_usage_cache always wires to Redis, like the CLI SSO session cache, and the SAML cache calls are async. * test(sso): drop docstring from the SAML proxy pair fixture * test(sso): cover the SAML login and metadata endpoints --- .../management_endpoints/sso/saml_sso.py | 14 +-- litellm/proxy/management_endpoints/ui_sso.py | 16 ++-- litellm/proxy/proxy_server.py | 17 +++- .../authorization/test_sso_blank_identity.py | 96 ++++++++++++++++++- .../management_endpoints/test_saml_sso.py | 30 ++++++ .../unit/proxy/test_redis_auth_cache_flag.py | 15 +++ 6 files changed, 165 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/management_endpoints/sso/saml_sso.py b/litellm/proxy/management_endpoints/sso/saml_sso.py index 292ba3c5988..4991436edac 100644 --- a/litellm/proxy/management_endpoints/sso/saml_sso.py +++ b/litellm/proxy/management_endpoints/sso/saml_sso.py @@ -60,7 +60,7 @@ _SAML_AUTHN_STATE_COOKIE: Final = "litellm_saml_authn" _SAML_IDP_SETTINGS_CACHE_PREFIX: Final = "saml_idp_settings" _SAML_AUTHN_REQUEST_CACHE_PREFIX: Final = "saml_authn_request" _SAML_CONSUMED_ASSERTION_CACHE_PREFIX: Final = "saml_consumed_assertion" -_SAML_AUTHN_REQUEST_TTL_SECONDS: Final = 600 +SAML_AUTHN_REQUEST_TTL_SECONDS: Final = 600 _SAML_IDP_METADATA_TTL_SECONDS: Final = 3600 _SAML_METADATA_FETCH_TIMEOUT_SECONDS: Final = 10 _SAML_MAX_POST_BYTES: Final = 5 * 1024 * 1024 @@ -161,7 +161,7 @@ class SAMLAuthHandler: ) cache_key: Final = f"{_SAML_IDP_SETTINGS_CACHE_PREFIX}:{hashlib.sha256(source.encode()).hexdigest()}" - cached: Final = cache.get_cache(key=cache_key) + cached: Final = await cache.async_get_cache(key=cache_key) if isinstance(cached, dict): return cast(dict[str, object], cached) # cast-ok: untyped python3-saml @@ -181,7 +181,7 @@ class SAMLAuthHandler: status_code=status.HTTP_502_BAD_GATEWAY, detail="Could not parse an IdP entityID/SSO URL/certificate from the SAML metadata.", ) - cache.set_cache(key=cache_key, value=idp_settings, ttl=_SAML_IDP_METADATA_TTL_SECONDS) + await cache.async_set_cache(key=cache_key, value=idp_settings, ttl=_SAML_IDP_METADATA_TTL_SECONDS) return idp_settings @staticmethod @@ -253,16 +253,16 @@ class SAMLAuthHandler: response: Final = RedirectResponse(url=redirect_url, status_code=303) request_id: Final = cast(str | None, auth.get_last_request_id()) # cast-ok: untyped python3-saml if request_id is not None: - cache.set_cache( + await cache.async_set_cache( key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", - ttl=_SAML_AUTHN_REQUEST_TTL_SECONDS, + ttl=SAML_AUTHN_REQUEST_TTL_SECONDS, ) secure: Final = SAMLAuthHandler._is_https(request) response.set_cookie( key=_SAML_AUTHN_STATE_COOKIE, value=request_id, - max_age=_SAML_AUTHN_REQUEST_TTL_SECONDS, + max_age=SAML_AUTHN_REQUEST_TTL_SECONDS, httponly=True, secure=secure, samesite="none" if secure else "lax", @@ -372,7 +372,7 @@ class SAMLAuthHandler: if in_response_to is not None: authn_key: Final = f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{in_response_to}" - if cache.get_cache(key=authn_key) is None: + if await cache.async_get_cache(key=authn_key) is None: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="SAML response references an unknown or already-used login request.", diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index c5b4940b384..4629103070f 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -1101,7 +1101,7 @@ async def google_login( general_settings, premium_user, prisma_client, - user_api_key_cache, + saml_sso_cache, user_custom_ui_sso_sign_in_handler, ) @@ -1169,7 +1169,7 @@ async def google_login( verbose_proxy_logger.info("Redirecting to SAML SSO login") return await SAMLAuthHandler.build_login_redirect( request=request, - cache=user_api_key_cache, + cache=saml_sso_cache, relay_state=return_to, ) @@ -2257,20 +2257,20 @@ async def auth_callback(request: Request, state: str | None = None): @router.get("/sso/saml/login", tags=["experimental"], include_in_schema=False) async def saml_login(request: Request, return_to: str | None = None): """SP-initiated SAML login. Redirects the user to the configured IdP.""" - from litellm.proxy.proxy_server import user_api_key_cache + from litellm.proxy.proxy_server import saml_sso_cache if is_admin_ui_disabled(): return admin_ui_disabled() - return await SAMLAuthHandler.build_login_redirect(request=request, cache=user_api_key_cache, relay_state=return_to) + return await SAMLAuthHandler.build_login_redirect(request=request, cache=saml_sso_cache, relay_state=return_to) @router.get("/sso/saml/metadata", tags=["experimental"], include_in_schema=False) async def saml_metadata(request: Request): """Service Provider metadata XML, for registering this proxy at the IdP.""" - from litellm.proxy.proxy_server import user_api_key_cache + from litellm.proxy.proxy_server import saml_sso_cache - metadata: Final = await SAMLAuthHandler.build_sp_metadata(request=request, cache=user_api_key_cache) + metadata: Final = await SAMLAuthHandler.build_sp_metadata(request=request, cache=saml_sso_cache) return Response(content=metadata, media_type="application/xml") @@ -2283,7 +2283,7 @@ async def saml_callback(request: Request): master_key, premium_user, prisma_client, - user_api_key_cache, + saml_sso_cache, ) if is_admin_ui_disabled(): @@ -2303,7 +2303,7 @@ async def saml_callback(request: Request): if "SAMLResponse" not in post_data: raise HTTPException(status_code=400, detail="Missing SAMLResponse in callback request.") - result: Final = await SAMLAuthHandler.handle_acs(request=request, cache=user_api_key_cache, post_data=post_data) + result: Final = await SAMLAuthHandler.handle_acs(request=request, cache=saml_sso_cache, post_data=post_data) await _raise_if_sso_exceeds_free_user_limit(premium_user, prisma_client) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 621dd0d8c04..4b4c6a13b1d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -714,6 +714,7 @@ from litellm.proxy.management_endpoints.router_settings_endpoints import ( from litellm.proxy.management_endpoints.session_endpoints import ( router as session_management_router, ) +from litellm.proxy.management_endpoints.sso.saml_sso import SAML_AUTHN_REQUEST_TTL_SECONDS from litellm.proxy.management_endpoints.tag_management_endpoints import ( router as tag_management_router, ) @@ -2830,6 +2831,7 @@ user_api_key_cache: UserApiKeyCache = UserApiKeyCache( ) spend_counter_cache: Final = DualCache(default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value) cli_sso_session_cache: Final = DualCache(default_in_memory_ttl=CLI_SSO_SESSION_TTL_SECONDS) +saml_sso_cache: Final = DualCache(default_in_memory_ttl=SAML_AUTHN_REQUEST_TTL_SECONDS) model_max_budget_limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=spend_counter_cache) litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter) redis_usage_cache: RedisCache | None = None # redis cache used for tracking spend, tpm/rpm limits @@ -5191,12 +5193,13 @@ def _attach_redis_usage_cache(redis_cache: RedisCache, enable_redis_auth_cache: """ Wires an established coordination Redis into the proxy-level caches that consume it directly: the spend counter cache, the CLI SSO login-session - cache, the cluster-wide config cache, and (only when opted in) the - virtual-key auth cache. + cache, the SAML SSO sign-in cache, the cluster-wide config cache, and (only + when opted in) the virtual-key auth cache. - The CLI SSO login-session cache is always backed by Redis when available so - that the browser SSO flow behind `lite login` survives landing on different - workers; it must not be gated behind enable_redis_auth_cache. + The CLI SSO login-session cache and the SAML SSO sign-in cache are always + backed by Redis when available so that the browser SSO flows behind + `lite login` and SAML sign-in survive landing on different workers; they + must not be gated behind enable_redis_auth_cache. """ spend_counter_cache.attach_redis_cache( redis_cache, @@ -5206,6 +5209,10 @@ def _attach_redis_usage_cache(redis_cache: RedisCache, enable_redis_auth_cache: redis_cache, default_redis_ttl=CLI_SSO_SESSION_TTL_SECONDS, ) + saml_sso_cache.attach_redis_cache( + redis_cache, + default_redis_ttl=SAML_AUTHN_REQUEST_TTL_SECONDS, + ) if enable_redis_auth_cache is True: user_api_key_cache.attach_redis_cache( redis_cache, diff --git a/tests/integration/authorization/test_sso_blank_identity.py b/tests/integration/authorization/test_sso_blank_identity.py index c2f0e29221d..6d9b3b50cc0 100644 --- a/tests/integration/authorization/test_sso_blank_identity.py +++ b/tests/integration/authorization/test_sso_blank_identity.py @@ -169,8 +169,9 @@ class SamlIdp: "" ) - def response(self, acs: str, request_id: str, name_id: str, email: str | None) -> str: + def response(self, acs: str, request_id: str | None, name_id: str, email: str | None) -> str: now: Final = datetime.now(timezone.utc) + in_response_to: Final = "" if request_id is None else f' InResponseTo="{request_id}"' issued: Final = (now - timedelta(seconds=1)).strftime("%Y-%m-%dT%H:%M:%SZ") not_before: Final = (now - timedelta(seconds=60)).strftime("%Y-%m-%dT%H:%M:%SZ") expires: Final = (now + timedelta(seconds=300)).strftime("%Y-%m-%dT%H:%M:%SZ") @@ -186,7 +187,7 @@ class SamlIdp: '' f"{name_id}" '' - f'' + f'' "" f'' f"{SAML_SP_ENTITY}" @@ -206,7 +207,7 @@ class SamlIdp: '' + f'Destination="{acs}"{in_response_to}>' f"{SAML_IDP_ENTITY}" '' f"{signed_text}" @@ -258,6 +259,54 @@ def _saml_sign_in(proxy: Gateway, idp: SamlIdp, name_id: str, email: str | None) return browser.post(acs, data={"SAMLResponse": idp.response(acs, request_id, name_id, email)}) +SAML_PAIR_BASE: Final = "http://saml-pair.integration.invalid" +SAML_PAIR_ACS: Final = f"{SAML_PAIR_BASE}/sso/saml/callback" + + +@dataclass(frozen=True, slots=True) +class SamlPair: + first: Gateway + second: Gateway + + +@pytest.fixture(scope="module") +def saml_pair( + saml_idp: SamlIdp, gateway_module: Gateway, tmp_path_factory: pytest.TempPathFactory +) -> Iterator[SamlPair]: + environment: Final = { + "SAML_IDP_METADATA_XML": saml_idp.metadata(), + "SAML_SP_ENTITY_ID": SAML_SP_ENTITY, + "SAML_ALLOW_UNSOLICITED": "true", + "PROXY_BASE_URL": SAML_PAIR_BASE, + } + config: Final = Path("tests/integration/coordination_redis_proxy_config.yaml") + with ( + owned_proxy_process( + gateway_module, tmp_path_factory.mktemp("saml-pair-first"), environment, config=config + ) as first, + owned_proxy_process( + gateway_module, tmp_path_factory.mktemp("saml-pair-second"), environment, config=config + ) as second, + ): + yield SamlPair(first.gateway, second.gateway) + + +def _saml_login_ticket(proxy: Gateway) -> str: + with _browser() as browser: + login: Final = browser.get(f"{_proxy_url(proxy)}/sso/key/generate", params={"return_to": "/ui/"}) + assert login.status_code == 303, f"{login.status_code} {login.text}" + return login.cookies[SAML_STATE_COOKIE] + + +def _post_saml_response(proxy: Gateway, saml_response: str, ticket: str | None) -> httpx.Response: + with _browser() as browser: + return browser.post( + f"{_proxy_url(proxy)}/sso/saml/callback", + data={"SAMLResponse": saml_response}, + headers={} if ticket is None else {"cookie": f"{SAML_STATE_COOKIE}={ticket}"}, + ) + + @pytest.fixture(scope="module") def graph() -> Iterator[Graph]: profiles: Final[dict[str, Reply]] = {} @@ -811,3 +860,44 @@ def test_saml_sign_in_with_an_email_signs_in_as_a_real_account( assert callback.status_code == 303 and session is not None, f"{callback.status_code} {callback.text}" assert session["user_id"] == expected, f"session issued for user_id={session['user_id']!r}" assert _blank_accounts() == (), _blank_accounts() + + +def test_saml_sign_in_started_on_one_proxy_completes_on_another( + saml_pair: SamlPair, saml_idp: SamlIdp, gateway_module: Gateway +) -> None: + email: Final = f"saml-pair-{uuid.uuid4().hex[:12]}@example.com" + with gateway_module.scenario() as scenario: + _forget_sign_in(scenario, email) + ticket: Final = _saml_login_ticket(saml_pair.first) + callback: Final = _post_saml_response( + saml_pair.second, saml_idp.response(SAML_PAIR_ACS, ticket, email, email), ticket + ) + assert callback.status_code == 303 and _session(callback) is not None, f"{callback.status_code} {callback.text}" + assert callback.headers["location"].startswith(f"{SAML_PAIR_BASE}/ui/"), callback.headers["location"] + + +def test_saml_response_replayed_on_another_proxy_is_refused( + saml_pair: SamlPair, saml_idp: SamlIdp, gateway_module: Gateway +) -> None: + email: Final = f"saml-pair-{uuid.uuid4().hex[:12]}@example.com" + with gateway_module.scenario() as scenario: + _forget_sign_in(scenario, email) + ticket: Final = _saml_login_ticket(saml_pair.first) + saml_response: Final = saml_idp.response(SAML_PAIR_ACS, ticket, email, email) + accepted: Final = _post_saml_response(saml_pair.second, saml_response, ticket) + assert accepted.status_code == 303, f"{accepted.status_code} {accepted.text}" + replayed: Final = _post_saml_response(saml_pair.first, saml_response, ticket) + assert replayed.status_code == 401 and _session(replayed) is None, f"{replayed.status_code} {replayed.text}" + assert "already been used" in string_value(JSON_OBJECT.validate_json(replayed.content)["detail"]), replayed.text + + +def test_idp_initiated_saml_sign_in_completes_on_any_proxy( + saml_pair: SamlPair, saml_idp: SamlIdp, gateway_module: Gateway +) -> None: + email: Final = f"saml-pair-{uuid.uuid4().hex[:12]}@example.com" + with gateway_module.scenario() as scenario: + _forget_sign_in(scenario, email) + callback: Final = _post_saml_response( + saml_pair.second, saml_idp.response(SAML_PAIR_ACS, None, email, email), None + ) + assert callback.status_code == 303 and _session(callback) is not None, f"{callback.status_code} {callback.text}" diff --git a/tests/unit/proxy/management_endpoints/test_saml_sso.py b/tests/unit/proxy/management_endpoints/test_saml_sso.py index 635e6332958..01f8fc929a3 100644 --- a/tests/unit/proxy/management_endpoints/test_saml_sso.py +++ b/tests/unit/proxy/management_endpoints/test_saml_sso.py @@ -10,6 +10,7 @@ makes a test fail. import base64 import datetime import time +from http.cookies import SimpleCookie import pytest from fastapi import HTTPException, Request @@ -637,6 +638,35 @@ async def test_idp_initiated_replay_rejected_across_workers(saml_env_idp_initiat assert exc.value.status_code == 401 +@pytest.mark.asyncio +async def test_sp_initiated_login_endpoint_ticket_is_accepted_by_another_worker(saml_env, monkeypatch): + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.management_endpoints.ui_sso import saml_login + + key_pem, cert_pem = saml_env + shared_store = InMemoryCache() + monkeypatch.setattr(proxy_server, "saml_sso_cache", _shared_cache(shared_store)) + + redirect = await saml_login(_fake_request()) + assert redirect.headers["location"].startswith(SSO_URL) + ticket = SimpleCookie(redirect.headers["set-cookie"])[_SAML_AUTHN_STATE_COOKIE].value + + resp = _build_signed_response(key_pem, cert_pem, in_response_to=ticket, email="carol@example.com") + result = await _acs(_b64(resp), _shared_cache(shared_store), cookies={_SAML_AUTHN_STATE_COOKIE: ticket}) + assert result.email == "carol@example.com" + + +@pytest.mark.asyncio +async def test_sp_metadata_endpoint_serves_acs_and_entity_id(saml_env): + from litellm.proxy.management_endpoints.ui_sso import saml_metadata + + response = await saml_metadata(_fake_request()) + body = bytes(response.body).decode() + assert response.media_type == "application/xml" + assert ACS in body + assert SP_ENTITY in body + + class _FakeChunkedRequest: def __init__(self, chunks, content_length=None): self._chunks = chunks diff --git a/tests/unit/proxy/test_redis_auth_cache_flag.py b/tests/unit/proxy/test_redis_auth_cache_flag.py index cb600bbb5fd..d414e17b078 100644 --- a/tests/unit/proxy/test_redis_auth_cache_flag.py +++ b/tests/unit/proxy/test_redis_auth_cache_flag.py @@ -73,6 +73,7 @@ def _patched_init_cache(litellm_settings: dict, cache_params: dict): patch.object(ps, "user_api_key_cache", fresh_user_cache), patch.object(ps, "spend_counter_cache", fresh_spend_cache), patch.object(ps, "cli_sso_session_cache", fresh_cli_sso_cache), + patch.object(ps, "saml_sso_cache", DualCache()), patch.object(ps, "litellm_config_cache", fresh_config_cache), patch.object(ps, "llm_router", None), # Cache is locally imported inside _init_cache: patch it at source. @@ -182,6 +183,7 @@ class TestRedisAuthCacheFlag: limiter_cache, ps.spend_counter_cache, ps.cli_sso_session_cache, + ps.saml_sso_cache, ps.user_api_key_cache, ps.litellm_config_cache, ) @@ -191,3 +193,16 @@ class TestRedisAuthCacheFlag: ps._attach_redis_usage_cache(fake_redis, enable_redis_auth_cache=False) assert limiter_cache.redis_cache is fake_redis assert ps.user_api_key_cache.redis_cache is None + + +@pytest.mark.parametrize( + "litellm_settings", + [{"enable_redis_auth_cache": True}, {"enable_redis_auth_cache": False}, {}], + ids=["flag-true", "flag-false", "flag-absent"], +) +def test_saml_sso_cache_always_gets_redis_regardless_of_flag(litellm_settings: dict[str, bool]) -> None: + with _patched_init_cache( + litellm_settings=litellm_settings, + cache_params={"type": "redis", "host": "localhost", "port": 6379}, + ): + assert ps.saml_sso_cache.redis_cache is not None From 5ed3629242ec572fcb8b5ee18fc766aa337b5a89 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Thu, 15 Oct 2026 19:05:22 -0700 Subject: [PATCH 3/8] fix(responses): write one spend-log row per native WebSocket turn (#46768) * fix(responses): write one spend-log row per native WebSocket turn Each response.create turn on a native Responses WebSocket connection logs and bills as its own call when its terminal event arrives, keyed by the resp_ id the client saw, instead of one row at socket close that sums every turn * test(responses): mock the websocket slot release and mark the per-turn start time naive * style(responses): keep the per-turn start time line within ruff format * refactor(responses): type websocket turn events as read-only mappings * fix(responses): settle the websocket budget reservation on the first turn only Every per-turn logging copy carried the connection's admission reservation, so two turns whose cost callbacks overlapped could both reconcile it. The first turn keeps the reservation and later turns charge their cost directly * fix(proxy): keep the legacy limiter's websocket slot until disconnect A websocket turn callback no longer decrements current_requests in the legacy limiter. The socket's slot is released once from arelease_max_parallel_requests_on_disconnect, which now dispatches to the legacy handler as well --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/cost_calculator.py | 18 ++ .../internal_call_metadata.py | 9 +- litellm/litellm_core_utils/litellm_logging.py | 51 ++-- .../proxy/hooks/parallel_request_limiter.py | 57 ++++- .../proxy/response_api_endpoints/endpoints.py | 2 + litellm/proxy/utils.py | 4 +- litellm/responses/streaming_iterator.py | 232 ++++++++++-------- litellm/types/utils.py | 1 + .../integration/_support/responses_ws_peer.py | 186 ++++++++++++++ .../test_openai_responses_websocket_wire.py | 226 +++-------------- .../test_responses_websocket_turn_rows.py | 114 +++++++++ .../test_litellm_logging.py | 52 ++++ .../hooks/test_parallel_request_limiter.py | 32 ++- .../response_api_endpoints/test_endpoints.py | 2 + .../test_responses_websocket_all_providers.py | 98 +++++++- tests/unit/test_cost_calculator.py | 20 ++ 16 files changed, 769 insertions(+), 335 deletions(-) create mode 100644 tests/integration/_support/responses_ws_peer.py create mode 100644 tests/integration/spend/test_responses_websocket_turn_rows.py diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 43ce71af7b8..0e315e97432 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -3008,6 +3008,7 @@ _RESPONSES_WS_BILLABLE_EVENT_TYPES: Final = frozenset({"response.completed", "re class _ResponsesWsEventResponse(LiteLLMBaseModel): + id: str | None = None usage: Mapping[str, object] | None = None service_tier: str | None = None @@ -3059,6 +3060,23 @@ class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor): ) return ResponsesWebSocketTokenUsageProcessor.combine_usage_objects(list(collected_usage_objects)) + @staticmethod + def response_id_from_responses_ws_results(results: Sequence[Mapping[str, object]]) -> str | None: + response_ids: Final = tuple( + response.id for _, response in _billable_responses_ws_events(results) if response.id + ) + return response_ids[-1] if response_ids else None + + @staticmethod + def create_logging_object(results: Sequence[Mapping[str, object]]) -> LiteLLMRealtimeStreamLoggingObject: + tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(results) + return LiteLLMRealtimeStreamLoggingObject( + id=ResponsesWebSocketTokenUsageProcessor.response_id_from_responses_ws_results(results), + usage=ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(results), + results=cast(OpenAIRealtimeStreamList, list(results)), # cast-ok: raw WS frames, unvalidated field + service_tier=next(iter(tier_partition)) if len(tier_partition) == 1 else None, + ) + _TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed" diff --git a/litellm/litellm_core_utils/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py index c5e57bad050..1200d799018 100644 --- a/litellm/litellm_core_utils/internal_call_metadata.py +++ b/litellm/litellm_core_utils/internal_call_metadata.py @@ -157,7 +157,10 @@ def sanitize_user_api_key_auth(auth: object) -> object: return auth -def _sanitized(parent_metadata: Mapping[str, object]) -> dict[str, object]: # mutable-ok: SDK metadata kwarg +def without_budget_reservation( + parent_metadata: Mapping[str, object], +) -> dict[str, object]: # mutable-ok: SDK metadata kwarg + """Parent metadata minus the admission reservation, also inside ``user_api_key_auth``.""" return { k: sanitize_user_api_key_auth(v) if k == _USER_API_KEY_AUTH_KEY else v for k, v in parent_metadata.items() @@ -176,7 +179,7 @@ def forwarded_internal_call_metadata( """ if not parent_metadata: return {} - return _sanitized(parent_metadata) | {INTERNAL_CALL_ORIGIN_METADATA_KEY: call_origin} + return without_budget_reservation(parent_metadata) | {INTERNAL_CALL_ORIGIN_METADATA_KEY: call_origin} def parent_session_kwargs(request_kwargs: Mapping[str, object] | None) -> Mapping[str, str]: @@ -202,4 +205,4 @@ def sanitized_forwardable_call_metadata( must not inherit per-request state such as its routing decision or logging payload. """ identity: Final = {k: v for k, v in parent_metadata.items() if k in FORWARDABLE_IDENTITY_METADATA_KEYS} - return _sanitized(identity) | {INTERNAL_CALL_ORIGIN_METADATA_KEY: call_origin} + return without_budget_reservation(identity) | {INTERNAL_CALL_ORIGIN_METADATA_KEY: call_origin} diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 7038cc531f5..e766a2be9c3 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -19,7 +19,7 @@ from types import MappingProxyType, TracebackType from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast from httpx import Response -from pydantic import BaseModel, JsonValue +from pydantic import BaseModel, JsonValue, TypeAdapter import litellm from litellm._internal_context import post_response_phase @@ -536,6 +536,10 @@ def _resolve_mantle_region_for_cost( return resolve_mantle_region(litellm_params or MappingProxyType({})) +_PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata"}) +_LITELLM_PARAMS: Final = TypeAdapter(dict[str, object]) + + def _provider_response_id(source: object) -> str | None: candidate: Final = source.get("id") if isinstance(source, dict) else getattr(source, "id", None) return candidate if isinstance(candidate, str) and candidate else None @@ -771,6 +775,36 @@ class Logging(LiteLLMLoggingBaseClass): self.on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None + def copy_for_new_call(self, litellm_call_id: str, start_time: datetime.datetime) -> "Logging": + """One turn of a long-lived connection: the connection's deployment, key and metadata under its own call id, + start time and chunk buffers, so the turn logs and bills as a request of its own.""" + connection_params: Final = _LITELLM_PARAMS.validate_python( + self.litellm_params # pyright: ignore[reportUnknownMemberType] # Logging.litellm_params is untyped + ) + litellm_params: Final[dict[str, object]] = { + **connection_params, + **{ + key: _LITELLM_PARAMS.validate_python(value) + for key, value in connection_params.items() + if key in _PER_CALL_LITELLM_PARAM_DICTS and isinstance(value, dict) + }, + "litellm_call_id": litellm_call_id, + } + new_call: Final = copy.copy(self) + new_call.litellm_call_id = litellm_call_id + new_call.start_time = start_time + new_call.litellm_params = litellm_params + new_call.model_call_details = { + **self.model_call_details, + "litellm_call_id": litellm_call_id, + "start_time": start_time, + "first_api_call_start_time": None, + "litellm_params": litellm_params, + } + new_call.streaming_chunks = [] + new_call.sync_streaming_chunks = [] + return new_call + def set_response_timing_metrics(self, timing_metrics: Mapping[str, float]) -> None: """Keep ``_response_ms`` / ``litellm_overhead_time_ms`` for a result that has no ``_hidden_params``.""" self.response_timing_metrics = dict(timing_metrics) @@ -2384,20 +2418,9 @@ class Logging(LiteLLMLoggingBaseClass): ) elif self.call_type == CallTypes.aresponses_websocket.value and isinstance(result, list): # pyright: ignore[reportUnknownMemberType] # Logging.call_type is untyped - combined_ws_usage: Final = ( - ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results( - results=result # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream - ) - ) - ws_tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier( + logging_result = ResponsesWebSocketTokenUsageProcessor.create_logging_object( results=result # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream ) - ws_service_tier: Final = next(iter(ws_tier_partition)) if len(ws_tier_partition) == 1 else None - logging_result = LiteLLMRealtimeStreamLoggingObject( - usage=combined_ws_usage, - results=result, # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream - service_tier=ws_service_tier, - ) elif self.relays_llm_passthrough and isinstance(result, Response): from litellm.utils import ProviderConfigManager @@ -6726,7 +6749,7 @@ def get_standard_logging_object_payload( else raw_usage_dict ) - id = response_obj.get("id", kwargs.get("litellm_call_id")) + id = response_obj.get("id") or kwargs.get("litellm_call_id") _model_id: Final = metadata.get("model_info", {}).get("id", "") _model_group: Final = metadata.get("model_group", "") diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index 19251b53f13..47ea50debc9 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -1,6 +1,6 @@ import asyncio import sys -from collections.abc import Callable +from collections.abc import Callable, Mapping from datetime import datetime, timedelta from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn @@ -23,7 +23,7 @@ from litellm.proxy.auth.budget_throttle import throttled_limit from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup -from litellm.types.utils import Usage +from litellm.types.utils import CallTypes, Usage if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -48,6 +48,14 @@ def _response_total_tokens(response_obj: object) -> int: return response_usage.total_tokens if isinstance(response_usage, Usage) else 0 +_CONNECTION_SCOPED_CALL_TYPES: Final = frozenset({CallTypes.aresponses_websocket.value}) + + +def _slots_released(kwargs: Mapping[str, object]) -> int: + """A websocket turn callback releases nothing; the socket's slot is released once at disconnect.""" + return 0 if kwargs.get("call_type") in _CONNECTION_SCOPED_CALL_TYPES else 1 + + class CacheObject(TypedDict): current_global_requests: dict | None request_count_api_key: dict | None @@ -509,6 +517,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): litellm_parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(kwargs=kwargs) try: self.print_verbose("INSIDE parallel request limiter ASYNC SUCCESS LOGGING") + released: Final = _slots_released(kwargs) # pyright: ignore[reportUnknownArgumentType] # legacy callback kwargs are untyped global_max_parallel_requests: Final = kwargs["litellm_params"]["metadata"].get( "global_max_parallel_requests", None @@ -540,7 +549,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): # decrement await self.internal_usage_cache.async_increment_cache( key=_key, - value=-1, + value=-released, local_only=True, litellm_parent_otel_span=litellm_parent_otel_span, ) @@ -568,7 +577,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - released, 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -609,7 +618,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - released, 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -635,7 +644,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - released, 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -661,7 +670,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - released, 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -687,7 +696,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - released, 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -709,6 +718,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): self.print_verbose("Inside Max Parallel Request Failure Hook") litellm_parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(kwargs=kwargs) _metadata: Final = kwargs["litellm_params"].get("metadata", {}) or {} + released: Final = _slots_released(kwargs) # pyright: ignore[reportUnknownArgumentType] # legacy callback kwargs are untyped global_max_parallel_requests: Final = _metadata.get("global_max_parallel_requests", None) user_api_key: Final = _metadata.get("user_api_key", None) self.print_verbose(f"user_api_key: [set={user_api_key is not None}]") @@ -736,7 +746,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): # decrement await self.internal_usage_cache.async_increment_cache( key=_key, - value=-1, + value=-released, local_only=True, litellm_parent_otel_span=litellm_parent_otel_span, ) @@ -758,7 +768,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): } new_val: Final = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - released, 0), "current_tpm": current["current_tpm"], "current_rpm": current["current_rpm"], } @@ -774,6 +784,33 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): verbose_proxy_logger.exception("Inside Parallel Request Limiter: An exception occurred - %s", e) @with_service_target("rate_limits") + async def async_release_max_parallel_requests_on_disconnect(self, user_api_key_dict: UserAPIKeyAuth) -> None: + """Release the slot ``async_pre_call_hook`` took for a connection whose per-turn callbacks released none.""" + from litellm.proxy.proxy_server import ( + general_settings, # pyright: ignore[reportUnknownVariableType] # module-level settings dict is untyped + ) + + now: Final = self._clock() + await self.async_log_success_event( # pyright: ignore[reportUnknownMemberType] # legacy callback signature is untyped + kwargs={ + "litellm_params": { + "metadata": { + "user_api_key": LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict), + "user_api_key_user_id": user_api_key_dict.user_id, + "user_api_key_team_id": user_api_key_dict.team_id, + "user_api_key_model_max_budget": user_api_key_dict.model_max_budget, # pyright: ignore[reportUnknownMemberType] # UserAPIKeyAuth.model_max_budget is a bare dict + "global_max_parallel_requests": general_settings.get( # pyright: ignore[reportUnknownMemberType] # module-level settings dict is untyped + "global_max_parallel_requests" + ), + } + }, + "user": user_api_key_dict.end_user_id, + }, + response_obj=None, + start_time=now, + end_time=now, + ) + async def get_internal_user_object( self, user_id: str, diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 4bfc4433815..b7f57d93a9b 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1905,6 +1905,8 @@ async def _responses_websocket_session( request_data=routed_data, ) await websocket.close(code=1011, reason="Internal server error") + finally: + await proxy_logging_obj.arelease_max_parallel_requests_on_disconnect(user_api_key_dict) @router.websocket("/v1/responses") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9788c4635c9..5d2ebe73cac 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -231,7 +231,7 @@ from litellm.proxy.hooks.cache_control_check import ( # noqa: F401, RUF100 # l ) from litellm.proxy.hooks.parallel_request_limiter import ( # noqa: F401, RUF100 # legacy module exports PROXY_MaxParallelRequestsHandler, - _PROXY_MaxParallelRequestsHandler, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export + _PROXY_MaxParallelRequestsHandler, # pyright: ignore[reportPrivateUsage] # backwards-compatible package export ) from litellm.proxy.hooks.parallel_request_limiter_v3 import ( # noqa: F401, RUF100 # legacy module exports PROXY_MaxParallelRequestsHandler_v3, @@ -4405,7 +4405,7 @@ class ProxyLogging: double-decrement under the limiter's in-memory fallback. """ limiter: Final = self.get_proxy_hook("parallel_request_limiter") - if not isinstance(limiter, PROXY_MaxParallelRequestsHandler_v3): + if not isinstance(limiter, (PROXY_MaxParallelRequestsHandler_v3, _PROXY_MaxParallelRequestsHandler)): return await limiter.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 7532a528248..7cf68332c07 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -7,8 +7,10 @@ import time import traceback import uuid from collections.abc import AsyncIterable, Awaitable, Callable, Coroutine, Iterable, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime from functools import lru_cache +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Protocol, cast, overload, runtime_checkable @@ -26,6 +28,7 @@ from litellm.constants import ( from litellm.exceptions import MidStreamFallbackError, RateLimitError from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import process_response_headers +from litellm.litellm_core_utils.internal_call_metadata import without_budget_reservation from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( @@ -1902,6 +1905,77 @@ _HISTORY_TERMINAL_EVENT_TYPES: Final = frozenset({"response.completed", "respons RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES: Final = frozenset({"input_text", "output_text", "text"}) _RESPONSES_WS_FAILURE_EVENT_TYPES: Final = frozenset({"error", "response.failed"}) +_PER_TURN_METADATA_KEYS: Final = frozenset({"metadata", "litellm_metadata"}) + +_RESPONSES_WS_TERMINAL_EVENT_TYPES: Final = _HISTORY_TERMINAL_EVENT_TYPES | _RESPONSES_WS_FAILURE_EVENT_TYPES + + +def _json_object_or_none(message: object) -> Mapping[str, object] | None: + if _is_json_object(message): + return message + if not isinstance(message, (str, bytes)): + return None + try: + parsed: Final = _load_json_value(message) + except json.JSONDecodeError: + return None + return parsed if _is_json_object(parsed) else None + + +def _user_texts_in_input_item(item: object) -> tuple[str, ...]: + if not _is_json_object(item) or item.get("type") != "message" or item.get("role") != "user": + return () + content: Final = item.get("content", []) + if isinstance(content, str): + return (content,) + return tuple( + text + for block in _json_array_or_empty(content) + if _is_json_object(block) + and block.get("type") == "input_text" + and isinstance(text := block.get("text"), str) + and text + ) + + +def _logged_user_messages(client_frame: object) -> tuple[Mapping[str, object], ...]: + frame: Final = _json_object_or_none(client_frame) + if frame is None or frame.get("type") != "response.create": + return () + input_items: Final = frame.get("input", []) + if isinstance(input_items, str): + return ({"role": "user", "content": input_items},) + texts: Final = chain.from_iterable(_user_texts_in_input_item(item) for item in _json_array_or_empty(input_items)) + return tuple({"role": "user", "content": text} for text in texts) + + +def _record_usage_for_failure(turn_logging_obj: LiteLLMLoggingObj, events: Sequence[Mapping[str, object]]) -> None: + from litellm.cost_calculator import ResponsesWebSocketTokenUsageProcessor + + logging_result: Final = ResponsesWebSocketTokenUsageProcessor.create_logging_object(events) + response_cost: Final = turn_logging_obj.response_cost_calculator(result=logging_result) or 0.0 + turn_logging_obj.record_partial_usage_for_failure(logging_result.usage, response_cost) + + +@dataclass(frozen=True, slots=True) +class _ResponsesWsTurn: + """One ``response.create`` on a native Responses socket, logged and billed as its own call once a terminal + event (``response.completed``, ``response.incomplete``, ``response.failed`` or ``error``) closes it.""" + + logging_obj: LiteLLMLoggingObj + events: tuple[Mapping[str, object], ...] = () + + def with_event(self, event: Mapping[str, object]) -> _ResponsesWsTurn: + return _ResponsesWsTurn(logging_obj=self.logging_obj, events=(*self.events, event)) + + def failure(self, model: str, custom_llm_provider: str) -> Exception | None: + failed_event: Final = next( + (event for event in self.events if event.get("type") in _RESPONSES_WS_FAILURE_EVENT_TYPES), None + ) + if failed_event is None: + return None + return _map_stream_error_to_exception(_ws_event_error(failed_event), model, custom_llm_provider) + _RESPONSES_WS_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"}) @@ -1997,8 +2071,8 @@ class ResponsesWebSocketStreaming: litellm_metadata: Final = self.request_data.get("litellm_metadata") self.litellm_metadata: dict[str, object] = litellm_metadata if _is_json_object(litellm_metadata) else {} self.custom_llm_provider: str | None = custom_llm_provider - self.messages: list[_MutableJsonObject] = [] - self.input_messages: list[dict[str, object]] = [] + self._open_turns: tuple[_ResponsesWsTurn, ...] = () + self._failure: Exception | None = None self.first_message = first_message self.guardrail_callbacks: Sequence[PresidioGuardrailCallback] = guardrail_callbacks or [] self.output_guardrail_callbacks: list[PresidioGuardrailCallback] = output_guardrail_callbacks or [] @@ -2008,105 +2082,69 @@ class ResponsesWebSocketStreaming: self.authorized_model: str | None = authorized_model self.request_defaults: ResponsesWebSocketRequestDefaults | None = request_defaults - def _should_store_event(self, event_obj: _MutableJsonObject) -> bool: - return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES - - def _store_event(self, event: str | bytes | dict[str, object]) -> None: - if isinstance(event, bytes): - event = event.decode("utf-8") - if isinstance(event, str): - try: - event_obj = _load_json_object(event) - except (json.JSONDecodeError, TypeError): - return - else: - event_obj = event - - if self._should_store_event(event_obj): - self.messages.append(event_obj) - - def _collect_input_from_client_event(self, message: object) -> None: - """Extract user input content from response.create for logging.""" - try: - if isinstance(message, str): - msg_obj = _load_json_object(message) - elif _is_json_object(message): - msg_obj = message - else: - return - - if msg_obj.get("type") != "response.create": - return - - input_items: Final = msg_obj.get("input", []) - if isinstance(input_items, str): - self.input_messages.append({"role": "user", "content": input_items}) - return - - if _is_json_array(input_items): - for item in input_items: - if not _is_json_object(item): - continue - if item.get("type") == "message" and item.get("role") == "user": - content = item.get("content", []) - if isinstance(content, str): - self.input_messages.append({"role": "user", "content": content}) - elif _is_json_array(content): - for c in content: - if _is_json_object(c) and c.get("type") == "input_text": - text = c.get("text", "") - if text: - self.input_messages.append({"role": "user", "content": text}) - except (json.JSONDecodeError, AttributeError, TypeError): - pass - - def _store_input(self, message: object) -> None: - self._collect_input_from_client_event(message) - if self.logging_obj: - self.logging_obj.pre_call(input=message, api_key="") - - def _failure_exception(self) -> Exception | None: - failed_event: Final = next( - (event for event in self.messages if event.get("type") in _RESPONSES_WS_FAILURE_EVENT_TYPES), None - ) - if failed_event is None: - return None - return _map_stream_error_to_exception( - _ws_event_error(failed_event), self.authorized_model or "", self.custom_llm_provider or "" + def _start_turn(self, client_frame: str | None) -> None: + start_time: Final = datetime.now() # noqa: DTZ005 # naive to match the logging pipeline start_time + turn_logging_obj: Final = self.logging_obj.copy_for_new_call( + litellm_call_id=str(uuid.uuid4()), start_time=start_time ) + self._detach_budget_reservation() + turn_logging_obj.model_call_details["messages"] = list(_logged_user_messages(client_frame)) + turn_logging_obj.pre_call(input=client_frame, api_key="") + self._open_turns = (*self._open_turns, _ResponsesWsTurn(logging_obj=turn_logging_obj)) - async def _log_messages(self) -> None: - if not self.logging_obj: + def _detach_budget_reservation(self) -> None: + """Only the first turn settles the connection's admission reservation; later turns charge their cost directly.""" + litellm_params: Final = _json_object_or_none( + self.logging_obj.litellm_params # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Logging.litellm_params is untyped + ) + if litellm_params is None: return - if self.input_messages: - self.logging_obj.model_call_details["messages"] = self.input_messages - if not self.messages: + self.logging_obj.litellm_params = { + key: without_budget_reservation(value) + if key in _PER_TURN_METADATA_KEYS and _is_json_object(value) + else value + for key, value in litellm_params.items() + } + + def _record_client_frame(self, client_frame: str) -> None: + frame: Final = _json_object_or_none(client_frame) + if frame is not None and frame.get("type") == "response.create": + self._start_turn(client_frame) + + def _record_backend_event(self, event: str) -> None: + event_obj: Final = _json_object_or_none(event) + if event_obj is None or event_obj.get("type") not in RESPONSES_WS_LOGGED_EVENT_TYPES: return - exception: Final = self._failure_exception() + if not self._open_turns: + self._start_turn(None) + turn: Final = self._open_turns[0].with_event(event_obj) + if event_obj.get("type") in _RESPONSES_WS_TERMINAL_EVENT_TYPES: + self._open_turns = self._open_turns[1:] + self._log_turn(turn) + return + self._open_turns = (turn, *self._open_turns[1:]) + + def _log_open_turns(self) -> None: + open_turns: Final = self._open_turns + self._open_turns = () + for turn in open_turns: + if turn.events: + self._log_turn(turn) + + def _log_turn(self, turn: _ResponsesWsTurn) -> None: + events: Final = list(turn.events) + exception: Final = turn.failure(self.authorized_model or "", self.custom_llm_provider or "") if exception is None: - asyncio.create_task(self.logging_obj.dispatch_success_handlers(self.messages, prefer_async_handlers=True)) + asyncio.create_task(turn.logging_obj.dispatch_success_handlers(events, prefer_async_handlers=True)) return - self._record_usage_for_failure() + if self._failure is None: + self._failure = exception + _record_usage_for_failure(turn.logging_obj, events) traceback_exception: Final = "".join(traceback.format_exception(exception)) asyncio.create_task( - self.logging_obj.dispatch_failure_handlers(exception, traceback_exception, prefer_async_handlers=True) + turn.logging_obj.dispatch_failure_handlers(exception, traceback_exception, prefer_async_handlers=True) ) - def _record_usage_for_failure(self) -> None: - from litellm.cost_calculator import ResponsesWebSocketTokenUsageProcessor - from litellm.types.utils import LiteLLMRealtimeStreamLoggingObject - - usage: Final = ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results( - self.messages - ) - tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(self.messages) - service_tier: Final = next(iter(tier_partition)) if len(tier_partition) == 1 else None - logging_result: Final = LiteLLMRealtimeStreamLoggingObject( - usage=usage, results=self.messages, service_tier=service_tier - ) - response_cost: Final = self.logging_obj.response_cost_calculator(result=logging_result) or 0.0 - self.logging_obj.record_partial_usage_for_failure(usage, response_cost) - def _wrap_response_event(self, response_str: str) -> str: try: event_obj: Final = _load_json_object(response_str) @@ -2164,7 +2202,7 @@ class ResponsesWebSocketStreaming: # Log the output-masked form so PII redacted by apply_to_output # guardrails does not appear in success logs. - self._store_event(wrapped_str) + self._record_backend_event(wrapped_str) await self.websocket.send_text(wrapped_str) @@ -2173,7 +2211,7 @@ class ResponsesWebSocketStreaming: except Exception as e: verbose_logger.exception("Error in responses WS backend_to_client: %s", e) finally: - await self._log_messages() + self._log_open_turns() def _enforce_authorized_model(self, msg_obj: _MutableJsonObject) -> bool: """ @@ -2518,8 +2556,7 @@ class ResponsesWebSocketStreaming: try: if self.first_message is not None and await self._enforce_or_reject_frame(self.first_message): masked_first: Final = await self._mask_response_create(self.first_message) - self._store_input(masked_first) - self._store_event(masked_first) + self._record_client_frame(masked_first) await self.backend_ws.send(masked_first) while True: @@ -2527,8 +2564,7 @@ class ResponsesWebSocketStreaming: if not await self._enforce_or_reject_frame(message): continue masked = await self._mask_response_create(message) - self._store_input(masked) - self._store_event(masked) + self._record_client_frame(masked) await self.backend_ws.send(masked) except Exception as e: @@ -2551,7 +2587,7 @@ class ResponsesWebSocketStreaming: await self.backend_ws.close() except Exception: pass - return self._failure_exception() + return self._failure # --------------------------------------------------------------------------- diff --git a/litellm/types/utils.py b/litellm/types/utils.py index a957e6957ee..0bcb732a0f5 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4587,6 +4587,7 @@ class LiteLLMRealtimeStreamLoggingObject(LiteLLMPydanticObjectBase): results: SkipValidation[OpenAIRealtimeStreamList] usage: Usage service_tier: str | None = None + id: str | None = None _hidden_params: dict = {} @property diff --git a/tests/integration/_support/responses_ws_peer.py b/tests/integration/_support/responses_ws_peer.py new file mode 100644 index 00000000000..1bf3e1fcbae --- /dev/null +++ b/tests/integration/_support/responses_ws_peer.py @@ -0,0 +1,186 @@ +"""A scripted provider-side Responses WebSocket peer. + +Answers every ``response.create`` frame with the events a script returns for that turn over +``wss://`` (the proxy only dials TLS provider sockets). Clients drive it through ``session``. +""" + +import asyncio +import itertools +import json +import ssl +import threading +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from typing import Final, TypeAlias + +import websockets +from integration._support.tls import server_context +from pydantic import JsonValue +from websockets.asyncio.server import ServerConnection, serve + +PROVIDER_MODEL: Final = "ws-peer-model" +USAGE: Final = {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7} +TERMINAL: Final = frozenset({"response.completed", "response.failed", "error"}) + +Script: TypeAlias = Callable[[str], tuple[dict[str, JsonValue], ...]] + + +@dataclass(frozen=True, slots=True) +class Credentials: + authorization: str | None + api_key: str | None + + +@dataclass(frozen=True, slots=True) +class Peer: + url: str + paths: SimpleQueue[str] + credentials: SimpleQueue[Credentials] + frames: SimpleQueue[dict[str, JsonValue]] + + +def events(response_id: str, text: str) -> tuple[dict[str, JsonValue], ...]: + message: Final = { + "type": "message", + "id": f"msg_{response_id}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + response: Final = {"id": response_id, "object": "response", "created_at": 1700000000, "model": PROVIDER_MODEL} + return ( + {"type": "response.created", "response": {**response, "status": "in_progress", "output": []}}, + { + "type": "response.output_text.delta", + "item_id": f"msg_{response_id}", + "output_index": 0, + "content_index": 0, + "delta": text, + }, + { + "type": "response.completed", + "response": {**response, "status": "completed", "output": [message], "usage": USAGE}, + }, + ) + + +def seven(response_id: str) -> tuple[dict[str, JsonValue], ...]: + return events(response_id, "seven") + + +async def _answer( + connection: ServerConnection, + paths: SimpleQueue[str], + credentials: SimpleQueue[Credentials], + frames: SimpleQueue[dict[str, JsonValue]], + script: Script, +) -> None: + request: Final = connection.request + paths.put(request.path if request is not None else "") + credentials.put( + Credentials(request.headers.get("Authorization"), request.headers.get("api-key")) + if request is not None + else Credentials(None, None) + ) + turns: Final = itertools.count(1) + async for raw in connection: + frame: Final = json.loads(raw) + frames.put(frame) + if frame.get("type") != "response.create": + continue + for event in script(f"resp_peer_{next(turns)}"): + await connection.send(json.dumps(event)) + + +async def _serve( + tls: ssl.SSLContext, + paths: SimpleQueue[str], + credentials: SimpleQueue[Credentials], + frames: SimpleQueue[dict[str, JsonValue]], + ports: SimpleQueue[int], + stop: asyncio.Event, + script: Script, +) -> None: + async with serve( + lambda connection: _answer(connection, paths, credentials, frames, script), "127.0.0.1", 0, ssl=tls + ) as server: + ports.put(next(iter(server.sockets)).getsockname()[1]) + await stop.wait() + + +@contextmanager +def responses_peer(cert: tuple[Path, Path], script: Script = seven) -> Iterator[Peer]: + loop: Final = asyncio.new_event_loop() + stop: Final = asyncio.Event() + paths: Final = SimpleQueue[str]() + credentials: Final = SimpleQueue[Credentials]() + frames: Final = SimpleQueue[dict[str, JsonValue]]() + ports: Final = SimpleQueue[int]() + thread: Final = threading.Thread( + target=loop.run_until_complete, + args=(_serve(server_context(*cert), paths, credentials, frames, ports, stop, script),), + daemon=True, + ) + thread.start() + try: + yield Peer(f"https://127.0.0.1:{ports.get(timeout=10)}/v1", paths, credentials, frames) + finally: + loop.call_soon_threadsafe(stop.set) + thread.join(timeout=10) + loop.close() + + +def create(model: str, text: str, previous_response_id: str | None = None) -> str: + return json.dumps( + { + "type": "response.create", + "model": model, + "store": True, + "input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": text}]}], + **({} if previous_response_id is None else {"previous_response_id": previous_response_id}), + } + ) + + +async def turn(connection: websockets.ClientConnection, frame: str) -> tuple[dict[str, JsonValue], ...]: + await connection.send(frame) + return await until_terminal(connection, ()) + + +async def until_terminal( + connection: websockets.ClientConnection, received: tuple[dict[str, JsonValue], ...] +) -> tuple[dict[str, JsonValue], ...]: + event: Final = json.loads(await asyncio.wait_for(connection.recv(), timeout=20)) + collected: Final = (*received, event) + if event.get("type") in TERMINAL or len(collected) >= 50: + return collected + return await until_terminal(connection, collected) + + +async def session( + proxy_url: str, key: str, model: str, texts: tuple[str, ...] +) -> tuple[tuple[dict[str, JsonValue], ...], ...]: + proxy: Final = proxy_url.rstrip("/").replace("http://", "ws://") + async with websockets.connect( + f"{proxy}/v1/responses?model={model}", + additional_headers={"Authorization": f"Bearer {key}"}, + open_timeout=10, + ) as connection: + first: Final = await turn(connection, create(model, texts[0])) + if len(texts) == 1: + return (first,) + previous: Final = str(first[-1]["response"]["id"]) + second: Final = await turn(connection, create(model, texts[1], previous)) + return (first, second) + + +def completed(received: tuple[dict[str, JsonValue], ...]) -> dict[str, JsonValue]: + assert received[-1]["type"] == "response.completed", [event.get("type") for event in received] + return received[-1]["response"] + + +def drain(queue: SimpleQueue[dict[str, JsonValue]]) -> tuple[dict[str, JsonValue], ...]: + return tuple(queue.get_nowait() for _ in range(queue.qsize())) diff --git a/tests/integration/providers/test_openai_responses_websocket_wire.py b/tests/integration/providers/test_openai_responses_websocket_wire.py index 27d328369d3..4f9a119db2f 100644 --- a/tests/integration/providers/test_openai_responses_websocket_wire.py +++ b/tests/integration/providers/test_openai_responses_websocket_wire.py @@ -1,197 +1,37 @@ import asyncio import hashlib -import itertools -import json -import ssl -import threading import uuid -from collections.abc import Callable, Iterator -from contextlib import contextmanager -from dataclasses import dataclass +from collections.abc import Iterator from pathlib import Path -from queue import SimpleQueue -from typing import Final, TypeAlias +from typing import Final import pytest -import websockets from integration._support.client import Gateway, eventually, gateway_from_environment, object_value, string_value from integration._support.database import read_rows from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.responses_stream import RATE_LIMIT_MESSAGE, openai_error_event -from integration._support.tls import server_context, write_self_signed_cert +from integration._support.responses_ws_peer import ( + PROVIDER_MODEL, + USAGE, + completed, + drain, + events, + responses_peer, + session, +) +from integration._support.tls import write_self_signed_cert from pydantic import JsonValue -from websockets.asyncio.server import ServerConnection, serve pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) - -PROVIDER_MODEL: Final = "ws-peer-model" -USAGE: Final = {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7} -TERMINAL: Final = frozenset({"response.completed", "response.failed", "error"}) PROXY_OPENAI_KEY: Final = "sk-proxy-environment-openai-key" PROXY_AZURE_KEY: Final = "proxy-environment-azure-key" -_Script: TypeAlias = Callable[[str], tuple[dict[str, JsonValue], ...]] - - -@dataclass(frozen=True, slots=True) -class Credentials: - authorization: str | None - api_key: str | None - - -@dataclass(frozen=True, slots=True) -class Peer: - url: str - paths: SimpleQueue[str] - credentials: SimpleQueue[Credentials] - frames: SimpleQueue[dict[str, JsonValue]] - - -def _events(response_id: str, text: str) -> tuple[dict[str, JsonValue], ...]: - message: Final = { - "type": "message", - "id": f"msg_{response_id}", - "status": "completed", - "role": "assistant", - "content": [{"type": "output_text", "text": text, "annotations": []}], - } - response: Final = {"id": response_id, "object": "response", "created_at": 1700000000, "model": PROVIDER_MODEL} - return ( - {"type": "response.created", "response": {**response, "status": "in_progress", "output": []}}, - { - "type": "response.output_text.delta", - "item_id": f"msg_{response_id}", - "output_index": 0, - "content_index": 0, - "delta": text, - }, - { - "type": "response.completed", - "response": {**response, "status": "completed", "output": [message], "usage": USAGE}, - }, - ) - - -def _seven(response_id: str) -> tuple[dict[str, JsonValue], ...]: - return _events(response_id, "seven") - def _top_level_rate_limit(response_id: str) -> tuple[dict[str, JsonValue], ...]: - created: Final = _events(response_id, "")[0] + created: Final = events(response_id, "")[0] return (created, dict(openai_error_event("rate_limit_exceeded", RATE_LIMIT_MESSAGE))) -async def _answer( - connection: ServerConnection, - paths: SimpleQueue[str], - credentials: SimpleQueue[Credentials], - frames: SimpleQueue[dict[str, JsonValue]], - script: _Script, -) -> None: - request: Final = connection.request - paths.put(request.path if request is not None else "") - credentials.put( - Credentials(request.headers.get("Authorization"), request.headers.get("api-key")) - if request is not None - else Credentials(None, None) - ) - turns: Final = itertools.count(1) - async for raw in connection: - frame: Final = json.loads(raw) - frames.put(frame) - if frame.get("type") != "response.create": - continue - for event in script(f"resp_peer_{next(turns)}"): - await connection.send(json.dumps(event)) - - -async def _serve( - tls: ssl.SSLContext, - paths: SimpleQueue[str], - credentials: SimpleQueue[Credentials], - frames: SimpleQueue[dict[str, JsonValue]], - ports: SimpleQueue[int], - stop: asyncio.Event, - script: _Script, -) -> None: - async with serve( - lambda connection: _answer(connection, paths, credentials, frames, script), "127.0.0.1", 0, ssl=tls - ) as server: - ports.put(next(iter(server.sockets)).getsockname()[1]) - await stop.wait() - - -@contextmanager -def responses_peer(cert: tuple[Path, Path], script: _Script = _seven) -> Iterator[Peer]: - loop: Final = asyncio.new_event_loop() - stop: Final = asyncio.Event() - paths: Final = SimpleQueue[str]() - credentials: Final = SimpleQueue[Credentials]() - frames: Final = SimpleQueue[dict[str, JsonValue]]() - ports: Final = SimpleQueue[int]() - thread: Final = threading.Thread( - target=loop.run_until_complete, - args=(_serve(server_context(*cert), paths, credentials, frames, ports, stop, script),), - daemon=True, - ) - thread.start() - try: - yield Peer(f"https://127.0.0.1:{ports.get(timeout=10)}/v1", paths, credentials, frames) - finally: - loop.call_soon_threadsafe(stop.set) - thread.join(timeout=10) - loop.close() - - -def _create(model: str, text: str, previous_response_id: str | None = None) -> str: - return json.dumps( - { - "type": "response.create", - "model": model, - "store": True, - "input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": text}]}], - **({} if previous_response_id is None else {"previous_response_id": previous_response_id}), - } - ) - - -async def _turn(connection: websockets.ClientConnection, frame: str) -> tuple[dict[str, JsonValue], ...]: - await connection.send(frame) - return await _until_terminal(connection, ()) - - -async def _until_terminal( - connection: websockets.ClientConnection, received: tuple[dict[str, JsonValue], ...] -) -> tuple[dict[str, JsonValue], ...]: - event: Final = json.loads(await asyncio.wait_for(connection.recv(), timeout=20)) - collected: Final = (*received, event) - if event.get("type") in TERMINAL or len(collected) >= 50: - return collected - return await _until_terminal(connection, collected) - - -async def _session( - proxy_url: str, key: str, model: str, texts: tuple[str, ...] -) -> tuple[tuple[dict[str, JsonValue], ...], ...]: - proxy: Final = proxy_url.rstrip("/").replace("http://", "ws://") - async with websockets.connect( - f"{proxy}/v1/responses?model={model}", - additional_headers={"Authorization": f"Bearer {key}"}, - open_timeout=10, - ) as connection: - first: Final = await _turn(connection, _create(model, texts[0])) - if len(texts) == 1: - return (first,) - previous: Final = str(first[-1]["response"]["id"]) - second: Final = await _turn(connection, _create(model, texts[1], previous)) - return (first, second) - - -def _completed(events: tuple[dict[str, JsonValue], ...]) -> dict[str, JsonValue]: - assert events[-1]["type"] == "response.completed", [event.get("type") for event in events] - return events[-1]["response"] - - @pytest.fixture(scope="module") def cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]: return write_self_signed_cert(tmp_path_factory.mktemp("responses-ws-cert")) @@ -208,10 +48,6 @@ def candidate(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) yield proxy -def _drain(queue: SimpleQueue[dict[str, JsonValue]]) -> tuple[dict[str, JsonValue], ...]: - return tuple(queue.get_nowait() for _ in range(queue.qsize())) - - def test_a_response_create_frame_streams_from_the_provider_socket_back_to_the_client( candidate: Gateway, cert: tuple[Path, Path] ) -> None: @@ -219,17 +55,17 @@ def test_a_response_create_frame_streams_from_the_provider_socket_back_to_the_cl model: Final = scenario.model(model=f"openai/{PROVIDER_MODEL}", api_base=peer.url) key: Final = scenario.key(models=[model]) text: Final = f"say seven {uuid.uuid4().hex}" - (events,) = asyncio.run(_session(str(candidate.client.base_url), key, model, (text,))) + (events,) = asyncio.run(session(str(candidate.client.base_url), key, model, (text,))) assert [event["type"] for event in events] == [ "response.created", "response.output_text.delta", "response.completed", ] - completed: Final = _completed(events) - assert completed["status"] == "completed" - assert completed["usage"] == USAGE + response: Final = completed(events) + assert response["status"] == "completed" + assert response["usage"] == USAGE assert peer.paths.get_nowait() == f"/v1/responses?model={PROVIDER_MODEL}" - (forwarded,) = _drain(peer.frames) + (forwarded,) = drain(peer.frames) assert forwarded["type"] == "response.create" assert forwarded["model"] == PROVIDER_MODEL assert forwarded["input"][0]["content"][0]["text"] == text @@ -242,12 +78,12 @@ def test_previous_response_id_from_the_first_turn_reaches_the_provider_as_its_ow model: Final = scenario.model(model=f"openai/{PROVIDER_MODEL}", api_base=peer.url) key: Final = scenario.key(models=[model]) texts: Final = (f"remember seven {uuid.uuid4().hex}", f"which number {uuid.uuid4().hex}") - first, second = asyncio.run(_session(str(candidate.client.base_url), key, model, texts)) - assert _completed(first)["status"] == "completed" - assert str(_completed(first)["id"]).startswith("resp_") and _completed(first)["id"] != "resp_peer_1" - assert _completed(second)["status"] == "completed" + first, second = asyncio.run(session(str(candidate.client.base_url), key, model, texts)) + assert completed(first)["status"] == "completed" + assert str(completed(first)["id"]).startswith("resp_") and completed(first)["id"] != "resp_peer_1" + assert completed(second)["status"] == "completed" assert peer.paths.qsize() == 1 - forwarded: Final = _drain(peer.frames) + forwarded: Final = drain(peer.frames) assert [frame["input"][0]["content"][0]["text"] for frame in forwarded] == list(texts) assert "previous_response_id" not in forwarded[0] assert forwarded[1]["previous_response_id"] == "resp_peer_1" @@ -260,10 +96,10 @@ def test_a_top_level_error_frame_reaches_the_client_and_logs_the_provider_code_a model: Final = scenario.model(model=f"openai/{PROVIDER_MODEL}", api_base=peer.url) key: Final = scenario.key(models=[model]) text: Final = f"say seven {uuid.uuid4().hex}" - (events,) = asyncio.run(_session(str(candidate.client.base_url), key, model, (text,))) + (events,) = asyncio.run(session(str(candidate.client.base_url), key, model, (text,))) assert [event["type"] for event in events] == ["response.created", "error"], events assert (events[-1]["code"], events[-1]["message"]) == ("rate_limit_exceeded", RATE_LIMIT_MESSAGE), events - (forwarded,) = _drain(peer.frames) + (forwarded,) = drain(peer.frames) assert forwarded["input"][0]["content"][0]["text"] == text (row,) = eventually( lambda: read_rows( @@ -280,7 +116,7 @@ def test_a_top_level_error_frame_reaches_the_client_and_logs_the_provider_code_a def _one_turn(candidate: Gateway, key: str, model: str) -> tuple[dict[str, JsonValue], ...]: - (events,) = asyncio.run(_session(str(candidate.client.base_url), key, model, (f"say seven {uuid.uuid4().hex}",))) + (events,) = asyncio.run(session(str(candidate.client.base_url), key, model, (f"say seven {uuid.uuid4().hex}",))) return events @@ -293,7 +129,7 @@ def test_a_bare_model_declared_as_openai_sends_its_deployment_key_to_the_provide model=PROVIDER_MODEL, custom_llm_provider="openai", api_base=peer.url, api_key=deployment_key ) events: Final = _one_turn(candidate, scenario.key(models=[model]), model) - assert _completed(events)["status"] == "completed" + assert completed(events)["status"] == "completed" assert peer.credentials.get_nowait().authorization == f"Bearer {deployment_key}" @@ -303,7 +139,7 @@ def test_an_openai_deployment_without_its_own_key_uses_the_proxy_openai_key( with responses_peer(cert) as peer, candidate.scenario() as scenario: model: Final = scenario.model(model=f"openai/{PROVIDER_MODEL}", api_base=peer.url, api_key=None) events: Final = _one_turn(candidate, scenario.key(models=[model]), model) - assert _completed(events)["status"] == "completed" + assert completed(events)["status"] == "completed" assert peer.credentials.get_nowait().authorization == f"Bearer {PROXY_OPENAI_KEY}" @@ -315,7 +151,7 @@ def test_an_azure_deployment_without_its_own_key_gets_the_azure_key_and_never_th model=f"azure/{PROVIDER_MODEL}", api_base=peer.url.removesuffix("/v1"), api_key=None ) events: Final = _one_turn(candidate, scenario.key(models=[model]), model) - assert _completed(events)["status"] == "completed" + assert completed(events)["status"] == "completed" credentials: Final = peer.credentials.get_nowait() assert credentials.api_key == PROXY_AZURE_KEY assert PROXY_OPENAI_KEY not in (credentials.authorization or "") @@ -329,8 +165,8 @@ def test_a_declared_openai_provider_wins_over_another_providers_model_prefix( model: Final = scenario.model(model=gateway_model, custom_llm_provider="openai", api_base=peer.url) key: Final = scenario.key(models=[model]) events: Final = _one_turn(candidate, key, model) - assert _completed(events)["status"] == "completed" - (forwarded,) = _drain(peer.frames) + assert completed(events)["status"] == "completed" + (forwarded,) = drain(peer.frames) assert forwarded["model"] == gateway_model (row,) = eventually( lambda: read_rows( diff --git a/tests/integration/spend/test_responses_websocket_turn_rows.py b/tests/integration/spend/test_responses_websocket_turn_rows.py new file mode 100644 index 00000000000..59b5925c6eb --- /dev/null +++ b/tests/integration/spend/test_responses_websocket_turn_rows.py @@ -0,0 +1,114 @@ +"""Every response.create turn on a native Responses WebSocket is its own spend row. + +OpenAI bills each turn as a request, so one socket carrying two turns has to leave two +rows, each keyed by the response id the client saw and priced from that turn's own usage, +instead of one row at socket close that sums both (LIT-8949). A turn the provider fails +leaves a failure row of its own and does not swallow the completed turn before it. +""" + +import asyncio +import hashlib +import uuid +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually, gateway_from_environment, object_value, string_value +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.responses_stream import RATE_LIMIT_MESSAGE, openai_error_event +from integration._support.responses_ws_peer import PROVIDER_MODEL, USAGE, completed, events, responses_peer, session +from integration._support.tls import write_self_signed_cert +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 180) + +INPUT_RATE: Final = 0.001 +OUTPUT_RATE: Final = 0.002 +TURN_SPEND: Final = USAGE["input_tokens"] * INPUT_RATE + USAGE["output_tokens"] * OUTPUT_RATE +SPEND_ROWS_FOR_KEY: Final = ( + 'SELECT request_id, spend, total_tokens, call_type, status, metadata FROM "LiteLLM_SpendLogs" WHERE api_key = %s' +) + + +def _second_turn_rate_limited(response_id: str) -> tuple[dict[str, JsonValue], ...]: + if response_id == "resp_peer_1": + return events(response_id, "seven") + return (events(response_id, "")[0], dict(openai_error_event("rate_limit_exceeded", RATE_LIMIT_MESSAGE))) + + +def _key_hash(key: str) -> str: + return hashlib.sha256(key.encode()).hexdigest() + + +@pytest.fixture(scope="module") +def cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]: + return write_self_signed_cert(tmp_path_factory.mktemp("responses-ws-rows-cert")) + + +@pytest.fixture(scope="module") +def candidate(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]: + with gateway_from_environment() as base: + with owned_proxy(base, tmp_path_factory.mktemp("responses-ws-rows"), {"SSL_CERT_FILE": str(cert[0])}) as proxy: + yield proxy + + +def test_two_turns_on_one_socket_leave_two_rows_keyed_by_the_response_ids_the_client_saw( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + with responses_peer(cert) as peer, candidate.scenario() as scenario: + model: Final = scenario.model( + model=f"openai/{PROVIDER_MODEL}", + api_base=peer.url, + input_cost_per_token=INPUT_RATE, + output_cost_per_token=OUTPUT_RATE, + ) + key: Final = scenario.key(models=[model]) + texts: Final = (f"first {uuid.uuid4().hex}", f"second {uuid.uuid4().hex}") + + first, second = asyncio.run(session(str(candidate.client.base_url), key, model, texts)) + + response_ids: Final = {str(completed(first)["id"]), str(completed(second)["id"])} + assert len(response_ids) == 2, response_ids + rows: Final = eventually( + lambda: read_rows(SPEND_ROWS_FOR_KEY, (_key_hash(key),)), + lambda found: len(found) >= 2, + seconds=70, + ) + assert {str(row["request_id"]) for row in rows} == response_ids, rows + assert [float(row["spend"]) for row in rows] == [pytest.approx(TURN_SPEND)] * 2, rows + assert {(row["call_type"], row["status"], row["total_tokens"]) for row in rows} == { + ("_aresponses_websocket", "success", USAGE["total_tokens"]) + }, rows + + +def test_a_turn_the_provider_rejects_leaves_its_own_failure_row_next_to_the_completed_turn( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + with responses_peer(cert, _second_turn_rate_limited) as peer, candidate.scenario() as scenario: + model: Final = scenario.model( + model=f"openai/{PROVIDER_MODEL}", + api_base=peer.url, + input_cost_per_token=INPUT_RATE, + output_cost_per_token=OUTPUT_RATE, + ) + key: Final = scenario.key(models=[model]) + texts: Final = (f"first {uuid.uuid4().hex}", f"second {uuid.uuid4().hex}") + + first, second = asyncio.run(session(str(candidate.client.base_url), key, model, texts)) + + assert second[-1]["type"] == "error", second + rows: Final = eventually( + lambda: read_rows(SPEND_ROWS_FOR_KEY, (_key_hash(key),)), + lambda found: len(found) >= 2, + seconds=70, + ) + by_status: Final = {str(row["status"]): row for row in rows} + assert set(by_status) == {"success", "failure"}, rows + assert str(by_status["success"]["request_id"]) == str(completed(first)["id"]), rows + assert float(by_status["success"]["spend"]) == pytest.approx(TURN_SPEND), rows + assert float(by_status["failure"]["spend"]) == 0.0, rows + error: Final = object_value(object_value(by_status["failure"]["metadata"])["error_information"]) + assert str(error["error_code"]) == "429", error + assert RATE_LIMIT_MESSAGE in string_value(error["error_message"]), error diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 14334690624..df6593d29b1 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -11880,3 +11880,55 @@ def test_masking_function_not_in_metadata_when_not_provided(): assert "_langfuse_masking_function" not in result assert result["metadata"]["some_key"] == "some_value" + + +def test_copy_for_new_call_logs_a_turn_under_its_own_call_id_against_the_connections_reservation() -> None: + reservation: Final[dict[str, object]] = {"reserved_cost": 0.5, "entries": []} + connection: Final = LitellmLogging( + model="gpt-5.4", + messages=[], + stream=False, + call_type=CallTypes.aresponses_websocket.value, + start_time=time.time(), + litellm_call_id="connection-call", + function_id="connection-call", + ) + connection.litellm_params = { + "litellm_call_id": "connection-call", + "metadata": {"user_api_key": "hashed-key", "user_api_key_budget_reservation": reservation}, + } + connection.model_call_details["litellm_params"] = connection.litellm_params + turn_start: Final = datetime.datetime(2026, 1, 1, 12, 0, 0) + + turn: Final = connection.copy_for_new_call(litellm_call_id="turn-1", start_time=turn_start) + + assert (turn.litellm_call_id, turn.start_time) == ("turn-1", turn_start) + assert (turn.model_call_details["litellm_call_id"], turn.model_call_details["start_time"]) == ("turn-1", turn_start) + assert turn.model_call_details["litellm_params"]["litellm_call_id"] == "turn-1" + assert turn.litellm_params["metadata"]["user_api_key"] == "hashed-key" + assert turn.litellm_params["metadata"]["user_api_key_budget_reservation"] is reservation + turn.litellm_params["metadata"]["turn_only"] = True + assert "turn_only" not in connection.litellm_params["metadata"] + assert (connection.litellm_call_id, connection.model_call_details["litellm_call_id"]) == ( + "connection-call", + "connection-call", + ) + + +@pytest.mark.parametrize("response_id, expected_row_id", [(None, "turn-7"), ("resp_9", "resp_9")]) +def test_a_websocket_turn_row_is_keyed_by_the_response_id_or_else_its_call_id( + logging_obj: LitellmLogging, response_id: str | None, expected_row_id: str +) -> None: + from litellm.litellm_core_utils.litellm_logging import get_standard_logging_object_payload + + now: Final = datetime_unit_test(2026, 1, 1, 12, 0, 0) + payload: Final = get_standard_logging_object_payload( + kwargs={"litellm_call_id": "turn-7", "litellm_params": {"metadata": {}}}, + init_response_obj=LiteLLMRealtimeStreamLoggingObject(id=response_id, usage=Usage(), results=[]), + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + assert payload is not None + assert payload["id"] == expected_row_id diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter.py b/tests/unit/proxy/hooks/test_parallel_request_limiter.py index 89cb6bd1ef2..f17cad56751 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter.py @@ -16,7 +16,7 @@ from litellm.proxy.hooks.parallel_request_limiter import ( PROXY_MaxParallelRequestsHandler, ) from litellm.proxy.utils import InternalUsageCache, hash_token -from litellm.types.utils import EmbeddingResponse, ModelResponse, TextCompletionResponse, Usage +from litellm.types.utils import CallTypes, EmbeddingResponse, ModelResponse, TextCompletionResponse, Usage FROZEN_INSTANT: Final = datetime(2026, 1, 31, 23, 59, 30) LAST_MICROSECOND_OF_JANUARY: Final = datetime(2026, 1, 31, 23, 59, 59, 999999) @@ -323,3 +323,33 @@ async def test_retry_after_is_the_seconds_until_the_next_minute_of_the_injected_ ) assert exc_info.value.headers["retry-after"] == "14.5" + + +@pytest.mark.asyncio +async def test_websocket_turn_callbacks_keep_the_sockets_slot_until_the_disconnect_release(): + api_key = hash_token("sk-websocket") + handler = PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()), clock=_frozen_clock) + bucket = f"{api_key}::{FROZEN_INSTANT.strftime('%Y-%m-%d-%H-%M')}::request_count" + await handler.internal_usage_cache.async_set_cache( + key=bucket, + value={"current_requests": 1, "current_tpm": 0, "current_rpm": 1}, + litellm_parent_otel_span=None, + ) + turn_kwargs = { + "call_type": CallTypes.aresponses_websocket.value, + "litellm_params": {"metadata": {"user_api_key": api_key}}, + "exception": RuntimeError("provider rejected the turn"), + } + + await handler.async_log_success_event( + kwargs=turn_kwargs, response_obj=None, start_time=FROZEN_INSTANT, end_time=FROZEN_INSTANT + ) + await handler.async_log_failure_event( + kwargs=turn_kwargs, response_obj=None, start_time=FROZEN_INSTANT, end_time=FROZEN_INSTANT + ) + after_two_turns = await handler.internal_usage_cache.async_get_cache(key=bucket, litellm_parent_otel_span=None) + assert after_two_turns["current_requests"] == 1, after_two_turns + + await handler.async_release_max_parallel_requests_on_disconnect(UserAPIKeyAuth(api_key=api_key)) + after_disconnect = await handler.internal_usage_cache.async_get_cache(key=bucket, litellm_parent_otel_span=None) + assert after_disconnect["current_requests"] == 0, after_disconnect diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 1af42947da7..0eaf619ba70 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -1952,6 +1952,7 @@ class TestResponsesWSFirstFrameModelAuth: proxy_logging_obj = MagicMock() proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + proxy_logging_obj.arelease_max_parallel_requests_on_disconnect = AsyncMock(return_value=None) user_api_key_dict = MagicMock() with ( @@ -2016,6 +2017,7 @@ class TestResponsesWSFirstFrameModelAuth: ) proxy_logging_obj = MagicMock() proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + proxy_logging_obj.arelease_max_parallel_requests_on_disconnect = AsyncMock(return_value=None) user_api_key_dict = MagicMock() with ( diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 3625bbb7c04..2b78aa98b04 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -1545,6 +1545,7 @@ class TestNativeWebSocketDeploymentDefaults: mock_logging = MagicMock() mock_logging.pre_call = MagicMock() mock_logging.dispatch_success_handlers = AsyncMock() + mock_logging.copy_for_new_call.return_value = mock_logging client_ws = MagicMock() client_ws.receive_text = AsyncMock(side_effect=RuntimeError("client closed")) @@ -1721,6 +1722,7 @@ class TestNativeWebSocketGuardrails: client_ws.send_text = AsyncMock() logging_obj = MagicMock() logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.copy_for_new_call.return_value = logging_obj delta_event = json.dumps( {"type": "response.output_text.delta", "delta": "alice@example.com"} @@ -1795,6 +1797,7 @@ class TestNativeWebSocketGuardrails: client_ws.send_text = AsyncMock() logging_obj = MagicMock() logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.copy_for_new_call.return_value = logging_obj done_events = [ json.dumps( @@ -2494,6 +2497,7 @@ class TestNativeWebSocketGuardrailMasking: ) logging_obj = MagicMock() logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.copy_for_new_call.return_value = logging_obj handler = _make_streaming( websocket=websocket, @@ -2550,6 +2554,7 @@ class TestNativeWebSocketGuardrailMasking: ) logging_obj = MagicMock() logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.copy_for_new_call.return_value = logging_obj handler = _make_streaming( websocket=websocket, @@ -2613,6 +2618,7 @@ class TestNativeWebSocketGuardrailMasking: ) logging_obj = MagicMock() logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.copy_for_new_call.return_value = logging_obj handler = _make_streaming( websocket=websocket, @@ -2676,6 +2682,7 @@ class TestNativeWebSocketGuardrailMasking: ) logging_obj = MagicMock() logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.copy_for_new_call.return_value = logging_obj handler = _make_streaming( websocket=websocket, @@ -3226,6 +3233,7 @@ class TestNativeWebSocketEncryptedContentAffinity: ) logging_obj = MagicMock() logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.copy_for_new_call.return_value = logging_obj handler = _make_streaming( websocket=websocket, backend_ws=backend_ws, @@ -3272,6 +3280,7 @@ class TestNativeWebSocketEncryptedContentAffinity: ) logging_obj = MagicMock() logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.copy_for_new_call.return_value = logging_obj handler = _make_streaming( websocket=websocket, backend_ws=backend_ws, @@ -3345,6 +3354,7 @@ class TestNativeWebSocketEncryptedContentAffinity: ) logging_obj = MagicMock() logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.copy_for_new_call.return_value = logging_obj logging_obj.dispatch_failure_handlers = AsyncMock() logging_obj.response_cost_calculator = MagicMock(return_value=0.0) handler = _make_streaming( @@ -3366,17 +3376,73 @@ class TestNativeWebSocketEncryptedContentAffinity: assert expected_message in str(exception) @pytest.mark.asyncio - async def test_backend_to_client_bills_completed_turns_before_a_failure(self): + async def test_backend_to_client_logs_each_turn_as_its_own_call(self, monkeypatch): import asyncio + import time from unittest.mock import AsyncMock import websockets.exceptions # noqa: F401 (lazy submodule must be importable) + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.utils import CallTypes + + pre_calls: list[tuple[str, object]] = [] + successes: list[tuple[str, list[object]]] = [] + failures: list[tuple[str, Exception]] = [] + partial_usage: list[tuple[str, int, float]] = [] + + def record_pre_call(self, input, api_key, model=None, additional_args=None): + pre_calls.append((self.litellm_call_id, self.model_call_details["messages"])) + + reservations: dict[str, tuple[object, object]] = {} + + def record_reservation(self): + metadata = self.litellm_params["metadata"] + reservations[self.litellm_call_id] = ( + metadata.get("user_api_key_budget_reservation"), + metadata["user_api_key_auth"].budget_reservation, + ) + + async def record_success(self, result, **kwargs): + record_reservation(self) + successes.append((self.litellm_call_id, [event["type"] for event in result])) + + async def record_failure(self, exception, traceback_exception, **kwargs): + record_reservation(self) + failures.append((self.litellm_call_id, exception)) + + def record_partial(self, usage, response_cost): + partial_usage.append((self.litellm_call_id, usage.total_tokens, response_cost)) + + monkeypatch.setattr(Logging, "pre_call", record_pre_call) + monkeypatch.setattr(Logging, "dispatch_success_handlers", record_success) + monkeypatch.setattr(Logging, "dispatch_failure_handlers", record_failure) + monkeypatch.setattr(Logging, "record_partial_usage_for_failure", record_partial) + monkeypatch.setattr(Logging, "response_cost_calculator", lambda self, result: 0.01) + connection_logging = Logging( + model="gpt-5.4", + messages=[], + stream=False, + call_type=CallTypes.aresponses_websocket.value, + start_time=time.time(), + litellm_call_id="connection-call", + function_id="connection-call", + ) + reservation = {"reserved_cost": 0.5, "entries": [], "finalized": False} + connection_logging.litellm_params = { + "metadata": { + "user_api_key_budget_reservation": reservation, + "user_api_key_auth": UserAPIKeyAuth(api_key="sk-hashed", budget_reservation=reservation), + } + } + websocket = MagicMock() websocket.send_text = AsyncMock() backend_ws = MagicMock() backend_ws.recv = AsyncMock( side_effect=[ + json.dumps({"type": "response.created", "response": {"id": "resp_1", "status": "in_progress"}}), json.dumps( { "type": "response.completed", @@ -3392,21 +3458,27 @@ class TestNativeWebSocketEncryptedContentAffinity: Exception("stop"), ] ) - logging_obj = MagicMock() - logging_obj.dispatch_success_handlers = AsyncMock() - logging_obj.dispatch_failure_handlers = AsyncMock() - logging_obj.response_cost_calculator = MagicMock(return_value=0.01) - handler = _make_streaming(websocket=websocket, backend_ws=backend_ws, logging_obj=logging_obj, request_data={}) + handler = _make_streaming( + websocket=websocket, backend_ws=backend_ws, logging_obj=connection_logging, request_data={} + ) + handler._record_client_frame(json.dumps({"type": "response.create", "input": "first question"})) + handler._record_client_frame(json.dumps({"type": "response.create", "input": "second question"})) await handler.backend_to_client() await asyncio.sleep(0) - logging_obj.record_partial_usage_for_failure.assert_called_once() - usage, response_cost = logging_obj.record_partial_usage_for_failure.call_args[0] - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15) - assert response_cost == 0.01 - logging_obj.dispatch_success_handlers.assert_not_awaited() - logging_obj.dispatch_failure_handlers.assert_awaited_once() + (first_turn, first_messages), (second_turn, second_messages) = pre_calls + assert len({first_turn, second_turn, "connection-call"}) == 3 + assert first_messages == [{"role": "user", "content": "first question"}] + assert second_messages == [{"role": "user", "content": "second question"}] + assert successes == [(first_turn, ["response.created", "response.completed"])] + ((failed_turn, exception),) = failures + assert failed_turn == second_turn + assert "bad turn" in str(exception) + assert partial_usage == [(second_turn, 0, 0.01)] + assert connection_logging.litellm_call_id == "connection-call" + assert reservations[first_turn] == (reservation, reservation) + assert reservations[second_turn] == (None, None) @pytest.mark.asyncio async def test_bidirectional_forward_returns_the_provider_failure(self): @@ -3450,6 +3522,7 @@ class TestNativeWebSocketEncryptedContentAffinity: backend_ws.close = AsyncMock() logging_obj = MagicMock() logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.copy_for_new_call.return_value = logging_obj logging_obj.dispatch_failure_handlers = AsyncMock() logging_obj.response_cost_calculator = MagicMock(return_value=0.0) handler = _make_streaming( @@ -3508,6 +3581,7 @@ class TestNativeWebSocketEncryptedContentAffinity: backend_ws.close = AsyncMock() logging_obj = MagicMock() logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.copy_for_new_call.return_value = logging_obj logging_obj.dispatch_failure_handlers = AsyncMock() handler = _make_streaming( websocket=websocket, diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 3db4432bd6c..cc2a1d62d52 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -9267,3 +9267,23 @@ def test_completion_cost_search_uses_search_tool_input_cost_per_query( assert map_cost == pytest.approx(expected_queries * litellm.model_cost["agentcore/search"]["input_cost_per_query"]) assert configured_cost == pytest.approx(expected_queries * 0.05) + + +def test_responses_ws_logging_object_carries_the_completed_responses_id_tier_and_usage() -> None: + events = [ + {"type": "response.created", "response": {"id": "resp_a", "status": "in_progress"}}, + { + "type": "response.completed", + "response": { + "id": "resp_a", + "service_tier": "priority", + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + }, + }, + ] + + logged = ResponsesWebSocketTokenUsageProcessor.create_logging_object(events) + + assert (logged.id, logged.service_tier) == ("resp_a", "priority") + assert (logged.usage.prompt_tokens, logged.usage.completion_tokens, logged.usage.total_tokens) == (10, 5, 15) + assert ResponsesWebSocketTokenUsageProcessor.create_logging_object([{"type": "error", "error": {}}]).id is None From de4d28653db0b10c0c3102dd574e035e3b7bee17 Mon Sep 17 00:00:00 2001 From: Emmanuel Acheampong Date: Thu, 15 Oct 2026 19:10:58 -0700 Subject: [PATCH 4/8] fix(crusoe): update base URL to current endpoint (#33121) Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/llms/openai_like/providers.json | 2 +- tests/unit/llms/crusoe/test_crusoe.py | 2 +- tests/unit/llms/openai_like/test_json_loader.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 1a987595da2..aa8a17c3824 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -108,7 +108,7 @@ "api_base_env": "AIHUBMIX_API_BASE" }, "crusoe": { - "base_url": "https://managed-inference-api-proxy.crusoecloud.com/v1", + "base_url": "https://api.inference.crusoecloud.com/v1", "api_key_env": "CRUSOE_API_KEY", "api_base_env": "CRUSOE_API_BASE", "param_mappings": { diff --git a/tests/unit/llms/crusoe/test_crusoe.py b/tests/unit/llms/crusoe/test_crusoe.py index 1fd2496cf3c..083941c3ab5 100644 --- a/tests/unit/llms/crusoe/test_crusoe.py +++ b/tests/unit/llms/crusoe/test_crusoe.py @@ -1,7 +1,7 @@ import os from unittest.mock import patch -CRUSOE_API_BASE = "https://managed-inference-api-proxy.crusoecloud.com/v1" +CRUSOE_API_BASE = "https://api.inference.crusoecloud.com/v1" def test_crusoe_json_registry(): diff --git a/tests/unit/llms/openai_like/test_json_loader.py b/tests/unit/llms/openai_like/test_json_loader.py index f6d412fe304..1a4835190cf 100644 --- a/tests/unit/llms/openai_like/test_json_loader.py +++ b/tests/unit/llms/openai_like/test_json_loader.py @@ -12,7 +12,7 @@ import pytest import litellm from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome -CRUSOE_API_BASE = "https://managed-inference-api-proxy.crusoecloud.com/v1" +CRUSOE_API_BASE = "https://api.inference.crusoecloud.com/v1" def test_crusoe_json_registry(): From f13d0430e852f66754b08039754d847ca79f8d17 Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Thu, 15 Oct 2026 19:12:07 -0700 Subject: [PATCH 5/8] test(e2e): attach perf case metrics and limits as JUnit properties (#46774) * test(e2e): attach perf case metrics and limits as JUnit properties The perf case only printed its metrics table, so a green run left no numbers behind. Each observed metric and its limit now land as entries on the test's JUnit testcase, named perf_[_p] and perf_[_p]_limit, so project-releaser can collect them * test(e2e): compare perf properties exactly instead of through pytest.approx basedpyright flags pytest.approx as partially unknown in the e2e lint step, so the harness test now uses metric values that are exact in binary floating point --- tests/e2e/perf/setup/observe.py | 32 ++++++++- tests/e2e/perf/test_perf_cases.py | 15 ++++- tests/e2e_harness/perf/setup/test_observe.py | 68 ++++++++++++++++++++ tests/e2e_harness/pytest.ini | 2 +- 4 files changed, 112 insertions(+), 5 deletions(-) create mode 100644 tests/e2e_harness/perf/setup/test_observe.py diff --git a/tests/e2e/perf/setup/observe.py b/tests/e2e/perf/setup/observe.py index af8295e8ffd..a322e5a7443 100644 --- a/tests/e2e/perf/setup/observe.py +++ b/tests/e2e/perf/setup/observe.py @@ -8,6 +8,7 @@ import statistics from collections import Counter from collections.abc import Iterable from dataclasses import dataclass +from itertools import chain from typing import Final from e2e_http import NoBody, get_external, post_external, unwrap @@ -203,10 +204,19 @@ def observe(metric: Metric, percentile_rank: int | None, observation: Observatio @dataclass(frozen=True, slots=True) class Row: - label: str + metric: Metric + percentile: int | None observed: float | None limit: float | None + @property + def label(self) -> str: + return f"{self.metric} percentile {self.percentile}" if self.percentile is not None else self.metric + + @property + def property_name(self) -> str: + return f"perf_{self.metric}" if self.percentile is None else f"perf_{self.metric}_p{self.percentile}" + @property def over(self) -> bool: return self.limit is not None and self.observed is not None and self.observed > self.limit @@ -240,8 +250,14 @@ def evaluate(case: PerformanceTestCase, observation: Observation) -> tuple[list[ rows: list[Row] = [] for metric, rank in keys: limit = limits.get((metric, rank)) - label = f"{metric} percentile {rank}" if rank is not None else metric - rows.append(Row(label=label, observed=observe(metric, rank, observation), limit=limit.max if limit else None)) + rows.append( + Row( + metric=metric, + percentile=rank, + observed=observe(metric, rank, observation), + limit=limit.max if limit else None, + ) + ) exceeded: Final = [ f"{row.label}: observed {row.observed:.4f} > limit {row.limit:.4f}" if row.over and row.observed is not None and row.limit is not None @@ -252,6 +268,16 @@ def evaluate(case: PerformanceTestCase, observation: Observation) -> tuple[list[ return rows, exceeded +def _row_properties(row: Row) -> tuple[tuple[str, float], ...]: + observed: Final = ((row.property_name, row.observed),) if row.observed is not None else () + limit: Final = ((f"{row.property_name}_limit", row.limit),) if row.limit is not None else () + return observed + limit + + +def junit_properties(rows: Iterable[Row]) -> tuple[tuple[str, float], ...]: + return tuple(chain.from_iterable(map(_row_properties, rows))) + + def render(case: PerformanceTestCase, observation: Observation, rows: list[Row]) -> str: by_endpoint: Final = Counter((record.endpoint, record.stream) for record in observation.records) failures: Final = [record for record in observation.records if not record.ok][:10] diff --git a/tests/e2e/perf/test_perf_cases.py b/tests/e2e/perf/test_perf_cases.py index a69c6bd63cc..d077cb9f8c7 100644 --- a/tests/e2e/perf/test_perf_cases.py +++ b/tests/e2e/perf/test_perf_cases.py @@ -17,6 +17,7 @@ from cases import CASES from setup.observe import ( Observation, evaluate, + junit_properties, read_guardrail_calls, read_mock_requests, render, @@ -48,7 +49,13 @@ def _spend_logged_share(proxy: ProxyClient, call_ids: tuple[str, ...]) -> float @pytest.mark.parametrize("case", CASES, ids=[case.name for case in CASES]) -def test_perf_case(case: PerformanceTestCase, proxy: ProxyClient, idp: Keycloak, resources: ResourceManager) -> None: +def test_perf_case( + case: PerformanceTestCase, + proxy: ProxyClient, + idp: Keycloak, + resources: ResourceManager, + request: pytest.FixtureRequest, +) -> None: plan: Final = build_plan(case) mock_urls: Final = {provider: mock_url(provider) for provider in MOCK_PROVIDERS} for url in (*mock_urls.values(), guardrail_url()): @@ -94,6 +101,12 @@ def test_perf_case(case: PerformanceTestCase, proxy: ProxyClient, idp: Keycloak, gateway_cpu_seconds=cpu.seconds, ) rows, exceeded_limits = evaluate(case, observation) + item: Final = request.node # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # pytest: untyped + assert isinstance(item, pytest.Item) + item.user_properties[:] = [ + entry for entry in item.user_properties if not entry[0].startswith("perf_") + ] # rebind-ok: a --reruns retry replaces the first attempt's numbers instead of appending a second set + item.user_properties.extend(junit_properties(rows)) print(render(case, observation, rows)) print( f" phases: setup {run_started - setup_started:.0f}s, locust {run_finished - run_started:.0f}s, " diff --git a/tests/e2e_harness/perf/setup/test_observe.py b/tests/e2e_harness/perf/setup/test_observe.py new file mode 100644 index 00000000000..998d7d9b87e --- /dev/null +++ b/tests/e2e_harness/perf/setup/test_observe.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +from dataclasses import replace +from typing import Final + +from cases.small_bodies_5rps import CASE +from memory_readings import RssCapture +from setup.datastore import DbCounters +from setup.observe import Observation, evaluate, junit_properties +from setup.performance_test_case import Limit +from setup.wire import WireRecord + +NO_RSS: Final = RssCapture(readings=(), failures=()) + + +def served(seq: int, total_s: float) -> WireRecord: + return WireRecord( + seq=seq, + marker=f"perfreq-{seq:06d}-n", + caller="key-a1", + endpoint="chat_completions", + stream=False, + blocked=False, + started_at=0.0, + status=200, + header_s=total_s, + first_byte_s=total_s, + first_token_s=None, + total_s=total_s, + text="", + ok=True, + mock_delay_s=1.0, + call_id=None, + error=None, + ) + + +def test_junit_properties_carry_each_observed_metric_and_each_limit_including_unobserved_ones() -> None: + case: Final = replace( + CASE, + limits=( + Limit("gateway_added_seconds", percentile=50, max=0.625), + Limit("redis_calls_per_request", max=35), + ), + ) + observation: Final = Observation( + records=(served(0, 1.25), served(1, 1.5), served(2, 1.75)), + mock_requests=(), + guardrail_calls=(), + rss_before=NO_RSS, + rss_after=NO_RSS, + wall_seconds=3.0, + db_before=DbCounters(commits=100), + db_after=DbCounters(commits=130), + ) + rows, _ = evaluate(case, observation) + + properties: Final = junit_properties(rows) + + assert len({name for name, _ in properties}) == len(properties) + assert dict(properties) == { + "perf_unexpected_status_count": 0.0, + "perf_gateway_added_seconds_p50": 0.5, + "perf_gateway_added_seconds_p50_limit": 0.625, + "perf_gateway_added_seconds_p99": 0.75, + "perf_postgres_calls_per_request": 10.0, + "perf_redis_calls_per_request_limit": 35.0, + } diff --git a/tests/e2e_harness/pytest.ini b/tests/e2e_harness/pytest.ini index 54d2630acc6..a29166d31bd 100644 --- a/tests/e2e_harness/pytest.ini +++ b/tests/e2e_harness/pytest.ini @@ -3,7 +3,7 @@ # (`from e2e_http import ...`, `from batch_cleanup import ...`) exactly as the suites # do, so tests/e2e and each suite folder that owns a module under test go on the path. addopts = --strict-markers --strict-config -p no:cacheprovider -pythonpath = ../e2e ../e2e/batches ../e2e/guardrails ../e2e/logging +pythonpath = ../e2e ../e2e/batches ../e2e/guardrails ../e2e/logging ../e2e/perf markers = covers(cell_id, *, exercised_on=()): coverage-registry cell(s) a test covers; exercised here only to prove the collector and the JUnit properties read it cli_determinism: drives the real claude CLI for several seconds From dfb949d4372eba86fab382a5038c56b86a0582b5 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 02:24:47 +0000 Subject: [PATCH 6/8] fix(terraform)!: send additional_drop_params to the proxy instead of deleting the keys it names (#46782) * fix(terraform): send additional_drop_params to the proxy instead of deleting the keys it names litellm_model treated additional_drop_params in additional_litellm_params as a provider-side delete list: it never sent the list and removed every key it named from litellm_params, so vertex_location vanished and requests still carried store. The list now goes to the proxy like any other JSON value, so the proxy drops those params per request. Approach credited to #38345. Fixes #38349 * fix(terraform)!: validate additional_drop_params and read it back from the proxy on refresh * docs(terraform): note that a non-list additional_drop_params now fails the apply * fix(terraform)!: check additional_drop_params at plan time and refresh listed keys and outside lists * fix(terraform)!: show every listed key an older provider deleted in the upgrade plan and leave outside drop lists alone Refresh no longer reads listed keys back from the proxy, which left a plan diff that never cleared when a listed key was removed from config. When the proxy has no stored list, the plan now shows the list and each listed key the config sets as changes. A drop list set outside Terraform on a model whose config has none stays out of the plan, and an empty value now fails the plan * fix(terraform): show a listed merge_reasoning_content_in_choices in the upgrade plan The refresh now clears merge_reasoning_content_in_choices when it is listed and the proxy has no stored list, since an older provider deleted it, and no longer clears litellm_credential_name, which the older provider always sent --- terraform/provider/CHANGELOG.md | 2 + terraform/provider/docs/resources/model.md | 14 +- .../examples/model_additional_params.tf | 21 +- terraform/provider/litellm/resource_model.go | 12 +- .../provider/litellm/resource_model_crud.go | 192 +++++++++----- .../provider/litellm/resource_model_test.go | 243 ++++++++++++++++++ terraform/provider/litellm/types.go | 1 + 7 files changed, 394 insertions(+), 91 deletions(-) diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index 4282485714a..597e0c6d4d8 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -51,6 +51,7 @@ longer signal it. ### Fixed +- **model**: `additional_drop_params` in `litellm_model.additional_litellm_params` is now sent to the proxy as the deployment's `additional_drop_params` list, the same list a `model_list` entry takes. Before, the provider never sent the list and instead deleted every key it named from `litellm_params`, so a `store` or `vertex_location` listed there was lost and requests still carried the parameters (#38349). The value must be a JSON list of parameter names such as `jsonencode(["store"])`, which `terraform plan` now checks, and `null` entries are skipped as before. When the configuration sets the list, refresh reads the stored list back from the proxy, so a list changed outside Terraform shows up as drift. On a model an older provider created, the proxy has no stored list, so the first plan after upgrading shows the list and every listed key the configuration sets as changes, since that apply writes them. A write-only secret such as `model_api_key_wo` is never shown in a plan - **user**: `litellm_user` updates no longer send `blocked`, which `/user/update` writes into the user table where no such column exists, so every update failed with a 400 (`Could not find field at upsertOneLiteLLM_UserTable.create.blocked`). Create still sends `blocked`, including `false`, so a `default_internal_user_params.blocked` on the proxy cannot block the auto-created key. Changing `blocked` after creation has no effect, since the proxy has no user-level block; use `litellm_key_block` instead - **model**: `litellm_model` refresh now reads the `{"data": [...]}` envelope `/model/info` returns, so `model_info` fields changed outside Terraform show up as drift instead of silently keeping the previous state - **key**: An update that changes `team_id` and fails because the key was already cascade-deleted along with its previous team now recovers by recreating the key under the new team, instead of aborting the apply. The key's absence is confirmed against the proxy first, so an unrelated failure still errors out, and a `team_id` change between two teams that both still exist stays a plain in-place update @@ -69,6 +70,7 @@ longer signal it. ### Changed +- **model** (breaking): a key named in `additional_litellm_params.additional_drop_params` is no longer deleted from the deployment's `litellm_params`. A configuration that also sets that key, such as `vertex_location`, stores it on the deployment at its next apply, and the proxy applies the drop list to requests instead. For `vertex_location` this changes routing. `/v1/messages` calls now go to the configured region, while chat completions drop a listed `vertex_location` and use the default region. To keep such a key off a deployment, remove its attribute and replace the model with `terraform apply -replace`, which gives it a new model id, since an update keeps a stored value the provider leaves out. `model`, `custom_llm_provider` and `merge_reasoning_content_in_choices` are always sent. A value that is not a JSON list of names, such as a bare `"store"` or an empty string, now fails the plan, and with it `terraform destroy`, instead of being stored as a string the proxy ignores. Rewrite it as `jsonencode([...])`, or remove the model without fixing it through `terraform destroy -refresh=false` - **key** (breaking): `model_max_budget` on `litellm_key` is now a JSON string of per-model budget objects (`jsonencode({"gpt-4o-mini" = {budget_limit = 50, time_period = "30d"}})`), matching `litellm_user`, `litellm_budget` and `litellm_tag`. The old `map(number)` form sent bare numbers to `/key/generate`, which the proxy rejects with a 500 (`'int' object is not iterable`), so every key with a non-empty `model_max_budget` failed to apply. Existing state upgrades automatically (schema version 1) and the attribute is refilled from the proxy on the next read; configurations still using the map form must be rewritten - **Versioning**: the provider is now published at the LiteLLM version, from the same commit as the proxy, on every LiteLLM release (dev, rc, stable). The `0.x` line ends at `0.4.0`; a `~> 0.4` constraint will not receive further releases, so re-pin to the LiteLLM version your proxy runs (for example `~> 1.99.0`). Existing `0.x` versions remain in the registry and keep verifying diff --git a/terraform/provider/docs/resources/model.md b/terraform/provider/docs/resources/model.md index b7c440b0e65..85c495bf1ca 100644 --- a/terraform/provider/docs/resources/model.md +++ b/terraform/provider/docs/resources/model.md @@ -183,16 +183,16 @@ The following arguments are supported: * Numeric strings are parsed first as integers; if integer parsing fails, parsed as floats (e.g., `"16384"` -> 16384, `"0.75"` -> 0.75) * JSON strings (starting with `[` or `{`) are parsed as JSON objects/arrays * Non-convertible strings remain strings - * Non-string map values (if supplied) are passed through unchanged. * The provider merges these keys into the `litellm_params` payload sent to the API. * Note: the remote API may not echo back all custom parameters; this provider preserves `additional_litellm_params` in state when present in configuration. - **Special parameter: `additional_drop_params`** - * When `additional_drop_params` is provided as a JSON array string, it specifies parameters to remove from the final `litellm_params` before sending to the API - * This allows you to override or remove built-in parameters if needed - * The `additional_drop_params` key itself is not included in the final parameters + **`additional_drop_params`** + * Give it a JSON array string, such as `jsonencode(["store"])`. The provider sends the list to the proxy as the deployment's `additional_drop_params`, and the proxy applies it the same as `additional_drop_params` in a `model_list` entry. `terraform plan` fails on any other value, an empty value, or an entry that is not a name, and `null` entries are skipped + * Other keys on the resource, such as `vertex_location`, are sent unchanged + * When the configuration sets the list, refresh reads the stored list back from the proxy, so a list changed outside Terraform shows up as drift. A model whose configuration has no list keeps a list set outside Terraform out of the plan, but the proxy clears it on the next update + * When the proxy has no stored list, as on a model an older provider created, the next plan shows the list and every listed key the configuration sets as changes, since the next apply writes them - Example showing booleans, integers, floats, strings, and parameter dropping: + Example showing booleans, integers, floats, strings, JSON, and dropped request parameters: ```hcl resource "litellm_model" "with_additional" { @@ -208,7 +208,7 @@ The following arguments are supported: "scale" = "0.75" # becomes float 0.75 "note" = "for testing" # stays string "complex_config" = "{\"nested\": {\"value\": 42}}" # parsed as JSON object - "additional_drop_params" = "[\"reasoningEffort\"]" # removes reasoningEffort parameter + "additional_drop_params" = "[\"reasoning_effort\"]" # proxy drops reasoning_effort from requests } } ``` diff --git a/terraform/provider/examples/model_additional_params.tf b/terraform/provider/examples/model_additional_params.tf index fb4981ec6fb..64c37993e04 100644 --- a/terraform/provider/examples/model_additional_params.tf +++ b/terraform/provider/examples/model_additional_params.tf @@ -9,8 +9,7 @@ provider "litellm" { # - Numeric strings will be parsed to integer (if possible) otherwise float. # - JSON strings (starting with [ or {) will be parsed as JSON objects/arrays. # - Non-convertible strings remain strings. -# - Non-string map values are passed through unchanged. -# - Use "additional_drop_params" as a JSON array to remove parameters from the final request. +# - "additional_drop_params" takes a JSON array of request parameters the proxy drops before calling the model provider. resource "litellm_model" "with_additional" { model_name = "custom-model" @@ -21,14 +20,12 @@ resource "litellm_model" "with_additional" { # Additional parameters not exposed as first-class arguments additional_litellm_params = { - "use_fine_tune" = "true" # becomes boolean true - "max_context" = "16384" # becomes integer 16384 - "temperature_scale" = "0.75" # becomes float 0.75 - "experimental_feature" = "enabled" # stays string "enabled" - "complex_config" = "{\"nested\": {\"value\": 42}}" # parsed as JSON object - "additional_drop_params" = "[\"reasoningEffort\"]" # removes reasoningEffort parameter - # You may also pass non-string values (they will be passed through unchanged) - # "raw_flag" = true + "use_fine_tune" = "true" # becomes boolean true + "max_context" = "16384" # becomes integer 16384 + "temperature_scale" = "0.75" # becomes float 0.75 + "experimental_feature" = "enabled" # stays string "enabled" + "complex_config" = "{\"nested\": {\"value\": 42}}" # parsed as JSON object + "additional_drop_params" = "[\"reasoning_effort\"]" # proxy drops reasoning_effort from requests } # Cost configuration (optional) @@ -47,9 +44,9 @@ resource "litellm_model" "azure_with_drop_params" { tier = "paid" mode = "completion" - # Drop the reasoningEffort parameter that might be automatically added + # Have the proxy drop reasoning_effort from requests before they reach Azure additional_litellm_params = { - "additional_drop_params" = "[\"reasoningEffort\"]" + "additional_drop_params" = "[\"reasoning_effort\"]" } input_cost_per_million_tokens = 0.25 diff --git a/terraform/provider/litellm/resource_model.go b/terraform/provider/litellm/resource_model.go index 1dc029b861c..e9582ad75fb 100644 --- a/terraform/provider/litellm/resource_model.go +++ b/terraform/provider/litellm/resource_model.go @@ -7,10 +7,11 @@ import ( func resourceLiteLLMModel() *schema.Resource { return &schema.Resource{ - Create: resourceLiteLLMModelCreate, - Read: resourceLiteLLMModelRead, - Update: resourceLiteLLMModelUpdate, - Delete: resourceLiteLLMModelDelete, + Create: resourceLiteLLMModelCreate, + Read: resourceLiteLLMModelRead, + Update: resourceLiteLLMModelUpdate, + Delete: resourceLiteLLMModelDelete, + CustomizeDiff: validateAdditionalDropParams, Importer: &schema.ResourceImporter{ StateContext: schema.ImportStatePassthroughContext, }, @@ -237,7 +238,8 @@ func resourceLiteLLMModel() *schema.Resource { Elem: &schema.Schema{ Type: schema.TypeString, }, - Description: "Additional parameters to pass to litellm_params beyond the standard ones", + DiffSuppressFunc: suppressEquivalentDropParams, + Description: "Additional parameters to pass to litellm_params beyond the standard ones", }, }, } diff --git a/terraform/provider/litellm/resource_model_crud.go b/terraform/provider/litellm/resource_model_crud.go index 3309dbe26dd..3db6f68760d 100644 --- a/terraform/provider/litellm/resource_model_crud.go +++ b/terraform/provider/litellm/resource_model_crud.go @@ -1,10 +1,13 @@ package litellm import ( + "context" "encoding/json" "fmt" "log" + "maps" "net/url" + "slices" "strconv" "strings" "time" @@ -68,6 +71,119 @@ func secretFromResource(d *schema.ResourceData, attr string) string { return d.Get(attr).(string) } +const additionalDropParamsKey = "additional_drop_params" + +func parseAdditionalDropParams(raw string) ([]string, error) { + var items []interface{} + if err := json.Unmarshal([]byte(raw), &items); err != nil { + return nil, fmt.Errorf("additional_litellm_params.%s must be a JSON list of parameter names, such as jsonencode([\"store\"]): %w", additionalDropParamsKey, err) + } + names := []string{} + for _, item := range items { + if item == nil { + continue + } + name, ok := item.(string) + if !ok { + return nil, fmt.Errorf("additional_litellm_params.%s entries must be parameter names, got %v", additionalDropParamsKey, item) + } + names = append(names, name) + } + return names, nil +} + +func validateAdditionalDropParams(_ context.Context, d *schema.ResourceDiff, _ interface{}) error { + params, _ := d.Get("additional_litellm_params").(map[string]interface{}) + raw, set := params[additionalDropParamsKey].(string) + if !set { + return nil + } + _, err := parseAdditionalDropParams(raw) + return err +} + +func suppressEquivalentDropParams(key, oldValue, newValue string, _ *schema.ResourceData) bool { + if key != "additional_litellm_params."+additionalDropParamsKey { + return false + } + oldNames, oldErr := parseAdditionalDropParams(oldValue) + newNames, newErr := parseAdditionalDropParams(newValue) + return oldErr == nil && newErr == nil && slices.Equal(oldNames, newNames) +} + +var attributeSentAsLiteLLMParam = map[string]string{ + "custom_llm_provider": "custom_llm_provider", + "tpm": "tpm", + "rpm": "rpm", + "input_cost_per_token": "input_cost_per_million_tokens", + "output_cost_per_token": "output_cost_per_million_tokens", + "api_key": "model_api_key", + "api_base": "model_api_base", + "api_version": "api_version", + "input_cost_per_pixel": "input_cost_per_pixel", + "output_cost_per_pixel": "output_cost_per_pixel", + "input_cost_per_second": "input_cost_per_second", + "output_cost_per_second": "output_cost_per_second", + "aws_access_key_id": "aws_access_key_id", + "aws_secret_access_key": "aws_secret_access_key", + "aws_region_name": "aws_region_name", + "aws_session_name": "aws_session_name", + "aws_role_name": "aws_role_name", + "vertex_project": "vertex_project", + "vertex_location": "vertex_location", + "vertex_credentials": "vertex_credentials", + "reasoning_effort": "reasoning_effort", + "thinking": "thinking_enabled", + "merge_reasoning_content_in_choices": "merge_reasoning_content_in_choices", +} + +func refreshAdditionalDropParams(d *schema.ResourceData, stored interface{}) { + params, _ := d.Get("additional_litellm_params").(map[string]interface{}) + raw, managed := params[additionalDropParamsKey].(string) + if !managed { + return + } + refreshed := maps.Clone(params) + if stored != nil { + encoded, _ := json.Marshal(stored) + refreshed[additionalDropParamsKey] = string(encoded) + d.Set("additional_litellm_params", refreshed) + return + } + delete(refreshed, additionalDropParamsKey) + listed, _ := parseAdditionalDropParams(raw) + for _, key := range listed { + delete(refreshed, key) + if attribute, ok := attributeSentAsLiteLLMParam[key]; ok { + d.Set(attribute, nil) + } + } + d.Set("additional_litellm_params", refreshed) +} + +func coerceAdditionalParamValue(strValue string) interface{} { + trimmedValue := strings.TrimSpace(strValue) + if strings.HasPrefix(trimmedValue, "[") || strings.HasPrefix(trimmedValue, "{") { + var parsedValue interface{} + if err := json.Unmarshal([]byte(strValue), &parsedValue); err == nil { + return parsedValue + } + } + switch strValue { + case "true": + return true + case "false": + return false + } + if intValue, err := strconv.Atoi(strValue); err == nil { + return intValue + } + if floatValue, err := strconv.ParseFloat(strValue, 64); err == nil { + return floatValue + } + return strValue +} + func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) error { client, ok := m.(*Client) if !ok { @@ -175,76 +291,17 @@ func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) e litellmParams["thinking"] = thinking } - // Add additional parameters if provided if additionalParams, ok := d.GetOk("additional_litellm_params"); ok { - var dropParams []string - for key, value := range additionalParams.(map[string]interface{}) { - // Convert string values to appropriate types where possible - if strValue, ok := value.(string); ok { - // Check if it's JSON (starts with [ or {) - trimmedValue := strings.TrimSpace(strValue) - if strings.HasPrefix(trimmedValue, "[") || strings.HasPrefix(trimmedValue, "{") { - var parsedValue interface{} - if err := json.Unmarshal([]byte(strValue), &parsedValue); err == nil { - // Successfully parsed JSON - if key == "additional_drop_params" { - // Handle drop params specially - if dropList, ok := parsedValue.([]interface{}); ok { - for _, item := range dropList { - if paramStr, ok := item.(string); ok { - dropParams = append(dropParams, paramStr) - } - } - } - continue // Don't add to litellmParams - } else { - litellmParams[key] = parsedValue - } - } else { - // Not valid JSON, apply existing conversion logic - if strValue == "true" { - litellmParams[key] = true - } else if strValue == "false" { - litellmParams[key] = false - } else { - // Try to convert numeric strings - if intValue, err := strconv.Atoi(strValue); err == nil { - litellmParams[key] = intValue - } else if floatValue, err := strconv.ParseFloat(strValue, 64); err == nil { - litellmParams[key] = floatValue - } else { - // Keep as string - litellmParams[key] = strValue - } - } - } - } else { - // Apply existing conversion logic for non-JSON strings - if strValue == "true" { - litellmParams[key] = true - } else if strValue == "false" { - litellmParams[key] = false - } else { - // Try to convert numeric strings - if intValue, err := strconv.Atoi(strValue); err == nil { - litellmParams[key] = intValue - } else if floatValue, err := strconv.ParseFloat(strValue, 64); err == nil { - litellmParams[key] = floatValue - } else { - // Keep as string - litellmParams[key] = strValue - } - } - } - } else { - litellmParams[key] = value + if key != additionalDropParamsKey { + litellmParams[key] = coerceAdditionalParamValue(value.(string)) + continue } - } - - // Apply drop params at the end - for _, paramToDrop := range dropParams { - delete(litellmParams, paramToDrop) + dropParams, err := parseAdditionalDropParams(value.(string)) + if err != nil { + return err + } + litellmParams[key] = dropParams } } @@ -419,6 +476,7 @@ func resourceLiteLLMModelRead(d *schema.ResourceData, m interface{}) error { if _, ok := d.GetOk("additional_litellm_params"); ok { d.Set("additional_litellm_params", d.Get("additional_litellm_params")) } + refreshAdditionalDropParams(d, modelResp.LiteLLMParams.AdditionalDropParams) return nil } diff --git a/terraform/provider/litellm/resource_model_test.go b/terraform/provider/litellm/resource_model_test.go index 22c226adfd3..4c2d5d27442 100644 --- a/terraform/provider/litellm/resource_model_test.go +++ b/terraform/provider/litellm/resource_model_test.go @@ -3,8 +3,11 @@ package litellm import ( "context" "encoding/json" + "maps" "net/http" "net/http/httptest" + "reflect" + "slices" "testing" "github.com/hashicorp/go-cty/cty" @@ -116,6 +119,246 @@ func TestResourceLiteLLMModelCreateOmitsUnsetDisplayName(t *testing.T) { } } +func TestResourceLiteLLMModelCreateSendsAdditionalLiteLLMParams(t *testing.T) { + var createPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/model/new": + if err := json.NewDecoder(r.Body).Decode(&createPayload); err != nil { + t.Errorf("failed to decode create payload: %v", err) + } + w.Write([]byte(modelInfoBody(""))) + case "/model/info": + w.Write([]byte(modelInfoBody(""))) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMModel().Schema, map[string]interface{}{ + "model_name": "vertex-claude-haiku", + "custom_llm_provider": "vertex_ai", + "base_model": "claude-haiku-4-5", + "vertex_project": "test-project", + "vertex_location": "eu", + "additional_litellm_params": map[string]interface{}{ + "additional_drop_params": `["store", "vertex_location"]`, + "tags": ` ["prod", "eu"]`, + "headers": `{"x-team": "search"}`, + "use_in_pass_through": "true", + "merge_reasoning_content_in_choices": "false", + "max_parallel_requests": "8", + "input_cost_per_token": "0.5", + "description": "[draft] haiku", + }, + }) + + if err := resourceLiteLLMModelCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + + litellmParams, ok := createPayload["litellm_params"].(map[string]interface{}) + if !ok { + t.Fatalf("expected litellm_params object in create payload, got %v", createPayload["litellm_params"]) + } + want := map[string]interface{}{ + "additional_drop_params": []interface{}{"store", "vertex_location"}, + "vertex_project": "test-project", + "vertex_location": "eu", + "tags": []interface{}{"prod", "eu"}, + "headers": map[string]interface{}{"x-team": "search"}, + "use_in_pass_through": true, + "merge_reasoning_content_in_choices": false, + "max_parallel_requests": float64(8), + "input_cost_per_token": 0.5, + "description": "[draft] haiku", + } + for key, wantValue := range want { + if got := litellmParams[key]; !reflect.DeepEqual(got, wantValue) { + t.Errorf("litellm_params[%q] = %#v, want %#v", key, got, wantValue) + } + } +} + +const terraformUnknownValue = "74D93920-ED26-11E3-AC10-0800200C9A66" + +func TestResourceLiteLLMModelChecksAdditionalDropParams(t *testing.T) { + cases := map[string]struct { + dropParams string + want []interface{} + wantErr bool + }{ + "null entry from an unset variable is skipped": {dropParams: `["store", null]`, want: []interface{}{"store"}}, + "number entry is rejected": {dropParams: `["store", 1]`, wantErr: true}, + "object entry is rejected": {dropParams: `["store", {"name": "temperature"}]`, wantErr: true}, + "object is rejected": {dropParams: `{"store": true}`, wantErr: true}, + "bare name is rejected": {dropParams: "store", wantErr: true}, + "empty value is rejected": {dropParams: "", wantErr: true}, + } + for name, tc := range cases { + t.Run(name, func(t *testing.T) { + var createPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/model/new": + if err := json.NewDecoder(r.Body).Decode(&createPayload); err != nil { + t.Errorf("failed to decode create payload: %v", err) + } + w.Write([]byte(modelInfoBody(""))) + case "/model/info": + w.Write([]byte(modelInfoBody(""))) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + res := resourceLiteLLMModel() + config := map[string]interface{}{ + "model_name": "vertex-claude-haiku", + "custom_llm_provider": "vertex_ai", + "base_model": "claude-haiku-4-5", + "additional_litellm_params": map[string]interface{}{"additional_drop_params": tc.dropParams}, + } + if _, err := res.Diff(context.Background(), nil, terraform.NewResourceConfigRaw(config), nil); (err != nil) != tc.wantErr { + t.Errorf("plan error = %v, want error %v", err, tc.wantErr) + } + + d := schema.TestResourceDataRaw(t, res.Schema, config) + err := resourceLiteLLMModelCreate(d, NewClient(srv.URL, "test-key", true)) + if tc.wantErr { + if err == nil { + t.Fatal("expected create to fail") + } + if createPayload != nil { + t.Errorf("expected no /model/new request, got %v", createPayload) + } + return + } + if err != nil { + t.Fatalf("create failed: %v", err) + } + litellmParams := createPayload["litellm_params"].(map[string]interface{}) + if got := litellmParams["additional_drop_params"]; !reflect.DeepEqual(got, tc.want) { + t.Errorf("additional_drop_params = %#v, want %#v", got, tc.want) + } + }) + } +} + +func TestResourceLiteLLMModelPlanAcceptsDropParamsUnknownUntilApply(t *testing.T) { + config := terraform.NewResourceConfigRaw(map[string]interface{}{ + "model_name": "vertex-claude-haiku", + "custom_llm_provider": "vertex_ai", + "additional_litellm_params": map[string]interface{}{"additional_drop_params": terraformUnknownValue}, + }) + if _, err := resourceLiteLLMModel().Diff(context.Background(), nil, config, nil); err != nil { + t.Errorf("plan failed on a drop list not known until apply: %v", err) + } +} + +func TestResourceLiteLLMModelPlanComparesAdditionalDropParamsWithProxy(t *testing.T) { + const ( + listKey = "additional_litellm_params.additional_drop_params" + parallelKey = "additional_litellm_params.max_parallel_requests" + ) + cases := map[string]struct { + configured string + configLocation string + stored interface{} + storedLocation string + extraConfig map[string]interface{} + wantChanged []string + }{ + "list never sent by an older provider": {configured: `["store"]`, configLocation: "eu", storedLocation: "eu", wantChanged: []string{listKey}}, + "listed keys an older provider deleted": { + configured: `["store", "vertex_location", "tpm", "max_parallel_requests"]`, + configLocation: "eu", + wantChanged: []string{listKey, "vertex_location", "tpm", parallelKey}, + }, + "listed flag an older provider deleted": { + configured: `["merge_reasoning_content_in_choices"]`, + configLocation: "eu", + storedLocation: "eu", + extraConfig: map[string]interface{}{"merge_reasoning_content_in_choices": true}, + wantChanged: []string{listKey, "merge_reasoning_content_in_choices"}, + }, + "listed key an older provider still sent": { + configured: `["litellm_credential_name"]`, + configLocation: "eu", + storedLocation: "eu", + extraConfig: map[string]interface{}{"litellm_credential_name": "vertex-cred"}, + wantChanged: []string{listKey}, + }, + "list changed in the Admin UI": {configured: `["store"]`, configLocation: "eu", stored: []string{"temperature"}, storedLocation: "eu", wantChanged: []string{listKey}}, + "list set outside Terraform": {configLocation: "eu", stored: []string{"store"}, storedLocation: "eu"}, + "same list written with spaces": {configured: `["store", "vertex_location"]`, configLocation: "eu", stored: []string{"store", "vertex_location"}, storedLocation: "eu"}, + "null entry the proxy never received": {configured: `["store", null]`, configLocation: "eu", stored: []string{"store"}, storedLocation: "eu"}, + "listed key left out of the config": {configured: `["vertex_location"]`, stored: []string{"vertex_location"}, storedLocation: "us-east5"}, + "no list on the model or in the proxy": {configLocation: "eu", storedLocation: "eu"}, + } + for name, tc := range cases { + t.Run(name, func(t *testing.T) { + litellmParams := map[string]interface{}{"model": "vertex_ai/claude-haiku-4-5", "custom_llm_provider": "vertex_ai"} + if tc.stored != nil { + litellmParams["additional_drop_params"] = tc.stored + } + if tc.storedLocation != "" { + litellmParams["vertex_location"] = tc.storedLocation + } + body, _ := json.Marshal(map[string]interface{}{ + "model_name": "vertex-claude-haiku", + "litellm_params": litellmParams, + "model_info": map[string]interface{}{"id": "model-123", "db_model": true, "base_model": "claude-haiku-4-5"}, + }) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write(body) + })) + defer srv.Close() + + res := resourceLiteLLMModel() + additionalParams := map[string]interface{}{"max_parallel_requests": "8"} + if tc.configured != "" { + additionalParams["additional_drop_params"] = tc.configured + } + config := map[string]interface{}{ + "model_name": "vertex-claude-haiku", + "custom_llm_provider": "vertex_ai", + "base_model": "claude-haiku-4-5", + "tpm": 1000, + "additional_litellm_params": additionalParams, + } + if tc.configLocation != "" { + config["vertex_location"] = tc.configLocation + } + maps.Copy(config, tc.extraConfig) + d := schema.TestResourceDataRaw(t, res.Schema, config) + d.SetId("model-123") + if err := resourceLiteLLMModelRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + + diff, err := res.Diff(context.Background(), d.State(), terraform.NewResourceConfigRaw(config), nil) + if err != nil { + t.Fatalf("diff failed: %v", err) + } + for _, key := range []string{listKey, "vertex_location", "tpm", parallelKey, "merge_reasoning_content_in_choices", "litellm_credential_name"} { + attr := (*terraform.ResourceAttrDiff)(nil) + if diff != nil { + attr = diff.Attributes[key] + } + got := attr != nil && attr.Old != attr.New + if want := slices.Contains(tc.wantChanged, key); got != want { + t.Errorf("plan shows %s change = %v, want %v", key, got, want) + } + } + }) + } +} + func TestResourceLiteLLMModelReadDisplayName(t *testing.T) { cases := map[string]struct { serverBody string diff --git a/terraform/provider/litellm/types.go b/terraform/provider/litellm/types.go index 3af4650de53..aff90cd64be 100644 --- a/terraform/provider/litellm/types.go +++ b/terraform/provider/litellm/types.go @@ -120,6 +120,7 @@ type LiteLLMParams struct { VertexProject string `json:"vertex_project,omitempty"` VertexLocation string `json:"vertex_location,omitempty"` VertexCredentials string `json:"vertex_credentials,omitempty"` + AdditionalDropParams interface{} `json:"additional_drop_params,omitempty"` } // ModelInfo represents information about a model. From dba86bec28cb483c9937470c5d950f5e03699d3b Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 02:32:24 +0000 Subject: [PATCH 7/8] fix(count_tokens): count Anthropic-format requests for Gemini the way /v1/messages inference converts them (#46767) * fix(count_tokens): convert Anthropic-format messages for the OpenAI counter the way inference does * test(count_tokens): check /v1/messages/count_tokens sends OpenAI the input /v1/messages inference sends * fix(count_tokens): build the Anthropic count request without a cast and sync dashboard API types * fix(count_tokens): count Anthropic-format requests for Gemini the way /v1/messages inference converts them * test(count_tokens): compare tools by value in the public api test * fix(count_tokens): count Anthropic bodies the way /v1/messages inference sends them to OpenAI Unwrap affinity-encoded reasoning, count the structured-output schema, flatten replayed web search results, join a list system prompt, and answer a message without a role with a 400 before calling OpenAI. messages_format is now internal to the Anthropic count route instead of a public /utils/token_counter field. * test(count_tokens): expect the image detail the adapter now sends after merging main * fix(count_tokens): fetch Gemini remote media without blocking the event loop * fix(count_tokens): type normalize_native_messages without the unimported Any On Python 3.10 to 3.13 the annotation is evaluated at import, so the proxy failed to start with NameError * chore(count_tokens): give the Gemini count helpers' list returns a reason * fix(count_tokens): normalize Anthropic messages for Gemini counts the way /v1/messages does * fix(count_tokens): count Gemini bodies after inference's history rewrites and answer converter errors with 400 * test(count_tokens): cover the Gemini count error paths and compaction with a text-block system --------- Co-authored-by: Mrinal Chanshetty --- .../pass_through/adapters/transformation.py | 2 +- litellm/llms/gemini/common_utils.py | 26 +- litellm/llms/vertex_ai/common_utils.py | 25 +- .../llms/vertex_ai/count_tokens/context.py | 114 ++++++- .../llms/vertex_ai/gemini/transformation.py | 8 +- .../test_vertex_gemini_count_tokens_wire.py | 87 +++++ .../llms/gemini/test_gemini_common_utils.py | 98 ++++-- .../vertex_ai/count_tokens/test_context.py | 315 +++++++++++++++++- .../vertex_ai/test_vertex_ai_common_utils.py | 91 ++++- 9 files changed, 703 insertions(+), 63 deletions(-) diff --git a/litellm/llms/anthropic/pass_through/adapters/transformation.py b/litellm/llms/anthropic/pass_through/adapters/transformation.py index 98cd16c6403..d1560c5b2ed 100644 --- a/litellm/llms/anthropic/pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/adapters/transformation.py @@ -465,7 +465,7 @@ class LiteLLMAnthropicMessagesAdapter: *, custom_llm_provider: str | None = None, preserve_midturn_system: bool = False, - ) -> list: + ) -> list[AllMessageValues]: new_messages: Final[list[AllMessageValues]] = [] replayable_messages: Final = strip_encrypted_reasoning_blocks_from_anthropic_messages(messages) leading_count: Final = next( diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 57ee0f33879..bac5824b981 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -504,25 +504,41 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): from litellm.llms.vertex_ai.count_tokens.context import ( VertexGeminiCountContextError, count_context_error_response, + gemini_count_contents, + inference_request_view, + provider_error_response, vertex_gemini_count_context, ) - count_context: Final = vertex_gemini_count_context(system=system, tools=tools) + request_view: Final = inference_request_view(messages or [], system, messages_format) + count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools) if isinstance(count_context, VertexGeminiCountContextError): return count_context_error_response(count_context, request_model, model_to_use) deployment = deployment or {} count_tokens_params_request: Final = copy.deepcopy(deployment.get("litellm_params", {})) + resolved_contents: Final = ( + contents + if contents is not None + else await gemini_count_contents( + request_view.messages, messages_format, model_to_use, "gemini", count_tokens_params_request + ) + ) + if isinstance(resolved_contents, VertexGeminiCountContextError): + return count_context_error_response(resolved_contents, request_model, model_to_use) count_tokens_params: Final = { "model": model_to_use, - "contents": contents, + "contents": resolved_contents, "system_instruction": count_context.system_instruction, "tools": count_context.tools, } count_tokens_params_request.update(count_tokens_params) - result: Final = await GoogleAIStudioTokenCounter().acount_tokens( - **count_tokens_params_request, - ) + try: + result: Final = await GoogleAIStudioTokenCounter().acount_tokens( + **count_tokens_params_request, + ) + except litellm.APIError as error: + return provider_error_response(error, request_model, model_to_use) if result is not None: return TokenCountResponse( diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 0e4cb856632..4c0026783f7 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1359,24 +1359,27 @@ class VertexAITokenCounter(BaseTokenCounter): from litellm.llms.vertex_ai.count_tokens.context import ( VertexGeminiCountContextError, count_context_error_response, + gemini_count_contents, + inference_request_view, + provider_error_response, vertex_gemini_count_context, ) from litellm.llms.vertex_ai.count_tokens.handler import VertexAITokenCounter - from litellm.llms.vertex_ai.gemini.transformation import ( - gemini_convert_messages_with_history, # pyright: ignore[reportPrivateUsage] # shared helper already used by gemini/chat, context_caching, and vertex_and_google_ai_studio_gemini - ) - count_context: Final = vertex_gemini_count_context(system=system, tools=tools) + request_view: Final = inference_request_view(messages or [], system, messages_format) + count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools) if isinstance(count_context, VertexGeminiCountContextError): return count_context_error_response(count_context, request_model, model_to_use) resolved_contents: Final = ( contents if contents is not None - else gemini_convert_messages_with_history( - messages=messages or [] # pyright: ignore[reportArgumentType] # preserve direct caller message handling + else await gemini_count_contents( + request_view.messages, messages_format, model_to_use, "vertex_ai", count_tokens_params_request ) ) + if isinstance(resolved_contents, VertexGeminiCountContextError): + return count_context_error_response(resolved_contents, request_model, model_to_use) count_tokens_params: Final = { "model": model_to_use, @@ -1390,15 +1393,7 @@ class VertexAITokenCounter(BaseTokenCounter): **count_tokens_params_request, ) except litellm.APIError as error: - return TokenCountResponse( - total_tokens=0, - request_model=request_model, - model_used=model_to_use, - tokenizer_type="gemini", - error=True, - error_message=error.message, - status_code=error.status_code, - ) + return provider_error_response(error, request_model, model_to_use) if "totalTokens" in result: return TokenCountResponse( diff --git a/litellm/llms/vertex_ai/count_tokens/context.py b/litellm/llms/vertex_ai/count_tokens/context.py index cb0bedf0238..3f9289c8244 100644 --- a/litellm/llms/vertex_ai/count_tokens/context.py +++ b/litellm/llms/vertex_ai/count_tokens/context.py @@ -1,12 +1,20 @@ -from collections.abc import Mapping +import copy +from collections.abc import Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import Final +from typing import TYPE_CHECKING, Final, Literal from pydantic import JsonValue, TypeAdapter, ValidationError +from typing_extensions import assert_never from litellm.types.llms.anthropic import AnthropicMessagesTextParam -from litellm.types.utils import TokenCountResponse +from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.vertex_ai import ContentType +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse + +if TYPE_CHECKING: + from litellm.exceptions import APIError + from litellm.llms.anthropic.pass_through.context_management.result import PolyfillResult @dataclass(frozen=True, slots=True) @@ -142,3 +150,103 @@ def count_context_error_response( error_message=error.error_message, status_code=400, ) + + +def provider_error_response(error: "APIError", request_model: str, model_used: str) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_used, + tokenizer_type="gemini", + error=True, + error_message=error.message, + status_code=error.status_code, + ) + + +@dataclass(frozen=True, slots=True) +class InferenceRequestView: + messages: Sequence[Mapping[str, object]] + system: object | None + + +def _client_compaction(messages: Sequence[Mapping[str, object]], system: object | None) -> "PolyfillResult | None": + from litellm.llms.anthropic.pass_through.context_management.editors.compact import ( + apply_client_compaction_block_history, + ) + + chat_messages: Final = [dict(message) for message in messages] + if system is None or isinstance(system, str): + return apply_client_compaction_block_history(messages=chat_messages, system=system) + try: + blocks: Final = _SYSTEM_BLOCKS.validate_python(system) + except ValidationError: + return None + return apply_client_compaction_block_history(messages=chat_messages, system=[dict(block) for block in blocks]) + + +def inference_request_view( + messages: Sequence[Mapping[str, object]], system: object | None, messages_format: TokenCountMessagesFormat +) -> InferenceRequestView: + """Apply the history rewrites /v1/messages runs before it converts a body for Gemini.""" + from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages + + match messages_format: + case "openai": + return InferenceRequestView(messages=messages, system=system) + case "anthropic": + normalized: Final = normalize_native_messages(messages) + compacted: Final = _client_compaction(normalized, system) + if compacted is None: + return InferenceRequestView(messages=normalized, system=system) + return InferenceRequestView(messages=compacted.messages, system=compacted.system) + case _: + assert_never(messages_format) + + +def _chat_messages( + messages: Sequence[Mapping[str, object]], + messages_format: TokenCountMessagesFormat, + model: str, + custom_llm_provider: Literal["vertex_ai", "gemini"], +) -> Sequence[AllMessageValues]: + from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter + + match messages_format: + case "anthropic": + return LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( + messages=list(messages), # pyright: ignore[reportArgumentType] # raw Anthropic body, read the way /v1/messages inference reads it + model=model, + custom_llm_provider=custom_llm_provider, + ) + case "openai": + return copy.deepcopy(list(messages)) # pyright: ignore[reportReturnType] # chat messages from the request, copied because the AI Studio conversion rewrites them in place + case _: + assert_never(messages_format) + + +async def gemini_count_contents( + messages: Sequence[Mapping[str, object]], + messages_format: TokenCountMessagesFormat, + model: str, + custom_llm_provider: Literal["vertex_ai", "gemini"], + litellm_params: Mapping[str, object], +) -> Sequence[ContentType] | VertexGeminiCountContextError: + from litellm.exceptions import BadRequestError + from litellm.litellm_core_utils.asyncify import asyncify + from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media + from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig + from litellm.llms.vertex_ai.gemini.transformation import should_inline_remote_media + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig + + config: Final = GoogleAIStudioGeminiConfig() if custom_llm_provider == "gemini" else VertexGeminiConfig() + try: + chat_messages: Final = await async_inline_remote_media( + list(_chat_messages(messages, messages_format, model, custom_llm_provider)), + should_inline=should_inline_remote_media(custom_llm_provider), + ) + return await asyncify(config.transform_messages)( + messages=chat_messages, model=model, litellm_params=dict(litellm_params) + ) + except BadRequestError as error: + return VertexGeminiCountContextError(error_message=error.message) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 57327731863..895d59965b7 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -6,7 +6,7 @@ Why separate file? Make it easy to see how transformation works import os import re -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Any, Final, Literal, cast from urllib.parse import quote @@ -1318,6 +1318,10 @@ def _vertex_inlines(media: RemoteMedia) -> bool: ) +def should_inline_remote_media(custom_llm_provider: str) -> Callable[[RemoteMedia], bool]: + return _ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines + + async def async_transform_request_body( gemini_api_key: str | None, messages: list[AllMessageValues], @@ -1360,7 +1364,7 @@ async def async_transform_request_body( ) inlined_messages: Final = await async_inline_remote_media( - messages, should_inline=_ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines + messages, should_inline=should_inline_remote_media(custom_llm_provider) ) if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages): diff --git a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py index 1a148b19f8d..7fef79845a1 100644 --- a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -167,6 +167,93 @@ def test_messages_and_count_tokens_send_matching_gemini_context( assert count_body == _count_tokens_request_from_generate_content(generated_body) +_LOOKUP_TOOL: Final[JsonValue] = { + "name": "lookup", + "description": "Look up a value", + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, +} +_CONTENT_BLOCK_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType( + { + "base64-pdf": { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the document."}, + { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}, + }, + ], + } + ], + }, + "base64-image": { + "messages": [ + { + "role": "user", + "content": [ + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}}, + {"type": "text", "text": "What is this?"}, + ], + } + ], + }, + "tool-turns": { + "tools": [_LOOKUP_TOOL], + "messages": [ + {"role": "user", "content": "Look up the value."}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Looking it up."}, + {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "value"}}, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}], + }, + ], + }, + } +) + + +def _ai_studio_deployment(scenario: Scenario, api_base: str) -> str: + return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key") + + +def _counted_contents(count_body: Mapping[str, JsonValue]) -> JsonValue: + if "contents" in count_body: + return count_body["contents"] + return _JSON_OBJECT.validate_python(count_body["generateContentRequest"])["contents"] + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +@pytest.mark.parametrize("case", sorted(_CONTENT_BLOCK_BODIES)) +def test_count_tokens_sends_the_contents_messages_inference_sends(gateway: Gateway, provider: str, case: str) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = ( + _deployment(gateway, scenario, wire.url) + if provider == "vertex_ai" + else _ai_studio_deployment(scenario, wire.url) + ) + body: Final = {"model": model, **_CONTENT_BLOCK_BODIES[case]} + generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16}) + counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body) + assert generated.status_code == 200, generated.text + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT, counted.text + requests: Final = peer.drain() + generated_body: Final = _json_body( + next(request for request in requests if request.target.endswith(":generateContent")) + ) + count_body: Final = _json_body(next(request for request in requests if request.target.endswith(":countTokens"))) + assert _counted_contents(count_body) == generated_body["contents"] + + def test_count_tokens_keeps_each_text_block_in_system_instruction(gateway: Gateway) -> None: peer: Final = _GeminiPeer() with wire_server(peer) as wire, gateway.scenario() as scenario: diff --git a/tests/unit/llms/gemini/test_gemini_common_utils.py b/tests/unit/llms/gemini/test_gemini_common_utils.py index 9a47d01ec46..e54023c6596 100644 --- a/tests/unit/llms/gemini/test_gemini_common_utils.py +++ b/tests/unit/llms/gemini/test_gemini_common_utils.py @@ -5,6 +5,7 @@ import httpx import pytest from pydantic import JsonValue, TypeAdapter +import litellm from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.common_utils import GeminiModelInfo, GoogleAIStudioTokenCounter from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler @@ -100,6 +101,58 @@ async def test_google_ai_studio_count_tokens_rejects_image_system_without_reques assert peer.request_body is None +@pytest.mark.asyncio +async def test_google_ai_studio_count_tokens_rejects_an_image_url_it_refuses_to_fetch_without_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", True) + peer: Final = _GeminiCountTokensPeer() + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + result: Final = await GoogleAIStudioTokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[ + { + "role": "user", + "content": [{"type": "image", "source": {"type": "url", "url": "http://127.0.0.1:9/image.png"}}], + } + ], + contents=None, + deployment={"litellm_params": {"api_key": "test-key", "api_base": "https://count-tokens.test"}}, + request_model="gemini/gemini-2.5-flash", + messages_format="anthropic", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "http://127.0.0.1:9/image.png" in (result.error_message or "") + assert peer.request_body is None + + +@pytest.mark.asyncio +async def test_google_ai_studio_count_tokens_answers_with_the_status_google_rejected_it_with() -> None: + rejection: Final = '{"error": {"code": 400, "message": "Invalid function name", "status": "INVALID_ARGUMENT"}}' + client: Final = AsyncHTTPHandler( + transport=httpx.MockTransport(lambda request: httpx.Response(400, text=rejection, request=request)) + ) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + result: Final = await GoogleAIStudioTokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=None, + contents=[{"role": "user", "parts": [{"text": "hello"}]}], + deployment={"litellm_params": {"api_key": "test-key", "api_base": "https://count-tokens.test"}}, + request_model="gemini/gemini-2.5-flash", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "Invalid function name" in (result.error_message or "") + + @pytest.mark.asyncio @pytest.mark.parametrize( ("system_instruction", "tools", "expected_body"), @@ -129,9 +182,7 @@ async def test_gemini_count_tokens_sends_optional_request_fields( expected_body: dict[str, JsonValue], ) -> None: peer: Final = _GeminiCountTokensPeer() - token_counter: Final = GeminiCountTokensHandler( - client=AsyncHTTPHandler(transport=httpx.MockTransport(peer)) - ) + token_counter: Final = GeminiCountTokensHandler(client=AsyncHTTPHandler(transport=httpx.MockTransport(peer))) result: Final = await token_counter.acount_tokens( contents=[{"role": "user", "parts": [{"text": "hello"}]}], @@ -177,18 +228,10 @@ class TestGeminiModelInfo: # Test edge cases where model names end with characters from "models/" # These would be incorrectly processed if using strip("models/") instead of replace("models/", "") models = [ - { - "name": "models/gemini-1.5-pro" - }, # ends with 'o' - would become "gemini-1.5-pr" with strip() - { - "name": "models/test-model" - }, # ends with 'l' - would become "gemini/test-mode" with strip() - { - "name": "models/custom-models" - }, # ends with 's' - would become "gemini/custom-model" with strip() - { - "name": "models/demo" - }, # ends with 'o' - would become "gemini/dem" with strip() + {"name": "models/gemini-1.5-pro"}, # ends with 'o' - would become "gemini-1.5-pr" with strip() + {"name": "models/test-model"}, # ends with 'l' - would become "gemini/test-mode" with strip() + {"name": "models/custom-models"}, # ends with 's' - would become "gemini/custom-model" with strip() + {"name": "models/demo"}, # ends with 'o' - would become "gemini/dem" with strip() ] result = gemini_model_info.process_model_name(models) @@ -239,16 +282,10 @@ class TestGoogleAIStudioTokenCounter: token_counter = GoogleAIStudioTokenCounter() # Test with gemini provider - should return True - assert ( - token_counter.should_use_token_counting_api(LlmProviders.GEMINI.value) - is True - ) + assert token_counter.should_use_token_counting_api(LlmProviders.GEMINI.value) is True # Test with other providers - should return False - assert ( - token_counter.should_use_token_counting_api(LlmProviders.OPENAI.value) - is False - ) + assert token_counter.should_use_token_counting_api(LlmProviders.OPENAI.value) is False assert token_counter.should_use_token_counting_api("anthropic") is False assert token_counter.should_use_token_counting_api("vertex_ai") is False @@ -320,9 +357,7 @@ class TestGoogleAIStudioTokenCounter: "functionResponse": { "id": "read_many_files-1757526647518-730a691aac11c", # This should be removed "name": "read_many_files", - "response": { - "output": "No files matching the criteria were found or all were skipped." - }, + "response": {"output": "No files matching the criteria were found or all were skipped."}, } } ], @@ -331,9 +366,7 @@ class TestGoogleAIStudioTokenCounter: ] # Clean the contents - cleaned_contents = token_counter._clean_contents_for_gemini_api( - contents_with_id - ) + cleaned_contents = token_counter._clean_contents_for_gemini_api(contents_with_id) # Verify the 'id' field was removed function_response = cleaned_contents[1]["parts"][0]["functionResponse"] @@ -342,8 +375,7 @@ class TestGoogleAIStudioTokenCounter: assert "response" in function_response assert function_response["name"] == "read_many_files" assert ( - function_response["response"]["output"] - == "No files matching the criteria were found or all were skipped." + function_response["response"]["output"] == "No files matching the criteria were found or all were skipped." ) def test_clean_contents_for_gemini_api_preserves_other_fields(self): @@ -359,9 +391,7 @@ class TestGoogleAIStudioTokenCounter: ] # Clean the contents - cleaned_contents = token_counter._clean_contents_for_gemini_api( - contents_without_function_response - ) + cleaned_contents = token_counter._clean_contents_for_gemini_api(contents_without_function_response) # Verify the contents are unchanged assert cleaned_contents == contents_without_function_response diff --git a/tests/unit/llms/vertex_ai/count_tokens/test_context.py b/tests/unit/llms/vertex_ai/count_tokens/test_context.py index 444920e1ecc..06dacedde67 100644 --- a/tests/unit/llms/vertex_ai/count_tokens/test_context.py +++ b/tests/unit/llms/vertex_ai/count_tokens/test_context.py @@ -1,8 +1,18 @@ -from typing import Final +import asyncio +import base64 +import time +import uuid +from collections.abc import Generator +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from threading import Thread +from typing import Final, Literal +from unittest.mock import patch import pytest -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter +import litellm from litellm.llms.anthropic.pass_through.adapters.transformation import ( LiteLLMAnthropicMessagesAdapter, truncate_tool_name, @@ -10,6 +20,8 @@ from litellm.llms.anthropic.pass_through.adapters.transformation import ( from litellm.llms.vertex_ai.count_tokens.context import ( VertexGeminiCountContext, VertexGeminiCountContextError, + gemini_count_contents, + inference_request_view, vertex_gemini_count_context, ) from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig @@ -193,3 +205,302 @@ def test_invalid_tools_return_an_error_value(tools: JsonValue) -> None: assert isinstance(context, VertexGeminiCountContextError) assert "tools" in context.error_message + + +_ANTHROPIC_CONVERSATION: Final = ( + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize both"}, + {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0="}}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}}, + ], + }, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "x"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}]}, +) + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_anthropic_documents_images_and_tool_turns_become_gemini_parts( + custom_llm_provider: Literal["vertex_ai", "gemini"], +) -> None: + contents: Final = await gemini_count_contents( + _ANTHROPIC_CONVERSATION, "anthropic", "gemini-2.5-flash", custom_llm_provider, {} + ) + + assert contents == [ + { + "role": "user", + "parts": [ + {"text": "Summarize both"}, + {"inline_data": {"data": "JVBERi0=", "mime_type": "application/pdf"}}, + {"inline_data": {"data": "aW1hZ2U=", "mime_type": "image/png"}}, + ], + }, + {"role": "model", "parts": [{"function_call": {"name": "lookup", "args": {"query": "x"}}}]}, + {"role": "user", "parts": [{"function_response": {"name": "lookup", "response": {"content": "42"}}}]}, + ] + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_openai_messages_keep_their_chat_conversion(custom_llm_provider: Literal["vertex_ai", "gemini"]) -> None: + contents: Final = await gemini_count_contents( + ({"role": "user", "content": "hello"}, {"role": "assistant", "content": "hi"}), + "openai", + "gemini-2.5-flash", + custom_llm_provider, + {}, + ) + + assert contents == [ + {"role": "user", "parts": [{"text": "hello"}]}, + {"role": "model", "parts": [{"text": "hi"}]}, + ] + + +_PNG_BYTES: Final = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" +) +_DOWNLOAD_SECONDS: Final = 0.5 + + +class _SlowPngHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + time.sleep(_DOWNLOAD_SECONDS) + self.send_response(200) + self.send_header("Content-Type", "image/png") + self.send_header("Content-Length", str(len(_PNG_BYTES))) + self.end_headers() + self.wfile.write(_PNG_BYTES) + + def log_message(self, format: str, *args: object) -> None: + return None + + +@contextmanager +def _slow_png_server() -> Generator[str, None, None]: + with ThreadingHTTPServer(("127.0.0.1", 0), _SlowPngHandler) as server: + thread: Final = Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}" + finally: + server.shutdown() + + +async def _event_loop_wake_delay(nap_seconds: float) -> float: + started: Final = time.monotonic() + await asyncio.sleep(nap_seconds) + return time.monotonic() - started + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_remote_image_downloads_without_blocking_the_event_loop( + custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", False) + with _slow_png_server() as base_url: + image_url: Final = f"{base_url}/{uuid.uuid4().hex}.png" + wake_delay, contents = await asyncio.gather( + _event_loop_wake_delay(0.05), + gemini_count_contents( + ( + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe it"}, + {"type": "image", "source": {"type": "url", "url": image_url}}, + ], + }, + ), + "anthropic", + "gemini-2.5-flash", + custom_llm_provider, + {}, + ), + ) + + assert contents == [ + { + "role": "user", + "parts": [ + {"text": "Describe it"}, + {"inline_data": {"data": base64.b64encode(_PNG_BYTES).decode(), "mime_type": "image/png"}}, + ], + } + ] + assert wake_delay < _DOWNLOAD_SECONDS / 2 + + +class _MissingImageHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + self.send_response(404) + self.send_header("Content-Length", "0") + self.end_headers() + + def log_message(self, format: str, *args: object) -> None: + return None + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_image_that_cannot_be_fetched_returns_an_error_value( + custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", False) + with ThreadingHTTPServer(("127.0.0.1", 0), _MissingImageHandler) as server: + Thread(target=server.serve_forever, daemon=True).start() + try: + contents = await gemini_count_contents( + ( + { + "role": "user", + "content": [ + { + "type": "image", + "source": { + "type": "url", + "url": f"http://127.0.0.1:{server.server_address[1]}/{uuid.uuid4().hex}.png", + }, + } + ], + }, + ), + "anthropic", + "gemini-2.5-flash", + custom_llm_provider, + {}, + ) + finally: + server.shutdown() + + assert isinstance(contents, VertexGeminiCountContextError) + + +@pytest.mark.parametrize( + "system", ("Be brief.", [{"type": "text", "text": "Be brief."}]), ids=("string-system", "text-block-system") +) +def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(system: JsonValue) -> None: + view: Final = inference_request_view( + ( + {"role": "user", "content": "first question " * 50}, + {"role": "assistant", "content": [{"type": "compaction", "content": "Summary of the earlier turns."}]}, + {"role": "user", "content": "latest question"}, + ), + system, + "anthropic", + ) + + assert view.messages == [{"role": "user", "content": "latest question"}] + assert "Summary of the earlier turns." in str(view.system) + assert "Be brief." in str(view.system) + + +def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> None: + snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." + view: Final = inference_request_view( + ( + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": snippet, + } + ], + }, + ], + }, + ), + None, + "anthropic", + ) + + assistant_blocks: Final = TypeAdapter(list[dict[str, str]]).validate_python(view.messages[1]["content"]) + assert [block["type"] for block in assistant_blocks] == ["text"] + assert snippet in assistant_blocks[0]["text"] + + +def test_openai_view_is_the_request_unchanged() -> None: + messages: Final = ({"role": "user", "content": [{"type": "compaction", "content": "kept"}]},) + + view: Final = inference_request_view(messages, "Be brief.", "openai") + + assert view.messages == messages + assert view.system == "Be brief." + + +async def test_gcs_media_type_lookup_uses_the_deployment_vertex_project() -> None: + with patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type", + return_value="image/png", + ) as lookup: + contents: Final = await gemini_count_contents( + ( + { + "role": "user", + "content": [ + {"type": "image", "source": {"type": "url", "url": "gs://bucket/image-without-extension"}} + ], + }, + ), + "anthropic", + "gemini-2.5-flash", + "vertex_ai", + {"vertex_project": "deployment-project"}, + ) + + assert lookup.call_args.kwargs.get("vertex_project") == "deployment-project" + assert not isinstance(contents, VertexGeminiCountContextError) + assert contents[0]["parts"][0] == { + "file_data": {"mime_type": "image/png", "file_uri": "gs://bucket/image-without-extension"} + } + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +def test_anthropic_replayed_web_search_result_is_counted_as_its_text( + custom_llm_provider: Literal["vertex_ai", "gemini"], +) -> None: + snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." + messages: Final = ( + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": snippet, + } + ], + }, + ], + }, + {"role": "user", "content": "Thanks."}, + ) + + view: Final = inference_request_view(messages, None, "anthropic") + contents: Final = asyncio.run( + gemini_count_contents(view.messages, "anthropic", "gemini-2.5-flash", custom_llm_provider, {}) + ) + + assert snippet in str(contents) diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py index 304424b2acc..8528724a24e 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -1,5 +1,6 @@ import time from collections.abc import Iterator +from contextlib import contextmanager from typing import Final from unittest.mock import AsyncMock, patch @@ -7,6 +8,7 @@ import httpx import pytest from pydantic import JsonValue, TypeAdapter +import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler @@ -24,16 +26,21 @@ from litellm.llms.vertex_ai.common_utils import ( vertex_request_labels_from_litellm_params, ) from litellm.llms.vertex_ai.count_tokens import handler as vertex_count_tokens_handler +from litellm.types.utils import TokenCountMessagesFormat @pytest.mark.asyncio -async def test_vertex_gemini_rejects_non_text_system_blocks_before_request() -> None: +@pytest.mark.parametrize("messages_format", ("openai", "anthropic")) +async def test_vertex_gemini_rejects_non_text_system_blocks_before_request( + messages_format: TokenCountMessagesFormat, +) -> None: result: Final = await VertexAITokenCounter().count_tokens( model_to_use="gemini-4-argon", messages=[{"role": "user", "content": "Count this"}], contents=None, request_model="vertex_ai/gemini-4-argon", system=[{"type": "image", "source": {"type": "base64", "data": "AAAA"}}], + messages_format=messages_format, ) assert result is not None @@ -42,6 +49,88 @@ async def test_vertex_gemini_rejects_non_text_system_blocks_before_request() -> assert "text blocks" in (result.error_message or "") +class _VertexCountTokensPeer: + def __init__(self, status_code: int) -> None: + self.status_code: Final = status_code + self.requests: tuple[httpx.Request, ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.requests = (*self.requests, request) + if self.status_code == 200: + return httpx.Response(200, json={"totalTokens": 37}, request=request) + return httpx.Response( + self.status_code, + json={"error": {"code": self.status_code, "message": "Invalid function name", "status": "INVALID_ARGUMENT"}}, + request=request, + ) + + +@contextmanager +def _vertex_count_tokens_peer(status_code: int) -> Iterator[_VertexCountTokensPeer]: + peer: Final = _VertexCountTokensPeer(status_code) + with ( + patch.object( + gemini_count_tokens_handler, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)), + ), + patch.object( + vertex_count_tokens_handler.VertexAITokenCounter, + "_ensure_access_token_async", + new=AsyncMock(return_value=("fake-token", "test-project")), + ), + ): + yield peer + + +_VERTEX_DEPLOYMENT: Final = {"litellm_params": {"vertex_project": "test-project", "vertex_location": "us-central1"}} + + +@pytest.mark.asyncio +async def test_vertex_gemini_count_answers_with_the_status_google_rejected_it_with() -> None: + with _vertex_count_tokens_peer(400) as peer: + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment=_VERTEX_DEPLOYMENT, + request_model="vertex_ai/gemini-2.5-flash", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "Invalid function name" in (result.error_message or "") + assert len(peer.requests) == 1 + + +@pytest.mark.asyncio +async def test_vertex_gemini_count_rejects_an_image_url_it_refuses_to_fetch_without_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", True) + with _vertex_count_tokens_peer(200) as peer: + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[ + { + "role": "user", + "content": [{"type": "image", "source": {"type": "url", "url": "http://127.0.0.1:9/image.png"}}], + } + ], + contents=None, + deployment=_VERTEX_DEPLOYMENT, + request_model="vertex_ai/gemini-2.5-flash", + messages_format="anthropic", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "http://127.0.0.1:9/image.png" in (result.error_message or "") + assert peer.requests == () + + @pytest.mark.parametrize("location", ["us", "eu", "us-central1", "europe-west1", "global"]) def test_validate_vertex_location_accepts_valid(location): assert validate_vertex_location(location) == location From 081cfa2a6264b9a4de8b647efefa9e72a3b296f0 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Thu, 15 Oct 2026 19:47:12 -0700 Subject: [PATCH 8/8] fix(proxy): reset end-user spend in committed pages so the budget cascade survives a large LiteLLM_EndUserTable (#46502) * fix(proxy): reset end-user spend in committed pages so the budget cascade survives large LiteLLM_EndUserTable Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy): collapse _enduser_ids_by_tier return to one line for ruff format Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): mark the end-user page IN list as bounded for the unbounded-IN check Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/common_utils/reset_budget_job.py | 110 +++++-- litellm/repositories/unit_of_work.py | 17 +- .../common_utils/test_reset_budget_job.py | 276 ++++++++++++------ tests/unit/repositories/test_unit_of_work.py | 24 +- 4 files changed, 315 insertions(+), 112 deletions(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 61330d12ac3..ccf2232e379 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -6,6 +6,7 @@ from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from enum import Enum +from itertools import chain from types import MappingProxyType from typing import Final, Generic, Literal, Protocol, TypeVar @@ -65,6 +66,7 @@ from litellm.repositories.team_repository import TeamRepository from litellm.repositories.unit_of_work import ( LinkedSpendResetWrites, budget_cascade_unit_of_work, + enduser_spend_reset_unit_of_work, spend_reset_unit_of_work, ) from litellm.repositories.verification_token_repository import ( @@ -222,8 +224,8 @@ def _budget_link_where( def _enduser_invalidation_where(budget_ids: Sequence[str]) -> dict[str, object]: """Customers whose cached spend a committed reset of these tiers invalidated. - Mirrors ``_queue_enduser_resets`` without its ``spend > 0`` filter, which - post-commit would match nobody. + ``_enduser_reset_where`` adds the ``spend > 0`` filter, which post-commit + would match nobody. """ linked: Final = _budget_link_where(budget_ids) default_budget_id: Final = litellm.max_end_user_budget_id @@ -251,26 +253,44 @@ def _queue_budget_linked_resets( writes.queue_spend_zero(where=_budget_link_where(plain_ids, extra)) -def _queue_enduser_resets(writes: LinkedSpendResetWrites, cascade: "_BudgetCascade") -> None: - """End users reset on the budget link like every other gated table, plus a - NULL-budget_id branch: rows created implicitly persist no link and ride the - default tier (litellm.max_end_user_budget_id). +def _enduser_reset_where(cascade: "_BudgetCascade") -> Mapping[str, object]: + """End users with spend to reset: linked to an expiring tier, or carrying + no link and riding the default tier (litellm.max_end_user_budget_id).""" + return {**_enduser_invalidation_where(cascade.budget_ids), **_SPENT_ROWS_WHERE} - Matching on the link rather than enumerating user ids keeps a statement's - bind count proportional to the expiring tiers instead of the customer - population, which past ~32,700 dependents exceeds PostgreSQL's per-statement - bind ceiling and wedges the cascade permanently (#40564). + +def _enduser_tier(row: _EndUserRow) -> str | None: + return row.budget_id if row.budget_id is not None else litellm.max_end_user_budget_id + + +def _enduser_ids_by_tier(rows: Sequence[_EndUserRow]) -> Mapping[str | None, tuple[str, ...]]: + tiers: Final = tuple(dict.fromkeys(_enduser_tier(row) for row in rows)) + return MappingProxyType({tier: tuple(row.user_id for row in rows if _enduser_tier(row) == tier) for tier in tiers}) + + +def _queue_enduser_page_resets( + writes: LinkedSpendResetWrites, rows: Sequence[_EndUserRow], cascade: "_BudgetCascade" +) -> None: + """The statements of ``_queue_budget_linked_resets`` for one page of end + users, keyed by their user ids so each statement touches at most one page. + + A statement matching the budget link reads the whole end-user table (no + index on budget_id) and updates every linked row at once, which past a few + million end users outlives the prisma client's 30s request timeout (#46404). + A page of ids stays under PostgreSQL's per-statement bind ceiling (#40564) + and the spend predicates still classify each row at write time. """ - _queue_budget_linked_resets(writes, cascade, extra=_SPENT_ROWS_WHERE) - default_budget_id: Final = litellm.max_end_user_budget_id - if default_budget_id is None or default_budget_id not in cascade.budget_ids: - return - cap: Final = cascade.rollover_caps.get(default_budget_id) + for tier, user_ids in _enduser_ids_by_tier(rows).items(): + _queue_enduser_tier_resets(writes, user_ids, cap=cascade.rollover_caps.get(tier) if tier is not None else None) + + +def _queue_enduser_tier_resets(writes: LinkedSpendResetWrites, user_ids: Sequence[str], cap: float | None) -> None: + page: Final = {"user_id": {"in": list(user_ids)}} # bounded-ok: one page of RESET_BUDGET_JOB_BATCH_SIZE ids if cap is None: - writes.queue_spend_zero(where={"budget_id": None, **_SPENT_ROWS_WHERE}) + writes.queue_spend_zero(where={**page, **_SPENT_ROWS_WHERE}) return - writes.queue_spend_zero(where={"budget_id": None, "spend": {"gt": 0, "lte": cap}}) - writes.queue_spend_decrement(where={"budget_id": None, "spend": {"gt": cap}}, amount=cap) + writes.queue_spend_zero(where={**page, "spend": {"gt": 0, "lte": cap}}) + writes.queue_spend_decrement(where={**page, "spend": {"gt": cap}}, amount=cap) @dataclass(frozen=True, slots=True) @@ -824,6 +844,49 @@ class ResetBudgetJob: ), ) + async def _reset_enduser_spend(self, cascade: _BudgetCascade) -> int: + """Reset the end users the expiring tiers gate, one committed page at a time. + + Runs before ``_commit_budget_cascade`` so a failure here leaves the + tiers due: the next tick resumes with the pages still carrying spend. + Each page invalidates its own caches right after its commit, so no + zeroed row sits behind a cached spend that still rejects requests. + Capless like ``_invalidate_enduser_caches``, for its reason. + """ + if not cascade.budget_ids: + return 0 + where: Final = _enduser_reset_where(cascade) + walk = _EndUserWalk() # rebind-ok: keyset cursor advances one committed page at a time + while walk.cursor is not None: + walk = await self._reset_enduser_page(cascade, where=where, cursor=walk.cursor, reached=walk.invalidated) + return walk.invalidated + + async def _reset_enduser_page( + self, cascade: _BudgetCascade, where: Mapping[str, object], cursor: str, reached: int + ) -> _EndUserWalk: + rows: Final = await self._fetch_enduser_page(where=where, cursor=cursor) + if not rows: + return _EndUserWalk(cursor=None, invalidated=reached) + await self._with_db_write_retry( + lambda: self._commit_enduser_page(rows, cascade), + reason="reset_budget_write_endusers_failure", + ) + await self._invalidate_caches( + counter_keys=tuple(_enduser_counter_key(row) for row in rows), + cache_keys=tuple(chain.from_iterable(_enduser_cache_keys(row) for row in rows)), + ) + walked: Final = reached + len(rows) + if len(rows) < RESET_BUDGET_JOB_BATCH_SIZE: + return _EndUserWalk(cursor=None, invalidated=walked) + return _EndUserWalk(cursor=rows[-1].user_id, invalidated=walked) + + async def _commit_enduser_page(self, rows: Sequence[_EndUserRow], cascade: _BudgetCascade) -> None: + async with ( + db_span("reset_budget_endusers", "LiteLLM_EndUserTable"), + enduser_spend_reset_unit_of_work(self._new_batch) as writes, + ): + _queue_enduser_page_resets(writes, rows, cascade) + async def _commit_budget_cascade(self, cascade: _BudgetCascade) -> None: """Zero the gated spend and advance ``budget_reset_at`` in one transaction. @@ -851,7 +914,6 @@ class ResetBudgetJob: _queue_budget_linked_resets(uow.tags, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.model_access_groups, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.projects, cascade, extra=_SPENT_ROWS_WHERE) - _queue_enduser_resets(uow.endusers, cascade) for budget_id, budget_reset_at in cascade.budget_resets: uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at) @@ -878,6 +940,7 @@ class ResetBudgetJob: return _BudgetCascadeFailed(cascade=_EMPTY_CASCADE, error=e) try: + await self._reset_enduser_spend(cascade) await self._commit_budget_cascade(cascade) except Exception as e: return _BudgetCascadeFailed(cascade=cascade, error=e) @@ -896,9 +959,10 @@ class ResetBudgetJob: """ Resets the spend a budget tier gates (end users, team members, keys, orgs, tags, model access groups) and advances the tier's - budget_reset_at, atomically. + budget_reset_at. End users commit in pages first, everything else + lands with the window advance in one transaction. - Caches are invalidated only after the transaction commits, so a failed + Caches are invalidated only after the commit they follow, so a failed run cannot leave a zeroed counter in front of an un-reset DB row. """ await _run_phase_in_chunks(self._reset_budget_for_litellm_budget_table_chunk) @@ -929,8 +993,8 @@ class ResetBudgetJob: case _BudgetCascadeFailed(cascade=cascade, error=error): verbose_proxy_logger.exception( "Failed to reset the budget table cascade (team member, enduser, org, tag and model access " - "group spend, plus budget_reset_at); nothing was committed and the budgets stay due for the " - "next run: %s", + "group spend, plus budget_reset_at); the budgets stay due for the next run, which resumes " + "with the end-user pages not committed yet: %s", error, exc_info=error, ) diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index 10f092c1cb2..95f4f5e81be 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -97,11 +97,13 @@ class SpendResetUnitOfWork: @dataclass(frozen=True, slots=True) class BudgetCascadeUnitOfWork: - """Every write a budget-tier reset performs, bound to one batch. + """Every write a budget-tier reset performs in its final batch. The dependent spend rows and the budget rows' ``budget_reset_at`` advance must land together: advancing the window without zeroing the spend it gates leaves the dependents pinned at their cap until the next window. + End users are the exception, paged and committed ahead of this batch by + ``enduser_spend_reset_unit_of_work`` because one tier can gate millions. """ team_memberships: LinkedSpendResetWrites @@ -110,7 +112,6 @@ class BudgetCascadeUnitOfWork: tags: LinkedSpendResetWrites model_access_groups: LinkedSpendResetWrites projects: LinkedSpendResetWrites - endusers: LinkedSpendResetWrites budgets: BudgetWindowWrites @@ -137,7 +138,17 @@ async def budget_cascade_unit_of_work( tags=LinkedSpendResetWrites(table=batch.litellm_tagtable), model_access_groups=LinkedSpendResetWrites(table=batch.litellm_modelaccessgroupbudgettable), projects=LinkedSpendResetWrites(table=batch.litellm_projecttable), - endusers=LinkedSpendResetWrites(table=batch.litellm_endusertable), budgets=BudgetWindowWrites(table=batch.litellm_budgettable), ) await batch.commit() + + +@asynccontextmanager +async def enduser_spend_reset_unit_of_work( + new_batch: Callable[[], PrismaBatch], +) -> AsyncGenerator[LinkedSpendResetWrites, None]: + """One page of end-user spend resets, committed on its own so a statement + never grows with the end-user population.""" + batch: Final = new_batch() + yield LinkedSpendResetWrites(table=batch.litellm_endusertable) + await batch.commit() diff --git a/tests/unit/proxy/common_utils/test_reset_budget_job.py b/tests/unit/proxy/common_utils/test_reset_budget_job.py index c17aed49014..70f7774e691 100644 --- a/tests/unit/proxy/common_utils/test_reset_budget_job.py +++ b/tests/unit/proxy/common_utils/test_reset_budget_job.py @@ -435,7 +435,7 @@ def test_reset_budget_for_team(reset_budget_job, mock_prisma_client): def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client): - """End-user spend is zeroed and the tier's window advances, in one batch.""" + """End-user spend is zeroed in its own committed page, then the tier's window advances.""" now = datetime.now(timezone.utc) test_budget = _budget_row(budget_id="test-budget-1", budget_duration="1d", budget_reset_at=now) @@ -446,11 +446,12 @@ def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client): "spend": 20.0, "litellm_budget_table": test_budget, "user_id": "test-enduser-1", + "budget_id": "test-budget-1", }, ) mock_prisma_client.data["budget"] = [test_budget] - mock_prisma_client.data["enduser"] = [test_enduser] + mock_prisma_client.db.litellm_endusertable.set_find_many_results([test_enduser]) asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) @@ -458,10 +459,13 @@ def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client): { "table": "enduser", "op": "update_many", - "where": {"budget_id": {"in": ["test-budget-1"]}, "spend": {"gt": 0}}, + "where": {"user_id": {"in": ["test-enduser-1"]}, "spend": {"gt": 0}}, "data": {"spend": 0}, } ] + assert [batcher.committed for batcher in mock_prisma_client.db.batchers] == [True, True] + assert [call["table"] for call in mock_prisma_client.db.batchers[0].calls] == ["enduser"] + assert "enduser" not in {call["table"] for call in mock_prisma_client.db.batchers[1].calls} budget_writes = _batch_writes(mock_prisma_client, "budget") assert len(budget_writes) == 1 @@ -470,6 +474,112 @@ def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client): assert set(budget_writes[0]["data"].keys()) == {"budget_reset_at"} +def test_enduser_reset_commits_in_pages_before_the_window_advances(reset_budget_job, mock_prisma_client): + """Regression for #46404: one statement over every end user a tier gates + reads the whole table and outlives the prisma client's 30s request timeout + once the table holds millions of rows, so the batch rolls back and + budget_reset_at never moves. The reset has to land one committed page at a + time, keyed by user id, with the window advance after the last page.""" + mock_prisma_client.data["budget"] = [_budget_row(budget_id="shared-tier", budget_duration="1d")] + population: Final = RESET_BUDGET_JOB_BATCH_SIZE * 2 + 3 + mock_prisma_client.db.litellm_endusertable.set_find_many_results( + [ + types.SimpleNamespace(spend=4.0, user_id=f"cust-{index:06d}", budget_id="shared-tier") + for index in range(population) + ] + ) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + batchers: Final = mock_prisma_client.db.batchers + assert [batcher.committed for batcher in batchers] == [True, True, True, True] + pages: Final = batchers[:3] + assert [[call["table"] for call in page.calls] for page in pages] == [["enduser"]] * 3 + assert [len(page.calls[0]["where"]["user_id"]["in"]) for page in pages] == [ + RESET_BUDGET_JOB_BATCH_SIZE, + RESET_BUDGET_JOB_BATCH_SIZE, + 3, + ] + assert [page.calls[0]["where"]["spend"] for page in pages] == [{"gt": 0}] * 3 + assert [page.calls[0]["data"] for page in pages] == [{"spend": 0}] * 3 + assert pages[0].calls[0]["where"]["user_id"]["in"][0] == "cust-000000" + assert pages[2].calls[0]["where"]["user_id"]["in"] == ["cust-001000", "cust-001001", "cust-001002"] + cascade: Final = batchers[3] + assert "enduser" not in {call["table"] for call in cascade.calls} + assert [call["where"] for call in cascade.calls if call["table"] == "budget"] == [{"budget_id": "shared-tier"}] + + +def test_enduser_page_failure_leaves_the_tier_due(reset_budget_job, mock_prisma_client): + """A page that fails after earlier pages committed must not advance the + window: the committed pages are at spend 0 and drop out of the next tick's + ``spend > 0`` read, so the next tick resumes where this one stopped.""" + mock_prisma_client.data["budget"] = [_budget_row(budget_id="shared-tier", budget_duration="1d")] + endusers: Final = mock_prisma_client.db.litellm_endusertable + endusers.set_find_many_results( + [ + types.SimpleNamespace(spend=4.0, user_id=f"cust-{index:06d}", budget_id="shared-tier") + for index in range(RESET_BUDGET_JOB_BATCH_SIZE + 3) + ] + ) + endusers.set_find_many_error(1, RuntimeError("connection reset while paging end users")) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + assert _batch_writes(mock_prisma_client, "budget") == [] + enduser_writes: Final = _batch_writes(mock_prisma_client, "enduser") + assert len(enduser_writes) == 1 + assert len(enduser_writes[0]["where"]["user_id"]["in"]) == RESET_BUDGET_JOB_BATCH_SIZE + + +def test_enduser_page_keeps_each_tier_s_rollover_statements_to_its_own_users( + rollover_enabled, reset_budget_job, mock_prisma_client, monkeypatch +): + """One page can hold end users of several expiring tiers. Each tier's + zero-then-decrement pair must name only that tier's user ids and its own + cap, or a 15.0 spend on a 20.0 tier would be carried as 5.0.""" + _make_counter_invalidation_job(monkeypatch) + mock_prisma_client.data["budget"] = [ + _budget_row(budget_id="tier-10", budget_duration="1d", max_budget=10.0), + _budget_row(budget_id="tier-20", budget_duration="1d", max_budget=20.0), + ] + mock_prisma_client.db.litellm_endusertable.set_find_many_results( + [ + types.SimpleNamespace(spend=15.0, user_id="cust-a", budget_id="tier-10"), + types.SimpleNamespace(spend=15.0, user_id="cust-b", budget_id="tier-20"), + types.SimpleNamespace(spend=25.0, user_id="cust-c", budget_id="tier-10"), + ] + ) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + assert _batch_writes(mock_prisma_client, "enduser") == [ + { + "table": "enduser", + "op": "update_many", + "where": {"user_id": {"in": ["cust-a", "cust-c"]}, "spend": {"gt": 0, "lte": 10.0}}, + "data": {"spend": 0}, + }, + { + "table": "enduser", + "op": "update_many", + "where": {"user_id": {"in": ["cust-a", "cust-c"]}, "spend": {"gt": 10.0}}, + "data": {"spend": {"decrement": 10.0}}, + }, + { + "table": "enduser", + "op": "update_many", + "where": {"user_id": {"in": ["cust-b"]}, "spend": {"gt": 0, "lte": 20.0}}, + "data": {"spend": 0}, + }, + { + "table": "enduser", + "op": "update_many", + "where": {"user_id": {"in": ["cust-b"]}, "spend": {"gt": 20.0}}, + "data": {"spend": {"decrement": 20.0}}, + }, + ] + + def test_reset_budget_all(reset_budget_job, mock_prisma_client): # Setup test data with timezone-aware datetime now = datetime.now(timezone.utc) @@ -529,6 +639,7 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): "spend": 20.0, "litellm_budget_table": test_budget, "user_id": "test-enduser-1", + "budget_id": "test-budget-1", }, ) @@ -536,7 +647,7 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): mock_prisma_client.data["user"] = [test_user] mock_prisma_client.data["team"] = [test_team] mock_prisma_client.data["budget"] = [test_budget] - mock_prisma_client.data["enduser"] = [test_enduser] + mock_prisma_client.db.litellm_endusertable.set_find_many_results([test_enduser]) # Run the test asyncio.run(reset_budget_job.reset_budget()) @@ -559,7 +670,7 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): { "table": "enduser", "op": "update_many", - "where": {"budget_id": {"in": ["test-budget-1"]}, "spend": {"gt": 0}}, + "where": {"user_id": {"in": ["test-enduser-1"]}, "spend": {"gt": 0}}, "data": {"spend": 0}, } ] @@ -579,7 +690,6 @@ _LINKED_TABLE_CASES = [ ("org", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), ("tag", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), ("model_access_group", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), - ("enduser", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), ] @@ -620,34 +730,33 @@ def _bind_count(where: Dict[str, Any]) -> int: @pytest.mark.parametrize("population", [3, 40_000], ids=["small", "over-pg-bind-ceiling"]) def test_enduser_reset_bind_count_does_not_scale_with_population(reset_budget_job, mock_prisma_client, population): - """Regression for #40564. + """Regression for #40564 and #46404. - Enumerating every dependent user id put one bind variable per customer into - a single prepared statement. Past PostgreSQL's ceiling the statement could - not be parsed at all, so the whole atomic cascade rolled back, - budget_reset_at never advanced, and the tier stayed due on every later tick - forever. Matching on the budget link keeps the statement the same size no - matter how many customers share a tier. + Enumerating every dependent user id in one statement put one bind variable + per end user into a single prepared statement, which past PostgreSQL's + ceiling could not be parsed at all. Matching on the budget link instead made + one statement read and update the whole end-user table, which past a few + million rows outlives the prisma client's 30s request timeout. Paging by + user id keeps every statement at one page of binds however many end users + share a tier, and the statements have to cover the whole population. """ budget = _budget_row(budget_id="shared-tier", budget_duration="1d") mock_prisma_client.data["budget"] = [budget] - mock_prisma_client.data["enduser"] = [ - types.SimpleNamespace( - spend=1.0, - litellm_budget_table=budget, - user_id=f"cust-{index:08d}", - budget_id="shared-tier", - ) - for index in range(population) - ] + mock_prisma_client.db.litellm_endusertable.set_find_many_results( + [ + types.SimpleNamespace(spend=1.0, user_id=f"cust-{index:08d}", budget_id="shared-tier") + for index in range(population) + ] + ) asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) writes = _batch_writes(mock_prisma_client, "enduser") - assert [_bind_count(write["where"]) for write in writes] == [2], ( - f"the cascade must not enumerate {population} user ids: past " - f"{_POSTGRES_MAX_BIND_VARIABLES} binds PostgreSQL refuses the statement, got {writes[:1]}" - ) + assert max(_bind_count(write["where"]) for write in writes) <= RESET_BUDGET_JOB_BATCH_SIZE + 1 + assert RESET_BUDGET_JOB_BATCH_SIZE + 1 < _POSTGRES_MAX_BIND_VARIABLES + assert sorted(user_id for write in writes for user_id in write["where"]["user_id"]["in"]) == [ + f"cust-{index:08d}" for index in range(population) + ] assert _batch_writes(mock_prisma_client, "budget")[0]["data"]["budget_reset_at"] is not None @@ -778,6 +887,7 @@ def test_reset_budget_resets_endusers_with_null_budget_id(reset_budget_job, mock "spend": 30.0, "litellm_budget_table": test_budget, "user_id": "enduser-explicit", + "budget_id": default_budget_id, }, ) @@ -811,40 +921,34 @@ def test_reset_budget_resets_endusers_with_null_budget_id(reset_budget_job, mock ) mock_prisma_client.data["budget"] = [test_budget] - mock_prisma_client.data["enduser"] = [enduser_with_budget] - - # Set up the DB mock for NULL-budget-id end users - mock_prisma_client.db.litellm_endusertable.set_find_many_results([enduser_no_budget_row]) + mock_prisma_client.db.litellm_endusertable.set_find_many_results([enduser_with_budget, enduser_no_budget_row]) asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - # Both end users are zeroed: the linked rows on the tier's budget_id, the - # implicit ones on the NULL branch that stands in for the default tier. + # Both end users are zeroed in one page: the linked row on the tier's + # budget_id and the implicit one whose NULL budget_id stands in for the + # default tier. assert _batch_writes(mock_prisma_client, "enduser") == [ { "table": "enduser", "op": "update_many", - "where": {"budget_id": {"in": [default_budget_id]}, "spend": {"gt": 0}}, - "data": {"spend": 0}, - }, - { - "table": "enduser", - "op": "update_many", - "where": {"budget_id": None, "spend": {"gt": 0}}, + "where": {"user_id": {"in": ["enduser-explicit", "enduser-implicit"]}, "spend": {"gt": 0}}, "data": {"spend": 0}, }, ] - # The post-commit invalidation walk covers both branches, so implicitly - # created customers on the default tier get their cached spend dropped too, - # and it is paged rather than reading the whole customer population. + # The paged reset and the post-commit invalidation walk both cover both + # branches, so implicitly created customers on the default tier are reset + # and get their cached spend dropped too. find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls - assert len(find_many_calls) == 1 - assert find_many_calls[0]["where"]["OR"] == [ - {"budget_id": {"in": [default_budget_id]}}, - {"budget_id": None}, - ] - assert find_many_calls[0]["take"] == RESET_BUDGET_JOB_BATCH_SIZE + assert len(find_many_calls) == 2 + assert find_many_calls[0]["where"]["spend"] == {"gt": 0} + for call in find_many_calls: + assert call["where"]["OR"] == [ + {"budget_id": {"in": [default_budget_id]}}, + {"budget_id": None}, + ] + assert call["take"] == RESET_BUDGET_JOB_BATCH_SIZE litellm.max_end_user_budget_id = None @@ -879,7 +983,8 @@ def test_reset_budget_skips_null_budget_id_endusers_when_default_not_configured( # ride a default tier that is not expiring, so their spend stays put. find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls assert [call["where"] for call in find_many_calls] == [ - {"budget_id": {"in": ["some-budget"]}, "user_id": {"gt": ""}} + {"budget_id": {"in": ["some-budget"]}, "spend": {"gt": 0}, "user_id": {"gt": ""}}, + {"budget_id": {"in": ["some-budget"]}, "user_id": {"gt": ""}}, ] litellm.max_end_user_budget_id = None @@ -919,7 +1024,8 @@ def test_reset_budget_skips_null_budget_id_endusers_when_default_not_in_reset_li # ride a default tier that is not expiring, so their spend stays put. find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls assert [call["where"] for call in find_many_calls] == [ - {"budget_id": {"in": ["other-budget"]}, "user_id": {"gt": ""}} + {"budget_id": {"in": ["other-budget"]}, "spend": {"gt": 0}, "user_id": {"gt": ""}}, + {"budget_id": {"in": ["other-budget"]}, "user_id": {"gt": ""}}, ] litellm.max_end_user_budget_id = None @@ -1671,7 +1777,8 @@ def test_enduser_invalidation_is_paged_and_batched(reset_budget_job, mock_prisma on a large enough deployment rather than a slow tick. Awaiting one cache call per customer makes the last customer wait out every customer ahead of it. Both regress silently, so pin the page size, the strictly advancing cursor, - and one batched call per page. + and one batched call per page. The paged reset ahead of the commit + invalidates its pages too, so the cache calls count both walks. """ counter_cache: Final = _make_counter_invalidation_job(monkeypatch) mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-1")] @@ -1685,7 +1792,9 @@ def test_enduser_invalidation_is_paged_and_batched(reset_budget_job, mock_prisma asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - reads: Final = mock_prisma_client.db.litellm_endusertable.find_many_calls + all_reads: Final = mock_prisma_client.db.litellm_endusertable.find_many_calls + reads: Final = [read for read in all_reads if "spend" not in read["where"]] + assert len(all_reads) == 6 assert [read["take"] for read in reads] == [RESET_BUDGET_JOB_BATCH_SIZE] * 3 assert [read["where"]["user_id"]["gt"] for read in reads] == [ "", @@ -1693,8 +1802,8 @@ def test_enduser_invalidation_is_paged_and_batched(reset_budget_job, mock_prisma f"cust-{RESET_BUDGET_JOB_BATCH_SIZE * 2 - 1:06d}", ] - assert counter_cache.async_delete_cache_keys.await_count == 3 - assert counter_cache.user_api_key_cache.async_delete_cache_keys.await_count == 3 + assert counter_cache.async_delete_cache_keys.await_count == 6 + assert counter_cache.user_api_key_cache.async_delete_cache_keys.await_count == 6 counter_cache.async_delete_cache.assert_not_called() invalidated: Final = { @@ -1727,7 +1836,7 @@ def test_enduser_invalidation_reports_a_page_read_failure_instead_of_a_clean_fin for i in range(RESET_BUDGET_JOB_BATCH_SIZE + 3) ] ) - endusers.set_find_many_error(1, RuntimeError("connection reset while paging customers")) + endusers.set_find_many_error(3, RuntimeError("connection reset while paging customers")) logging_obj: Final = RecordingProxyLogging() job: Final = ResetBudgetJob(proxy_logging_obj=logging_obj, prisma_client=mock_prisma_client) @@ -1985,19 +2094,22 @@ def test_budget_reset_at_is_not_advanced_when_the_cascade_fails(db_factory, monk def test_budget_cascade_writes_land_in_a_single_transaction(reset_budget_job, mock_prisma_client, monkeypatch): - """Dependent spend and the budget_reset_at advance ride one batch.""" + """Dependent spend and the budget_reset_at advance ride one batch, after + the end-user pages have committed on their own.""" _make_counter_invalidation_job(monkeypatch) now = datetime.now(timezone.utc) budget = _budget_row(budget_id="budget-1", budget_duration="7d") mock_prisma_client.data["budget"] = [budget] - mock_prisma_client.data["enduser"] = [ - type("EndUser", (), {"spend": 5.0, "litellm_budget_table": budget, "user_id": "enduser-1", "budget_id": "budget-1"}) - ] + mock_prisma_client.db.litellm_endusertable.set_find_many_results( + [type("EndUser", (), {"spend": 5.0, "user_id": "enduser-1", "budget_id": "budget-1"})] + ) asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - assert len(mock_prisma_client.db.batchers) == 1, "the cascade must not be split across transactions" - batcher = mock_prisma_client.db.batchers[0] + assert len(mock_prisma_client.db.batchers) == 2, "one end-user page, then the cascade" + enduser_page, batcher = mock_prisma_client.db.batchers + assert enduser_page.committed is True + assert {(call["table"], call["op"]) for call in enduser_page.calls} == {("enduser", "update_many")} assert batcher.committed is True assert {(call["table"], call["op"]) for call in batcher.calls} == { ("team_membership", "update_many"), @@ -2006,7 +2118,6 @@ def test_budget_cascade_writes_land_in_a_single_transaction(reset_budget_job, mo ("tag", "update_many"), ("model_access_group", "update_many"), ("project", "update_many"), - ("enduser", "update_many"), ("budget", "update_many"), } budget_write = next(call for call in batcher.calls if call["table"] == "budget") @@ -3288,13 +3399,9 @@ def test_budget_cascade_carries_enduser_overage_when_rollover_enabled( _make_counter_invalidation_job(monkeypatch) budget = _budget_row(budget_id="budget-roll", budget_duration="1d", max_budget=10.0) mock_prisma_client.data["budget"] = [budget] - mock_prisma_client.data["enduser"] = [ - type( - "EndUser", - (), - {"spend": 15.0, "litellm_budget_table": budget, "user_id": "enduser-roll", "budget_id": "budget-roll"}, - ) - ] + mock_prisma_client.db.litellm_endusertable.set_find_many_results( + [type("EndUser", (), {"spend": 15.0, "user_id": "enduser-roll", "budget_id": "budget-roll"})] + ) asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) @@ -3302,13 +3409,13 @@ def test_budget_cascade_carries_enduser_overage_when_rollover_enabled( assert { "table": "enduser", "op": "update_many", - "where": {"budget_id": "budget-roll", "spend": {"gt": 10.0}}, + "where": {"user_id": {"in": ["enduser-roll"]}, "spend": {"gt": 10.0}}, "data": {"spend": {"decrement": 10.0}}, } in enduser_writes assert { "table": "enduser", "op": "update_many", - "where": {"budget_id": "budget-roll", "spend": {"gt": 0, "lte": 10.0}}, + "where": {"user_id": {"in": ["enduser-roll"]}, "spend": {"gt": 0, "lte": 10.0}}, "data": {"spend": 0}, } in enduser_writes @@ -3377,13 +3484,9 @@ def test_cascade_rollover_writes_survive_sequential_execution( {"user_id": "member-1", "team_id": "team-1", "spend": 15.0, "budget_id": "budget-roll"}, ) mock_prisma_client.db.litellm_teammembership.set_find_many_results([membership]) - mock_prisma_client.data["enduser"] = [ - type( - "EndUser", - (), - {"spend": 15.0, "litellm_budget_table": budget, "user_id": "enduser-roll", "budget_id": "budget-roll"}, - ) - ] + mock_prisma_client.db.litellm_endusertable.set_find_many_results( + [type("EndUser", (), {"spend": 15.0, "user_id": "enduser-roll", "budget_id": "budget-roll"})] + ) asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) @@ -3702,8 +3805,9 @@ def _wire_cascade_reads_for_test(prisma_client, endusers=()): tables an awaitable find_many so the reads resolve instead of falling into the job's warn-and-continue path. - End users are read by the post-commit invalidation walk rather than by - ``get_data``, so callers that care about customers pass them here. + End users are read by the paged reset and the post-commit invalidation + walk rather than by ``get_data``, so callers that care about them pass + them here. """ for table in ( "litellm_teammembership", @@ -4010,6 +4114,7 @@ async def test_reset_budget_endusers_are_zeroed_with_the_budget_window_advance(f prisma_client.get_data.side_effect = get_data_mock prisma_client.update_data = AsyncMock() batch_calls: Final = _wire_batcher_for_test(prisma_client) + _wire_cascade_reads_for_test(prisma_client, endusers=endusers) proxy_logging_obj: Final = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() @@ -4020,11 +4125,14 @@ async def test_reset_budget_endusers_are_zeroed_with_the_budget_window_advance(f await _await_tasks_spawned_by(job.reset_budget_for_litellm_budget_table()) - assert prisma_client.db.batch_.call_count == 1, "the cascade must be one transaction" + assert prisma_client.db.batch_.call_count == 2, "one end-user page, then the cascade in one transaction" enduser_writes: Final = [c for c in batch_calls if c["table"] == "enduser"] assert len(enduser_writes) == 1 - assert enduser_writes[0]["where"] == {"budget_id": {"in": ["budget1"]}, "spend": {"gt": 0}} + assert enduser_writes[0]["where"] == { + "user_id": {"in": ["user1", "user2", "user3", "user4", "user5", "user6"]}, + "spend": {"gt": 0}, + } assert enduser_writes[0]["data"] == {"spend": 0} budget_writes: Final = [c for c in batch_calls if c["table"] == "budget"] @@ -4098,7 +4206,7 @@ async def test_reset_budget_continues_other_categories_on_failure(frozen_reset_c assert len([c for c in batch_calls if c["table"] == "team_membership"]) == 1 enduser_writes: Final = [c for c in batch_calls if c["table"] == "enduser"] assert len(enduser_writes) == 1 - assert enduser_writes[0]["where"] == {"budget_id": {"in": ["budget1"]}, "spend": {"gt": 0}} + assert enduser_writes[0]["where"] == {"user_id": {"in": ["user1"]}, "spend": {"gt": 0}} assert enduser_writes[0]["data"] == {"spend": 0} key_writes: Final = [c for c in batch_calls if c["table"] == "key" and c["op"] == "update"] @@ -4576,7 +4684,7 @@ async def test_service_logger_endusers_success(): enduser_writes = [c for c in batch_calls if c["table"] == "enduser"] assert len(enduser_writes) == 1 - assert enduser_writes[0]["where"] == {"budget_id": {"in": ["budget1"]}, "spend": {"gt": 0}} + assert enduser_writes[0]["where"] == {"user_id": {"in": ["user1", "user2"]}, "spend": {"gt": 0}} proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_called_once() ( diff --git a/tests/unit/repositories/test_unit_of_work.py b/tests/unit/repositories/test_unit_of_work.py index 9eff248b917..a08d44b9b4d 100644 --- a/tests/unit/repositories/test_unit_of_work.py +++ b/tests/unit/repositories/test_unit_of_work.py @@ -7,6 +7,7 @@ import pytest from litellm.repositories.unit_of_work import ( LinkedSpendResetWrites, budget_cascade_unit_of_work, + enduser_spend_reset_unit_of_work, spend_reset_unit_of_work, ) @@ -96,7 +97,6 @@ async def test_budget_cascade_dependents_and_window_advance_share_one_batch(): uow.tags.queue_spend_zero(where=linked) uow.model_access_groups.queue_spend_zero(where=linked) uow.projects.queue_spend_zero(where=linked) - uow.endusers.queue_spend_zero(where={"user_id": {"in": ["enduser-1"]}}) uow.budgets.queue_window_advance(budget_id="budget-1", budget_reset_at=reset_at) assert batch.commit_count == 0 @@ -108,11 +108,31 @@ async def test_budget_cascade_dependents_and_window_advance_share_one_batch(): ("litellm_tagtable.update_many", linked, {"spend": 0}), ("litellm_modelaccessgroupbudgettable.update_many", linked, {"spend": 0}), ("litellm_projecttable.update_many", linked, {"spend": 0}), - ("litellm_endusertable.update_many", {"user_id": {"in": ["enduser-1"]}}, {"spend": 0}), ("litellm_budgettable.update_many", {"budget_id": "budget-1"}, {"budget_reset_at": reset_at}), ] +async def test_enduser_page_commits_on_its_own_batch(): + """One tier can gate millions of end users, so a page of them lands in its + own transaction instead of riding the cascade's.""" + batch = FakeBatch() + + async with enduser_spend_reset_unit_of_work(lambda: batch) as writes: + writes.queue_spend_zero(where={"user_id": {"in": ["enduser-1", "enduser-2"]}, "spend": {"gt": 0}}) + writes.queue_spend_decrement(where={"user_id": {"in": ["enduser-3"]}, "spend": {"gt": 5.0}}, amount=5.0) + assert batch.commit_count == 0 + + assert batch.commit_count == 1 + assert batch.calls == [ + ( + "litellm_endusertable.update_many", + {"user_id": {"in": ["enduser-1", "enduser-2"]}, "spend": {"gt": 0}}, + {"spend": 0}, + ), + ("litellm_endusertable.update_many", {"user_id": {"in": ["enduser-3"]}, "spend": {"gt": 5.0}}, {"spend": {"decrement": 5.0}}), + ] + + async def test_budget_window_advance_tolerates_a_tier_deleted_mid_chunk(): """A tier deleted between the read and the commit must not abort the batch: ``update`` raises P2025 on a missing row and takes every other write in the