scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-19 04:54:37 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 70

Merge origin/main, keeping main's per-turn WebSocket logging from #46768

Browse source
This commit is contained in:
mateo-berri 2026-10-15 19:50:23 -07:00
commit ee55d9c2b6
80 changed files with 3374 additions and 2680 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/cost_calculator.py
View file

@ -3004,7 +3004,7 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor):
)
_RESPONSES_WS_BILLABLE_EVENT_TYPES: Final = frozenset({"response.completed", "response.incomplete", "response.failed"})
_RESPONSES_WS_BILLABLE_EVENT_TYPES: Final = frozenset({"response.completed", "response.incomplete"})
class _ResponsesWsEventResponse(LiteLLMBaseModel):
@ -3063,20 +3063,17 @@ class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor):
@staticmethod
def response_id_from_responses_ws_results(results: Sequence[Mapping[str, object]]) -> str | None:
response_ids: Final = tuple(
event.response.id
for event in (_ResponsesWsEvent.model_validate(result) for result in results)
if event.response is not None and event.response.id
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)
raw_events: Final = cast(OpenAIRealtimeStreamList, list(results)) # cast-ok: raw WS frames, unvalidated field
return LiteLLMRealtimeStreamLoggingObject(
id=ResponsesWebSocketTokenUsageProcessor.response_id_from_responses_ws_results(results),
usage=ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(results),
results=raw_events,
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,
)

9
litellm/litellm_core_utils/internal_call_metadata.py
View file

@ -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}

43
litellm/litellm_core_utils/litellm_logging.py
View file

@ -540,34 +540,11 @@ _PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata"
_LITELLM_PARAMS: Final = TypeAdapter(dict[str, object])
def _proxy_server_request_with_body(
litellm_params: Mapping[str, object], request_body: Mapping[str, object]
) -> Mapping[str, object]:
proxy_server_request: Final = litellm_params.get("proxy_server_request")
return (
{"proxy_server_request": {**_LITELLM_PARAMS.validate_python(proxy_server_request), "body": dict(request_body)}}
if isinstance(proxy_server_request, Mapping)
else {}
)
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
def _logged_response_id(
response: Mapping[str, object], litellm_call_id: object, falls_back_on_empty_id: bool
) -> str | None:
call_id: Final = litellm_call_id if isinstance(litellm_call_id, str) else None
if falls_back_on_empty_id:
return _provider_response_id(response) or call_id
if "id" not in response:
return call_id
response_id: Final = response["id"]
return None if response_id is None else str(response_id)
def mask_api_base_credentials(api_base: str) -> str:
if "key=" not in api_base:
return api_base
@ -798,18 +775,19 @@ 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, request_body: Mapping[str, object]
) -> "Logging":
connection_params: Final = _LITELLM_PARAMS.validate_python(self.litellm_params)
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: {**value}
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)
},
**_proxy_server_request_with_body(connection_params, request_body),
"litellm_call_id": litellm_call_id,
}
new_call: Final = copy.copy(self)
@ -6771,12 +6749,7 @@ def get_standard_logging_object_payload(
else raw_usage_dict
)
litellm_call_id: Final[object] = kwargs.get("litellm_call_id")
id = _logged_response_id(
response_obj,
litellm_call_id,
falls_back_on_empty_id=isinstance(init_response_obj, LiteLLMRealtimeStreamLoggingObject),
)
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", "")

4
litellm/llms/anthropic/count_tokens/token_counter.py
View file

@ -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.

2
litellm/llms/anthropic/pass_through/adapters/transformation.py
View file

@ -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(

15
litellm/llms/anthropic/pass_through/messages/utils.py
View file

@ -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(

12
litellm/llms/anthropic/pass_through/responses_adapters/transformation.py
View file

@ -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")

4
litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py
View file

@ -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.

4
litellm/llms/base_llm/base_utils.py
View file

@ -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

4
litellm/llms/bedrock/count_tokens/bedrock_token_counter.py
View file

@ -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

14
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -157,10 +157,7 @@ from litellm.types.llms.openai import (
from litellm.types.realtime import RealtimeQueryParams
from litellm.types.rerank import RerankResponse
from litellm.types.responses.main import DeleteResponseResult
from litellm.types.responses.streaming_websocket import (
ResponsesWebSocketRequestDefaults,
ResponsesWebSocketTurnFailureHook,
)
from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import (
CallTypes,
@ -6809,9 +6806,8 @@ class BaseLLMHTTPHandler:
custom_llm_provider: str | None = None,
first_message: str | None = None,
request_defaults: ResponsesWebSocketRequestDefaults | None = None,
on_turn_failure: ResponsesWebSocketTurnFailureHook | None = None,
**kwargs: object,
) -> None:
) -> Exception | None:
"""
Handles Responses API WebSocket mode.
@ -6845,7 +6841,7 @@ class BaseLLMHTTPHandler:
**kwargs,
)
await handler.run()
return
return None
import websockets
from websockets.asyncio.client import ClientConnection
@ -6966,9 +6962,8 @@ class BaseLLMHTTPHandler:
authorized_model=model,
custom_llm_provider=custom_llm_provider,
request_defaults=request_defaults,
on_turn_failure=on_turn_failure,
)
await streaming.bidirectional_forward()
return await streaming.bidirectional_forward()
except websockets.exceptions.InvalidStatus as e:
verbose_logger.exception("Error connecting to responses WS backend: %s", e)
@ -6982,6 +6977,7 @@ class BaseLLMHTTPHandler:
pass
else:
raise Exception(f"Unexpected error while closing WebSocket: {close_error}")
return None
def image_edit_handler(
self,

32
litellm/llms/gemini/common_utils.py
View file

@ -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
@ -502,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(

6
litellm/llms/openai/responses/count_tokens/handler.py
View file

@ -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)

78
litellm/llms/openai/responses/count_tokens/token_counter.py
View file

@ -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:

88
litellm/llms/openai/responses/count_tokens/transformation.py
View file

@ -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))

2
litellm/llms/openai_like/providers.json
View file

@ -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": {

31
litellm/llms/vertex_ai/common_utils.py
View file

@ -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
@ -1357,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,
@ -1388,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(

114
litellm/llms/vertex_ai/count_tokens/context.py
View file

@ -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)

8
litellm/llms/vertex_ai/gemini/transformation.py
View file

@ -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):

17
litellm/proxy/anthropic_endpoints/endpoints.py
View file

@ -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": <number>}
"""
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):

110
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -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,
)

110
litellm/proxy/hooks/parallel_request_limiter.py
View file

@ -1,8 +1,6 @@
import asyncio
import sys
from collections.abc import Callable
from contextvars import ContextVar
from dataclasses import dataclass
from collections.abc import Callable, Mapping
from datetime import datetime, timedelta
from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn
@ -50,15 +48,12 @@ def _response_total_tokens(response_obj: object) -> int:
return response_usage.total_tokens if isinstance(response_usage, Usage) else 0
@dataclass(frozen=True, slots=True)
class _WebSocketConnectionSlot:
request_count_key: str | None
holds_global_slot: bool
_CONNECTION_SCOPED_CALL_TYPES: Final = frozenset({CallTypes.aresponses_websocket.value})
_websocket_connection_slot: Final[ContextVar[_WebSocketConnectionSlot | None]] = ContextVar(
"litellm_legacy_rate_limiter_websocket_connection_slot", default=None
)
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):
@ -76,7 +71,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
self.internal_usage_cache = internal_usage_cache
self._clock: Final = clock
def print_verbose(self, print_statement: object) -> None:
def print_verbose(self, print_statement) -> None:
try:
verbose_proxy_logger.debug(print_statement)
if litellm.set_verbose:
@ -512,41 +507,6 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
) # don't block execution for cache updates
)
if call_type == CallTypes.aresponses_websocket.value:
_websocket_connection_slot.set(
_WebSocketConnectionSlot(
request_count_key=(f"{api_key}::{precise_minute}::request_count" if api_key is not None else None),
holds_global_slot=global_max_parallel_requests is not None,
)
)
@with_service_target("rate_limits")
async def async_release_websocket_connection_slot(self) -> None:
connection: Final = _websocket_connection_slot.get()
if connection is None:
return
_websocket_connection_slot.set(None)
if connection.holds_global_slot:
await self.internal_usage_cache.async_increment_cache(
key="global_max_parallel_requests",
value=-1,
local_only=True,
litellm_parent_otel_span=None,
)
if connection.request_count_key is None:
return
current: Final = await self.internal_usage_cache.async_get_cache(
key=connection.request_count_key,
litellm_parent_otel_span=None,
)
if current is None:
return
await self.internal_usage_cache.async_set_cache(
connection.request_count_key,
{**current, "current_requests": max(current["current_requests"] - 1, 0)},
ttl=60,
litellm_parent_otel_span=None,
)
@with_service_target("rate_limits")
async def async_log_success_event(self, kwargs, response_obj: object, start_time, end_time) -> None:
@ -557,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
@ -568,7 +529,6 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
"user_api_key_model_max_budget", None
)
user_api_key_end_user_id: Final = kwargs.get("user")
requests_released: Final = 0 if _websocket_connection_slot.get() is not None else 1
user_api_key_metadata: Final = kwargs["litellm_params"]["metadata"].get("user_api_key_metadata", {}) or {}
user_api_key_team_metadata = kwargs["litellm_params"]["metadata"].get("user_api_key_team_metadata", None)
@ -589,7 +549,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
# decrement
await self.internal_usage_cache.async_increment_cache(
key=_key,
value=-requests_released,
value=-released,
local_only=True,
litellm_parent_otel_span=litellm_parent_otel_span,
)
@ -611,13 +571,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
key=request_count_api_key,
litellm_parent_otel_span=litellm_parent_otel_span,
) or {
"current_requests": requests_released,
"current_requests": 1,
"current_tpm": 0,
"current_rpm": 0,
}
new_val = {
"current_requests": max(current["current_requests"] - requests_released, 0),
"current_requests": max(current["current_requests"] - released, 0),
"current_tpm": current["current_tpm"] + total_tokens,
"current_rpm": current["current_rpm"],
}
@ -652,13 +612,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
key=request_count_api_key,
litellm_parent_otel_span=litellm_parent_otel_span,
) or {
"current_requests": requests_released,
"current_requests": 1,
"current_tpm": 0,
"current_rpm": 0,
}
new_val = {
"current_requests": max(current["current_requests"] - requests_released, 0),
"current_requests": max(current["current_requests"] - released, 0),
"current_tpm": current["current_tpm"] + total_tokens,
"current_rpm": current["current_rpm"],
}
@ -678,13 +638,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
key=request_count_api_key,
litellm_parent_otel_span=litellm_parent_otel_span,
) or {
"current_requests": requests_released,
"current_requests": 1,
"current_tpm": total_tokens,
"current_rpm": 1,
}
new_val = {
"current_requests": max(current["current_requests"] - requests_released, 0),
"current_requests": max(current["current_requests"] - released, 0),
"current_tpm": current["current_tpm"] + total_tokens,
"current_rpm": current["current_rpm"],
}
@ -704,13 +664,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
key=request_count_api_key,
litellm_parent_otel_span=litellm_parent_otel_span,
) or {
"current_requests": requests_released,
"current_requests": 1,
"current_tpm": total_tokens,
"current_rpm": 1,
}
new_val = {
"current_requests": max(current["current_requests"] - requests_released, 0),
"current_requests": max(current["current_requests"] - released, 0),
"current_tpm": current["current_tpm"] + total_tokens,
"current_rpm": current["current_rpm"],
}
@ -730,13 +690,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
key=request_count_api_key,
litellm_parent_otel_span=litellm_parent_otel_span,
) or {
"current_requests": requests_released,
"current_requests": 1,
"current_tpm": total_tokens,
"current_rpm": 1,
}
new_val = {
"current_requests": max(current["current_requests"] - requests_released, 0),
"current_requests": max(current["current_requests"] - released, 0),
"current_tpm": current["current_tpm"] + total_tokens,
"current_rpm": current["current_rpm"],
}
@ -758,10 +718,11 @@ 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}]")
if user_api_key is None or _websocket_connection_slot.get() is not None:
if user_api_key is None:
return
## decrement call count if call failed
@ -785,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,
)
@ -807,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"],
}
@ -823,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,

130
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -10,7 +10,7 @@ import itertools
import logging
import os
import uuid
from collections.abc import AsyncGenerator, Awaitable, Callable, Iterable, Mapping, Sequence, Set
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence, Set
from contextlib import asynccontextmanager
from contextvars import ContextVar
from dataclasses import dataclass, field
@ -42,7 +42,6 @@ from litellm.caching.redis_batch import (
)
from litellm.caching.redis_cache import log_redis_failure
from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.cost_calculator import ResponsesWebSocketTokenUsageProcessor
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_str_from_messages,
@ -695,7 +694,6 @@ class RequestRateLimiterStash:
reserved_tokens: int = 0
reserved_model: RateLimitedModel | None = None
reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset)
reserved_window_identities: frozenset[ReservationWindowIdentity] = field(default_factory=frozenset)
ptu_ceiling: PTUTeamCeiling | None = None
ptu_reserved_tokens: int = 0
ptu_reserved_window_identities: frozenset[ReservationWindowIdentity] = field(default_factory=frozenset)
@ -712,7 +710,6 @@ class RequestRateLimiterStash:
batch_enqueued_reservation: BatchEnqueuedTokenReservation | None = None
batch_tpd_refund_ops: tuple[ReservationAwareIncrementOperation, ...] = ()
reservation_released: bool = False
websocket_connection: bool = False
tpm_limited_tags: frozenset[str] = field(default_factory=frozenset)
@ -781,21 +778,10 @@ def get_or_create_request_stash() -> RequestRateLimiterStash:
return stash
def _tpm_limited_tags(call_stash: RequestRateLimiterStash | None, call_tags: Iterable[object]) -> frozenset[str]:
if call_stash is not None:
return call_stash.tpm_limited_tags
connection: Final = get_request_stash()
if connection is None or not connection.websocket_connection:
return frozenset()
return connection.tpm_limited_tags.intersection(call_tags)
def claim_request_stash_for_data(data: dict) -> RequestRateLimiterStash:
from litellm.proxy.common_request_processing import request_litellm_call_id
stash: Final = get_or_create_request_stash()
owner_call_id: Final = request_litellm_call_id(_REQUEST_RATE_LIMIT_DATA.validate_python(data))
if owner_call_id is not None:
owner_call_id: Final = data.get("litellm_call_id")
if isinstance(owner_call_id, str):
stash.owner_litellm_call_id = owner_call_id
return stash
@ -831,21 +817,6 @@ def _call_id_from_callback_kwargs(kwargs: object) -> str | None:
return call_id if isinstance(call_id, str) else None
_RESPONSES_WS_TURN_EVENTS: Final = TypeAdapter(tuple[Mapping[str, object], ...])
_REQUEST_TAGS: Final = TypeAdapter(tuple[object, ...])
def _unmodeled_response_usage(kwargs: Mapping[str, object], response_obj: object) -> Usage | None:
combined_usage: Final = kwargs.get("combined_usage_object")
if isinstance(combined_usage, Usage):
return combined_usage
if kwargs.get("call_type") != CallTypes.aresponses_websocket.value or not isinstance(response_obj, list):
return None
return ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(
_RESPONSES_WS_TURN_EVENTS.validate_python(response_obj)
)
def _parse_output_cap_value(raw_value: object) -> int | None:
if isinstance(raw_value, bool) or not isinstance(raw_value, (int, float, str)):
return None
@ -4192,7 +4163,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
verbose_proxy_logger.debug("Inside Rate Limit Pre-Call Hook")
stash: Final = claim_request_stash_for_data(data)
stash.websocket_connection = call_type == CallTypes.aresponses_websocket.value
#########################################################
# Check if the call type has a specific rate limiter
@ -4398,7 +4368,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if d["key"] not in (PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY)
and (d.get("rate_limit") or {}).get("tokens_per_unit") is not None
)
stash.reserved_window_identities = frozenset(tpm_response.get("reservation_windows", frozenset()))
tpm_reservation_scopes = tuple( # rebind-ok: record successful reservation scopes
stash.reserved_scopes
)
@ -4977,33 +4946,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
return tuple((*itpm_ops, *otpm_ops))
async def _charge_websocket_turn_usage(
self,
kwargs: Mapping[str, object],
response_obj: object,
parent_otel_span: Span | None,
) -> None:
"""
A failed turn still consumed the usage its failure event reports; count it toward TPM the way a
success does. Project ITPM/OTPM already charged the turn's ``response.create`` frame estimate.
"""
connection: Final = get_request_stash()
if connection is None or not connection.websocket_connection:
return
pipeline_operations: Final = self._build_success_event_pipeline_operations(
kwargs=dict(kwargs),
response_obj=response_obj,
rate_limit_type=self.get_rate_limit_type(),
)
if pipeline_operations:
await self.async_increment_tokens_with_ttl_preservation(
pipeline_operations=pipeline_operations,
parent_otel_span=parent_otel_span,
)
def _collect_tpm_scope_targets(
self,
standard_logging_metadata: Mapping[str, Any],
standard_logging_metadata: dict[str, Any],
kwargs: object,
model_group: str | None,
tpm_limited_tags: Set[str] = frozenset(),
@ -5189,8 +5134,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# Get total tokens from response. Responses LiteLLM does not model
# (e.g. pass-through, whose usage is reported by the upstream rather
# than parsed out of the body) carry their usage in
# ``combined_usage_object`` instead, and a native Responses WebSocket
# turn in its raw event list; either would otherwise never charge
# ``combined_usage_object`` instead, and would otherwise never charge
# the TPM window.
_usage: Usage | dict[str, object] | None = None
if isinstance(
@ -5204,7 +5148,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
):
_usage = getattr(response_obj, "usage", None)
else:
_usage = _unmodeled_response_usage(kwargs, response_obj)
_combined_usage: Final = kwargs.get("combined_usage_object")
if isinstance(_combined_usage, Usage):
_usage = _combined_usage
total_tokens = self._get_total_tokens_from_usage(usage=_usage, rate_limit_type=rate_limit_type)
if total_tokens == 0:
total_tokens = self._aggregate_only_total_tokens(usage=_usage)
@ -5236,9 +5182,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
targets: Final = self._collect_tpm_scope_targets(
standard_logging_metadata=standard_logging_metadata,
kwargs=kwargs,
tpm_limited_tags=_tpm_limited_tags(
stash, _REQUEST_TAGS.validate_python(standard_logging_object.get("request_tags") or ())
),
tpm_limited_tags=stash.tpm_limited_tags if stash is not None else frozenset(),
model_group=reconcile_model.group if reconcile_model is not None else None,
)
targets.extend(
@ -5398,8 +5342,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
pipeline_operations=io_token_operations,
parent_otel_span=litellm_parent_otel_span,
)
if stash is not None:
stash.reservation_released = True
except Exception as e:
verbose_proxy_logger.exception("Error in rate limit success event: %s", e)
@ -5503,9 +5445,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
@with_service_target("rate_limits")
async def async_log_failure_event(
self, kwargs: object, response_obj: object, start_time: object, end_time: object
) -> None:
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None:
"""
On failure: decrement max_parallel_requests and refund the upfront
TPM reservation only against the scopes the reservation actually
@ -5519,15 +5459,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
try:
failure_kwargs: Final = dict(_REQUEST_RATE_LIMIT_DATA.validate_python(kwargs))
litellm_parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(failure_kwargs)
litellm_parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(kwargs)
pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = []
stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(failure_kwargs))
stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs))
await self._release_stashed_parallel_slot(stash, litellm_parent_otel_span, in_logging_callback=True)
if stash is None:
await self._charge_websocket_turn_usage(failure_kwargs, response_obj, litellm_parent_otel_span)
# Skip the reservation refund if async_post_call_failure_hook
# already released it (proxy-level rejection that also bubbles up
@ -5538,7 +5475,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if stash is None or stash.reservation_released
else (stash.reserved_tokens, stash.itpm_reserved_tokens, stash.otpm_reserved_tokens)
)
tpm_actual, itpm_actual, otpm_actual = self._recovered_partial_usage_tokens(failure_kwargs)
tpm_actual, itpm_actual, otpm_actual = self._recovered_partial_usage_tokens(kwargs)
if stash is not None and reserved_tokens > 0:
verbose_proxy_logger.debug(
@ -5643,41 +5580,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
TTL prunes it. The stashed acquisition's presence (not the key
object's current max_parallel_requests configuration, which can
change mid-request) decides whether there is anything to release.
A native Responses WebSocket connection logs its turns under their own
call ids, so no callback of its own ever settles the TPM, ITPM or OTPM
reservation it made at pre-call; its disconnect refunds that estimate.
Every other request keeps its reservation for the success or failure
event that settles it, even when the client went away mid-stream.
"""
stash: Final = get_request_stash()
await self._release_stashed_parallel_slot(stash, None)
if stash is not None and stash.websocket_connection:
await self._refund_unsettled_reservation(stash)
async def _refund_unsettled_reservation(self, stash: RequestRateLimiterStash) -> None:
if stash.reservation_released:
return
stash.reservation_released = True
await self._refund_reserved_tokens(
scopes=tuple(scope for scope in stash.reserved_scopes if scope[0] != PTU_TEAM_DESCRIPTOR_KEY),
amount=stash.reserved_tokens,
reservation_windows=stash.reserved_window_identities,
)
await self._refund_reserved_tokens(
scopes=tuple(scope for scope in stash.reserved_scopes if scope[0] == PTU_TEAM_DESCRIPTOR_KEY),
amount=stash.ptu_reserved_tokens,
reservation_windows=stash.ptu_reserved_window_identities,
)
await self._refund_reserved_tokens(
scopes=tuple(stash.itpm_reserved_scopes),
amount=stash.itpm_reserved_tokens,
reservation_windows=stash.itpm_reserved_window_identities,
)
await self._refund_reserved_tokens(
scopes=tuple(stash.otpm_reserved_scopes),
amount=stash.otpm_reserved_tokens,
reservation_windows=stash.otpm_reserved_window_identities,
)
await self._release_stashed_parallel_slot(get_request_stash(), None)
@with_service_target("rate_limits")
async def async_post_call_success_hook(self, data: dict, user_api_key_dict: UserAPIKeyAuth, response) -> None:
@ -5757,13 +5661,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
both this hook and async_log_failure_event end up running in the same
flow, only the first release/refund applies. A mid-stream failure
relayed here with recovered partial usage settles the reservation at
that usage instead of refunding it. A native Responses WebSocket
connection keeps its slot and reservation through a failed turn, and
its disconnect releases them.
that usage instead of refunding it.
"""
try:
stash: Final = get_request_stash()
if stash is None or stash.websocket_connection:
if stash is None:
return
await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)

53
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -4,8 +4,6 @@ from collections.abc import Callable, Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Final, Protocol, cast
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.batches.batch_utils import batch_cost_is_final
@ -139,12 +137,6 @@ class _ProxyDBLogger(CustomLogger):
if self.spend_event_producer is None or not is_offloadable_success(response_obj):
await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time)
return
_claim_budget_reservation(
budget_reservation=_get_budget_reservation_from_metadata(
metadata=get_litellm_metadata_from_kwargs(kwargs=dict(kwargs))
),
litellm_call_id=_string_call_id(_CALLBACK_DATA.validate_python(kwargs)),
)
event: Final = build_spend_event(
kwargs,
response_obj,
@ -174,18 +166,14 @@ class _ProxyDBLogger(CustomLogger):
user_api_key_dict: UserAPIKeyAuth,
traceback_str: str | None = None,
) -> None:
budget_reservation: Final = _claim_budget_reservation(
budget_reservation=user_api_key_dict.budget_reservation,
litellm_call_id=_string_call_id(_CALLBACK_DATA.validate_python(request_data)),
)
try:
await _release_budget_reservation(budget_reservation=budget_reservation)
await _release_budget_reservation(budget_reservation=user_api_key_dict.budget_reservation)
except Exception:
verbose_proxy_logger.exception("Failed to release budget reservation during failure handling")
try:
await _invalidate_budget_reservation_counters(budget_reservation=budget_reservation)
if budget_reservation is not None:
budget_reservation["finalized"] = True
await _invalidate_budget_reservation_counters(budget_reservation=user_api_key_dict.budget_reservation)
if user_api_key_dict.budget_reservation is not None:
user_api_key_dict.budget_reservation["finalized"] = True
except Exception:
verbose_proxy_logger.exception(
"Failed to invalidate budget reservation counters after failure release failed"
@ -357,10 +345,7 @@ class _ProxyDBLogger(CustomLogger):
resolve_missing_key_identity=str(kwargs.get("call_type")) not in _CAPTURED_IDENTITY_CALL_TYPES,
)
_write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata)
budget_reservation: Final = _claim_budget_reservation(
budget_reservation=_get_budget_reservation_from_metadata(metadata=metadata),
litellm_call_id=_string_call_id(_CALLBACK_DATA.validate_python(kwargs)),
)
budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata)
if (
isinstance(completion_response, LiteLLMBatch)
and kwargs.get("call_type") == CallTypes.aretrieve_batch.value
@ -486,7 +471,11 @@ class _ProxyDBLogger(CustomLogger):
await _release_budget_reservation(budget_reservation=budget_reservation)
# Non-model call types (health checks, afile_delete) have no model or standard_logging_object.
# Use .get() for "stream" to avoid KeyError on health checks.
if sl_object is None and (not kwargs.get("model") or kwargs.get("call_type") == "_arealtime"):
# WS session wrappers (_aresponses_websocket, _arealtime) also reach here with
# result=None; their per-turn costs are tracked on the inner aresponses/realtime calls.
if sl_object is None and (
not kwargs.get("model") or kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime")
):
verbose_proxy_logger.warning(
"Cost tracking - skipping, no standard_logging_object for call_type=%s",
kwargs.get("call_type", "unknown"),
@ -632,7 +621,7 @@ class _ProxyDBLogger(CustomLogger):
ProxyDBLogger: Final = _ProxyDBLogger
def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: Mapping[str, object]) -> None:
def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: dict) -> None:
patch = {k: v for k, v in metadata.items() if (k.startswith("user_api_key") or k == "tags") and v is not None}
if not patch:
return
@ -708,29 +697,13 @@ def _metadata_keys(metadata: object) -> tuple[str, ...]:
return tuple(sorted(str(key) for key in metadata))
def _get_budget_reservation_from_metadata(metadata: Mapping[str, object]) -> dict[str, object] | None:
def _get_budget_reservation_from_metadata(metadata: dict) -> dict | None:
return budget_reservation_from_metadata(metadata)
_CALLBACK_DATA: Final = TypeAdapter(Mapping[str, object])
def _string_call_id(callback_data: Mapping[str, object]) -> str | None:
call_id: Final = callback_data.get("litellm_call_id")
return call_id if isinstance(call_id, str) else None
def _claim_budget_reservation(
budget_reservation: dict[str, object] | None, litellm_call_id: str | None
) -> dict[str, object] | None:
from litellm.proxy.spend_tracking.budget_reservation import claim_budget_reservation
return claim_budget_reservation(budget_reservation=budget_reservation, litellm_call_id=litellm_call_id)
def _get_request_tags_for_cost_tracking(
sl_object: StandardLoggingPayload | None,
metadata: Mapping[str, object],
metadata: dict,
) -> list[str] | None:
if sl_object is not None:
request_tags: Final = sl_object.get("request_tags", None)

14
litellm/proxy/management_endpoints/sso/saml_sso.py
View file

@ -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.",

16
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -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)

33
litellm/proxy/proxy_server.py
View file

@ -173,6 +173,7 @@ from litellm.types.utils import (
ModelResponseStream,
StreamingChoices,
TextCompletionResponse,
TokenCountMessagesFormat,
TokenCountResponse,
echoed_cost_map_pricing_fields,
is_server_derived_pricing_key,
@ -713,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,
)
@ -2829,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
@ -5190,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,
@ -5205,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,
@ -14415,6 +14423,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 +14438,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 +14482,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 +14549,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

52
litellm/proxy/response_api_endpoints/endpoints.py
View file

@ -57,7 +57,6 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio
request_tags_from_request_data,
)
from litellm.proxy.route_llm_request import raise_if_required_body_param_missing
from litellm.proxy.spend_tracking.budget_reservation import share_budget_reservation_across_calls
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.openai import (
REASONING_EFFORT,
@ -65,7 +64,6 @@ from litellm.types.llms.openai import (
ResponsesAPIResponse,
)
from litellm.types.responses.main import DeleteResponseResult
from litellm.types.responses.streaming_websocket import ResponsesWebSocketTurnFailureHook
from litellm.types.utils import TokenCountResponse
if TYPE_CHECKING:
@ -1732,42 +1730,6 @@ def _routing_hints_from_first_ws_frame(first_message: str) -> Mapping[str, objec
return MappingProxyType({key: value for key, value in hints.items() if value is not None})
_RESPONSES_WS_FIRST_FRAME_KEYS: Final = frozenset({"first_message", "input", "previous_response_id"})
_RESPONSES_WS_PER_TURN_DICT_KEYS: Final = frozenset({"metadata", "litellm_metadata"})
def _responses_ws_turn_request_data(
connection_request_data: Mapping[str, object], turn_logging_obj: "LiteLLMLoggingObj"
) -> dict[str, object]:
turn_proxy_server_request: Final = turn_logging_obj.litellm_params.get("proxy_server_request")
return {
**{key: value for key, value in connection_request_data.items() if key not in _RESPONSES_WS_FIRST_FRAME_KEYS},
**{
key: {**value}
for key, value in connection_request_data.items()
if key in _RESPONSES_WS_PER_TURN_DICT_KEYS and isinstance(value, dict)
},
**({"proxy_server_request": turn_proxy_server_request} if isinstance(turn_proxy_server_request, dict) else {}),
"litellm_call_id": turn_logging_obj.litellm_call_id,
"litellm_logging_obj": turn_logging_obj,
}
def _responses_ws_turn_failure_hook(
proxy_logging_obj: "ProxyLogging",
user_api_key_dict: UserAPIKeyAuth,
connection_request_data: Mapping[str, object],
) -> ResponsesWebSocketTurnFailureHook:
async def book_turn_failure(turn_logging_obj: "LiteLLMLoggingObj", exception: Exception) -> None:
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=exception,
request_data=_responses_ws_turn_request_data(connection_request_data, turn_logging_obj),
)
return book_turn_failure
def _responses_ws_failure_frame(failure: Exception) -> str:
raw_status: Final = getattr(failure, "status_code", None)
status: Final = raw_status if isinstance(raw_status, int) and not isinstance(raw_status, bool) else 500
@ -1915,12 +1877,8 @@ async def _responses_websocket_session(
await websocket.close(code=1008, reason="Pre-call error")
return
share_budget_reservation_across_calls(user_api_key_dict.budget_reservation)
routed_data: Final = dict(
data,
user_api_key_dict=user_api_key_dict,
responses_ws_turn_failure_hook=_responses_ws_turn_failure_hook(proxy_logging_obj, user_api_key_dict, data),
**_routing_hints_from_first_ws_frame(first_message),
data, user_api_key_dict=user_api_key_dict, **_routing_hints_from_first_ws_frame(first_message)
)
# Phase 2: route to upstream provider
try:
@ -1930,7 +1888,13 @@ async def _responses_websocket_session(
llm_router=llm_router,
user_model=user_model,
)
await llm_call
failure: Final = await llm_call
if isinstance(failure, Exception):
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=failure,
request_data=routed_data,
)
except Exception as e:
verbose_proxy_logger.exception("Responses WebSocket error")
with contextlib.suppress(Exception):

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

@ -378,30 +378,6 @@ async def reconcile_budget_reservation(
return pending
_SETTLING_CALL_ID: Final = "settling_call_id"
def share_budget_reservation_across_calls(budget_reservation: dict[str, object] | None) -> None:
if budget_reservation is not None:
budget_reservation.setdefault(_SETTLING_CALL_ID, None)
def claim_budget_reservation(
budget_reservation: dict[str, object] | None, litellm_call_id: str | None
) -> dict[str, object] | None:
"""Hand a shared reservation to the first call that settles it.
A Responses WebSocket connection shares its one reservation with every turn it logs, each under its own
call id. Two turns settling it at once would both move the counter from the estimate to their own cost,
so the first claimant keeps the reservation and every later call sees none and charges its cost directly.
A reservation no connection shared belongs to its one call and is returned unchanged."""
if budget_reservation is None or litellm_call_id is None or _SETTLING_CALL_ID not in budget_reservation:
return budget_reservation
if budget_reservation[_SETTLING_CALL_ID] is None:
budget_reservation[_SETTLING_CALL_ID] = litellm_call_id
return budget_reservation if budget_reservation[_SETTLING_CALL_ID] == litellm_call_id else None
def stamp_budget_reservation_actual_cost(budget_reservation: dict | None, actual_cost: float | None) -> None:
"""Record that every reserved counter now holds ``actual_cost``, once the adjustments handed back by
``reconcile_budget_reservation(apply_consistent=False)`` have been written."""

7
litellm/proxy/utils.py
View file

@ -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,10 +4405,7 @@ class ProxyLogging:
double-decrement under the limiter's in-memory fallback.
"""
limiter: Final = self.get_proxy_hook("parallel_request_limiter")
if isinstance(limiter, PROXY_MaxParallelRequestsHandler):
await limiter.async_release_websocket_connection_slot()
return
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)

17
litellm/repositories/unit_of_work.py
View file

@ -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()

13
litellm/responses/main.py
View file

@ -62,10 +62,7 @@ from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.litellm_core_utils.get_llm_provider_logic import is_openai_compatible_endpoint
from litellm.llms.openai.data_residency import infer_openai_data_residency
from litellm.types.responses.main import *
from litellm.types.responses.streaming_websocket import (
ResponsesWebSocketRequestDefaults,
ResponsesWebSocketTurnFailureHook,
)
from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import all_litellm_params
from litellm.utils import (
@ -2478,13 +2475,12 @@ async def _aresponses_websocket(
api_base: str | None = None,
api_key: str | None = None,
timeout: float | None = None,
responses_ws_turn_failure_hook: ResponsesWebSocketTurnFailureHook | None = None,
**kwargs,
) -> None:
) -> Exception | None:
"""
Private function to handle the Responses API WebSocket mode.
For PROXY use only. ``responses_ws_turn_failure_hook`` books each failed turn of a native socket.
For PROXY use only. Returns the provider failure that ended the connection, if any.
Resolves the LLM provider from ``model``, looks up the matching
``BaseResponsesAPIConfig``, and hands off to
@ -2554,7 +2550,7 @@ async def _aresponses_websocket(
deployment_kwargs: Final = {k: v for k, v in kwargs.items() if k not in _RESPONSES_WS_ROUTING_HINT_KEYS}
first_message: Final = kwargs.get("first_message")
await base_llm_http_handler.async_responses_websocket(
return await base_llm_http_handler.async_responses_websocket(
model=resolved_model,
websocket=websocket,
logging_obj=litellm_logging_obj,
@ -2571,6 +2567,5 @@ async def _aresponses_websocket(
litellm_metadata=_build_litellm_metadata_for_ws(kwargs),
custom_llm_provider=_custom_llm_provider,
request_defaults=_build_responses_websocket_request_defaults(deployment_kwargs),
on_turn_failure=responses_ws_turn_failure_hook,
**remaining_kwargs,
)

166
litellm/responses/streaming_iterator.py
View file

@ -28,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 (
@ -58,7 +59,6 @@ if TYPE_CHECKING:
ResponsesBackendWebSocket,
ResponsesClientWebSocket,
ResponsesWebSocketRequestDefaults,
ResponsesWebSocketTurnFailureHook,
)
from litellm.types.router import LiteLLM_Params
@ -1905,11 +1905,10 @@ _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
_RESPONSES_WS_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"})
def _json_object_or_none(message: object) -> Mapping[str, object] | None:
if _is_json_object(message):
@ -1939,38 +1938,47 @@ def _user_texts_in_input_item(item: object) -> tuple[str, ...]:
)
def _logged_user_messages(client_frame: object) -> tuple[dict[str, object], ...]:
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 ()
nested_response: Final = frame.get("response")
input_items: Final = frame.get(
"input", nested_response.get("input", []) if _is_json_object(nested_response) else []
)
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 _ws_event_response_id(event: object) -> str | None:
response: Final = event.get("response") if _is_json_object(event) else None
response_id: Final = response.get("id") if _is_json_object(response) else None
return response_id if isinstance(response_id, str) else None
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[_MutableJsonObject, ...] = ()
events: tuple[Mapping[str, object], ...] = ()
@property
def response_id(self) -> str | None:
return next(filter(None, map(_ws_event_response_id, self.events)), None)
def with_event(self, event: _MutableJsonObject) -> _ResponsesWsTurn:
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"})
def _ws_event_error(event: Mapping[str, object]) -> object:
if event.get("type") == "error":
@ -2009,31 +2017,6 @@ def _restore_wrapped_ids_in_response_create(msg_obj: Mapping[str, object]) -> di
return {**msg_obj, **top_fields, **restored_nested}
_RESPONSE_CREATE_MASKED_FIELDS: Final = ("input", "instructions")
def _with_masked_fields(client: Mapping[str, object], forwarded: Mapping[str, object]) -> Mapping[str, object]:
masked: Final = {
key: forwarded[key] for key in _RESPONSE_CREATE_MASKED_FIELDS if key in client and key in forwarded
}
return {**client, **masked}
def _client_frame_with_masked_text(client_frame: str, forwarded_frame: str) -> str:
client: Final = _json_object_or_none(client_frame)
forwarded: Final = _json_object_or_none(forwarded_frame)
if client is None or forwarded is None:
return client_frame
client_nested: Final = client.get("response")
forwarded_nested: Final = forwarded.get("response")
nested: Final = (
{"response": _with_masked_fields(client_nested, forwarded_nested)}
if _is_json_object(client_nested) and _is_json_object(forwarded_nested)
else {}
)
return json.dumps({**_with_masked_fields(client, forwarded), **nested})
def _wrap_output_item_encrypted_content(
event_obj: Mapping[str, object], litellm_metadata: Mapping[str, object]
) -> dict[str, object] | None:
@ -2079,7 +2062,6 @@ class ResponsesWebSocketStreaming:
authorized_model: str | None = None,
custom_llm_provider: str | None = None,
request_defaults: ResponsesWebSocketRequestDefaults | None = None,
on_turn_failure: ResponsesWebSocketTurnFailureHook | None = None,
):
self.websocket = websocket
self.backend_ws = backend_ws
@ -2089,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.on_turn_failure: ResponsesWebSocketTurnFailureHook | None = on_turn_failure
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 []
@ -2101,20 +2083,28 @@ class ResponsesWebSocketStreaming:
self.request_defaults: ResponsesWebSocketRequestDefaults | None = request_defaults
def _start_turn(self, client_frame: str | None) -> None:
turn_start: Final = datetime.now() # noqa: DTZ005 # logging pipeline uses naive datetimes
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=turn_start,
request_body=_json_object_or_none(client_frame) or {},
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))
def _logged_client_frame(self, client_frame: str, forwarded_frame: str) -> str:
if not self.guardrail_callbacks or forwarded_frame is client_frame:
return client_frame
return _client_frame_with_masked_text(client_frame, forwarded_frame)
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
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)
@ -2127,34 +2117,12 @@ class ResponsesWebSocketStreaming:
return
if not self._open_turns:
self._start_turn(None)
index: Final = self._turn_index_for_event(event_obj)
turn: Final = self._open_turns[index].with_event(event_obj)
other_turns: Final = (*self._open_turns[:index], *self._open_turns[index + 1 :])
turn: Final = self._open_turns[0].with_event(event_obj)
if event_obj.get("type") in _RESPONSES_WS_TERMINAL_EVENT_TYPES:
self._open_turns = other_turns
self._open_turns = self._open_turns[1:]
self._log_turn(turn)
return
self._open_turns = (*self._open_turns[:index], turn, *self._open_turns[index + 1 :])
def _turn_index_for_event(self, event_obj: Mapping[str, object]) -> int:
"""
An event naming a response id belongs to the turn whose ``response.created``
carried it. The backend answers turns strictly in order, so an id-less event,
a rejected turn's ``error`` included, goes to the oldest turn still waiting for
its ``response.created``. The one exception is an id-less ``error`` while a turn
is streaming: the backend rejects a ``response.create`` it cannot parse as soon
as it reads it, so that error belongs to the newest pending turn.
"""
response_id: Final = _ws_event_response_id(event_obj)
turn_ids: Final = tuple(turn.response_id for turn in self._open_turns)
if response_id is not None and response_id in turn_ids:
return turn_ids.index(response_id)
pending: Final = tuple(index for index, turn_id in enumerate(turn_ids) if turn_id is None)
if not pending:
return 0
if response_id is None and event_obj.get("type") == "error" and len(pending) < len(turn_ids):
return pending[-1]
return pending[0]
self._open_turns = (turn, *self._open_turns[1:])
def _log_open_turns(self) -> None:
open_turns: Final = self._open_turns
@ -2164,35 +2132,18 @@ class ResponsesWebSocketStreaming:
self._log_turn(turn)
def _log_turn(self, turn: _ResponsesWsTurn) -> None:
_spawn_logging_task(
asyncio.get_running_loop(), self._dispatch_turn(turn), task_name="Responses WS turn logging"
)
async def _dispatch_turn(self, turn: _ResponsesWsTurn) -> None:
events: Final = list(turn.events)
failed_event: Final = next(
(event for event in events if event.get("type") in _RESPONSES_WS_FAILURE_EVENT_TYPES), None
)
if failed_event is None:
await turn.logging_obj.dispatch_success_handlers(events, prefer_async_handlers=True)
exception: Final = turn.failure(self.authorized_model or "", self.custom_llm_provider or "")
if exception is None:
asyncio.create_task(turn.logging_obj.dispatch_success_handlers(events, prefer_async_handlers=True))
return
exception: Final = _map_stream_error_to_exception(
_ws_event_error(failed_event), self.authorized_model or "", self.custom_llm_provider or ""
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(
turn.logging_obj.dispatch_failure_handlers(exception, traceback_exception, prefer_async_handlers=True)
)
self._record_usage_for_failure(turn.logging_obj, events)
await turn.logging_obj.dispatch_failure_handlers(
exception, "".join(traceback.format_exception(exception)), prefer_async_handlers=True
)
if self.on_turn_failure is not None:
await self.on_turn_failure(turn.logging_obj, exception)
@staticmethod
def _record_usage_for_failure(turn_logging_obj: LiteLLMLoggingObj, events: Sequence[_MutableJsonObject]) -> 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)
def _wrap_response_event(self, response_str: str) -> str:
try:
@ -2605,7 +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._record_client_frame(self._logged_client_frame(self.first_message, masked_first))
self._record_client_frame(masked_first)
await self.backend_ws.send(masked_first)
while True:
@ -2613,13 +2564,13 @@ class ResponsesWebSocketStreaming:
if not await self._enforce_or_reject_frame(message):
continue
masked = await self._mask_response_create(message)
self._record_client_frame(self._logged_client_frame(message, masked))
self._record_client_frame(masked)
await self.backend_ws.send(masked)
except Exception as e:
verbose_logger.debug("Responses WS client_to_backend ended: %s", e)
async def bidirectional_forward(self) -> None:
async def bidirectional_forward(self) -> Exception | None:
forward_task: Final = asyncio.create_task(self.backend_to_client())
try:
await self.client_to_backend()
@ -2636,6 +2587,7 @@ class ResponsesWebSocketStreaming:
await self.backend_ws.close()
except Exception:
pass
return self._failure
# ---------------------------------------------------------------------------

3
litellm/router.py
View file

@ -8579,7 +8579,8 @@ class Router:
Track remaining tpm/rpm quota for model in model_list
"""
try:
if kwargs.get("call_type") == "_arealtime":
# WS session wrappers fire with result=None; per-turn costs tracked by inner calls.
if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"):
return
if is_batch_retrieve_call_type(kwargs.get("call_type")):
return

3
litellm/router_strategy/budget_limiter.py
View file

@ -502,7 +502,8 @@ class RouterBudgetLimiting(CustomLogger):
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
"""Original method now uses helper functions"""
verbose_router_logger.debug("in RouterBudgetLimiting.async_log_success_event")
if kwargs.get("call_type") == "_arealtime":
# WS session wrappers fire with result=None; per-turn costs tracked by inner calls.
if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"):
return
standard_logging_payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None)
if standard_logging_payload is None:

11
litellm/types/responses/streaming_websocket.py
View file

@ -2,13 +2,10 @@ from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass
from typing import TYPE_CHECKING, Protocol
from typing import Protocol
from litellm.types.guardrails import PresidioPerRequestConfig
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging
class ResponsesClientWebSocket(Protocol):
"""Client-facing websocket surface used by the Responses API websocket handlers."""
@ -46,12 +43,6 @@ class PresidioGuardrailCallback(Protocol):
) -> str: ...
class ResponsesWebSocketTurnFailureHook(Protocol):
"""Books one failed ``response.create`` turn of a native Responses API socket from that turn's logging object."""
async def __call__(self, turn_logging_obj: Logging, exception: Exception) -> None: ...
@dataclass(frozen=True, slots=True)
class ResponsesWebSocketRequestDefaults:
"""Deployment-level request parameters merged into every ``response.create`` frame relayed over a native websocket."""

4
litellm/types/utils.py
View file

@ -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

2
terraform/provider/CHANGELOG.md
View file

@ -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

14
terraform/provider/docs/resources/model.md
View file

@ -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
}
}
```

21
terraform/provider/examples/model_additional_params.tf
View file

@ -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

12
terraform/provider/litellm/resource_model.go
View file

@ -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",
},
},
}

192
terraform/provider/litellm/resource_model_crud.go
View file

@ -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
}

243
terraform/provider/litellm/resource_model_test.go
View file

@ -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

1
terraform/provider/litellm/types.go
View file

@ -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.

32
tests/e2e/perf/setup/observe.py
View file

@ -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]

15
tests/e2e/perf/test_perf_cases.py
View file

@ -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, "

68
tests/e2e_harness/perf/setup/test_observe.py Normal file
View file

@ -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,
}

2
tests/e2e_harness/pytest.ini
View file

@ -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

186
tests/integration/_support/responses_ws_peer.py Normal file
View file

@ -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()))

96
tests/integration/authorization/test_sso_blank_identity.py
View file

@ -169,8 +169,9 @@ class SamlIdp:
"</IDPSSODescriptor></EntityDescriptor>"
)
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:
'<saml:Subject><saml:NameID Format="urn:oasis:names:tc:SAML:1.1:nameid-format:unspecified">'
f"{name_id}</saml:NameID>"
'<saml:SubjectConfirmation Method="urn:oasis:names:tc:SAML:2.0:cm:bearer">'
f'<saml:SubjectConfirmationData InResponseTo="{request_id}" NotOnOrAfter="{expires}" Recipient="{acs}"/>'
f'<saml:SubjectConfirmationData{in_response_to} NotOnOrAfter="{expires}" Recipient="{acs}"/>'
"</saml:SubjectConfirmation></saml:Subject>"
f'<saml:Conditions NotBefore="{not_before}" NotOnOrAfter="{expires}">'
f"<saml:AudienceRestriction><saml:Audience>{SAML_SP_ENTITY}</saml:Audience>"
@ -206,7 +207,7 @@ class SamlIdp:
'<samlp:Response xmlns:samlp="urn:oasis:names:tc:SAML:2.0:protocol" '
'xmlns:saml="urn:oasis:names:tc:SAML:2.0:assertion" '
f'ID="_response_{uuid.uuid4().hex}" Version="2.0" IssueInstant="{issued}" '
f'Destination="{acs}" InResponseTo="{request_id}">'
f'Destination="{acs}"{in_response_to}>'
f"<saml:Issuer>{SAML_IDP_ENTITY}</saml:Issuer>"
'<samlp:Status><samlp:StatusCode Value="urn:oasis:names:tc:SAML:2.0:status:Success"/></samlp:Status>'
f"{signed_text}</samlp:Response>"
@ -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}"

192
tests/integration/providers/test_openai_count_tokens_tools_wire.py
View file

@ -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."

226
tests/integration/providers/test_openai_responses_websocket_wire.py
View file

@ -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(

87
tests/integration/providers/test_vertex_gemini_count_tokens_wire.py
View file

@ -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:

7
tests/integration/spend/test_cache_and_quota.py
View file

@ -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": "",
}

114
tests/integration/spend/test_responses_websocket_turn_rows.py Normal file
View file

@ -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

172
tests/unit/litellm_core_utils/test_litellm_logging.py
View file

@ -39,7 +39,6 @@ from litellm.litellm_core_utils.litellm_logging import (
StandardLoggingPayloadSetup,
_extract_response_obj_and_hidden_params,
_get_status_fields,
get_standard_logging_object_payload,
set_callbacks,
)
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
@ -2051,34 +2050,34 @@ async def test_arealtime_marks_litellm_params_async(monkeypatch):
@pytest.mark.asyncio
async def test_aresponses_websocket_writes_no_connection_level_success_log(
async def test_aresponses_websocket_hands_back_the_provider_failure_without_a_success_log(
monkeypatch: pytest.MonkeyPatch,
):
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.responses.main import base_llm_http_handler
logged_call_ids = []
success_events = []
class CaptureLogger(CustomLogger):
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
logged_call_ids.append(kwargs.get("litellm_call_id"))
success_events.append(response_obj)
monkeypatch.setattr(litellm, "callbacks", [CaptureLogger()])
monkeypatch.setattr(litellm, "failure_callback", [])
monkeypatch.setattr(litellm, "_async_failure_callback", [])
monkeypatch.setattr(litellm, "success_callback", [])
monkeypatch.setattr(litellm, "_async_success_callback", [])
with patch.object( # test-quality-ok: the provider socket is the seam; each turn logs inside it, so the wrapper must not
base_llm_http_handler, "async_responses_websocket", AsyncMock(return_value=None)
failure = litellm.BadRequestError(message="invalid_encrypted_content", model="gpt-4o", llm_provider="openai")
with patch.object( # test-quality-ok: the provider socket is the seam; how the wrapper treats the relay's outcome is under test
base_llm_http_handler, "async_responses_websocket", AsyncMock(return_value=failure)
):
await litellm._aresponses_websocket(
model="openai/gpt-4o", websocket=MagicMock(), api_key="sk-test", litellm_call_id="socket-call"
)
outcome = await litellm._aresponses_websocket(model="openai/gpt-4o", websocket=MagicMock(), api_key="sk-test")
await asyncio.sleep(0)
with contextlib.suppress(asyncio.TimeoutError):
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0)
assert "socket-call" not in logged_call_ids
assert outcome is failure
assert success_events == []
@pytest.mark.asyncio
@ -8354,107 +8353,6 @@ def test_normalize_logging_result_prices_responses_websocket_at_returned_service
assert priority_http_cost > default_http_cost
@pytest.mark.parametrize(
("completed_response_id", "expected_request_id"),
(("resp_wrapped_turn_2", "resp_wrapped_turn_2"), (None, "turn-2-call")),
)
def test_responses_websocket_turn_is_logged_under_the_response_id_the_client_saw(
completed_response_id: str | None, expected_request_id: str
):
events = [
{"type": "response.created", "response": {"id": completed_response_id}},
{
"type": "response.completed",
"response": {
"id": completed_response_id,
"usage": {"input_tokens": 100, "output_tokens": 40, "total_tokens": 140},
},
},
]
logging_obj = _responses_ws_logging_obj()
now = datetime.datetime.now()
payload = litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload(
kwargs={"litellm_call_id": "turn-2-call", "model": "gpt-4o", "messages": []},
init_response_obj=logging_obj.normalize_logging_result(result=events),
start_time=now,
end_time=now,
logging_obj=logging_obj,
status="success",
)
assert payload is not None
assert payload["id"] == expected_request_id
@pytest.mark.parametrize(
("response_obj", "expected_id"),
(({"id": ""}, ""), ({"id": "chatcmpl-1"}, "chatcmpl-1"), ({}, "http-call")),
)
def test_non_websocket_payload_id_is_the_response_id_field_as_returned(
logging_obj: LitellmLogging, response_obj: dict[str, str], expected_id: str
):
now = datetime.datetime.now()
payload = get_standard_logging_object_payload(
kwargs={"litellm_call_id": "http-call", "model": "gpt-4o", "messages": []},
init_response_obj=response_obj,
start_time=now,
end_time=now,
logging_obj=logging_obj,
status="success",
)
assert payload is not None
assert payload["id"] == expected_id
def test_copy_for_new_call_gives_each_call_its_own_id_clock_and_metadata():
reservation = {"reserved_cost": 0.5, "finalized": False}
connection = _responses_ws_logging_obj()
connection_start = connection.start_time
connection_request = {"url": "ws://proxy/v1/responses", "body": {"first_message": "turn 1 frame"}}
connection.litellm_params = {
"metadata": {"user_api_key_alias": "codex", "budget_reservation": reservation},
"litellm_metadata": {"model_info": {"id": "dep-1"}},
"litellm_call_id": connection.litellm_call_id,
"proxy_server_request": connection_request,
}
connection.model_call_details = {
**connection.model_call_details,
"litellm_params": connection.litellm_params,
"first_api_call_start_time": connection_start,
}
turn_start = datetime.datetime(2026, 10, 9, 17, 24, 47)
turn_frame = {"type": "response.create", "input": "turn 2 frame"}
turn = connection.copy_for_new_call(litellm_call_id="turn-1-call", start_time=turn_start, request_body=turn_frame)
turn.litellm_params["metadata"]["status"] = "failure"
turn.litellm_params["litellm_metadata"]["model_group"] = "gpt-4o"
turn.model_call_details["standard_logging_object"] = {"id": "resp_1"}
assert (turn.litellm_call_id, turn.start_time) == ("turn-1-call", turn_start)
assert turn.model_call_details["litellm_call_id"] == "turn-1-call"
assert turn.model_call_details["start_time"] == turn_start
assert turn.model_call_details["litellm_params"]["litellm_call_id"] == "turn-1-call"
assert turn.model_call_details["first_api_call_start_time"] is None
assert turn.litellm_params["metadata"]["budget_reservation"] is reservation
assert turn.model_call_details["litellm_params"]["proxy_server_request"] == {
"url": "ws://proxy/v1/responses",
"body": turn_frame,
}
assert connection.litellm_params["proxy_server_request"] == {
"url": "ws://proxy/v1/responses",
"body": {"first_message": "turn 1 frame"},
}
assert (connection.litellm_call_id, connection.start_time) == ("responses-ws-usage-test", connection_start)
assert connection.litellm_params["litellm_call_id"] == "responses-ws-usage-test"
assert "status" not in connection.litellm_params["metadata"]
assert "model_group" not in connection.litellm_params["litellm_metadata"]
assert "standard_logging_object" not in connection.model_call_details
assert connection.model_call_details["first_api_call_start_time"] == connection_start
def test_get_standard_logging_object_payload_reads_overhead_from_logging_obj_for_dict_results(logging_obj):
"""LIT-5466: /v1/messages returns a plain dict with no _hidden_params, so the overhead
recorded on the logging object must reach hidden_params.litellm_overhead_time_ms (SpendLogs)."""
@ -11982,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

2
tests/unit/llms/crusoe/test_crusoe.py
View file

@ -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():

98
tests/unit/llms/gemini/test_gemini_common_utils.py
View file

@ -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

207
tests/unit/llms/openai/responses/count_tokens/test_token_counter.py
View file

@ -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,

108
tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py
View file

@ -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."

2
tests/unit/llms/openai_like/test_json_loader.py
View file

@ -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():

315
tests/unit/llms/vertex_ai/count_tokens/test_context.py
View file

@ -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)

91
tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py
View file

@ -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

36
tests/unit/proxy/anthropic_endpoints/test_endpoints.py
View file

@ -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,
):

276
tests/unit/proxy/common_utils/test_reset_budget_job.py
View file

@ -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()
(

144
tests/unit/proxy/hooks/test_parallel_request_limiter.py
View file

@ -2,9 +2,8 @@
Unit Tests for the max parallel request limiter v1 for the proxy
"""
import asyncio
import itertools
from collections.abc import Awaitable, Callable, Iterator, Mapping
from collections.abc import Callable, Iterator, Mapping
from datetime import datetime
from typing import Final
@ -13,11 +12,10 @@ import pytest
from litellm.caching.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.hooks.parallel_request_limiter import (
PROXY_MaxParallelRequestsHandler,
)
from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token
from litellm.proxy.utils import InternalUsageCache, hash_token
from litellm.types.utils import CallTypes, EmbeddingResponse, ModelResponse, TextCompletionResponse, Usage
FROZEN_INSTANT: Final = datetime(2026, 1, 31, 23, 59, 30)
@ -327,123 +325,31 @@ async def test_retry_after_is_the_seconds_until_the_next_minute_of_the_injected_
assert exc_info.value.headers["retry-after"] == "14.5"
async def _log_websocket_turn_success(
handler: PROXY_MaxParallelRequestsHandler, metadata: Mapping[str, object]
) -> None:
@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={"litellm_params": {"metadata": dict(metadata)}},
response_obj=None,
start_time=FROZEN_INSTANT,
end_time=FROZEN_INSTANT,
kwargs=turn_kwargs, response_obj=None, start_time=FROZEN_INSTANT, end_time=FROZEN_INSTANT
)
async def _log_websocket_turn_failure(
handler: PROXY_MaxParallelRequestsHandler, metadata: Mapping[str, object]
) -> None:
await handler.async_log_failure_event(
kwargs={"litellm_params": {"metadata": dict(metadata)}, "exception": Exception("upstream turn failed")},
response_obj=None,
start_time=FROZEN_INSTANT,
end_time=FROZEN_INSTANT,
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
def _proxy_logging_with(handler: PROXY_MaxParallelRequestsHandler) -> ProxyLogging:
proxy_logging: Final = ProxyLogging(user_api_key_cache=UserApiKeyCache())
proxy_logging.proxy_hook_mapping["parallel_request_limiter"] = handler
return proxy_logging
async def _requests_admitted(
handler: PROXY_MaxParallelRequestsHandler, session: UserAPIKeyAuth, metadata: Mapping[str, object]
) -> int:
for admitted in range(3):
try:
await handler.async_pre_call_hook(
user_api_key_dict=session,
cache=DualCache(),
data={"metadata": dict(metadata)},
call_type=CallTypes.acompletion.value,
)
except ProxyRateLimitError:
return admitted
await asyncio.sleep(0)
return 3
@pytest.mark.parametrize("log_turn", [_log_websocket_turn_success, _log_websocket_turn_failure])
@pytest.mark.parametrize(
"max_parallel_requests,limit_metadata,admitted_while_open",
[(1, {}, 0), (None, {"global_max_parallel_requests": 2}, 1)],
ids=["key_limit", "global_limit"],
)
@pytest.mark.asyncio
async def test_a_websocket_connection_holds_its_parallel_slot_through_its_turns_until_it_closes(
log_turn: Callable[[PROXY_MaxParallelRequestsHandler, Mapping[str, object]], Awaitable[None]],
max_parallel_requests: int | None,
limit_metadata: Mapping[str, object],
admitted_while_open: int,
):
handler: Final = PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(DualCache()), clock=_frozen_clock
)
api_key: Final = hash_token("sk-websocket-slot")
session: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=max_parallel_requests)
turn_logged: Final = asyncio.Event()
socket_closed: Final = asyncio.Event()
async def websocket_connection() -> None:
await handler.async_pre_call_hook(
user_api_key_dict=session,
cache=DualCache(),
data={"metadata": dict(limit_metadata)},
call_type=CallTypes.aresponses_websocket.value,
)
await asyncio.sleep(0)
await log_turn(handler, {"user_api_key": api_key, "user_api_key_model_max_budget": {}, **limit_metadata})
turn_logged.set()
await socket_closed.wait()
await _proxy_logging_with(handler).arelease_max_parallel_requests_on_disconnect(session)
connection: Final = asyncio.create_task(websocket_connection())
await turn_logged.wait()
assert await _requests_admitted(handler, session, limit_metadata) == admitted_while_open
socket_closed.set()
await connection
assert await _requests_admitted(handler, session, limit_metadata) == 1
@pytest.mark.asyncio
async def test_a_websocket_turn_logged_in_a_later_minute_leaves_no_slot_held_after_the_socket_closes():
handler: Final = PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(DualCache()), clock=_clock_rolling_over_after_first_read()
)
api_key: Final = hash_token("sk-websocket-rollover")
session: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1)
async def websocket_connection() -> None:
await handler.async_pre_call_hook(
user_api_key_dict=session, cache=DualCache(), data={}, call_type=CallTypes.aresponses_websocket.value
)
await asyncio.sleep(0)
await _log_websocket_turn_success(handler, {"user_api_key": api_key, "user_api_key_model_max_budget": {}})
await handler.async_release_websocket_connection_slot()
await asyncio.create_task(websocket_connection())
assert await _requests_admitted(handler, session, {}) == 1
@pytest.mark.asyncio
async def test_a_disconnect_release_outside_a_websocket_connection_leaves_the_request_slot_held():
handler: Final = PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(DualCache()), clock=_frozen_clock
)
session: Final = UserAPIKeyAuth(api_key=hash_token("sk-http-slot"), max_parallel_requests=1)
assert await _requests_admitted(handler, session, {}) == 1
await _proxy_logging_with(handler).arelease_max_parallel_requests_on_disconnect(session)
assert await _requests_admitted(handler, session, {}) == 0
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

388
tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -24,7 +24,6 @@ from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.integrations.custom_secret_manager import CustomSecretManager
from litellm.litellm_core_utils import ptu_pricing
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
@ -54,7 +53,6 @@ from litellm.types.mcp import MCPPreCallRequestObject, MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
from litellm.types.utils import (
CallTypes,
EmbeddingResponse,
ModelResponse,
PromptTokensDetailsWrapper,
@ -4246,82 +4244,6 @@ async def test_release_on_disconnect_works_when_key_config_changed_v3():
) == 0
async def _admit_with_tpm_reservation(
handler: _PROXY_MaxParallelRequestsHandler,
local_cache: DualCache,
user_api_key_dict: UserAPIKeyAuth,
call_type: str = "completion",
) -> int:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={
"model": "gpt-4o-mini",
"messages": [{"role": "user", "content": "hello"}],
"max_tokens": 50,
"litellm_call_id": "owner-call-id",
},
call_type=call_type,
)
stash = get_request_stash()
assert stash is not None and stash.reserved_tokens > 0
return stash.reserved_tokens
@pytest.mark.asyncio
async def test_release_on_disconnect_refunds_an_unsettled_tpm_reservation_v3(monkeypatch):
"""
A native Responses WebSocket connection reserves tokens at pre-call and logs its
turns under their own call ids, so no callback ever settles the connection's
reservation. The disconnect release must hand the estimate back.
"""
monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False)
_api_key = hash_token("sk-ws-unsettled-reservation")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2)
tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens")
reserved = await _admit_with_tpm_reservation(
handler, local_cache, user_api_key_dict, CallTypes.aresponses_websocket.value
)
assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == reserved
await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict)
assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 0
assert get_request_stash().reservation_released is True
@pytest.mark.asyncio
async def test_release_on_disconnect_keeps_a_settled_tpm_reservation_v3(monkeypatch):
"""
When the owner's success event already settled the reservation to the actual
usage, the disconnect release must not refund the estimate a second time.
"""
monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False)
_api_key = hash_token("sk-ws-settled-reservation")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2)
tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens")
await _admit_with_tpm_reservation(handler, local_cache, user_api_key_dict, CallTypes.aresponses_websocket.value)
await handler.async_log_success_event(
kwargs={
"litellm_call_id": "owner-call-id",
"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}},
},
response_obj=ModelResponse(usage=Usage(prompt_tokens=7, completion_tokens=3, total_tokens=10)),
start_time=datetime.now(),
end_time=datetime.now(),
)
assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 10
await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict)
assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 10
@pytest.mark.asyncio
async def test_post_call_failure_hook_releases_parallel_slot_v3():
"""
@ -4381,108 +4303,6 @@ async def test_post_call_failure_hook_releases_parallel_slot_v3():
)
@pytest.mark.asyncio
async def test_post_call_failure_hook_for_another_call_keeps_the_owner_slot_v3():
"""
A native Responses WebSocket holds one parallel slot for the whole socket and books each failed turn
under that turn's own call id. A failed turn must not free the slot while the socket is still open.
"""
_api_key = hash_token("sk-ws-turn-failure")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1)
counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests"
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={"model": "gpt-4o-mini", "litellm_call_id": "socket-call"},
call_type=CallTypes.aresponses_websocket.value,
)
await handler.async_post_call_failure_hook(
request_data={"model": "gpt-4o-mini", "litellm_call_id": "turn-2-call"},
original_exception=Exception("previous_response_not_found"),
user_api_key_dict=user_api_key_dict,
)
in_flight_after_turn_failure: Final = handler._gauge_in_flight_from_cache_value(
await local_cache.async_get_cache(key=counter_key)
)
await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict)
assert in_flight_after_turn_failure == 1
assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0
def _pass_through_logging_obj(litellm_call_id: str) -> Logging:
return Logging(
model="gpt-4o-mini",
messages=[],
stream=False,
call_type="pass_through_endpoint",
start_time=datetime.now(),
litellm_call_id=litellm_call_id,
function_id="pass-through",
)
async def _admit_pass_through_with_client_call_id(
handler: _PROXY_MaxParallelRequestsHandler, local_cache: DualCache, user_api_key_dict: UserAPIKeyAuth
) -> Logging:
"""
A pass-through body may carry its own litellm_call_id. The proxy books the request under the id of the
logging object it put in the same data, so that id, not the client's, must own the parallel slot.
"""
logging_obj = _pass_through_logging_obj("proxy-call")
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={"litellm_call_id": "client-chosen-id", "litellm_logging_obj": logging_obj},
call_type="pass_through_endpoint",
)
return logging_obj
@pytest.mark.asyncio
async def test_post_call_failure_hook_releases_the_slot_the_proxy_call_id_owns_v3(monkeypatch: pytest.MonkeyPatch):
_api_key = hash_token("sk-passthrough-client-id-failure")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1)
counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests"
logging_obj = await _admit_pass_through_with_client_call_id(handler, local_cache, user_api_key_dict)
monkeypatch.setattr(litellm, "callbacks", [handler])
await ProxyLogging(user_api_key_cache=UserApiKeyCache()).post_call_failure_hook(
request_data={"litellm_call_id": "client-chosen-id", "litellm_logging_obj": logging_obj},
original_exception=Exception("rejected by a pre-call hook"),
user_api_key_dict=user_api_key_dict,
)
assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0
@pytest.mark.asyncio
async def test_success_event_releases_the_slot_the_proxy_call_id_owns_v3():
_api_key = hash_token("sk-passthrough-client-id-success")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1)
counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests"
await _admit_pass_through_with_client_call_id(handler, local_cache, user_api_key_dict)
await handler.async_log_success_event(
kwargs={
"litellm_call_id": "proxy-call",
"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}},
},
response_obj=ModelResponse(usage=Usage(prompt_tokens=5, completion_tokens=5, total_tokens=10)),
start_time=datetime.now(),
end_time=datetime.now(),
)
assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0
@pytest.mark.asyncio
async def test_success_event_releases_parallel_slot_v3(monkeypatch):
"""
@ -5232,45 +5052,6 @@ async def test_tag_429_names_the_tag_that_is_over_its_limit(monkeypatch):
assert "cell-blocked" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_only_tagged_websocket_turns_count_toward_the_connection_tag_tpm_v3(monkeypatch):
monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60")
monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", "false")
_request_stash.set(None)
resolver, _ = _static_tag_limits({"cell-1": TagRateLimit(rpm_limit=None, tpm_limit=1_000)})
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache), tag_rate_limit_resolver=resolver
)
user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-ws-tag-tpm"))
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={"model": "gpt-4o-mini", "metadata": {"tags": ["cell-1"]}, "litellm_call_id": "socket-call"},
call_type=CallTypes.aresponses_websocket.value,
)
tokens_before_turns: Final = int(await local_cache.async_get_cache("{tag:cell-1}:tokens") or 0)
calls_in_the_connection = (("turn-1-call", ["cell-1"]), ("nested-untagged-call", []), ("turn-2-call", ["cell-1"]))
for call_id, request_tags in calls_in_the_connection:
await handler.async_log_success_event(
kwargs={
"litellm_call_id": call_id,
"standard_logging_object": {
"metadata": {"user_api_key_hash": user_api_key_dict.api_key},
"request_tags": request_tags,
},
"model": "gpt-4o-mini",
},
response_obj=ModelResponse(usage=Usage(prompt_tokens=60, completion_tokens=60, total_tokens=120)),
start_time=datetime.now(),
end_time=datetime.now(),
)
await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict)
assert await local_cache.async_get_cache("{tag:cell-1}:tokens") == tokens_before_turns + 240
@pytest.mark.asyncio
async def test_tag_object_tpm_limit_enforced_v3(monkeypatch):
monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60")
@ -8523,7 +8304,7 @@ async def test_success_settles_the_scope_the_reservation_was_taken_on():
async def _reserve_a_ptu_minute(
cache: DualCache, call_id: str, clock: list[datetime] | None = None, call_type: str = "acompletion"
cache: DualCache, call_id: str, clock: list[datetime] | None = None
) -> tuple[_PROXY_MaxParallelRequestsHandler, str]:
resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=2000, ratio=4.0)
handler = _PROXY_MaxParallelRequestsHandler(
@ -8536,7 +8317,7 @@ async def _reserve_a_ptu_minute(
user_api_key_dict=key,
cache=cache,
data={**_ptu_request(), "litellm_call_id": call_id},
call_type=call_type,
call_type="acompletion",
)
stash = get_request_stash()
assert stash is not None and stash.ptu_reserved_tokens > 0
@ -8580,18 +8361,6 @@ async def test_a_proxy_side_rejection_refunds_the_whole_normalized_ptu_reservati
assert int(await cache.async_get_cache(key=ptu_key) or 0) == 0
@pytest.mark.asyncio
async def test_a_websocket_disconnect_refunds_the_whole_normalized_ptu_reservation():
cache = DualCache()
handler, ptu_key = await _reserve_a_ptu_minute(cache, "ptu-ws", call_type=CallTypes.aresponses_websocket.value)
await handler.async_release_max_parallel_requests_on_disconnect(
UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t")
)
assert int(await cache.async_get_cache(key=ptu_key) or 0) == 0
async def _settle_a_ptu_request(handler: _PROXY_MaxParallelRequestsHandler, call_id: str, outcome: str) -> None:
team_metadata = {"metadata": {"user_api_key_hash": hash_token("sk-ptu"), "user_api_key_team_id": "t"}}
if outcome == "success":
@ -8797,156 +8566,3 @@ async def test_a2a_url_target_owns_invocation_fee_and_request_limit(
await _rpm_request(limiter, cache, auth, "a2a/cheap")
assert denied.value.status_code == 429
assert "expensive" in str(denied.value.detail)
def _websocket_turn_kwargs(call_id: str, api_key_hash: str, usage: Usage) -> dict[str, object]:
return {
"litellm_call_id": call_id,
"model": "gpt-4o-mini",
"combined_usage_object": usage,
"litellm_params": {"metadata": {"user_api_key": api_key_hash}},
"standard_logging_object": {"metadata": {"user_api_key_hash": api_key_hash}},
}
@pytest.mark.asyncio
async def test_release_on_disconnect_keeps_a_plain_request_tpm_reservation_v3(monkeypatch):
"""
A cancelled HTTP stream keeps its reservation for the success or failure event
that still settles it; only a WebSocket connection, whose turns log under their
own call ids, hands its estimate back at disconnect.
"""
monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False)
_api_key = hash_token("sk-http-cancelled-stream")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2)
tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens")
reserved = await _admit_with_tpm_reservation(handler, local_cache, user_api_key_dict, "completion")
await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict)
assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == reserved
assert get_request_stash().reservation_released is False
@pytest.mark.asyncio
async def test_release_on_disconnect_skips_a_tpm_window_that_rolled_v3(monkeypatch):
"""
The TPM window can roll while a WebSocket connection stays open. The disconnect
refund belongs to the window the connection reserved against, so it must leave
the window that is live now untouched instead of handing it tokens it never took.
"""
monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False)
_api_key = hash_token("sk-ws-rolled-window")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2)
tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens")
window_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="window")
await _admit_with_tpm_reservation(handler, local_cache, user_api_key_dict, CallTypes.aresponses_websocket.value)
reserved_window_start = await local_cache.async_get_cache(key=window_key)
assert reserved_window_start is not None
await local_cache.async_set_cache(key=window_key, value=int(reserved_window_start) + handler.window_size)
await local_cache.async_set_cache(key=tokens_key, value=500)
await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict)
assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 500
assert get_request_stash().reservation_released is True
@pytest.mark.asyncio
async def test_websocket_turn_failure_counts_its_usage_toward_tpm_v3(monkeypatch):
"""
A turn that ends in ``response.failed`` still consumed the usage the event
reports. It logs under its own call id, so the failure event owns no stash; it
must still charge the connection's TPM scopes the way a successful turn does.
"""
monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False)
_api_key = hash_token("sk-ws-failed-turn")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2)
tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens")
reserved = await _admit_with_tpm_reservation(
handler, local_cache, user_api_key_dict, CallTypes.aresponses_websocket.value
)
await handler.async_log_failure_event(
kwargs=_websocket_turn_kwargs(
"turn-1", _api_key, Usage(prompt_tokens=12, completion_tokens=5, total_tokens=17)
),
response_obj=None,
start_time=datetime.now(),
end_time=datetime.now(),
)
assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == reserved + 17
assert get_request_stash().reservation_released is False
@pytest.mark.asyncio
async def test_websocket_turn_success_counts_its_usage_toward_tpm_v3(monkeypatch):
"""
A successful native Responses WebSocket turn is logged with the raw event list
as its response, not a modeled response object. Its usage must still reach the
key's TPM window right away, while the socket stays open.
"""
monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False)
_api_key = hash_token("sk-ws-successful-turn")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2)
tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens")
reserved = await _admit_with_tpm_reservation(
handler, local_cache, user_api_key_dict, CallTypes.aresponses_websocket.value
)
turn_kwargs = {
key: value
for key, value in _websocket_turn_kwargs("turn-1", _api_key, Usage()).items()
if key != "combined_usage_object"
}
await handler.async_log_success_event(
kwargs={**turn_kwargs, "call_type": CallTypes.aresponses_websocket.value},
response_obj=[
{"type": "response.created", "response": {"id": "resp_turn_1"}},
{
"type": "response.completed",
"response": {"id": "resp_turn_1", "usage": {"input_tokens": 12, "output_tokens": 5, "total_tokens": 17}},
},
],
start_time=datetime.now(),
end_time=datetime.now(),
)
assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == reserved + 17
assert get_request_stash().reservation_released is False
@pytest.mark.asyncio
async def test_failure_event_without_a_stash_charges_nothing_outside_a_websocket_connection_v3(monkeypatch):
"""
A nested call made inside a plain request (an LLM-judge guardrail) logs under a
fresh call id too. Its failure owns no stash and must not touch the request's
counters: only a WebSocket connection's turns are charged that way.
"""
monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False)
_api_key = hash_token("sk-nested-failed-call")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2)
tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens")
reserved = await _admit_with_tpm_reservation(handler, local_cache, user_api_key_dict, "completion")
await handler.async_log_failure_event(
kwargs=_websocket_turn_kwargs(
"nested-call-id", _api_key, Usage(prompt_tokens=12, completion_tokens=5, total_tokens=17)
),
response_obj=None,
start_time=datetime.now(),
end_time=datetime.now(),
)
assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == reserved

166
tests/unit/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -2737,50 +2737,6 @@ async def test_async_log_success_event_hands_the_sidecar_a_compact_event_and_ski
assert event.response_cost == 0.0125
@pytest.mark.asyncio
async def test_collector_events_let_only_the_first_turn_settle_a_shared_reservation(tmp_path):
from litellm.litellm_core_utils.core_helpers import (
budget_reservation_from_metadata,
get_litellm_metadata_from_kwargs,
)
from litellm.proxy.spend_tracking.budget_reservation import (
claim_budget_reservation,
share_budget_reservation_across_calls,
)
from litellm.proxy.spend_tracking.spend_event import spend_event_callback_args
handler = _RecordingHandler()
consumer = SpendEventConsumer(handler)
address = UnixAddress(path=str(tmp_path / "spend.sock"))
server = await consumer.serve(address)
producer = SpendEventProducer(
address=address, on_unavailable="fallback", buffer_size=10, connect_timeout=1.0, fallback=_no_fallback
)
logger = ProxyDBLogger(producer)
connection_kwargs = _offload_kwargs()
share_budget_reservation_across_calls(
budget_reservation_from_metadata(get_litellm_metadata_from_kwargs(connection_kwargs))
)
for turn_call_id in ("turn-1", "turn-2", "turn-3"):
await logger.async_log_success_event(
{**connection_kwargs, "litellm_call_id": turn_call_id}, _offload_response(), datetime.now(), datetime.now()
)
await producer.close(drain_timeout=5.0)
server.close()
assert await consumer.drain(timeout=5.0) == 0
sidecar_kwargs = [spend_event_callback_args(decode_spend_event(line)).kwargs for line in handler.lines]
assert [
claim_budget_reservation(
budget_reservation=budget_reservation_from_metadata(get_litellm_metadata_from_kwargs(kwargs)),
litellm_call_id=kwargs["litellm_call_id"],
)
is not None
for kwargs in sidecar_kwargs
] == [True, False, False]
@pytest.mark.asyncio
async def test_async_log_success_event_keeps_batch_retrieves_in_process():
producer = SpendEventProducer(
@ -3064,125 +3020,3 @@ async def test_track_cost_callback_enqueue_emits_no_service_span(): # test-qual
assert update_database.await_count == 1, "the spend enqueue itself must still run"
assert success_hook.await_count == 0, [call.kwargs for call in success_hook.await_args_list]
def _shared_reservation_turn_kwargs(call_id: str, reservation: dict) -> dict:
return {
"model": "gpt-5.4-mini",
"litellm_call_id": call_id,
"litellm_params": {
"metadata": {"user_api_key": "test-key-hash", "user_api_key_budget_reservation": reservation}
},
"stream": False,
"standard_logging_object": {"response_cost": 0.25, "request_tags": None},
}
async def _reservation_handed_to_the_counters(kwargs) -> dict | None:
"""The reservation the callback hands the spend counters for one logged call."""
logger = ProxyDBLogger()
with (
patch( # test-quality-ok: callback imports proxy_logging_obj off proxy_server in its body, no seam
"litellm.proxy.proxy_server.proxy_logging_obj"
) as mock_proxy_logging,
patch( # test-quality-ok: the arguments to this call are the boundary under test
"litellm.proxy.hooks.proxy_track_cost_callback._update_database_and_spend_counters",
new=AsyncMock(),
) as mock_update,
patch( # test-quality-ok: llm_router is a proxy_server global the callback reads lazily, no seam
"litellm.proxy.proxy_server.llm_router", new=_FakeDeploymentLookup({})
),
):
mock_proxy_logging.failed_tracking_alert = AsyncMock()
mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock()
mock_proxy_logging.db_spend_update_writer = MagicMock()
mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock()
await logger._PROXY_track_cost_callback(
kwargs=kwargs,
completion_response=None,
start_time=datetime.now(),
end_time=datetime.now(),
)
return mock_update.await_args.kwargs["budget_reservation"]
@pytest.mark.asyncio
async def test_track_cost_callback_hands_a_shared_reservation_to_the_first_turn_only():
"""Every turn of a Responses WebSocket connection carries the connection's one reservation under its
own call id. The first turn settles it; a later turn reconciling the same estimate again would
under-count the key, so it must charge its cost without the reservation."""
reservation = {"reserved_cost": 0.5, "entries": [], "settling_call_id": None}
first_turn = await _reservation_handed_to_the_counters(_shared_reservation_turn_kwargs("turn-1", reservation))
second_turn = await _reservation_handed_to_the_counters(_shared_reservation_turn_kwargs("turn-2", reservation))
assert first_turn is reservation
assert second_turn is None
@pytest.mark.asyncio
async def test_async_post_call_failure_hook_releases_a_shared_reservation_for_the_first_turn_only():
logger = ProxyDBLogger()
budget_reservation = {"reserved_cost": 0.5, "entries": [], "settling_call_id": None}
user_api_key_dict = UserAPIKeyAuth(
api_key="test_api_key",
request_route="/custom/route",
budget_reservation=budget_reservation,
)
with (
patch(
"litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation",
new_callable=AsyncMock,
) as mock_release_budget_reservation,
patch(
"litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
new_callable=AsyncMock,
),
):
for call_id in ("turn-1", "turn-2"):
await logger.async_post_call_failure_hook(
request_data={"litellm_call_id": call_id},
original_exception=Exception("Test exception"),
user_api_key_dict=user_api_key_dict,
)
assert mock_release_budget_reservation.await_count == 1
assert (
mock_release_budget_reservation.await_args.kwargs["budget_reservation"]
is user_api_key_dict.budget_reservation
)
assert user_api_key_dict.budget_reservation["settling_call_id"] == "turn-1"
@pytest.mark.asyncio
@pytest.mark.parametrize("client_call_id", [12345, ["a", "b"], {"x": 1}])
async def test_async_post_call_failure_hook_writes_the_failure_row_for_a_non_string_body_call_id(client_call_id):
"""A rejected request's body is the client's, so its litellm_call_id can be any JSON value. The failure
row must still be written and the reservation released."""
logger = ProxyDBLogger()
user_api_key_dict = UserAPIKeyAuth(
api_key="test_api_key", budget_reservation={"reserved_cost": 0.5, "entries": []}
)
with (
patch(
"litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation",
new_callable=AsyncMock,
) as mock_release_budget_reservation,
patch(
"litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
new_callable=AsyncMock,
) as mock_update_database,
):
await logger.async_post_call_failure_hook(
request_data={"model": "gpt-5.4-mini", "litellm_call_id": client_call_id},
original_exception=Exception("Authentication Error"),
user_api_key_dict=user_api_key_dict,
)
released = mock_release_budget_reservation.await_args.kwargs["budget_reservation"]
assert released is user_api_key_dict.budget_reservation
assert mock_update_database.call_args.kwargs["kwargs"]["litellm_params"]["metadata"]["status"] == "failure"

97
tests/unit/proxy/hooks/test_tpm_concurrent.py
View file

@ -42,7 +42,7 @@ from litellm.types.llms.openai import (
ResponsesAPIResponse,
)
from litellm.types.rerank import RerankResponse
from litellm.types.utils import CallTypes, ModelResponse, PromptTokensDetailsWrapper, Usage
from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage
@pytest.fixture
@ -3727,98 +3727,3 @@ async def test_the_project_itpm_reservation_counts_the_request_off_the_event_loo
if __name__ == "__main__":
pytest.main([__file__, "-v", "-s"])
@pytest.mark.asyncio
async def test_websocket_turn_usage_is_not_counted_twice_toward_project_itpm_and_otpm(rate_limiter):
"""
A native Responses WebSocket connection reserves project ITPM/OTPM at pre-call, and
each response.create frame is charged its own estimate before it is forwarded. The
turns that frame starts log under their own call ids, success and failure alike, and
must not count toward the project counters a second time. The disconnect hands back
only the connection's estimate.
"""
handler, cache = rate_limiter
api_key = hash_token("sk-ws-project-turns")
user_api_key_dict = UserAPIKeyAuth(
api_key=api_key,
project_id="proj-ws-turns",
project_metadata={
"model_itpm_limit": {"bedrock_mantle/claude-opus": 1000},
"model_otpm_limit": {"bedrock_mantle/claude-opus": 500},
},
)
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=cache,
data={
"model": "bedrock_mantle/claude-opus",
"messages": [{"role": "user", "content": "hello"}],
"max_tokens": 50,
"litellm_call_id": "connection-call-id",
},
call_type=CallTypes.aresponses_websocket.value,
)
stash = get_request_stash()
assert stash is not None and stash.itpm_reserved_tokens > 0 and stash.otpm_reserved_tokens > 0
(itpm_scope,) = stash.itpm_reserved_scopes
(otpm_scope,) = stash.otpm_reserved_scopes
itpm_key = handler.create_rate_limit_keys(itpm_scope[0], itpm_scope[1], "tokens")
otpm_key = handler.create_rate_limit_keys(otpm_scope[0], otpm_scope[1], "tokens")
async def counters() -> tuple[int, int]:
return (
int(await cache.async_get_cache(key=itpm_key) or 0),
int(await cache.async_get_cache(key=otpm_key) or 0),
)
def turn_kwargs(call_id: str, usage: Usage | None = None) -> dict[str, object]:
return {
"litellm_call_id": call_id,
"model": "bedrock_mantle/claude-opus",
"combined_usage_object": usage,
"litellm_params": {"metadata": {"user_api_key": api_key}},
"standard_logging_object": {
"metadata": {"user_api_key_hash": api_key, "user_api_key_project_id": "proj-ws-turns"}
},
}
async def charge_frame() -> None:
await handler.enforce_project_io_token_quota_for_frame(
user_api_key_dict=user_api_key_dict,
requested_model="bedrock_mantle/claude-opus",
estimated_input_tokens=9,
estimated_output_tokens=40,
)
await charge_frame()
reserved_and_one_frame = (stash.itpm_reserved_tokens + 9, stash.otpm_reserved_tokens + 40)
assert await counters() == reserved_and_one_frame
await handler.async_log_success_event(
kwargs={**turn_kwargs("turn-1"), "call_type": CallTypes.aresponses_websocket.value},
response_obj=[
{"type": "response.created", "response": {"id": "resp_1"}},
{
"type": "response.completed",
"response": {"id": "resp_1", "usage": {"input_tokens": 30, "output_tokens": 7, "total_tokens": 37}},
},
],
start_time=datetime.now(),
end_time=datetime.now(),
)
assert await counters() == reserved_and_one_frame
await charge_frame()
await handler.async_log_failure_event(
kwargs=turn_kwargs("turn-2", Usage(prompt_tokens=10, completion_tokens=3, total_tokens=13)),
response_obj=None,
start_time=datetime.now(),
end_time=datetime.now(),
)
assert await counters() == (stash.itpm_reserved_tokens + 18, stash.otpm_reserved_tokens + 80)
await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict=user_api_key_dict)
assert await counters() == (18, 80)
assert stash.reservation_released is True

30
tests/unit/proxy/management_endpoints/test_saml_sso.py
View file

@ -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

137
tests/unit/proxy/response_api_endpoints/test_endpoints.py
View file

@ -7,7 +7,7 @@ import contextlib
import dataclasses
import json
import unittest
from collections.abc import Coroutine, Mapping
from collections.abc import Mapping
from typing import Any, Final, Literal, cast
from unittest.mock import AsyncMock, MagicMock, patch
@ -1642,6 +1642,29 @@ class TestResponsesWSStreamingFirstMessage:
class TestWSSessionCostTracking:
@pytest.mark.asyncio
async def test_router_budget_limiter_skips_aresponses_websocket_call_type(self):
"""
RouterBudgetLimiting.async_log_success_event must not raise when
call_type='_aresponses_websocket', even when standard_logging_object is None.
Per-turn costs are tracked by individual aresponses calls inside the session;
the outer session wrapper fires with result=None.
"""
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
limiter = RouterBudgetLimiting.__new__(RouterBudgetLimiting)
kwargs = {
"call_type": "_aresponses_websocket",
"standard_logging_object": None,
"litellm_params": {"custom_llm_provider": "vertex_ai"},
}
await limiter.async_log_success_event(
kwargs=kwargs,
response_obj=None,
start_time=None,
end_time=None,
)
@pytest.mark.asyncio
async def test_router_budget_limiter_skips_arealtime_call_type(self):
"""Same guard applies to _arealtime WS session wrappers."""
@ -1899,7 +1922,8 @@ class TestResponsesWSFirstFrameModelAuth:
ws.close.assert_not_awaited()
@pytest.mark.asyncio
async def test_endpoint_books_each_failed_turn_as_its_own_failed_request(self):
@pytest.mark.parametrize("provider_rejected", [True, False])
async def test_endpoint_books_a_provider_rejected_connection_as_a_failed_request(self, provider_rejected: bool):
from litellm.proxy.response_api_endpoints.endpoints import (
responses_websocket_endpoint,
)
@ -1911,39 +1935,24 @@ class TestResponsesWSFirstFrameModelAuth:
ws.url = "ws://testserver/v1/responses"
ws.accept = AsyncMock()
ws.receive_text = AsyncMock(
return_value=json.dumps({"type": "response.create", "model": "gpt-4o-mini", "input": "first turn"})
return_value=json.dumps({"type": "response.create", "model": "gpt-4o-mini", "input": []})
)
ws.close = AsyncMock()
connection_metadata = {"user_api_key_alias": "codex", "model_group": "gpt-4o-mini"}
processor = MagicMock()
processor.common_processing_pre_call_logic = AsyncMock(
return_value=(
{
"model": "gpt-4o-mini",
"metadata": connection_metadata,
"litellm_metadata": {},
"litellm_call_id": "connection-call",
"proxy_server_request": {"url": "ws://testserver/v1/responses", "body": {"input": "first turn"}},
},
MagicMock(),
)
return_value=({"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock())
)
turn_proxy_server_request = {"url": "ws://testserver/v1/responses", "body": {"input": "second turn"}}
turn_logging_obj = MagicMock()
turn_logging_obj.litellm_call_id = "turn-2-call"
turn_logging_obj.litellm_params = {"proxy_server_request": turn_proxy_server_request}
failure = litellm.BadRequestError(
message="previous_response_not_found", model="gpt-4o-mini", llm_provider="openai"
message="invalid_encrypted_content", model="gpt-4o-mini", llm_provider="openai"
)
async def relay_with_one_failed_turn():
routed = mock_route_request.await_args.kwargs["data"]
await routed["responses_ws_turn_failure_hook"](turn_logging_obj, failure)
async def fake_llm_call():
return failure if provider_rejected else None
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()
proxy_logging_obj.arelease_max_parallel_requests_on_disconnect = AsyncMock(return_value=None)
user_api_key_dict = MagicMock()
with (
@ -1951,15 +1960,15 @@ class TestResponsesWSFirstFrameModelAuth:
"litellm.proxy.response_api_endpoints.endpoints._enforce_responses_ws_first_frame_model_auth",
new_callable=AsyncMock,
),
patch( # test-quality-ok: the pre-call processor needs a live proxy; how the endpoint books a failed turn is under test
patch( # test-quality-ok: the pre-call processor needs a live proxy; what the endpoint does with the relay's outcome is under test
"litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing",
return_value=processor,
),
patch( # test-quality-ok: routing is the seam that hands the turn failure hook to the relay
patch( # test-quality-ok: routing is the seam that hands back the relay's outcome
"litellm.proxy.route_llm_request.route_request",
new_callable=AsyncMock,
side_effect=lambda **_: relay_with_one_failed_turn(),
) as mock_route_request,
return_value=fake_llm_call(),
),
patch( # test-quality-ok: the failure hook is the proxy's only path to a failed spend log row
"litellm.proxy.proxy_server.proxy_logging_obj",
proxy_logging_obj,
@ -1972,79 +1981,14 @@ class TestResponsesWSFirstFrameModelAuth:
)
ws.close.assert_not_awaited()
if not provider_rejected:
proxy_logging_obj.post_call_failure_hook.assert_not_awaited()
return
proxy_logging_obj.post_call_failure_hook.assert_awaited_once()
booked = proxy_logging_obj.post_call_failure_hook.await_args.kwargs
assert booked["original_exception"] is failure
assert booked["user_api_key_dict"] is user_api_key_dict
turn_request = booked["request_data"]
assert turn_request["litellm_call_id"] == "turn-2-call"
assert turn_request["litellm_logging_obj"] is turn_logging_obj
assert turn_request["model"] == "gpt-4o-mini"
assert turn_request["metadata"] == connection_metadata
assert turn_request["metadata"] is not connection_metadata
assert "input" not in turn_request
assert "first_message" not in turn_request
assert turn_request["proxy_server_request"] == turn_proxy_server_request
proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once_with(user_api_key_dict)
@pytest.mark.asyncio
async def test_endpoint_hands_the_connection_reservation_to_its_first_settling_turn_only(self):
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.response_api_endpoints.endpoints import (
responses_websocket_endpoint,
)
from litellm.proxy.spend_tracking.budget_reservation import claim_budget_reservation
ws = MagicMock()
ws.headers = {}
ws.query_params = {}
ws.scope = {"headers": []}
ws.url = "ws://testserver/v1/responses"
ws.accept = AsyncMock()
ws.receive_text = AsyncMock(
return_value=json.dumps({"type": "response.create", "model": "gpt-4o-mini", "input": "first turn"})
)
ws.close = AsyncMock()
user_api_key_dict = UserAPIKeyAuth(
api_key="hashed-token", budget_reservation={"reserved_cost": 0.5, "entries": []}
)
processor = MagicMock()
processor.common_processing_pre_call_logic = AsyncMock(
return_value=({"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock())
)
proxy_logging_obj = MagicMock()
proxy_logging_obj.arelease_max_parallel_requests_on_disconnect = AsyncMock()
async def relay() -> None:
return None
def route(**_: object) -> Coroutine[object, object, None]:
return relay()
with (
patch( # test-quality-ok: first-frame model auth needs a live router and key table and has its own tests above
"litellm.proxy.response_api_endpoints.endpoints._enforce_responses_ws_first_frame_model_auth",
new_callable=AsyncMock,
),
patch( # test-quality-ok: the pre-call processor needs a live proxy; the reservation it leaves is under test
"litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing",
return_value=processor,
),
patch( # test-quality-ok: routing starts the relay, which needs a live upstream socket
"litellm.proxy.route_llm_request.route_request",
new_callable=AsyncMock,
side_effect=route,
),
patch( # test-quality-ok: the disconnect release reads proxy_logging_obj off proxy_server
"litellm.proxy.proxy_server.proxy_logging_obj",
proxy_logging_obj,
),
):
await responses_websocket_endpoint(websocket=ws, model=None, user_api_key_dict=user_api_key_dict)
reservation = user_api_key_dict.budget_reservation
assert claim_budget_reservation(budget_reservation=reservation, litellm_call_id="turn-1") is reservation
assert claim_budget_reservation(budget_reservation=reservation, litellm_call_id="turn-2") is None
assert booked["request_data"]["model"] == "gpt-4o-mini"
@pytest.mark.asyncio
async def test_endpoint_sends_an_error_frame_when_routing_rejects_the_connection(self):
@ -2111,7 +2055,6 @@ class TestResponsesWSFirstFrameModelAuth:
assert booked["original_exception"] is rejection
assert booked["user_api_key_dict"] is user_api_key_dict
assert booked["request_data"]["model"] == "gpt-4o-mini"
proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once_with(user_api_key_dict)
@pytest.mark.asyncio
async def test_reruns_model_auth_for_first_frame_model(self):

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

@ -1,5 +1,4 @@
import asyncio
import copy
import threading
import time
from collections.abc import Mapping, Sequence
@ -41,13 +40,11 @@ from litellm.proxy.common_utils.user_api_key_cache import (
)
from litellm.proxy.spend_tracking.budget_reservation import (
_get_model_access_group_budget_counters,
claim_budget_reservation,
estimate_request_max_cost,
get_budget_window_start,
invalidate_budget_reservation_counters,
release_budget_reservation,
release_budget_reservation_on_cancel,
share_budget_reservation_across_calls,
reserve_budget_for_request,
)
from litellm.proxy.spend_tracking.input_tokens import (
@ -3723,77 +3720,3 @@ async def test_unreserved_model_access_group_is_charged_alongside_a_reserved_one
assert counter_cache.in_memory_cache.get_cache(
key=model_access_group_spend_counter_key("starter")
) == pytest.approx(4.2)
def test_a_reservation_no_connection_shared_is_left_as_its_call_made_it():
reservation: dict[str, object] = {
"reserved_cost": 0.5,
"entries": [{"counter_key": "spend:key:k", "reserved_cost": 0.5}],
}
made_by_its_call = copy.deepcopy(reservation)
claims = [claim_budget_reservation(budget_reservation=reservation, litellm_call_id=call) for call in ("a", "b")]
assert claims == [reservation, reservation]
assert reservation == made_by_its_call
@pytest.mark.asyncio
async def test_should_charge_both_turns_when_two_calls_settle_one_reservation_at_once(
spend_counter_state,
):
"""A Responses WebSocket connection's turns share its reservation under their own call ids. Only the
first claim keeps it; the other turn adds its cost on top instead of settling the same estimate again,
which would leave the key counter short by the estimate."""
counter_cache, key_cache = spend_counter_state
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
valid_token = UserAPIKeyAuth(
token="key-budget-shared-turns",
spend=0.0,
max_budget=1.0,
)
with patch(
"litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost",
return_value=0.6,
):
reservation = await reserve_budget_for_request(
request_body=_request_body(),
route="/responses",
llm_router=None,
valid_token=valid_token,
team_object=None,
user_object=None,
prisma_client=None,
user_api_key_cache=key_cache,
proxy_logging_obj=proxy_logging_obj,
)
from litellm.proxy.proxy_server import increment_spend_counters
share_budget_reservation_across_calls(reservation)
first_turn = claim_budget_reservation(budget_reservation=reservation, litellm_call_id="turn-1")
second_turn = claim_budget_reservation(budget_reservation=reservation, litellm_call_id="turn-2")
assert first_turn is reservation
assert second_turn is None
assert claim_budget_reservation(budget_reservation=reservation, litellm_call_id="turn-1") is reservation
await asyncio.gather(
increment_spend_counters(
token="key-budget-shared-turns",
team_id="team-without-budget",
user_id=None,
response_cost=0.2,
budget_reservation=first_turn,
),
increment_spend_counters(
token="key-budget-shared-turns",
team_id="team-without-budget",
user_id=None,
response_cost=0.3,
budget_reservation=second_turn,
),
)
assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-budget-shared-turns") == pytest.approx(0.5)
assert counter_cache.in_memory_cache.get_cache(key="spend:team:team-without-budget") == pytest.approx(0.5)

159
tests/unit/proxy/test_proxy_token_counter.py
View file

@ -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

15
tests/unit/proxy/test_redis_auth_cache_flag.py
View file

@ -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

24
tests/unit/repositories/test_unit_of_work.py
View file

@ -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

666
tests/unit/responses/test_responses_websocket_all_providers.py
View file

@ -8,7 +8,7 @@ Tests that:
"""
import json
from collections.abc import Callable, Mapping
from collections.abc import Mapping
from typing import Final
from unittest.mock import MagicMock
@ -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(
@ -1899,34 +1902,6 @@ def _make_streaming(**kwargs):
return ResponsesWebSocketStreaming(**kwargs)
class _TurnLoggingRecorder:
"""Stands in for a socket's connection-level Logging object and keeps every per-turn copy it hands out."""
def __init__(self) -> None:
self.turns: list[MagicMock] = []
self.connection = MagicMock()
self.connection.copy_for_new_call = MagicMock(side_effect=self._copy_for_new_call)
def _copy_for_new_call(self, litellm_call_id: str, start_time: object, request_body: object) -> MagicMock:
from unittest.mock import AsyncMock
turn = MagicMock()
turn.litellm_call_id = litellm_call_id
turn.start_time = start_time
turn.request_body = request_body
turn.model_call_details = {}
turn.dispatch_success_handlers = AsyncMock()
turn.dispatch_failure_handlers = AsyncMock()
turn.response_cost_calculator = MagicMock(return_value=0.0)
self.turns.append(turn)
return turn
def logged_turns(self) -> int:
return sum(
1 for turn in self.turns if turn.dispatch_success_handlers.await_count or turn.dispatch_failure_handlers.await_count
)
class TestNativeWebSocketGuardrailMasking:
"""Exercises the input/output PII masking hooks on ResponsesWebSocketStreaming."""
@ -2522,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,
@ -2578,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,
@ -2641,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,
@ -2704,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,
@ -3252,11 +3231,13 @@ class TestNativeWebSocketEncryptedContentAffinity:
Exception("stop"),
]
)
recorder = _TurnLoggingRecorder()
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,
logging_obj=recorder.connection,
logging_obj=logging_obj,
request_data={"litellm_metadata": dict(_AFFINITY_METADATA)},
custom_llm_provider="openai",
)
@ -3275,7 +3256,7 @@ class TestNativeWebSocketEncryptedContentAffinity:
)
assert completed["response"]["output"][0]["encrypted_content"] == wrapped_content
await asyncio.sleep(0)
logged = recorder.turns[0].dispatch_success_handlers.await_args[0][0]
logged = logging_obj.dispatch_success_handlers.await_args[0][0]
assert logged[0]["response"]["id"] == completed["response"]["id"]
@pytest.mark.asyncio
@ -3299,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,
@ -3370,11 +3352,15 @@ class TestNativeWebSocketEncryptedContentAffinity:
Exception("stop"),
]
)
recorder = _TurnLoggingRecorder()
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(
websocket=websocket,
backend_ws=backend_ws,
logging_obj=recorder.connection,
logging_obj=logging_obj,
request_data={},
authorized_model="gpt-5.6",
custom_llm_provider="openai",
@ -3383,461 +3369,227 @@ class TestNativeWebSocketEncryptedContentAffinity:
await handler.backend_to_client()
await asyncio.sleep(0)
(turn,) = recorder.turns
turn.dispatch_success_handlers.assert_not_awaited()
turn.dispatch_failure_handlers.assert_awaited_once()
exception = turn.dispatch_failure_handlers.await_args[0][0]
logging_obj.dispatch_success_handlers.assert_not_awaited()
logging_obj.dispatch_failure_handlers.assert_awaited_once()
exception = logging_obj.dispatch_failure_handlers.await_args[0][0]
assert exception.status_code == expected_status
assert expected_message in str(exception)
def _completed_turn(response_id: str, input_tokens: int, output_tokens: int) -> list[dict[str, object]]:
return [
{"type": "response.created", "response": {"id": response_id, "status": "in_progress"}},
{"type": "response.output_text.delta", "delta": "hi"},
{
"type": "response.completed",
"response": {
"id": response_id,
"status": "completed",
"output": [],
"usage": {
"input_tokens": input_tokens,
"output_tokens": output_tokens,
"total_tokens": input_tokens + output_tokens,
},
},
},
]
def _response_create(text: str) -> str:
return json.dumps(
{
"type": "response.create",
"model": "gpt-5.4-mini",
"input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": text}]}],
}
)
class _ScriptedBackend:
"""Answers each forwarded response.create with the next scripted turn of upstream events."""
def __init__(self, turns: list[list[dict[str, object]]]) -> None:
@pytest.mark.asyncio
async def test_backend_to_client_logs_each_turn_as_its_own_call(self, monkeypatch):
import asyncio
self.turns = list(turns)
self.events: asyncio.Queue[str] = asyncio.Queue()
self.sent: list[str] = []
async def send(self, frame: str) -> None:
self.sent.append(frame)
if json.loads(frame).get("type") == "response.create" and self.turns:
for event in self.turns.pop(0):
self.events.put_nowait(json.dumps(event))
async def recv(self, decode: bool = False) -> str:
return await self.events.get()
async def close(self) -> None:
return None
async def _until(condition) -> bool:
import asyncio
for _ in range(200):
if condition():
return True
await asyncio.sleep(0)
return False
class _ScriptedClient:
"""Sends each frame only after the previous turn's log was dispatched, then hangs up."""
def __init__(self, frames: list[str], logged_turns: Callable[[], int]) -> None:
self.frames = list(frames)
self.logged_turns = logged_turns
self.sent_text: list[str] = []
self.logged_turns_before_each_frame: list[int] = []
self.logged_turns_before_hang_up: int | None = None
async def receive_text(self) -> str:
sent_so_far = len(self.logged_turns_before_each_frame)
await _until(lambda: self.logged_turns() >= sent_so_far)
if self.frames:
self.logged_turns_before_each_frame.append(self.logged_turns())
return self.frames.pop(0)
self.logged_turns_before_hang_up = self.logged_turns()
raise Exception("client hung up")
async def send_text(self, text: str) -> None:
self.sent_text.append(text)
class TestNativeWebSocketPerTurnLogging:
"""One native Responses socket carries many response.create turns, and each turn is its own logged call."""
@pytest.mark.asyncio
async def test_each_turn_is_logged_as_its_own_call_while_the_socket_stays_open(self):
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
from litellm.responses.utils import ResponsesAPIRequestUtils
recorder = _TurnLoggingRecorder()
client = _ScriptedClient([_response_create(text) for text in ("one", "two", "three")], recorder.logged_turns)
backend = _ScriptedBackend(
[_completed_turn("resp_1", 10, 1), _completed_turn("resp_2", 20, 2), _completed_turn("resp_3", 30, 3)]
)
handler = _make_streaming(
websocket=client, backend_ws=backend, logging_obj=recorder.connection, request_data={}
)
await handler.bidirectional_forward()
assert client.logged_turns_before_each_frame == [0, 1, 2]
assert client.logged_turns_before_hang_up == 3
assert len({turn.litellm_call_id for turn in recorder.turns}) == 3
assert recorder.connection.litellm_call_id not in {turn.litellm_call_id for turn in recorder.turns}
logged_events = [turn.dispatch_success_handlers.await_args[0][0] for turn in recorder.turns]
sent_frames = [json.loads(text) for text in client.sent_text]
ids_the_client_saw = [frame["response"]["id"] for frame in sent_frames if frame["type"] == "response.completed"]
assert [events[-1]["response"]["id"] for events in logged_events] == ids_the_client_saw
assert [
ResponsesAPIRequestUtils.decode_responses_api_response_id(response_id)["response_id"]
for response_id in ids_the_client_saw
] == ["resp_1", "resp_2", "resp_3"]
assert [len(events) for events in logged_events] == [2, 2, 2]
assert [events[-1]["response"]["usage"]["input_tokens"] for events in logged_events] == [10, 20, 30]
assert [turn.model_call_details["messages"] for turn in recorder.turns] == [
[{"role": "user", "content": "one"}],
[{"role": "user", "content": "two"}],
[{"role": "user", "content": "three"}],
]
assert [turn.request_body for turn in recorder.turns] == [
json.loads(_response_create(text)) for text in ("one", "two", "three")
]
for turn in recorder.turns:
turn.dispatch_failure_handlers.assert_not_awaited()
@pytest.mark.parametrize(
("guardrail_callbacks", "logged_text"),
[([], "mail alice@example.com"), ([_FakeWSGuardrail()], "mail <EMAIL_ADDRESS_1>")],
)
@pytest.mark.asyncio
async def test_each_turn_logs_the_frame_the_client_sent_not_the_frame_sent_upstream(
self, guardrail_callbacks: list[_FakeWSGuardrail], logged_text: str
):
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
from litellm.responses.utils import ResponsesAPIRequestUtils
wrapped_previous = ResponsesAPIRequestUtils._build_responses_api_response_id(
custom_llm_provider="openai", model_id="deployment-1", response_id="resp_upstream_prev"
)
client_frame = {
"type": "response.create",
"model": "client-model-group",
"input": "mail alice@example.com",
"previous_response_id": wrapped_previous,
}
recorder = _TurnLoggingRecorder()
client = _ScriptedClient([json.dumps(client_frame)], recorder.logged_turns)
backend = _ScriptedBackend([_completed_turn("resp_1", 10, 1)])
handler = _make_streaming(
websocket=client,
backend_ws=backend,
logging_obj=recorder.connection,
request_data={},
authorized_model="gpt-5.4-mini-deployment",
request_defaults=_deployment_defaults(),
guardrail_callbacks=guardrail_callbacks,
)
await handler.bidirectional_forward()
assert json.loads(backend.sent[0]) == {
"type": "response.create",
"model": "gpt-5.4-mini-deployment",
"input": logged_text,
"previous_response_id": "resp_upstream_prev",
"reasoning": {"effort": "high"},
"service_tier": "priority",
"provider_default": "configured",
}
assert [turn.request_body for turn in recorder.turns] == [{**client_frame, "input": logged_text}]
assert recorder.turns[0].model_call_details["messages"] == [{"role": "user", "content": logged_text}]
@pytest.mark.asyncio
async def test_a_turn_whose_input_is_nested_under_response_logs_its_user_message(self):
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
recorder = _TurnLoggingRecorder()
nested_frame = json.dumps(
{"type": "response.create", "response": {"model": "gpt-5.4-mini", "input": "nested"}}
)
client = _ScriptedClient([nested_frame], recorder.logged_turns)
backend = _ScriptedBackend([_completed_turn("resp_1", 10, 1)])
handler = _make_streaming(
websocket=client, backend_ws=backend, logging_obj=recorder.connection, request_data={}
)
await handler.bidirectional_forward()
assert [turn.model_call_details["messages"] for turn in recorder.turns] == [
[{"role": "user", "content": "nested"}]
]
@pytest.mark.asyncio
async def test_an_overlapped_turn_rejected_at_parse_time_fails_alone_and_the_streaming_turn_keeps_its_usage(
self,
):
"""OpenAI rejects an unparseable response.create with an id-less error while the earlier turn is
still streaming. That error is the second turn's failure, and the first turn's completed event with
its usage stays on the first turn."""
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
recorder = _TurnLoggingRecorder()
handler = _make_streaming(
websocket=MagicMock(),
backend_ws=MagicMock(),
logging_obj=recorder.connection,
request_data={},
authorized_model="gpt-5.4-mini",
custom_llm_provider="openai",
)
created, _delta, completed = _completed_turn("resp_1", 300, 23)
rejection = {
"type": "error",
"status": None,
"error": {"type": "invalid_request_error", "code": None, "message": "Invalid 'max_output_tokens'"},
}
handler._record_client_frame(_response_create("first"))
handler._record_client_frame(_response_create("second"))
for event in (created, rejection, completed):
handler._record_backend_event(json.dumps(event))
assert await _until(lambda: recorder.logged_turns() == 2)
first, second = recorder.turns
first.dispatch_failure_handlers.assert_not_awaited()
first.dispatch_success_handlers.assert_awaited_once()
assert [event["type"] for event in first.dispatch_success_handlers.await_args[0][0]] == [
"response.created",
"response.completed",
]
second.dispatch_success_handlers.assert_not_awaited()
second.dispatch_failure_handlers.assert_awaited_once()
assert "max_output_tokens" in str(second.dispatch_failure_handlers.await_args[0][0])
assert second.record_partial_usage_for_failure.call_args[0][0].total_tokens == 0
@pytest.mark.asyncio
async def test_a_rejected_turn_answered_in_order_fails_and_the_next_turn_keeps_its_response(self):
"""OpenAI answers a response.create it parses but rejects (previous_response_not_found) in order,
before it starts the next frame's response. With both turns still waiting, the id-less error is the
older turn's failure and the next response.created belongs to the newer turn."""
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
recorder = _TurnLoggingRecorder()
handler = _make_streaming(
websocket=MagicMock(),
backend_ws=MagicMock(),
logging_obj=recorder.connection,
request_data={},
authorized_model="gpt-5.4-mini",
custom_llm_provider="openai",
)
created, _delta, completed = _completed_turn("resp_2", 40, 4)
rejection = {
"type": "error",
"error": {
"type": "invalid_request_error",
"code": "previous_response_not_found",
"message": "Previous response with id 'resp_missing' not found.",
},
}
handler._record_client_frame(_response_create("first"))
handler._record_client_frame(_response_create("second"))
for event in (rejection, created, completed):
handler._record_backend_event(json.dumps(event))
assert await _until(lambda: recorder.logged_turns() == 2)
first, second = recorder.turns
first.dispatch_success_handlers.assert_not_awaited()
assert "resp_missing" in str(first.dispatch_failure_handlers.await_args[0][0])
second.dispatch_failure_handlers.assert_not_awaited()
assert [event["type"] for event in second.dispatch_success_handlers.await_args[0][0]] == [
"response.created",
"response.completed",
]
@pytest.mark.asyncio
async def test_a_failed_turn_is_booked_alone_and_leaves_the_other_turns_successful(self):
import time
from unittest.mock import AsyncMock
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
recorder = _TurnLoggingRecorder()
on_turn_failure = AsyncMock()
client = _ScriptedClient([_response_create(text) for text in ("one", "two", "three")], recorder.logged_turns)
backend = _ScriptedBackend(
[
_completed_turn("resp_1", 10, 1),
[
{
"type": "error",
"status": 400,
"error": {
"type": "invalid_request_error",
"code": "previous_response_not_found",
"message": "Previous response with id 'resp_x' not found.",
},
}
],
_completed_turn("resp_3", 30, 3),
]
)
handler = _make_streaming(
websocket=client,
backend_ws=backend,
logging_obj=recorder.connection,
request_data={},
authorized_model="gpt-5.4-mini",
custom_llm_provider="openai",
on_turn_failure=on_turn_failure,
)
await handler.bidirectional_forward()
first, failed, third = recorder.turns
for successful in (first, third):
successful.dispatch_success_handlers.assert_awaited_once()
successful.dispatch_failure_handlers.assert_not_awaited()
failed.dispatch_success_handlers.assert_not_awaited()
failed.dispatch_failure_handlers.assert_awaited_once()
on_turn_failure.assert_awaited_once()
booked_turn, booked_exception = on_turn_failure.await_args[0]
assert booked_turn is failed
assert booked_exception is failed.dispatch_failure_handlers.await_args[0][0]
assert booked_exception.status_code == 400
assert "not found" in str(booked_exception)
@pytest.mark.asyncio
async def test_a_failed_turn_bills_the_usage_its_failure_event_reports(self):
from datetime import datetime
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import CallTypes
logged: list[dict[str, object]] = []
pre_calls: list[tuple[str, object]] = []
successes: list[tuple[str, list[object]]] = []
failures: list[tuple[str, Exception]] = []
partial_usage: list[tuple[str, int, float]] = []
class _SpendRows(CustomLogger):
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
logged.append(kwargs["standard_logging_object"])
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"]))
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
logged.append(kwargs["standard_logging_object"])
reservations: dict[str, tuple[object, object]] = {}
spend_rows = _SpendRows()
connection = Logging(
model="gpt-4o",
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=datetime.now(), # noqa: DTZ005 # logging pipeline uses naive datetimes
litellm_call_id="ws-connection",
function_id="ws-connection",
dynamic_async_success_callbacks=[spend_rows],
dynamic_async_failure_callbacks=[spend_rows],
start_time=time.time(),
litellm_call_id="connection-call",
function_id="connection-call",
)
connection.update_environment_variables(
litellm_params={"metadata": {}, "proxy_server_request": {}},
optional_params={},
model="gpt-4o",
custom_llm_provider="openai",
)
failed_usage = {"input_tokens": 7, "output_tokens": 2, "total_tokens": 9}
client = _ScriptedClient([_response_create(text) for text in ("one", "two")], lambda: len(logged))
backend = _ScriptedBackend(
[
_completed_turn("resp_1", 10, 5),
[
{"type": "response.created", "response": {"id": "resp_2", "status": "in_progress"}},
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.failed",
"type": "response.completed",
"response": {
"id": "resp_2",
"status": "failed",
"error": {"code": "server_error", "message": "upstream blew up"},
"usage": failed_usage,
"id": "resp_1",
"status": "completed",
"output": [],
"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
},
},
],
}
),
json.dumps({"type": "error", "error": {"type": "invalid_request_error", "message": "bad turn"}}),
Exception("stop"),
]
)
handler = _make_streaming(
websocket=client,
backend_ws=backend,
logging_obj=connection,
request_data={},
authorized_model="gpt-4o",
custom_llm_provider="openai",
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.bidirectional_forward()
await _until(lambda: len(logged) == 2)
await handler.backend_to_client()
await asyncio.sleep(0)
completed, failed = logged
assert (completed["status"], failed["status"]) == ("success", "failure")
assert (completed["prompt_tokens"], completed["completion_tokens"]) == (10, 5)
assert (failed["prompt_tokens"], failed["completion_tokens"]) == (7, 2)
http_failed_cost = litellm.completion_cost(
completion_response=ResponsesAPIResponse(
id="resp_2",
created_at=0,
model="gpt-4o",
object="response",
output=[],
status="failed",
usage=ResponseAPIUsage(**failed_usage),
),
model="gpt-4o",
call_type=CallTypes.aresponses.value,
custom_llm_provider="openai",
)
assert http_failed_cost > 0
assert failed["response_cost"] == pytest.approx(http_failed_cost)
(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_a_turn_still_in_flight_when_the_socket_closes_is_logged_with_what_it_received(self):
async def test_bidirectional_forward_returns_the_provider_failure(self):
import asyncio
from unittest.mock import AsyncMock
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
recorder = _TurnLoggingRecorder()
client = _ScriptedClient([_response_create(text) for text in ("one", "two", "three")], recorder.logged_turns)
backend = _ScriptedBackend(
[
_completed_turn("resp_1", 10, 1),
[{"type": "response.created", "response": {"id": "resp_2", "status": "in_progress"}}],
]
)
backend_drained = asyncio.Event()
backend_events = [
json.dumps({"type": "response.created", "response": {"id": "resp_1", "status": "in_progress"}}),
json.dumps(
{
"type": "error",
"status": 400,
"error": {
"type": "invalid_request_error",
"code": "invalid_encrypted_content",
"message": "could not be verified",
},
}
),
]
async def recv(decode=False):
if backend_events:
return backend_events.pop(0)
backend_drained.set()
raise Exception("stop")
async def receive_text():
await backend_drained.wait()
raise Exception("client gone")
websocket = MagicMock()
websocket.send_text = AsyncMock()
websocket.receive_text = receive_text
backend_ws = MagicMock()
backend_ws.recv = recv
backend_ws.send = AsyncMock()
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(
websocket=client, backend_ws=backend, logging_obj=recorder.connection, request_data={}
websocket=websocket,
backend_ws=backend_ws,
logging_obj=logging_obj,
request_data={},
authorized_model="gpt-5.6",
custom_llm_provider="openai",
)
await handler.bidirectional_forward()
await _until(lambda: recorder.turns[1].dispatch_success_handlers.await_count == 1)
failure = await handler.bidirectional_forward()
first, in_flight, unanswered = recorder.turns
first.dispatch_success_handlers.assert_awaited_once()
in_flight.dispatch_success_handlers.assert_awaited_once()
assert [event["type"] for event in in_flight.dispatch_success_handlers.await_args[0][0]] == ["response.created"]
unanswered.dispatch_success_handlers.assert_not_awaited()
unanswered.dispatch_failure_handlers.assert_not_awaited()
assert isinstance(failure, Exception)
assert failure.status_code == 400
assert "could not be verified" in str(failure)
@pytest.mark.asyncio
async def test_bidirectional_forward_returns_none_after_a_completed_turn(self):
import asyncio
from unittest.mock import AsyncMock
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
backend_drained = asyncio.Event()
backend_events = [
json.dumps(
{
"type": "response.completed",
"response": {
"id": "resp_1",
"status": "completed",
"output": [],
"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
},
}
),
]
async def recv(decode=False):
if backend_events:
return backend_events.pop(0)
backend_drained.set()
raise Exception("stop")
async def receive_text():
await backend_drained.wait()
raise Exception("client gone")
websocket = MagicMock()
websocket.send_text = AsyncMock()
websocket.receive_text = receive_text
backend_ws = MagicMock()
backend_ws.recv = recv
backend_ws.send = AsyncMock()
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,
backend_ws=backend_ws,
logging_obj=logging_obj,
request_data={},
authorized_model="gpt-5.6",
custom_llm_provider="openai",
)
assert await handler.bidirectional_forward() is None

4
tests/unit/router_strategy/test_budget_limiter.py
View file

@ -59,9 +59,7 @@ async def _log_success(limiter: RouterBudgetLimiting, kwargs: dict[str, object])
@pytest.mark.asyncio
@pytest.mark.parametrize(
"call_type", ["aresponses", "_aresponses_websocket", "anthropic_messages", "aembedding", "arerank"]
)
@pytest.mark.parametrize("call_type", ["aresponses", "anthropic_messages", "aembedding", "arerank"])
async def test_provider_spend_tracked_when_litellm_params_omits_provider(disable_budget_sync, call_type):
"""Non-chat surfaces carry the provider only on the standard logging payload."""
limiter = RouterBudgetLimiting(

20
tests/unit/test_cost_calculator.py
View file

@ -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

2
tests/unit/test_count_tokens_public_api.py
View file

@ -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():

46
tests/unit/test_router/test_router_helper_utils.py
View file

@ -482,52 +482,6 @@ async def test_deployment_callback_on_success(sync_mode):
assert tpm_key is not None
@pytest.mark.asyncio
async def test_deployment_callback_on_success_counts_a_responses_websocket_turn():
"""Each native Responses WebSocket turn logs under its own call, so the deployment's
tpm and rpm counters move once per turn, while a realtime session still counts nothing."""
import time
model_list = [
{
"model_name": "gpt-5-mini",
"litellm_params": {"model": "gpt-5-mini", "api_key": "sk-fake", "rpm": 100},
"model_info": {"id": "ws-deployment"},
}
]
router = Router(model_list=model_list)
standard_logging_payload = create_standard_logging_payload()
standard_logging_payload["total_tokens"] = 17
def turn_kwargs(call_type: str) -> dict:
return {
"call_type": call_type,
"litellm_params": {
"metadata": {"model_group": "gpt-5-mini"},
"model_info": {"id": "ws-deployment"},
},
"standard_logging_object": standard_logging_payload,
}
tpm_key = await router.deployment_callback_on_success(
kwargs=turn_kwargs("_aresponses_websocket"),
completion_response=None,
start_time=time.time(),
end_time=time.time(),
)
assert tpm_key is not None
assert await router.cache.async_get_cache(key=tpm_key) == 17
realtime_key = await router.deployment_callback_on_success(
kwargs=turn_kwargs("_arealtime"),
completion_response=None,
start_time=time.time(),
end_time=time.time(),
)
assert realtime_key is None
assert await router.cache.async_get_cache(key=tpm_key) == 17
@pytest.mark.asyncio
async def test_deployment_callback_on_success_tracks_tpm_for_io_deployment():
"""