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

merge(main): integrate latest main

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Moe Khalil 2026-10-31 20:29:09 +00:00
commit 327a7fdfc8
77 changed files with 24928 additions and 17792 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
.github/PULL_REQUEST_TEMPLATE/general.md vendored
View file

@ -4,9 +4,14 @@
Drop every section you have nothing to put in, heading included: a bare "## Relevant issues" or
"## Affected release" with nothing under it must not appear in the final description -->
<!-- Two or three plain sentences a teammate can read in 10 seconds, written the way you'd explain it in Slack
## Handwritten 1-3 sentence description
<!-- The PR author rewrites this in their own words before asking for review. Being able to sum up the PR in a few
plain sentences proves the author understands what it does, and human-written technical English is far easier
for a reviewer to parse than AI-written English
An AI agent may draft a starting point here, but the author replaces it with their own handwritten version
Say what the PR does for users and why, then anything that keeps it working over time (a sync job, a config someone owns)
No headings, bullets, tables, or code names here. Everything else goes inside the Details block below -->
No bullets, tables, or code names here. Everything else goes inside the Details block below -->
<details>
<summary>Details</summary>

2
.github/workflows/create-rc-branch.yml vendored
View file

@ -2,7 +2,7 @@ name: Create RC Branch
on:
schedule:
- cron: "0 3 * * 5"
- cron: "0 3 * * 1"
timezone: "America/Los_Angeles"
workflow_dispatch:

3
litellm/constants.py
View file

@ -432,6 +432,7 @@ WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123
DEEPGRAM_DEFAULT_API_BASE: Final = "https://api.deepgram.com/v1"
NADIR_DEFAULT_API_BASE: Final = "https://api.getnadir.com/v1"
DEEPGRAM_LISTEN_DEFAULT_MODEL: Final = "nova-3"
DEEPGRAM_SPEAK_DEFAULT_MODEL: Final = "aura-asteria-en"
BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_opening_message"
BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed"
@ -675,7 +676,7 @@ request_timeout: float = (
DEFAULT_REQUEST_TIMEOUT_SECONDS if _request_timeout_from_env is None else _request_timeout_from_env
)
request_timeout_explicitly_set: bool = _request_timeout_from_env is not None
DEFAULT_A2A_AGENT_TIMEOUT: Final[float] = get_env_float("DEFAULT_A2A_AGENT_TIMEOUT", 6000.0) # 10 minutes
DEFAULT_A2A_AGENT_TIMEOUT: Final[float] = get_env_float("DEFAULT_A2A_AGENT_TIMEOUT", 6000.0)
A2A_CALL_TYPES: Final = frozenset({"asend_message", "send_message", "asend_message_streaming"})
AGENT_KILL_SWITCH_TIMEOUT_SECONDS: Final = 10.0
AGENT_KILL_SWITCH_RESPONSE_BODY_MAX_CHARS: Final = 2000

50
litellm/integrations/anthropic_cache_control_hook.py
View file

@ -13,7 +13,7 @@ import copy
import itertools
import os
import re
from collections.abc import Callable, Container, Iterable, Mapping, Sequence
from collections.abc import Callable, Container, Iterable, Iterator, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypeVar, cast
from urllib.parse import urlparse
@ -445,22 +445,41 @@ def _as_object_list(value: object | None) -> list[object] | None:
return _validated_object_list(value)
def _tool_call_carries_cache_breakpoint(tool_call: object, message: object, answered_tool_ids: Container[str]) -> bool:
def _server_tool_call_id(tool_call: object) -> str | None:
tool_call_id: Final = _attribute_or_key(tool_call, "id")
if isinstance(tool_call_id, str) and tool_call_id.startswith(ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX):
return tool_call_id
return None
def _tool_call_carries_cache_breakpoint(tool_call: object, client_tool_call_ids: Container[str]) -> bool:
if _attribute_or_key(tool_call, "cache_control") is None:
return False
server_tool_call_id: Final = _server_tool_call_id(tool_call)
return server_tool_call_id is None or server_tool_call_id in client_tool_call_ids
return not tool_call_is_rebuilt_as_server_tool_use(
_attribute_or_key(tool_call, "id"), _attribute_or_key(message, "provider_specific_fields"), answered_tool_ids
def _server_tool_call_ids_without_a_kept_result(message: object) -> Iterator[str]:
provider_specific_fields: Final = _attribute_or_key(message, "provider_specific_fields")
tool_calls: Final = _as_object_list(_attribute_or_key(message, "tool_calls")) or ()
return (
tool_call_id
for tool_call in tool_calls
if (tool_call_id := _server_tool_call_id(tool_call)) is not None
and not tool_call_is_rebuilt_as_server_tool_use(tool_call_id, provider_specific_fields)
)
def _tool_call_ids_answered_by_client_when_a_server_tool_call_is_marked(
def _client_tool_call_ids_when_a_server_tool_call_is_marked(
messages: Sequence[object], model: str | None = None, custom_llm_provider: str | None = None
) -> frozenset[str]:
if not _marks_a_server_tool_call(messages):
return frozenset()
if model is not None and not _builds_anthropic_messages(model, custom_llm_provider):
return tool_call_ids_answered_in(messages)
server_calls_sent_as_client_calls: Final = itertools.chain.from_iterable(
map(_server_tool_call_ids_without_a_kept_result, messages)
)
return tool_call_ids_answered_in(messages) | frozenset(server_calls_sent_as_client_calls)
return tool_call_ids_answered_by_client(messages, _keeps_mid_conversation_system(model, custom_llm_provider))
@ -517,12 +536,7 @@ def _model_info_flags_mid_conversation_system(model: str, custom_llm_provider: s
def _is_marked_server_tool_call(tool_call: object) -> bool:
tool_call_id: Final = _attribute_or_key(tool_call, "id")
return (
isinstance(tool_call_id, str)
and tool_call_id.startswith(ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX)
and _attribute_or_key(tool_call, "cache_control") is not None
)
return _server_tool_call_id(tool_call) is not None and _attribute_or_key(tool_call, "cache_control") is not None
def _tool_carries_cache_breakpoint(tool: object) -> bool:
@ -668,7 +682,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
reserved_blocks: Final = AnthropicCacheControlHook._blocks_reserved_outside_messages(
remaining_points, external_breakpoints, openai_dialect
)
answered_tool_ids: Final = _tool_call_ids_answered_by_client_when_a_server_tool_call_is_marked(
answered_tool_ids: Final = _client_tool_call_ids_when_a_server_tool_call_is_marked(
processed_messages, model, custom_llm_provider
)
breakpoints_before: Final = AnthropicCacheControlHook.count_request_cache_breakpoints(
@ -757,7 +771,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
)
message_list: Final = tuple(messages)
client_answers: Final = (
_tool_call_ids_answered_by_client_when_a_server_tool_call_is_marked(message_list)
_client_tool_call_ids_when_a_server_tool_call_is_marked(message_list)
if answered_tool_ids is None
else answered_tool_ids
)
@ -849,7 +863,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
so earlier points win when slots are scarce.
"""
client_answers: Final = (
_tool_call_ids_answered_by_client_when_a_server_tool_call_is_marked(messages)
_client_tool_call_ids_when_a_server_tool_call_is_marked(messages)
if answered_tool_ids is None
else answered_tool_ids
)
@ -939,11 +953,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
content_count: Final = sum(1 for block in content if _carries_cache_breakpoint(block)) if content else 0
tool_calls: Final = _as_object_list(_attribute_or_key(message, "tool_calls"))
tool_call_count: Final = (
sum(
1
for tool_call in tool_calls
if _tool_call_carries_cache_breakpoint(tool_call, message, answered_tool_ids)
)
sum(1 for tool_call in tool_calls if _tool_call_carries_cache_breakpoint(tool_call, answered_tool_ids))
if tool_calls
else 0
)

9
litellm/litellm_core_utils/litellm_logging.py
View file

@ -3273,7 +3273,14 @@ class Logging(LiteLLMLoggingBaseClass):
start_time = self.start_time
if end_time is None:
end_time = datetime.datetime.now()
if self.completion_start_time is None:
# A request re-dispatched to a fallback reuses this logging object but runs in
# its own window, so a first-token time from the previous attempt cannot belong
# to it. Keeping it makes the attempt report a negative TTFT.
first_token_seconds: Final = _timestamp_seconds(self.completion_start_time)
attempt_start_seconds: Final = _timestamp_seconds(start_time)
if first_token_seconds is None or (
attempt_start_seconds is not None and first_token_seconds < attempt_start_seconds
):
self.completion_start_time = end_time
self.model_call_details["completion_start_time"] = self.completion_start_time

32
litellm/litellm_core_utils/streaming_handler.py
View file

@ -349,11 +349,14 @@ class CustomStreamWrapper:
make_call: Callable | None = None,
_response_headers: dict | httpx.Headers | None = None,
count_prompt_tokens: Callable[[], int] | None = None,
is_mock: bool = False,
):
self.model = model
self.make_call = make_call
self.count_prompt_tokens = count_prompt_tokens
self.custom_llm_provider = custom_llm_provider
self.is_mock = is_mock
self.uses_openai_chat_stream = False
self.logging_obj: LiteLLMLoggingObject = logging_obj
self.completion_stream = completion_stream
self.sent_first_chunk = False
@ -376,6 +379,7 @@ class CustomStreamWrapper:
self._upstream_provider: str | None = None
self.received_finish_reason: str | None = None
self.intermittent_finish_reason: str | None = None # finish reasons that show up mid-stream
self.received_any_finish_reason = False # a later chunk can reset intermittent_finish_reason
self.special_tokens = [
"<|assistant|>",
"<|system|>",
@ -1597,6 +1601,7 @@ class CustomStreamWrapper:
if response_obj["is_finished"]:
self.received_finish_reason = response_obj["finish_reason"]
else: # openai / azure chat model
self.uses_openai_chat_stream = True
if isinstance(chunk, ChatCompletionChunk):
self._openai_handler_receives_chat_chunks = True
if any(str(choice.finish_reason) == "abort" for choice in chunk.choices or ()):
@ -1617,6 +1622,7 @@ class CustomStreamWrapper:
return _ProviderChunkEarlyReturn(None)
completion_obj["content"] = response_obj["text"]
self.intermittent_finish_reason = response_obj.get("finish_reason", None)
self.received_any_finish_reason |= response_obj["finish_reason"] is not None
if response_obj["is_finished"]:
if response_obj["finish_reason"] == "error":
raise Exception(
@ -1996,6 +2002,22 @@ class CustomStreamWrapper:
model_response.choices[0].delta.content = close_tag
return model_response
def _should_fail_on_clean_eof(self) -> bool:
return (
not self.is_mock
and self.uses_openai_chat_stream
and self.received_finish_reason is None
and self.intermittent_finish_reason is None
and not self.received_any_finish_reason
)
@staticmethod
def _stream_eof_error() -> httpx.RemoteProtocolError:
return httpx.RemoteProtocolError(
"Stream ended before a finish_reason was received",
request=httpx.Request(method="POST", url="https://docs.litellm.ai/docs"),
)
def _final_chunk_for_logging(self, processed_chunk: "ModelResponseStream") -> "ModelResponseStream | None":
if self.received_finish_reason is not None or self.intermittent_finish_reason is not None:
return processed_chunk
@ -2181,6 +2203,13 @@ class CustomStreamWrapper:
return response
except StopIteration:
if self._should_fail_on_clean_eof():
error: Final = self._stream_eof_error()
eof_traceback: Final = traceback.format_exc()
self._record_partial_usage_for_failure()
threading.Thread(target=self.logging_obj.failure_handler, args=(error, eof_traceback)).start()
self.close()
self._handle_stream_fallback_error(error)
if self.sent_last_chunk is True:
complete_streaming_response: Final = self._build_complete_streaming_response()
response = self._new_usage_chunk()
@ -2384,6 +2413,9 @@ class CustomStreamWrapper:
continue
return usage_free_chunk
except (StopAsyncIteration, StopIteration):
if self._should_fail_on_clean_eof():
await self.aclose()
self._log_stream_failure_and_raise(self._stream_eof_error())
return await self._finalize_completed_stream(cache_hit=cache_hit)
except httpx.TimeoutException as e: # if httpx read timeout error occues
traceback_exception = traceback.format_exc()

9
litellm/llms/a2a/common_utils.py
View file

@ -160,8 +160,11 @@ async def resolve_a2a_hop_auth_header(
litellm_params: Mapping[str, object],
custom_llm_provider: object,
resolve_entra_header: AgentAuthHeaderResolver = resolve_azure_ai_agent_auth_header,
agent_extra_headers: Mapping[str, str] | None = None,
) -> Mapping[str, str] | None:
"""Entra credentials authenticate the A2A hop only; a completion-bridge agent hands them to the model provider it bridges to."""
if not a2a_hop_uses_entra(litellm_params, custom_llm_provider):
if a2a_hop_uses_entra(litellm_params, custom_llm_provider):
return await resolve_entra_header(litellm_params)
if custom_llm_provider or any(name.lower() == "authorization" for name in (agent_extra_headers or {})):
return None
return await resolve_entra_header(litellm_params)
api_key: Final = litellm_params.get("api_key")
return {"Authorization": f"Bearer {api_key}"} if isinstance(api_key, str) and api_key else None

3
litellm/llms/anthropic/messages/adapters/handler.py
View file

@ -118,9 +118,10 @@ def _extract_proxy_litellm_metadata(
def _requested_router(kwargs: Mapping[str, object]) -> "Router | None":
from litellm.router import Router
from litellm.router_python_core import PythonCore
requested: Final = kwargs.get("litellm_router")
return requested if isinstance(requested, Router) else None
return requested if isinstance(requested, (Router, PythonCore)) else None # pyright: ignore[reportReturnType] # the router passes its core, which carries Router's members
def request_safeguards_evaluator(

27
litellm/llms/bedrock/chat/agentcore/transformation.py
View file

@ -544,6 +544,17 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
return AgentCoreParsedResponse(content=content, usage=usage_data, final_message=final_message)
@staticmethod
def _message_stop_chunk(model: str) -> ModelResponseStream:
chunk: Final = ModelResponseStream(
id=f"chatcmpl-{uuid.uuid4()}",
created=0,
model=model,
object="chat.completion.chunk",
)
chunk.choices = [StreamingChoices(finish_reason="stop", index=0, delta=Delta())]
return chunk
def _stream_agentcore_response_sync(
self,
response: httpx.Response,
@ -553,6 +564,8 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
Internal sync generator that parses SSE and yields ModelResponse chunks.
"""
buffer = ""
sent_finish = False # rebind-ok: flips when a finish chunk goes out
message_stopped = False # rebind-ok: flips on messageStop, the only end some runtimes send
for text_chunk in response.iter_text():
buffer += text_chunk
@ -624,7 +637,9 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
total_tokens=usage_data.get("totalTokens", 0),
),
)
sent_finish = True
yield chunk
message_stopped |= "messageStop" in event_payload
# Process final message
if "message" in data_obj and isinstance(data_obj["message"], dict):
@ -641,12 +656,16 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
delta=Delta(),
)
]
sent_finish = True
yield chunk
except json.JSONDecodeError:
verbose_logger.debug("Skipping non-JSON SSE line: %s", line[:100])
continue
if message_stopped and not sent_finish:
yield self._message_stop_chunk(model)
def get_sync_custom_stream_wrapper(
self,
model: str,
@ -775,6 +794,8 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
Internal async generator that parses SSE and yields ModelResponse chunks.
"""
buffer = ""
sent_finish = False # rebind-ok: flips when a finish chunk goes out
message_stopped = False # rebind-ok: flips on messageStop, the only end some runtimes send
async for text_chunk in response.aiter_text():
buffer += text_chunk
@ -846,7 +867,9 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
total_tokens=usage_data.get("totalTokens", 0),
),
)
sent_finish = True
yield chunk
message_stopped |= "messageStop" in event_payload
# Process final message
if "message" in data_obj and isinstance(data_obj["message"], dict):
@ -863,12 +886,16 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
delta=Delta(),
)
]
sent_finish = True
yield chunk
except json.JSONDecodeError:
verbose_logger.debug("Skipping non-JSON SSE line: %s", line[:100])
continue
if message_stopped and not sent_finish:
yield self._message_stop_chunk(model)
async def get_async_custom_stream_wrapper(
self,
model: str,

77
litellm/llms/deepgram/common_utils.py
View file

@ -1,16 +1,21 @@
import json
import math
import re
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from urllib.parse import parse_qs, urlparse
from urllib.parse import parse_qs, urlparse, urlsplit, urlunsplit
import httpx
import litellm
from litellm.constants import DEEPGRAM_DEFAULT_API_BASE, DEEPGRAM_LISTEN_DEFAULT_MODEL
from litellm.constants import DEEPGRAM_DEFAULT_API_BASE, DEEPGRAM_LISTEN_DEFAULT_MODEL, DEEPGRAM_SPEAK_DEFAULT_MODEL
from litellm.litellm_core_utils.url_utils import join_api_base
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.passthrough_endpoints.pass_through_endpoints import WebSocketSessionUsage
from litellm.types.utils import LlmProviders
from litellm.utils import ModelInfoLookupMiss, lookup_model_info_for_registration
_WEBSOCKET_SCHEMES: Final = MappingProxyType({"https": "wss", "http": "ws", "wss": "wss", "ws": "ws"})
DEEPGRAM_LISTEN_CALLBACK_PARAMS: Final = frozenset({"callback", "callback_method"})
@ -199,3 +204,71 @@ def deepgram_listen_transcript(websocket_messages: Sequence[Mapping[str, object]
if frame.get("type") == "Results"
if (transcript := _final_transcript(frame)) is not None
)
_AURA_GENERATION: Final = re.compile(r"aura(?:-\d+)?(?=-[a-z])")
def deepgram_speak_websocket_target(api_base: str | None, query_string: str) -> str:
speak_url: Final = urlsplit(join_api_base(api_base or DEEPGRAM_DEFAULT_API_BASE, "speak"))
scheme: Final = _WEBSOCKET_SCHEMES.get(speak_url.scheme, speak_url.scheme)
params: Final = _first_occurrences(query_string)
query: Final = params if params.get("model") else params.remove("model").add("model", DEEPGRAM_SPEAK_DEFAULT_MODEL)
return urlunsplit((scheme, speak_url.netloc, speak_url.path, str(query), ""))
def deepgram_speak_model(upstream_url: str) -> str:
models: Final = parse_qs(urlparse(upstream_url).query).get("model")
return models[0] if models else DEEPGRAM_SPEAK_DEFAULT_MODEL
def deepgram_speak_requested_model(query_string: str) -> str:
"""Read off the upstream target, so the voice a key is authorized for is the voice Deepgram receives"""
return deepgram_speak_model(deepgram_speak_websocket_target(None, query_string))
def deepgram_speak_registry_key(upstream_url: str) -> str:
"""Deepgram bills every voice of an Aura generation at one rate, so ``aura-2-thalia-en`` is priced at the
``deepgram/aura-2`` row"""
model: Final = deepgram_speak_model(upstream_url)
generation: Final = _AURA_GENERATION.match(model)
return f"{LlmProviders.DEEPGRAM.value}/{generation.group() if generation else model}"
def deepgram_speak_character_rate(upstream_url: str) -> float | None:
"""Only an exact registry hit counts, so a voice is never billed at a neighbouring row's rate"""
registry_key: Final = deepgram_speak_registry_key(upstream_url)
model_info: Final = lookup_model_info_for_registration(registry_key, LlmProviders.DEEPGRAM.value)
if isinstance(model_info, ModelInfoLookupMiss) or model_info["key"] != registry_key:
return None
return model_info.get("input_cost_per_character")
@dataclass(frozen=True, slots=True)
class DeepgramSpeakMeter:
"""Holds the rate checked when the socket opened, so a cost-map reload mid-session can't change or drop the price"""
upstream_url: str
character_rate: float
def measure(self, client_text: str) -> int:
"""Counts a Speak message's text as it is sent. Deepgram also speaks, and bills, unflushed text on its own (a
long buffer, Close, a close frame), so counting only flushed text would let a client get audio free, while
counting on send can only over-bill text Deepgram discards, such as a Clear or an oversized Speak. Never
raises, so a malformed frame can't end the session"""
try:
message: Final[object] = json.loads(client_text)
except (ValueError, RecursionError):
return 0
match message:
case {"type": "Speak", "text": str() as text}:
return len(text)
case _:
return 0
def price(self, total: int) -> WebSocketSessionUsage:
return WebSocketSessionUsage(
model=deepgram_speak_model(self.upstream_url),
custom_llm_provider=LlmProviders.DEEPGRAM.value,
response_cost=total * self.character_rate,
)

5
litellm/llms/langgraph/chat/sse_iterator.py
View file

@ -41,8 +41,9 @@ class LangGraphSSEStreamIterator:
return self
def __aiter__(self):
"""Initialize async iteration."""
self.async_line_iterator = aiter_sse_lines(self.response.aiter_text())
"""Initialize async iteration once; CustomStreamWrapper re-enters `async for` on every chunk."""
if self.async_line_iterator is None:
self.async_line_iterator = aiter_sse_lines(self.response.aiter_text())
return self
def _parse_sse_line(self, line: str) -> ModelResponseStream | None:

20
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -40,7 +40,11 @@ from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, TypeVar, Uni
from openai.types.responses.response_output_refusal import ResponseOutputRefusal
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
from typing_extensions import (
ReadOnly,
TypedDict,
TypeIs, # noqa: TID251 # TypeIs reaches typing only on 3.13
)
from litellm._logging import verbose_proxy_logger
from litellm.completion_extras.litellm_responses_transformation.transformation import (
@ -100,6 +104,7 @@ from litellm.types.llms.openai import (
OutputTextDeltaEvent,
OutputTextDoneEvent,
ResponseAPIUsage,
ResponseCompactionItem,
ResponseCompletedEvent,
ResponsesAPIOptionalRequestParams,
ResponsesAPIResponse,
@ -712,6 +717,15 @@ def _next_stream_sequence_number(responses_so_far: Sequence[object] | None) -> i
return next((n + 1 for n in sequence_numbers if isinstance(n, int)), 0)
def _is_responses_input(
input_data: object,
) -> TypeIs["str | ResponseInputParam"]: # guard-ok: str, or a list of Mapping items
if isinstance(input_data, str):
return True
items: Final = cast(list[object], input_data) # cast-ok: isinstance narrows to list[Unknown] only
return isinstance(input_data, list) and all(isinstance(item, (Mapping, ResponseCompactionItem)) for item in items)
class OpenAIResponsesHandler(BaseTranslation):
"""
Handler for processing OpenAI Responses API with guardrails.
@ -733,8 +747,8 @@ class OpenAIResponsesHandler(BaseTranslation):
Transforms `input` (string or ResponseInputParam) and optional
`instructions` into chat completion messages.
"""
input_data: Final = data.get("input")
if input_data is None:
input_data: Final[object] = data.get("input")
if not _is_responses_input(input_data):
return None
messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
input=input_data,

2
litellm/main.py
View file

@ -1237,6 +1237,7 @@ def mock_completion(
custom_llm_provider="openai",
logging_obj=logging,
count_prompt_tokens=count_prompt_tokens,
is_mock=True,
)
return CustomStreamWrapper(
completion_stream=mock_completion_streaming_obj(
@ -1246,6 +1247,7 @@ def mock_completion(
custom_llm_provider="openai",
logging_obj=logging,
count_prompt_tokens=count_prompt_tokens,
is_mock=True,
)
if isinstance(mock_response, litellm.MockException):
raise mock_response

84
litellm/proxy/agent_endpoints/a2a_endpoints.py
View file

@ -11,7 +11,7 @@ The A2A SDK can point to LiteLLM's URL and invoke agents registered with LiteLLM
"""
import json
from collections.abc import AsyncGenerator, Mapping
from collections.abc import AsyncGenerator, Callable, Coroutine, Mapping
from copy import deepcopy
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Protocol
@ -19,7 +19,8 @@ from urllib.parse import urlparse
from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import TypeAdapter, ValidationError
from fastapi.routing import APIRoute
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_proxy_logger
@ -27,7 +28,7 @@ from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header
from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy.a2a.version_convert import (
A2AVersion,
normalize_agent_card,
@ -42,6 +43,7 @@ from litellm.proxy.agent_endpoints.databricks_oauth import (
)
from litellm.proxy.agent_endpoints.utils import merge_agent_headers
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.openai_error_payload import status_typed_proxy_exception
from litellm.proxy.common_utils.sse_keepalive import (
SSE_COMMENT_PING,
coerce_keepalive_interval,
@ -172,13 +174,16 @@ def _caller_identity_headers(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str,
async def _resolve_backend_auth(
litellm_params: dict[str, object],
custom_llm_provider: object,
agent_extra_headers: Mapping[str, str] | None = None,
) -> AgentBackendAuth:
databricks_auth: Final = await resolve_databricks_backend_auth(litellm_params)
if has_databricks_oauth(litellm_params):
return databricks_auth
return AgentBackendAuth(
litellm_params=databricks_auth.litellm_params,
header=await resolve_a2a_hop_auth_header(litellm_params, custom_llm_provider),
header=await resolve_a2a_hop_auth_header(
litellm_params, custom_llm_provider, agent_extra_headers=agent_extra_headers
),
)
@ -218,6 +223,57 @@ def _jsonrpc_error(
)
class _JsonRpcRequestId(BaseModel):
model_config = ConfigDict(frozen=True)
id: str | int | None = None
async def _jsonrpc_request_id(request: Request) -> str | int | None:
try:
return _JsonRpcRequestId.model_validate_json(await request.body()).id
except ValidationError:
return None
async def _jsonrpc_refusal(request: Request, refusal: ProxyException | HTTPException) -> JSONResponse:
from litellm.proxy.proxy_server import (
_close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; a route that answers an error itself must stamp the OTel server span like the global ProxyException handler does
)
error: Final = refusal if isinstance(refusal, ProxyException) else status_typed_proxy_exception(refusal)
_close_dangling_otel_server_span(request, int(error.code), exc=refusal)
return JSONResponse(
content={
"jsonrpc": "2.0",
"id": await _jsonrpc_request_id(request),
"error": {
"code": -32000,
"message": error.message,
"data": {"type": error.type, "param": error.param, "code": error.code},
},
},
status_code=int(error.code),
headers=error.headers if isinstance(refusal, ProxyException) else refusal.headers,
)
class _JsonRpcRefusalRoute(APIRoute):
def get_route_handler(self) -> Callable[[Request], Coroutine[object, object, Response]]:
handler: Final = super().get_route_handler()
async def answer_refusals_as_jsonrpc(request: Request) -> Response:
try:
return await handler(request)
except (ProxyException, HTTPException) as refusal:
return await _jsonrpc_refusal(request, refusal)
return answer_refusals_as_jsonrpc
jsonrpc_router: Final = APIRouter(route_class=_JsonRpcRefusalRoute)
async def _get_agent(agent_id: str) -> "AgentResponse | None":
"""Look up an agent by ID or name. Returns None if not found."""
from litellm.proxy.common_utils.registry_read_through import (
@ -650,17 +706,17 @@ async def get_agent_card(
raise HTTPException(status_code=500, detail=str(e))
@router.post(
@jsonrpc_router.post(
"/a2a/{agent_id}",
tags=["[beta] A2A Agents"],
dependencies=[Depends(user_api_key_auth)],
)
@router.post(
@jsonrpc_router.post(
"/a2a/{agent_id}/message/send",
tags=["[beta] A2A Agents"],
dependencies=[Depends(user_api_key_auth)],
)
@router.post(
@jsonrpc_router.post(
"/v1/a2a/{agent_id}/message/send",
tags=["[beta] A2A Agents"],
dependencies=[Depends(user_api_key_auth)],
@ -838,14 +894,15 @@ async def invoke_agent_a2a(
if header_name:
dynamic_headers[header_name] = val
backend_auth: Final = await _resolve_backend_auth(litellm_params, custom_llm_provider)
merged_agent_headers: Final = merge_agent_headers(
dynamic_headers=dynamic_headers or None,
static_headers=static_headers or None,
)
backend_auth: Final = await _resolve_backend_auth(litellm_params, custom_llm_provider, merged_agent_headers)
agent_extra_headers: Final = _forwarding_headers(
caller_identity=caller_identity,
request_data=data,
agent_extra_headers=merge_agent_headers(
dynamic_headers=dynamic_headers or None,
static_headers=static_headers or None,
),
agent_extra_headers=merged_agent_headers,
backend_auth_header=backend_auth.header,
)
forwarded_litellm_params: Final = backend_auth.litellm_params
@ -1057,3 +1114,6 @@ async def invoke_agent_a2a(
if isinstance(e, litellm.BadRequestError):
return _jsonrpc_error(body.get("id"), -32602, e.message, 400)
return _jsonrpc_error(body.get("id"), -32603, f"Internal error: {e}", 500)
router.include_router(jsonrpc_router)

35
litellm/proxy/auth/auth_checks.py
View file

@ -1086,13 +1086,31 @@ def model_max_budget_of(holder: ModelMaxBudgetHolder) -> Mapping[str, object] |
def route_skips_budget_checks(route: str) -> bool:
return route not in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES and (
route in MODEL_DISCOVERY_ROUTES or not RouteChecks.is_llm_api_route(route=route)
if route in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES:
return False
if route in MODEL_DISCOVERY_ROUTES:
return True
if RouteChecks.is_llm_api_route(route=route):
return False
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( # noqa: PLC0415 # pass-through endpoints import the auth stack
InitPassThroughEndpointHelpers,
)
return not InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=route)
def request_skips_budget_checks(route: str, model: str | list[str] | None, llm_router: Router | None) -> bool:
return route_skips_budget_checks(route=route) or is_model_cost_zero(model=model, llm_router=llm_router)
def request_skips_budget_checks(
route: str,
model: str | list[str] | None,
llm_router: Router | None,
*,
request_body: Mapping[str, object] | None = None,
) -> bool:
if route_skips_budget_checks(route=route):
return True
if request_body is not None and is_mcp_discovery_request(route=route, request_body=request_body):
return True
return is_model_cost_zero(model=model, llm_router=llm_router)
async def common_checks(
@ -1142,10 +1160,11 @@ async def common_checks(
team_id=valid_token.team_id if valid_token is not None else None,
)
skip_all_budget_checks: Final = (
skip_budget_checks
or route_skips_budget_checks(route=route)
or is_mcp_discovery_request(route=route, request_body=_typed_request_body(request_body))
skip_all_budget_checks: Final = skip_budget_checks or request_skips_budget_checks(
route=route,
model=_model,
llm_router=llm_router,
request_body=_typed_request_body(request_body),
)
membership_user_id: Final = (

32
litellm/proxy/auth/user_api_key_auth.py
View file

@ -80,7 +80,6 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
get_team_object,
get_user_object,
get_user_role,
is_mcp_discovery_request,
is_model_cost_zero,
is_user_proxy_admin,
is_valid_fallback_model,
@ -88,6 +87,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
key_access_group_grants_model,
key_model_aliases_for_auth_check,
model_max_budget_of,
request_skips_budget_checks,
resolve_and_validate_end_user_id,
resolve_default_end_user_budget,
route_skips_budget_checks,
@ -2497,15 +2497,9 @@ async def user_api_key_auth_builder(
llm_router=llm_router,
team_id=valid_token.team_id,
)
skip_budget_checks = is_mcp_discovery_request(route=route, request_body=request_data)
if not skip_budget_checks and model is not None and llm_router is not None:
from litellm.proxy.auth.auth_checks import is_model_cost_zero
skip_budget_checks = ( # rebind-ok: pre-existing rebinding on a rename-only line
is_model_cost_zero(model=model, llm_router=llm_router)
)
if skip_budget_checks:
verbose_proxy_logger.info("Skipping all budget checks for zero-cost model: %s", model)
skip_budget_checks = request_skips_budget_checks(
route=route, model=model, llm_router=llm_router, request_body=request_data
)
# Fetch project object for JWT path if project_id is set
_jwt_project_obj = None
@ -3019,15 +3013,9 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e
llm_router=llm_router,
team_id=valid_token.team_id,
)
skip_budget_checks = is_mcp_discovery_request(route=route, request_body=request_data)
if not skip_budget_checks and model is not None and llm_router is not None:
from litellm.proxy.auth.auth_checks import is_model_cost_zero
skip_budget_checks = is_model_cost_zero( # rebind-ok: pre-existing rebinding on a rename-only line
model=model, llm_router=llm_router
)
if skip_budget_checks:
verbose_proxy_logger.info("Skipping all budget checks for zero-cost model: %s", model)
skip_budget_checks: Final = request_skips_budget_checks(
route=route, model=model, llm_router=llm_router, request_body=request_data
)
# Check 3. Check if user is in their team budget
if not skip_budget_checks and valid_token.team_member_spend is not None:
@ -4088,11 +4076,7 @@ def _should_skip_budget_checks(
llm_router=llm_router,
team_id=team_id,
)
if is_mcp_discovery_request(route=route, request_body=request_data):
return True
if model is not None and llm_router is not None:
return is_model_cost_zero(model=model, llm_router=llm_router)
return False
return request_skips_budget_checks(route=route, model=model, llm_router=llm_router, request_body=request_data)
def _resolve_request_principal(request: Request, valid_token: UserAPIKeyAuth) -> Principal:

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

@ -64,11 +64,16 @@ from litellm.llms.azure.passthrough.transformation import (
from litellm.llms.bedrock.passthrough.transformation import is_bedrock_streaming_endpoint
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.llms.deepgram.common_utils import (
DeepgramSpeakMeter,
deepgram_listen_callback_params,
deepgram_listen_is_priced,
deepgram_listen_registry_key,
deepgram_listen_requested_model,
deepgram_listen_websocket_target,
deepgram_speak_character_rate,
deepgram_speak_registry_key,
deepgram_speak_requested_model,
deepgram_speak_websocket_target,
)
from litellm.llms.fal_ai.cost_calculator import fal_ai_passthrough_cost, fal_ai_queue_base
from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_group_in_path
@ -129,6 +134,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
YieldsToConfiguredPassThroughRoute,
create_pass_through_route,
create_websocket_passthrough_route,
truncated_close_reason,
websocket_passthrough_request,
)
from litellm.proxy.utils import ProxyLogging as ProxyLoggingType
@ -145,6 +151,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY,
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY,
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
WebSocketSessionMeter,
)
from litellm.types.passthrough_endpoints.tinyfish import (
TINYFISH_AUTHENTICATED_RUN_FIELDS,
@ -3460,6 +3467,7 @@ class _WebsocketRelay(Protocol):
forward_headers: bool,
endpoint: str,
accept_websocket: bool,
meter: WebSocketSessionMeter | None = None,
) -> None: ...
@ -3621,6 +3629,54 @@ async def deepgram_listen_websocket_route(
)
async def deepgram_speak_user_api_key_auth(websocket: WebSocket) -> UserAPIKeyAuth:
return await user_api_key_auth_websocket_for_model(
websocket, model=deepgram_speak_requested_model(websocket.url.query)
)
@router.websocket("/deepgram/v1/speak")
@router.websocket("/deepgram/speak")
async def deepgram_speak_websocket_route(
websocket: WebSocket,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(deepgram_speak_user_api_key_auth)],
relay: Annotated[_WebsocketRelay, Depends(_websocket_relay)],
) -> None:
deepgram_api_key: Final = passthrough_endpoint_router.get_credentials(
custom_llm_provider=litellm.LlmProviders.DEEPGRAM.value,
region_name=None,
)
if deepgram_api_key is None:
await websocket.close(code=1011, reason=_DEEPGRAM_WS_MISSING_KEY_REASON)
return
await websocket.accept(subprotocol=_negotiated_websocket_subprotocol(websocket))
target: Final = deepgram_speak_websocket_target(
api_base=get_secret_str("DEEPGRAM_API_BASE"),
query_string=websocket.url.query,
)
character_rate: Final = deepgram_speak_character_rate(target)
if character_rate is None:
await websocket.close(
code=1008,
reason=truncated_close_reason(
_DEEPGRAM_WS_UNPRICED_REASON.format(registry_key=deepgram_speak_registry_key(target))
),
)
return
await relay(
websocket=websocket,
target=target,
custom_headers={"Authorization": f"Token {deepgram_api_key}"},
user_api_key_dict=user_api_key_dict,
forward_headers=False,
endpoint=websocket.url.path,
accept_websocket=False,
meter=DeepgramSpeakMeter(target, character_rate),
)
_OPENAI_BETA_HEADER: Final = "openai-beta"
_CLIENT_OPENAI_BETA_RESOURCES: Final = frozenset({"agents", "vaults", "chatkit"})
@ -4474,10 +4530,16 @@ async def watsonx_proxy_route(
if provider_config is None:
raise HTTPException(status_code=404, detail="Watsonx passthrough config not found")
watsonx_credentials: Final = passthrough_endpoint_router.get_proxy_wide_deployment_credentials(
custom_llm_provider="watsonx"
)
api_key: Final = watsonx_credentials.api_key if watsonx_credentials is not None else None
api_base: Final = watsonx_credentials.api_base if watsonx_credentials is not None else None
# Get complete URL with version parameter
complete_url, _ = provider_config.get_complete_url(
api_base=None,
api_key=None,
api_base=api_base,
api_key=api_key,
model="",
endpoint=endpoint,
request_query_params=None,
@ -4491,8 +4553,8 @@ async def watsonx_proxy_route(
messages=[],
optional_params={},
litellm_params={},
api_key=None,
api_base=None,
api_key=api_key,
api_base=api_base,
)
# Check for streaming

4
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -2584,7 +2584,7 @@ def _json_object_frame(frame: str | bytes) -> dict[str, object] | None:
return decoded if isinstance(decoded, dict) else None
def _truncated_close_reason(reason: str) -> str:
def truncated_close_reason(reason: str) -> str:
"""
Fit a close reason inside the byte budget a WebSocket close frame allows, without splitting a character
"""
@ -2958,7 +2958,7 @@ async def websocket_passthrough_request(
if upstream_close is not None and _client_socket_is_open(websocket):
await websocket.close(
code=upstream_close.code,
reason=_truncated_close_reason(upstream_close.reason),
reason=truncated_close_reason(upstream_close.reason),
)
end_time: Final = datetime.now()

13
litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py
View file

@ -132,6 +132,19 @@ class PassthroughEndpointRouter:
api_base=deployment_credentials.api_base if deployment_credentials is not None else None,
)
def get_proxy_wide_deployment_credentials(self, custom_llm_provider: str) -> _PassThroughCredentials | None:
deployment_credentials: Final = self._get_deployment_credentials(
custom_llm_provider=custom_llm_provider,
region_name=None,
include_team_deployments=False,
)
if deployment_credentials is None:
return None
api_key: Final = _resolve_oidc_reference(deployment_credentials.api_key)
if not api_key:
return None
return _PassThroughCredentials(api_key=api_key, api_base=deployment_credentials.api_base)
def has_credentials(self, custom_llm_provider: str) -> bool:
if (
self._get_deployment_credentials(

64
litellm/proxy/proxy_server.py
View file

@ -846,6 +846,8 @@ from litellm.proxy.telemetry.middleware import TelemetryMiddleware
from litellm.proxy.telemetry.runtime import TelemetryRuntime
from litellm.proxy.telemetry.settings import describe_errors as describe_telemetry_errors
from litellm.proxy.telemetry.settings import load_settings as load_telemetry_settings
from litellm.router import router_arguments
from litellm.router_config import build_router_config
try:
from litellm.proxy.enterprise_billing.billing_metrics import (
@ -7407,20 +7409,24 @@ class ProxyConfig:
verbose_proxy_logger.warning(
"Key '%s' is not a valid argument for Router.__init__(). Ignoring this key.", k
)
router = litellm.Router(
**router_params,
assistants_config=assistants_config,
search_tools=search_tools,
router_general_settings=RouterGeneralSettings(
async_only_mode=True # only init async clients
),
ignore_invalid_deployments=True, # don't raise an error if a deployment is invalid
fallback_access_check=router_fallback_access_check,
fallback_budget_check=RouterFallbackBudgetCheck(
is_enforced=router_fallback_budget_check.is_enforced,
team_member_model_budget_limiter=model_max_budget_limiter,
),
auto_router_capability_limit=_license_check.auto_router_capability_limit,
router = litellm.Router.from_config(
build_router_config(
router_arguments(
**router_params,
assistants_config=assistants_config,
search_tools=search_tools,
router_general_settings=RouterGeneralSettings(
async_only_mode=True # only init async clients
),
ignore_invalid_deployments=True, # don't raise an error if a deployment is invalid
fallback_access_check=router_fallback_access_check,
fallback_budget_check=RouterFallbackBudgetCheck(
is_enforced=router_fallback_budget_check.is_enforced,
team_member_model_budget_limiter=model_max_budget_limiter,
),
auto_router_capability_limit=_license_check.auto_router_capability_limit,
)
)
)
if redis_usage_cache is not None and router.health().store.redis_cache is None:
@ -7910,19 +7916,23 @@ class ProxyConfig:
# Router can function with model_list=[] if search_tools are configured
if len(_model_list) > 0 or search_tools:
verbose_proxy_logger.debug("_model_list: %s", _model_list)
llm_router = litellm.Router(
model_list=_model_list,
router_general_settings=RouterGeneralSettings(
async_only_mode=True # only init async clients
),
search_tools=search_tools,
ignore_invalid_deployments=True,
fallback_access_check=router_fallback_access_check,
fallback_budget_check=RouterFallbackBudgetCheck(
is_enforced=router_fallback_budget_check.is_enforced,
team_member_model_budget_limiter=model_max_budget_limiter,
),
auto_router_capability_limit=_license_check.auto_router_capability_limit,
llm_router = litellm.Router.from_config(
build_router_config(
router_arguments(
model_list=_model_list,
router_general_settings=RouterGeneralSettings(
async_only_mode=True # only init async clients
),
search_tools=search_tools,
ignore_invalid_deployments=True,
fallback_access_check=router_fallback_access_check,
fallback_budget_check=RouterFallbackBudgetCheck(
is_enforced=router_fallback_budget_check.is_enforced,
team_member_model_budget_limiter=model_max_budget_limiter,
),
auto_router_capability_limit=_license_check.auto_router_capability_limit,
)
)
)
verbose_proxy_logger.debug("updated llm_router: %s", llm_router)
else:

6
litellm/proxy/route_llm_request.py
View file

@ -18,7 +18,9 @@ from litellm.litellm_core_utils.core_helpers import normalize_drop_params
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.routing_common_utils import is_proxy_admin_request
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.router import router_arguments
from litellm.router_api import router_init_args
from litellm.router_config import build_router_config
# Client-supplied params that make the router or the call path fabricate a
# failure or a delay instead of calling the provider. The ``mock_testing_*``
@ -98,7 +100,9 @@ async def _route_user_config_request(data: dict, route_type: str):
filtered_config: Final = {k: v for k, v in router_config.items() if k in valid_args}
logging_obj: Final = _request_logging_obj(JSON_OBJECT_ADAPTER.validate_python(data), route_type)
cached_router: Final = user_config_routers.get(filtered_config, lambda: litellm.Router(**filtered_config))
cached_router: Final = user_config_routers.get(
filtered_config, lambda: litellm.Router.from_config(build_router_config(router_arguments(**filtered_config)))
)
cached_router.observe(logging_obj)
return await getattr(cached_router, f"{route_type}")(**{**data, "litellm_logging_obj": logging_obj})

5
litellm/responses/streaming_iterator.py
View file

@ -233,6 +233,11 @@ def spawn_logging_task(
async def drain_logging_tasks(timeout: float = 5.0) -> None:
# A task that finished as its loop closed never ran its discard callback, and a
# task from another loop cannot be awaited here: forget both, wait on the rest.
running_loop: Final = asyncio.get_running_loop()
stale: Final = {task for task in _PENDING_LOGGING_TASKS if task.done() or task.get_loop() is not running_loop}
_PENDING_LOGGING_TASKS.difference_update(stale)
if not _PENDING_LOGGING_TASKS:
return
_, pending = await asyncio.wait(tuple(_PENDING_LOGGING_TASKS), timeout=timeout)

17811
litellm/router.py
View file

File diff suppressed because it is too large Load diff

44
litellm/router_api.py
View file

@ -42,6 +42,8 @@ from dataclasses import dataclass, field
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias, overload
from litellm.router_config import detached
if TYPE_CHECKING:
from typing import Any # noqa: TID251 # only to spell the cores' untyped (Unknown) containers and one Any type
@ -176,9 +178,9 @@ def deployment_usable_by_team(deployment: "Mapping[str, object] | Deployment", t
Whether ``team_id`` may use ``deployment``: a team-scoped deployment (``model_info.team_id``
set) only by that team, any other deployment by everyone. Needs no router instance.
"""
from litellm.router import Router
from litellm.router_python_core import PythonCore
return Router._core_deployment_usable_by_team(deployment, team_id)
return PythonCore._core_deployment_usable_by_team(deployment, team_id)
def deployment_id(model_group: str, litellm_params: Mapping[str, object]) -> str:
@ -200,9 +202,9 @@ def deployment_id(model_group: str, litellm_params: Mapping[str, object]) -> str
def router_init_args() -> tuple[str, ...]:
"""The arguments ``Router.__init__`` accepts (was ``Router.get_valid_args``). Needs no router instance."""
from litellm.router import Router
from litellm.router_python_core import PythonCore
return tuple(Router.get_valid_args())
return tuple(PythonCore.get_valid_args())
def mark_advisor_orchestration_failure(exception: BaseException) -> None:
@ -266,9 +268,9 @@ def deployment_metadata(request_kwargs: Mapping[str, object]) -> Mapping[str, ob
def _json_default_stable_id(obj: object) -> str:
from litellm.router import Router
from litellm.router_python_core import PythonCore
return Router._json_default_stable_id(obj)
return PythonCore._json_default_stable_id(obj)
class RouterRegistry:
@ -529,10 +531,10 @@ class RouterRegistry:
``should_include_deployment``), or, without ``model_name``, whether the team may use it at
all (was ``deployment_usable_by_team``).
"""
from litellm.router import Router
from litellm.router_python_core import PythonCore
if model_name is None:
return Router._core_deployment_usable_by_team(deployment, team_id)
return PythonCore._core_deployment_usable_by_team(deployment, team_id)
return self._router._core_should_include_deployment(model_name=model_name, model=deployment, team_id=team_id)
@overload
@ -1053,6 +1055,12 @@ SETTING_FIELDS: Final[tuple[str, ...]] = (
"cache_responses",
)
# Runtime settings kept as plain router attributes rather than ``settings()`` fields. A write
# to one (``router.num_retries = 2``) is a changeset too: ``ChangeSet(assign={...})``.
PLAIN_SETTING_ATTRIBUTES: Final[frozenset[str]] = frozenset({"num_retries", "max_fallbacks", "retry_after", "timeout"})
# Former attributes whose value lives under another name, assignable through a changeset.
ASSIGNABLE_ALIASES: Final[Mapping[str, str]] = MappingProxyType({"model_group_alias": "_core_model_group_alias"})
_MISSING: Final = object()
@ -1168,16 +1176,16 @@ class RouterHealth:
@staticmethod
def is_blocked(deployment: "Deployment") -> bool:
"""Whether an admin paused this deployment (was ``is_deployment_blocked``)."""
from litellm.router import Router
from litellm.router_python_core import PythonCore
return Router._core_is_deployment_blocked(deployment)
return PythonCore._core_is_deployment_blocked(deployment)
@staticmethod
def all_blocked(deployments: "Sequence[DeploymentTypedDict]") -> bool:
"""Whether every deployment in the list is paused (was ``are_all_deployments_blocked``)."""
from litellm.router import Router
from litellm.router_python_core import PythonCore
return Router._core_are_all_deployments_blocked(deployments)
return PythonCore._core_are_all_deployments_blocked(deployments)
def fully_blocked_model_names(self) -> set[str]: # mutable-ok: the core's type, a fresh set
"""Groups whose every deployment is paused (was ``get_fully_blocked_model_names``)."""
@ -1296,13 +1304,17 @@ class RoutingObservers:
def apply_changes(router: "Router", changes: ChangeSet) -> ApplyResult:
if changes.refresh_model_info:
raise TypeError("ChangeSet(refresh_model_info=True) needs aapply()")
unknown: Final = sorted(set(changes.assign) - set(SETTING_FIELDS))
unknown: Final = sorted(set(changes.assign) - {*SETTING_FIELDS, *PLAIN_SETTING_ATTRIBUTES, *ASSIGNABLE_ALIASES})
if unknown:
raise ValueError(f"not assignable router settings: {unknown}")
for name, value in changes.assign.items():
setattr(router, "_core_" + name, value)
setattr(
router,
name if name in PLAIN_SETTING_ATTRIBUTES else ASSIGNABLE_ALIASES.get(name, "_core_" + name),
detached(value),
)
if changes.settings:
router._core_update_settings(**changes.settings)
router._core_update_settings(**{name: detached(value) for name, value in changes.settings.items()})
upserted: Final = [router._core_upsert_deployment(deployment=d) for d in changes.upsert]
deleted: Final = [router._core_delete_deployment(id=i) for i in changes.delete]
if changes.redis_cache is not None:
@ -1388,7 +1400,7 @@ class RoutingContext:
return self._router.observe()
def belongs_to(self, router: "Router") -> bool:
return self._router is router
return self._router is router or router.__dict__.get("_core") is self._router
@property
def explicit_allowed_fails(self) -> int | None:

157
litellm/router_config.py Normal file
View file

@ -0,0 +1,157 @@
"""
The router's configuration, built outside the router.
``RouterConfig`` is static: the constructor arguments, frozen for the router's life.
``RouterDynamicConfig`` is a snapshot of the process settings the router reads while it serves
(``litellm.num_retries``, ``litellm.model_cost``, the fallback lists, ...). A call reads one snapshot
at its start. This module is the only router code that reads ``litellm.*`` settings.
"""
import itertools
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
import litellm
if TYPE_CHECKING:
from litellm.caching.caching import Cache
from litellm.litellm_core_utils.logging_callback_manager import LoggingCallbackManager
@dataclass(frozen=True, slots=True)
class RouterConfig:
"""The arguments the router was built with, exactly as passed (unset ones are absent)."""
arguments: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({}))
def get(self, name: str, default: object = None) -> object:
return self.arguments.get(name, default)
@dataclass(frozen=True, slots=True)
class RouterDynamicConfig:
"""Process settings the router reads per call. Values are the live objects, not copies."""
version: int
num_retries: int | None
max_fallbacks: int | None
router_max_fallbacks: int
allowed_fails: int
request_timeout: float
fallbacks: Sequence[object] | None
default_fallbacks: Sequence[str] | None
context_window_fallbacks: Sequence[object] | None
content_policy_fallbacks: Sequence[object] | None
drop_params: bool | None
provider_list: Sequence[str]
expose_router_debug_in_errors: bool
redact_user_api_key_info: bool | None
model_alias_map: Mapping[str, str]
model_cost: Mapping[str, object]
known_custom_logger_compatible_callbacks: Sequence[str]
exception_types: Sequence[type[BaseException]]
callbacks: Sequence[object]
success_callback: Sequence[object]
failure_callback: Sequence[object]
input_callback: Sequence[object]
callback_manager: "LoggingCallbackManager"
cache: "Cache | None"
_versions: Final = itertools.count(1)
def detached(value: object) -> object:
"""``value`` with its plain containers (dict, list, tuple, set) copied, all the way down, so the
caller's own objects can't change the router afterwards. Other objects (clients, caches,
callbacks, pydantic models) are kept as they are."""
if isinstance(value, dict):
return {key: detached(item) for key, item in value.items()}
if isinstance(value, list):
return [detached(item) for item in value]
if isinstance(value, tuple) and type(value) is tuple:
return tuple(detached(item) for item in value)
if isinstance(value, set):
return {detached(item) for item in value}
return value
def build_router_config(arguments: Mapping[str, object]) -> RouterConfig:
"""The constructor arguments, detached from the caller's objects: settings are static."""
return RouterConfig(arguments=MappingProxyType({name: detached(value) for name, value in arguments.items()}))
def current_dynamic_config() -> RouterDynamicConfig:
"""A snapshot of the process settings as they are right now."""
return RouterDynamicConfig(
version=next(_versions),
num_retries=litellm.num_retries,
max_fallbacks=litellm.max_fallbacks,
router_max_fallbacks=litellm.ROUTER_MAX_FALLBACKS,
allowed_fails=litellm.allowed_fails,
request_timeout=litellm.request_timeout,
fallbacks=litellm.fallbacks,
default_fallbacks=litellm.default_fallbacks,
context_window_fallbacks=litellm.context_window_fallbacks,
content_policy_fallbacks=litellm.content_policy_fallbacks,
drop_params=litellm.drop_params,
provider_list=litellm.provider_list,
expose_router_debug_in_errors=litellm.expose_router_debug_in_errors,
redact_user_api_key_info=litellm.redact_user_api_key_info,
model_alias_map=litellm.model_alias_map,
model_cost=litellm.model_cost,
known_custom_logger_compatible_callbacks=litellm._known_custom_logger_compatible_callbacks,
exception_types=litellm.LITELLM_EXCEPTION_TYPES,
callbacks=litellm.callbacks,
success_callback=litellm.success_callback,
failure_callback=litellm.failure_callback,
input_callback=litellm.input_callback,
callback_manager=litellm.logging_callback_manager,
cache=litellm.cache,
)
_LIVE_SOURCES: Final = MappingProxyType(
{
"num_retries": "num_retries",
"max_fallbacks": "max_fallbacks",
"router_max_fallbacks": "ROUTER_MAX_FALLBACKS",
"allowed_fails": "allowed_fails",
"request_timeout": "request_timeout",
"fallbacks": "fallbacks",
"default_fallbacks": "default_fallbacks",
"context_window_fallbacks": "context_window_fallbacks",
"content_policy_fallbacks": "content_policy_fallbacks",
"drop_params": "drop_params",
"provider_list": "provider_list",
"expose_router_debug_in_errors": "expose_router_debug_in_errors",
"redact_user_api_key_info": "redact_user_api_key_info",
"model_alias_map": "model_alias_map",
"model_cost": "model_cost",
"known_custom_logger_compatible_callbacks": "_known_custom_logger_compatible_callbacks",
"exception_types": "LITELLM_EXCEPTION_TYPES",
"callbacks": "callbacks",
"success_callback": "success_callback",
"failure_callback": "failure_callback",
"input_callback": "input_callback",
"callback_manager": "logging_callback_manager",
"cache": "cache",
}
)
class LiveDynamicConfig:
"""``RouterDynamicConfig``'s fields read from the process settings at access time, with no snapshot."""
__slots__ = ()
def __getattr__(self, name: str) -> object:
source: Final = _LIVE_SOURCES.get(name)
if source is None:
raise AttributeError(name)
return getattr(litellm, source)
live_dynamic_config: Final = LiveDynamicConfig()

56
litellm/router_core.py Normal file
View file

@ -0,0 +1,56 @@
"""
``RouterCore``: the interface ``litellm.Router`` delegates to, and ``make_core``, which picks one.
``PythonCore`` (``litellm/router_python_core.py``) is today's routing logic. The Rust core
(litellm-rust/crates/router-core) implements the same operations. The core is picked once, when
the router is built, from its ``RouterConfig``.
"""
from collections.abc import Mapping
from typing import TYPE_CHECKING, Protocol, runtime_checkable
from litellm.router_config import RouterConfig
if TYPE_CHECKING:
from litellm.router_api import (
ApplyResult,
ChangeSet,
RouterHealth,
RouterRegistry,
RouterSettings,
RoutingObservers,
)
from litellm.router_python_core import PythonCore
from litellm.router_python_host import PythonHost
from litellm.types.router import CustomRoutingStrategyBase
@runtime_checkable
class RouterCore(Protocol):
async def call(self, op: str, request: Mapping[str, object]) -> object: ...
def registry(self) -> "RouterRegistry": ...
def settings(self) -> "RouterSettings": ...
def health(self) -> "RouterHealth": ...
async def select(
self,
model: str,
request_kwargs: dict[str, object] | None = None, # mutable-ok: the selection cores take a dict
) -> object: ...
def apply(self, changes: "ChangeSet") -> "ApplyResult": ...
def observe(self) -> "RoutingObservers": ...
def register_plugin(self, plugin: "CustomRoutingStrategyBase") -> None: ...
def retire(self) -> None: ...
def make_core(config: RouterConfig, host: "PythonHost", core_class: "type[PythonCore] | None" = None) -> RouterCore:
from litellm.router_python_core import PythonCore
return (core_class or PythonCore)(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the arguments were bound against PythonCore's own signature

30
litellm/router_deprecated.py
View file

@ -31,7 +31,7 @@ from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Generic, TypeVar, overload
from litellm.router_api import SETTING_FIELDS, ChangeSet
from litellm.router_api import ASSIGNABLE_ALIASES, SETTING_FIELDS, ChangeSet
if TYPE_CHECKING:
from typing import Any # noqa: TID251 # only to repeat two old signatures that had Any, unchanged
@ -65,7 +65,7 @@ _INTERNAL_PREFIX: Final = "litellm.router"
# Router's own attribute hooks: a write or delete through an old name passes through them first.
_ATTRIBUTE_HOOKS: Final = frozenset({"__setattr__", "__delattr__"})
_ATTRIBUTE_HOOKS: Final = frozenset({"__getattribute__", "__setattr__", "__delattr__"})
def _note(name: str, replacement: str) -> None:
@ -88,9 +88,10 @@ def _note(name: str, replacement: str) -> None:
def _as_router(obj: "DeprecatedRouterMembers") -> "Router":
"""The Router this mixin is part of (it is mixed into ``litellm.Router`` only)."""
from litellm.router import Router
from litellm.router_python_core import PythonCore
if isinstance(obj, Router):
return obj
if isinstance(obj, (Router, PythonCore)):
return obj # pyright: ignore[reportReturnType] # a PythonCore carries every member the mixin reads
raise TypeError(f"{type(obj).__name__} is not a litellm.Router; DeprecatedRouterMembers is mixed into Router only")
@ -134,7 +135,7 @@ class _DeprecatedAttribute(Generic[_T]):
def __set__(self, obj: "DeprecatedRouterMembers", value: _T) -> None:
_note(self.name, self.replacement)
if self.name in SETTING_FIELDS:
if self.name in SETTING_FIELDS or self.name in ASSIGNABLE_ALIASES:
_as_router(obj).apply(ChangeSet(assign={self.name: value}))
else:
setattr(obj, "_core_" + self.name, value)
@ -169,6 +170,13 @@ _OVERRIDES_RUNNING: Final[contextvars.ContextVar[frozenset[tuple[int, str]]]] =
)
def _override_key(router: object, name: str) -> tuple[int, str]:
"""The running-override key for ``router``: a Router and its core share one, so an old member
entered on the Router is seen by the override routing that runs on the core."""
core: Final = getattr(router, "__dict__", {}).get("_core")
return (id(router if core is None else core), name)
def _enter_override(name: str, override: Callable[..., object]) -> Callable[..., object]:
"""
A direct call of a subclass override (``router.<name>()``) marks it running, so its
@ -184,7 +192,7 @@ def _enter_override(name: str, override: Callable[..., object]) -> Callable[...,
*args: object,
**kwargs: object, # kwargs-ok: passes the override's arguments through unchanged
) -> object:
token: Final = _OVERRIDES_RUNNING.set(_OVERRIDES_RUNNING.get() | {(id(router), name)})
token: Final = _OVERRIDES_RUNNING.set(_OVERRIDES_RUNNING.get() | {_override_key(router, name)})
try:
return await async_override(router, *args, **kwargs)
finally:
@ -198,7 +206,7 @@ def _enter_override(name: str, override: Callable[..., object]) -> Callable[...,
*args: object,
**kwargs: object, # kwargs-ok: passes the override's arguments through unchanged
) -> object:
token: Final = _OVERRIDES_RUNNING.set(_OVERRIDES_RUNNING.get() | {(id(router), name)})
token: Final = _OVERRIDES_RUNNING.set(_OVERRIDES_RUNNING.get() | {_override_key(router, name)})
try:
return override(router, *args, **kwargs)
finally:
@ -236,7 +244,7 @@ def _route_override(cls: type, name: str, override: Callable[..., object]) -> Ca
*args: object,
**kwargs: object, # kwargs-ok: passes the override's arguments through unchanged
) -> object:
key: Final = (id(router), name)
key: Final = _override_key(router, name)
running: Final = _OVERRIDES_RUNNING.get()
if key in running:
parent_async: Final[Callable[..., Awaitable[object]]] = getattr(super(cls, router), core)
@ -255,7 +263,7 @@ def _route_override(cls: type, name: str, override: Callable[..., object]) -> Ca
*args: object,
**kwargs: object, # kwargs-ok: passes the override's arguments through unchanged
) -> object:
key: Final = (id(router), name)
key: Final = _override_key(router, name)
running: Final = _OVERRIDES_RUNNING.get()
if key in running:
return parent(router)(*args, **kwargs)
@ -502,9 +510,9 @@ class DeprecatedRouterMembers:
@staticmethod
def deployment_usable_by_team(model: "Mapping[str, object] | Deployment", team_id: str | None) -> bool:
_note("deployment_usable_by_team", "registry().includes(deployment, team_id=...)")
from litellm.router import Router
from litellm.router_python_core import PythonCore
return Router._core_deployment_usable_by_team(model, team_id)
return PythonCore._core_deployment_usable_by_team(model, team_id)
_deployment_usable_by_team = deployment_usable_by_team

17669
litellm/router_python_core.py Normal file
View file

File diff suppressed because it is too large Load diff

135
litellm/router_python_host.py Normal file
View file

@ -0,0 +1,135 @@
"""
``PythonHost``: what the router core reaches outside itself, served from the litellm package.
The core never touches ``litellm.*`` directly. Inference calls (``invoke``), backoff sleeps,
callback registration and the process-wide cache all go through the host, and the process settings
come from ``dynamic()`` (``litellm/router_config.py``). A Rust core gets the same operations.
"""
import asyncio
from collections.abc import Awaitable, Callable, Sequence
from typing import (
TYPE_CHECKING,
Final,
Literal,
TypeAlias,
cast, # noqa: TID251 # untyped litellm globals: the live config view and the callback lists
)
from typing_extensions import overload
import litellm
from litellm.router_config import RouterDynamicConfig, live_dynamic_config
if TYPE_CHECKING:
from pydantic import BaseModel
from litellm.caching.caching import Cache
from litellm.litellm_core_utils.logging_callback_manager import LoggingCallbackManager
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper
from litellm.types.llms.openai import HttpxBinaryResponseContent, OpenAIFileObject
from litellm.types.rerank import RerankResponse
from litellm.types.utils import (
AdapterCompletionStreamWrapper,
EmbeddingResponse,
ImageResponse,
LiteLLMBatch,
ModelInfo,
ModelResponse,
TextCompletionResponse,
TranscriptionResponse,
)
CallbackList: TypeAlias = Literal[
"success_callback", "_async_success_callback", "failure_callback", "_async_failure_callback", "input_callback"
]
Endpoint: TypeAlias = Callable[..., object]
class PythonHost:
__slots__ = ()
# The endpoints whose results the router core reads, typed as litellm declares them; any other
# op gets the untyped Endpoint.
@overload
def invoke(self, op: Literal["acompletion"]) -> Callable[..., Awaitable["ModelResponse | CustomStreamWrapper"]]: ...
@overload
def invoke(
self, op: Literal["atext_completion"]
) -> Callable[..., Awaitable["TextCompletionResponse | TextCompletionStreamWrapper"]]: ...
@overload
def invoke(
self, op: Literal["aadapter_completion"]
) -> Callable[..., Awaitable["BaseModel | AdapterCompletionStreamWrapper | None"]]: ...
@overload
def invoke(self, op: Literal["aembedding"]) -> Callable[..., Awaitable["EmbeddingResponse"]]: ...
@overload
def invoke(self, op: Literal["aimage_generation"]) -> Callable[..., Awaitable["ImageResponse"]]: ...
@overload
def invoke(self, op: Literal["atranscription"]) -> Callable[..., Awaitable["TranscriptionResponse"]]: ...
@overload
def invoke(self, op: Literal["aspeech"]) -> Callable[..., Awaitable["HttpxBinaryResponseContent"]]: ...
@overload
def invoke(self, op: Literal["arerank"]) -> Callable[..., Awaitable["RerankResponse"]]: ...
@overload
def invoke(self, op: Literal["acreate_file"]) -> Callable[..., Awaitable["OpenAIFileObject"]]: ...
@overload
def invoke(self, op: Literal["acreate_batch"]) -> Callable[..., Awaitable["LiteLLMBatch"]]: ...
@overload
def invoke(self, op: Literal["aretrieve_batch"]) -> Callable[..., Awaitable["LiteLLMBatch"]]: ...
@overload
def invoke(self, op: Literal["acancel_batch"]) -> Callable[..., Awaitable["LiteLLMBatch"]]: ...
@overload
def invoke(self, op: Literal["alist_batches"]) -> Callable[..., Awaitable[object]]: ...
@overload
def invoke(self, op: Literal["get_llm_provider"]) -> Callable[..., "tuple[str, str, str | None, str | None]"]: ...
@overload
def invoke(self, op: Literal["get_model_info"]) -> Callable[..., "ModelInfo"]: ...
@overload
def invoke(self, op: Literal["get_supported_openai_params"]) -> Callable[..., "Sequence[str] | None"]: ...
@overload
def invoke(self, op: Literal["token_counter"]) -> Callable[..., int]: ...
@overload
def invoke(self, op: Literal["calculate_retry_after"]) -> Callable[..., float]: ...
@overload
def invoke(self, op: str) -> Endpoint: ...
def invoke(self, op: str) -> Endpoint:
"""The litellm function named ``op`` (``acompletion``, ``get_model_info``, ...), looked up at call time."""
endpoint: Final[object] = getattr(litellm, op, None)
if endpoint is None:
raise AttributeError(f"litellm has no endpoint {op!r}")
if not callable(endpoint):
raise TypeError(f"litellm.{op} is not an endpoint")
return endpoint
def sleep(self, seconds: float) -> Awaitable[None]:
return asyncio.sleep(seconds)
def dynamic(self) -> RouterDynamicConfig:
"""The process settings, read live on each access (``current_dynamic_config()`` takes a snapshot)."""
return cast(RouterDynamicConfig, live_dynamic_config) # cast-ok: the live view reads the same fields
@property
def callbacks(self) -> "LoggingCallbackManager":
return litellm.logging_callback_manager
def callback_list(self, name: CallbackList) -> list[object]: # mutable-ok: the process callback list itself
value: Final = getattr(litellm, name)
return cast(list[object], value) # cast-ok: the process callback list is an untyped litellm global
def set_callback_list(self, name: CallbackList, value: list[object]) -> None: # mutable-ok: assigned as given
setattr(litellm, name, value)
def global_cache(self) -> "Cache | None":
return litellm.cache
def set_global_cache(self, cache: "Cache | None") -> None:
litellm.cache = cache
def suppress_debug_info(self) -> None:
litellm.suppress_debug_info = True
python_host: Final = PythonHost()

27
litellm/scheduler.py
View file

@ -1,15 +1,20 @@
import asyncio
import enum
import heapq
import logging as _logging
import time
from collections.abc import Awaitable, Callable, Sequence
from typing import Final, TypeAlias
from typing import Protocol as _Protocol
from pydantic import TypeAdapter
from litellm import print_verbose
from litellm._internal_context import with_service_target
from litellm._logging import verbose_logger as _verbose_logger
from litellm.caching.caching import DualCache, RedisCache
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError as _RedisCircuitBreakerOpenError
from litellm.caching.redis_cache import log_redis_failure as _log_redis_failure
from litellm.constants import DEFAULT_IN_MEMORY_TTL, DEFAULT_POLLING_INTERVAL
from litellm.exceptions import Timeout
from litellm.types.llms.base import LiteLLMBaseModel
@ -30,6 +35,26 @@ class FlowItem(LiteLLMBaseModel):
model_name: str
class _QueueCache(_Protocol):
async def async_get_cache(self, key: str, *, local_only: bool = False, raise_on_error: bool = False) -> object: ...
async def async_set_cache(self, key: str, value: object, *, local_only: bool = False) -> object: ...
async def _read_queue(cache: _QueueCache, redis_cache: _QueueCache | None, key: str) -> object:
if redis_cache is None:
return await cache.async_get_cache(key=key, local_only=True)
from redis.exceptions import RedisError as _RedisError
try:
response: Final = await redis_cache.async_get_cache(key=key, raise_on_error=True)
except (_RedisError, _RedisCircuitBreakerOpenError) as error:
_log_redis_failure(_verbose_logger, _logging.ERROR, "Scheduler queue using local snapshot", error)
return await cache.async_get_cache(key=key, local_only=True)
await cache.async_set_cache(key=key, value=[] if response is None else response, local_only=True)
return response
class Scheduler:
cache: DualCache
@ -153,7 +178,7 @@ class Scheduler:
"""
if self.cache is not None:
_cache_key: Final = f"{SchedulerCacheKeys.queue.value}:{model_name}"
response: Final = await self.cache.async_get_cache(key=_cache_key)
response: Final = await _read_queue(self.cache, self.cache.redis_cache, _cache_key)
if not isinstance(response, list):
return []
return _QUEUE_ENTRIES.validate_python(response)

1
tests/code_coverage_tests/recursive_detector.py
View file

@ -3,6 +3,7 @@ import os
IGNORE_FUNCTIONS = [
"_json_cost", # bounded at depth 32 and consumed under byte/node limits.
"detached", # router_config: copies only the plain containers of the router's own constructor arguments, once per build or changeset
"_format_type",
"remove_additional_properties",
"remove_strict_from_schema",

54
tests/code_coverage_tests/router_litellm_globals_check.py Normal file
View file

@ -0,0 +1,54 @@
"""
The router and its cores never read or write ``litellm.*``.
Process settings reach the router through ``RouterDynamicConfig`` (litellm/router_config.py) and
everything else outside it (inference calls, callbacks, the global cache) through ``PythonHost``
(litellm/router_python_host.py). Those two modules are the only router code allowed to touch the
``litellm`` package namespace. This check fails on ``import litellm`` (or ``import litellm.<x>``,
which binds the same name) and on any ``litellm.<attr>`` access in the files below.
Run: python tests/code_coverage_tests/router_litellm_globals_check.py
"""
import ast
import os
import sys
ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
CHECKED = (
"litellm/router.py",
"litellm/router_api.py",
"litellm/router_core.py",
"litellm/router_deprecated.py",
"litellm/router_python_core.py",
)
def violations(rel: str) -> list[str]:
tree = ast.parse(open(os.path.join(ROOT, rel)).read(), rel)
found = []
for node in ast.walk(tree):
if isinstance(node, ast.Import):
found.extend(
f"{rel}:{node.lineno}: import {alias.name}"
for alias in node.names
if alias.asname is None and (alias.name == "litellm" or alias.name.startswith("litellm."))
)
elif isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name) and node.value.id == "litellm":
found.append(f"{rel}:{node.lineno}: litellm.{node.attr}")
return found
def main() -> int:
found = [line for rel in CHECKED for line in violations(rel)]
for line in found:
print(line)
if found:
print(f"{len(found)} litellm.* uses in router modules; read settings from RouterDynamicConfig and use PythonHost")
return 1
print("router modules use no litellm.* globals")
return 0
if __name__ == "__main__":
sys.exit(main())

8
tests/code_coverage_tests/router_surface_check.py
View file

@ -38,11 +38,15 @@ ALLOWLIST = os.path.join(os.path.dirname(__file__), "router_surface_allowlist.js
ROUTER_FILES = {
"litellm/router.py",
"litellm/router_api.py",
"litellm/router_config.py",
"litellm/router_core.py",
"litellm/router_python_core.py",
"litellm/router_python_host.py",
"litellm/router_deprecated.py",
"litellm/router_host.py",
}
ROUTER_CLASS_FILES = ["litellm/router.py", "litellm/router_deprecated.py"]
ROUTER_CLASS_NAMES = {"Router", "DeprecatedRouterMembers"}
ROUTER_CLASS_FILES = ["litellm/router.py", "litellm/router_python_core.py", "litellm/router_deprecated.py"]
ROUTER_CLASS_NAMES = {"Router", "PythonCore", "DeprecatedRouterMembers"}
SCAN_ROOTS = ["litellm", "enterprise"]
SKIP_DIRS = {

5
tests/integration/_support/anthropic_refusal.py
View file

@ -29,6 +29,7 @@ DOCS_STREAMING_REFUSAL_URL: Final = (
)
# Copied verbatim from DOCS_REFUSAL_URL and DOCS_STREAMING_REFUSAL_URL on 2026-10-12.
DOCS_REFUSAL_EXPLANATION: Final = "This request was declined because it could enable cyber harm."
DOCS_REFUSAL: Final[Event] = {
"id": "msg_01XFUDYJgAACzvnptvVoYEL",
"type": "message",
@ -39,7 +40,7 @@ DOCS_REFUSAL: Final[Event] = {
"stop_details": {
"type": "refusal",
"category": "cyber",
"explanation": "This request was declined because it could enable cyber harm.",
"explanation": DOCS_REFUSAL_EXPLANATION,
},
"usage": {"input_tokens": 412, "output_tokens": 0},
}
@ -50,7 +51,7 @@ DOCS_PARTIAL_REFUSAL: Final[Event] = {
"stop_details": {
"type": "refusal",
"category": "cyber",
"explanation": "This request was declined because it could enable cyber harm.",
"explanation": DOCS_REFUSAL_EXPLANATION,
},
}

98
tests/integration/authorization/test_agent_access.py
View file

@ -1,7 +1,8 @@
import itertools
import json
import uuid
from collections.abc import Callable
from typing import Final
from typing import Final, Literal
import pytest
from integration._support.client import Gateway, Scenario
@ -18,6 +19,22 @@ class _JsonRpcResult(BaseModel):
result: dict[str, JsonValue]
class _JsonRpcErrorBody(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
code: int
message: str
data: dict[str, JsonValue]
class _JsonRpcError(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
jsonrpc: Literal["2.0"]
id: str | int | None
error: _JsonRpcErrorBody
_STREAM_EVENTS: Final = (
{"kind": "status-update", "taskId": "t1", "contextId": "c1", "status": {"state": "working"}, "final": False},
{
@ -163,6 +180,14 @@ def test_agent_permissions_on_keys_teams_and_access_groups_gate_a2a_send_and_car
separators=(",", ":"),
)
def jsonrpc_refusal(text: str, identity: str) -> bool:
error: Final = _JsonRpcError.model_validate_json(text)
return (error.id, error.error.code, error.error.message) == (
marker,
-32000,
f"Agent '{identity}' is not allowed for your key/team. Contact proxy admin for access.",
)
spellings: Final = (
("send by id", lambda agent, name: f"/a2a/{agent}", send_payload, "message/send"),
("send by name", lambda agent, name: f"/a2a/{name}", send_payload, "message/send"),
@ -224,13 +249,13 @@ def test_agent_permissions_on_keys_teams_and_access_groups_gate_a2a_send_and_car
for label, key in callers.items():
for spelling, status, text in attempt(key, denied, denied_name):
refused: Final = refusal(denied_name if "by name" in spelling else denied)
assert (status, text) == (403, refused), f"{label} {spelling}: {status} {text}"
identity: Final = denied_name if "by name" in spelling else denied
assert status == 403 and jsonrpc_refusal(text, identity), f"{label} {spelling}: {status} {text}"
for spelling, payload in denied_tasks:
response: Final = gateway.client.post(
f"/a2a/{denied}", headers={"Authorization": f"Bearer {key}"}, json=payload
)
assert (response.status_code, response.text) == (403, refusal(denied)), (
assert response.status_code == 403 and jsonrpc_refusal(response.text, denied), (
f"{label} {spelling}: {response.status_code} {response.text}"
)
assert card(key, denied) == (403, refusal(denied)), label
@ -278,13 +303,72 @@ def test_agent_permissions_on_keys_teams_and_access_groups_gate_a2a_send_and_car
emptied: Final = gateway.request("PUT", f"/v1/access_group/{group_id}", {"access_agent_ids": []})
assert emptied.status_code == 200, emptied.text
emptied_outcomes: Final = attempt(callers["access group"], allowed, allowed_name)
assert emptied_outcomes == tuple(
(spelling, 403, refusal(allowed_name if "by name" in spelling else allowed))
for spelling, _, _ in emptied_outcomes
assert all(
status == 403 and jsonrpc_refusal(text, allowed_name if "by name" in spelling else allowed)
for spelling, status, text in emptied_outcomes
), emptied_outcomes
assert wire.drain() == (), emptied_outcomes
def test_a2a_refusals_answer_as_jsonrpc_errors_that_echo_the_callers_request_id(gateway: Gateway) -> None:
marker: Final = "a2arefuse" + uuid.uuid4().hex[:12]
def upstream(request: Request) -> Reply:
return Reply(body=json.dumps(_peer_card(wire.url, marker)).encode())
with wire_server(upstream) as wire, gateway.scenario() as scenario:
agent: Final = _register_agent(gateway, scenario, marker, wire.url)
refused_callers: Final = (
(
"over budget",
scenario.key(object_permission={"agents": [agent]}, max_budget=0.0001, spend=1.0),
422,
"Budget has been exceeded!",
"budget_exceeded",
),
(
"agent not granted",
scenario.key(object_permission={"agents": [marker + "-other"]}),
403,
f"Agent '{agent}' is not allowed for your key/team. Contact proxy admin for access.",
None,
),
("unknown key", "sk-" + uuid.uuid4().hex, 401, "Authentication Error", None),
)
message: Final = {
"kind": "message",
"role": "user",
"messageId": marker + "-in",
"parts": [{"kind": "text", "text": "ping"}],
}
requests: Final = tuple(
(path, method, request_id)
for path, (method, request_id) in itertools.product(
(f"/a2a/{agent}", f"/a2a/{agent}/message/send", f"/v1/a2a/{agent}/message/send"),
(("message/send", marker), ("message/stream", 7), ("SendMessage", "l10")),
)
)
for label, key, status, message_start, error_type in refused_callers:
for path, method, request_id in requests:
response = gateway.client.post(
path,
headers={"Authorization": f"Bearer {key}", "Accept": "application/json, text/event-stream"},
json={"jsonrpc": "2.0", "id": request_id, "method": method, "params": {"message": message}},
)
context = f"{label} {path} {method}: {response.status_code} {response.text}"
assert response.status_code == status, context
error = _JsonRpcError.model_validate_json(response.text)
assert (error.id, error.error.code) == (request_id, -32000), context
assert error.error.message.startswith(message_start), context
assert error_type is None or error.error.data["type"] == error_type, context
unparseable: Final = gateway.client.post(
f"/a2a/{agent}", headers={"Authorization": "Bearer sk-" + uuid.uuid4().hex}, content=b"not json"
)
assert unparseable.status_code == 400, unparseable.text
assert _JsonRpcError.model_validate_json(unparseable.text).id is None, unparseable.text
assert wire.drain() == ()
def test_legacy_agent_json_card_path_applies_agent_permissions_to_non_admin_keys(gateway: Gateway) -> None:
pytest.skip(
"BUG: GET /a2a/{agent_id}/.well-known/agent.json returns 401 'Only proxy admin can be used' for non-admin keys because agent_inference_routes lists only agent-card.json"

9
tests/integration/authorization/test_agent_least_privilege.py
View file

@ -221,8 +221,13 @@ def test_agent_access_default_deny_denies_ungranted_key_but_keeps_explicit_grant
denied_sends: Final = tuple(
_send(candidate, agent_id, key_a, f"denied-{index}-{uuid.uuid4().hex}") for index in range(8)
)
_assert_denial(denied_sends[0], agent_id)
assert all(response.json() == denied_sends[0].json() for response in denied_sends), [
assert denied_sends[0].status_code == 403, denied_sends[0].text
assert denied_sends[0].json()["error"] == {
"code": -32000,
"message": f"Agent '{agent_id}' is not allowed for your key/team. Contact proxy admin for access.",
"data": denied_sends[0].json()["error"]["data"],
}, denied_sends[0].text
assert all(response.json()["error"] == denied_sends[0].json()["error"] for response in denied_sends), [
response.text for response in denied_sends
]
denied_completion: Final = _chat_completion(candidate, agent_name, key_a)

99
tests/integration/compatibility/test_a2a_wire_versions.py
View file

@ -348,6 +348,105 @@ _TRANSPORT_HEADERS: Final = frozenset(
)
def _a2a_auth_requests(
gateway: Gateway, method: str, agent_options: dict[str, JsonValue], client_headers: dict[str, str]
) -> tuple[Request, ...]:
marker: Final = "a2akey" + uuid.uuid4().hex[:12]
def upstream(request: Request) -> Reply:
if request.method == "GET":
return Reply(body=json.dumps(_peer_card(wire.url, marker)).encode())
body: Final = json.loads(request.body)
if body["method"] == "message/stream":
frames: Final = tuple(
f"data: {json.dumps({'jsonrpc': '2.0', 'id': body['id'], 'result': event})}\n\n".encode()
for event in _STREAM_EVENTS
)
return Reply(content_type="text/event-stream", chunks=frames)
if body["method"] == "tasks/get":
return Reply(
body=json.dumps(
{
"jsonrpc": "2.0",
"id": body["id"],
"result": {"kind": "task", "id": "t1", "contextId": "c1", "status": {"state": "completed"}},
}
).encode()
)
return _message_reply(body, "key test")
with wire_server(upstream) as wire, gateway.scenario() as scenario:
identity: Final = _register_agent(
gateway,
scenario,
{"agent_name": marker, "agent_card_params": _peer_card(wire.url, marker, "1.0"), **agent_options},
)
params: Final = {"id": "t1"} if method == "tasks/get" else {"message": _MESSAGE_V10}
headers: Final = {
"Authorization": f"Bearer {gateway.key}",
"a2a-version": "1.0",
**{name.replace("{agent}", marker): value for name, value in client_headers.items()},
}
response: Final = gateway.client.post(
f"/a2a/{identity}/message/send",
headers=headers,
json={"jsonrpc": "2.0", "id": marker, "method": method, "params": params},
)
assert response.status_code == 200, response.text
assert "error" not in response.text, response.text
return wire.drain()
@pytest.mark.parametrize(
("method", "observed_method"),
(("SendMessage", "GET"), ("SendMessage", "POST"), ("SendStreamingMessage", "POST"), ("tasks/get", "POST")),
ids=("card", "send", "stream", "task"),
)
def test_regression_db_agent_api_key_is_bearer_on_each_a2a_hop(
gateway: Gateway, method: str, observed_method: str
) -> None:
key: Final = "synthetic-backend-" + uuid.uuid4().hex
requests: Final = _a2a_auth_requests(gateway, method, {"litellm_params": {"api_key": key}}, {})
observed: Final = tuple(request for request in requests if request.method == observed_method)
assert len(observed) == 1, requests
assert observed[0].headers.get("authorization") == f"Bearer {key}", observed[0].headers
@pytest.mark.parametrize("header_name", ("Authorization", "aUtHoRiZaTiOn"))
def test_guard_static_authorization_overrides_api_key(gateway: Gateway, header_name: str) -> None:
requests: Final = _a2a_auth_requests(
gateway,
"SendMessage",
{"litellm_params": {"api_key": "backend-key"}, "static_headers": {header_name: "Bearer static-key"}},
{},
)
assert tuple(request.method for request in requests) == ("GET", "POST"), requests
assert tuple(request.headers.get("authorization") for request in requests) == ("Bearer static-key",) * 2
@pytest.mark.parametrize("forwarding", ("configured", "convention"))
def test_guard_forwarded_authorization_arrives_with_api_key(gateway: Gateway, forwarding: str) -> None:
options: Final = {
"litellm_params": {"api_key": "backend-key"},
**({"extra_headers": ["authorization"]} if forwarding == "configured" else {}),
}
client_headers: Final = (
{"x-litellm-api-key": gateway.key, "Authorization": "Bearer client-agent-key"}
if forwarding == "configured"
else {"x-a2a-{agent}-authorization": "Bearer client-agent-key"}
)
requests: Final = _a2a_auth_requests(gateway, "SendMessage", options, client_headers)
assert tuple(request.method for request in requests) == ("GET", "POST"), requests
assert tuple(request.headers.get("authorization") for request in requests) == ("Bearer client-agent-key",) * 2
def test_guard_no_agent_key_or_proxy_caller_key_is_sent(gateway: Gateway) -> None:
requests: Final = _a2a_auth_requests(gateway, "SendMessage", {}, {})
assert tuple(request.method for request in requests) == ("GET", "POST"), requests
assert tuple(request.headers.get("authorization") for request in requests) == (None, None)
assert all(gateway.key not in value for request in requests for value in request.headers.values()), requests
@pytest.mark.parametrize(
("version", "method", "message"),
(("0.3", "message/stream", _MESSAGE_V03), ("1.0", "SendStreamingMessage", _MESSAGE_V10)),

23
tests/integration/observability/test_otel_team_capture_message_content.py
View file

@ -303,6 +303,23 @@ def _with_content(url: str, since: int, sent: Sent) -> tuple[Span, ...]:
)
def _is_root(span: Span) -> bool:
return span["kind"] == SERVER and span["name"].startswith("POST")
def _with_root(url: str, since: int, sent: Sent) -> tuple[Span, ...]:
"""The copy of ``sent`` at ``url`` once its model span, the marker and its server span arrived."""
return eventually(
lambda: _of_request(url, since, sent),
lambda spans: (
bool(_model_spans(spans))
and any(_carries(span, sent.marker) for span in spans)
and any(_is_root(span) for span in spans)
),
seconds=40,
)
def _arrived(
url: str, since: int, sent: Sent, expected: frozenset[str], credential: Credential | None = None
) -> tuple[Span, ...]:
@ -626,11 +643,7 @@ def test_stalled_tenant_during_a_mixed_burst_keeps_serving_and_each_teams_policy
configure_sink(rig.sinks.tenant, paused=False)
assert [result.status for result in sent] == [200] * 30, [result.text for result in sent if result.status != 200]
for (key, _, _), result in zip(plan, sent, strict=True):
roots = [
span
for span in _with_content(rig.sinks.operator, cursors.operator, result)
if span["kind"] == SERVER and span["name"].startswith("POST ")
]
roots = [span for span in _with_root(rig.sinks.operator, cursors.operator, result) if _is_root(span)]
assert len(roots) == 1, roots
if key == restricted:
_assert_redacted_twin(rig, result, cursors, rig.sinks.tenant, cursors.tenant)

13
tests/integration/providers/test_anthropic_refusal_stop_details_wire.py
View file

@ -18,6 +18,7 @@ from integration._support.anthropic_refusal import (
DECLARED_PARTIAL_JSON,
DOCS_PARTIAL_TEXT,
DOCS_REFUSAL,
DOCS_REFUSAL_EXPLANATION,
DOCS_STOP_DETAILS,
ENVELOPED_PARTIAL_JSON,
FIXTURES,
@ -651,8 +652,7 @@ def test_family_messages_stream_refusal_keeps_frames_and_logs_the_stop_details(r
model: Final = _deployment(scenario, provider, rig)
key: Final = scenario.key()
events: Final = _messages_stream(rig, key, model, refusal_prompt("refusal", _marker()))
native: Final = {"stop_reason": "refusal"} if provider == "snowflake" else REFUSAL_FIELDS
assert _native_deltas(events) == (native,), events
assert _native_deltas(events) == (REFUSAL_FIELDS,), events
assert _relayed_in(events) == (), events
assert _paths(rig.wire) == (_STREAM_TARGETS[provider][1],)
logged: Final = _logged_response(rig, key)
@ -783,10 +783,10 @@ def test_responses_cache_hit_logs_the_same_stop_details_as_the_miss(rig: Rig) ->
served_from_cache: Final = tuple(event.get("cache_hit") is True for event in logged)
assert (len(_paths(rig.wire)), served_from_cache.count(True)) == (1, 1), (served_from_cache, logged)
assert relayed(first) == ((_RESPONSE_FIELDS, REFUSAL_FIELDS),), first
assert tuple(relayed(event["response"]) for event in logged) == (((_CHAT_FIELDS, REFUSAL_FIELDS),),) * 2, (
served_from_cache,
logged,
)
assert tuple(relayed(event["response"]) for event in logged) == (
((_CHAT_FIELDS, REFUSAL_FIELDS),),
relayed(first),
), (served_from_cache, logged)
def test_passthrough_is_not_served_from_the_response_cache(rig: Rig) -> None:
@ -1310,6 +1310,7 @@ def test_chat_refusal_body_serializes_the_stop_details_beside_the_existing_messa
"message": {
"content": content,
"role": "assistant",
"refusal": DOCS_REFUSAL_EXPLANATION,
"provider_specific_fields": {**_ANTHROPIC_MESSAGE_FIELDS, **REFUSAL_FIELDS},
},
}

36
tests/integration/providers/test_cache_control_tool_call_marks_wire.py
View file

@ -354,7 +354,7 @@ _WEB_RESULTS: Final[dict[str, JsonValue]] = {
),
pytest.param(
{},
[ASK_LABEL, tool_use_label(CITIES[0]), tool_use_label(CITIES[1]), "assistant:tool_use:srvtoolu_web"],
[SYSTEM_LABEL, ASK_LABEL, tool_use_label(CITIES[0]), tool_use_label(CITIES[1])],
id="server-tool-without-result",
),
),
@ -413,16 +413,14 @@ _GEMINI_CACHE_WRITE: Final = [
@pytest.mark.parametrize(
("calls", "expected"),
"calls",
(
pytest.param(
marked_calls(), [("POST", "/models/gemini-3.8-flash:generateContent")], id="tool-calls-fill-the-cap"
),
pytest.param([tool_call(city) for city in CITIES], _GEMINI_CACHE_WRITE, id="unmarked-tool-calls"),
pytest.param(marked_calls(), id="tool-calls-fill-the-cap"),
pytest.param([tool_call(city) for city in CITIES], id="unmarked-tool-calls"),
),
)
def test_gemini_context_cache_follows_the_cap_census(
gateway: Gateway, calls: list[JsonValue], expected: list[tuple[str, str]]
def test_gemini_context_cache_takes_the_system_message_with_the_marked_ask(
gateway: Gateway, calls: list[JsonValue]
) -> None:
marker: Final = new_marker()
long_system: Final[JsonValue] = {"role": "system", "content": "lorem " * (2 * get_prompt_cache_min_tokens(_GEMINI))}
@ -436,8 +434,18 @@ def test_gemini_context_cache_follows_the_cap_census(
)
status, _, text = post_chat(gateway, chat_body(model, messages))
assert status == 200, text
targets: Final = [(request.method, request.target.split("?")[0]) for request in wire.drain()]
assert targets == expected
received: Final = wire.drain()
targets: Final = [(request.method, request.target.split("?")[0]) for request in received]
cache_writes: Final = [
request.body for request in received if request.method == "POST" and "cachedContents" in request.target
]
generate: Final = _JSON_OBJECT.validate_json(received[-1].body)
assert (
targets,
[b"lorem lorem" in body for body in cache_writes],
generate.get("cachedContent"),
"system_instruction" in generate,
) == (_GEMINI_CACHE_WRITE, [True], "cachedContents/census", False), (cache_writes, generate)
def _openai_marks(request: Request) -> list[str]:
@ -747,11 +755,11 @@ def test_request_level_points_skip_injection_when_client_marks_fill_the_cap(gate
_OTHER_CALL_RESULT: Final[dict[str, JsonValue]] = {
"web_search_results": [{"type": "web_search_tool_result", "tool_use_id": "srvtoolu_other", "content": []}]
}
_SERVER_MARK_COUNTED: Final = [
_UNRUN_SERVER_MARK_UNCOUNTED: Final = [
SYSTEM_LABEL,
ASK_LABEL,
tool_use_label(CITIES[0]),
tool_use_label(CITIES[1]),
"assistant:tool_use:srvtoolu_web",
]
@ -767,7 +775,7 @@ _SERVER_MARK_COUNTED: Final = [
pytest.param(_OTHER_CALL_RESULT, id="result-for-another-call"),
),
)
def test_malformed_provider_specific_fields_count_the_server_tool_call_mark(
def test_malformed_provider_specific_fields_leave_the_unrun_server_tool_call_mark_uncounted(
gateway: Gateway, fields: JsonValue
) -> None:
marker: Final = new_marker()
@ -779,7 +787,7 @@ def test_malformed_provider_specific_fields_count_the_server_tool_call_mark(
)
assert status == 200, text
labels: Final = anthropic_labels(_only_request(wire))
assert labels == _SERVER_MARK_COUNTED
assert labels == _UNRUN_SERVER_MARK_UNCOUNTED
@pytest.mark.parametrize(

81
tests/integration/providers/test_watsonx_passthrough_router_credentials_wire.py Normal file
View file

@ -0,0 +1,81 @@
import json
import uuid
from pathlib import Path
from typing import Final
import yaml
from integration._support.client import Gateway, JsonValue, object_value
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
_DEPLOYMENT_KEY: Final = "synthetic-watsonx-deployment-key"
_IAM_TOKEN: Final = "synthetic-watsonx-iam-token"
_WATSONX_ENV: Final = (
"WATSONX_API_KEY",
"WATSONX_APIKEY",
"WX_API_KEY",
"WATSONX_API_BASE",
"WATSONX_URL",
"WX_URL",
"WML_URL",
"WATSONX_TOKEN",
"WATSONX_ZENAPIKEY",
)
def _config(directory: Path, api_base: str) -> Path:
stock: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
config: Final[dict[str, JsonValue]] = {
**stock,
"model_list": [
{
"model_name": "granite",
"litellm_params": {
"model": "watsonx/ibm/granite-3-8b-instruct",
"api_key": _DEPLOYMENT_KEY,
"api_base": api_base,
"project_id": "synthetic-project",
"use_in_pass_through": True,
},
}
],
}
path: Final = directory / "watsonx-passthrough.yaml"
path.write_text(yaml.safe_dump(config))
return path
def test_watsonx_passthrough_uses_the_flagged_deployment_key_and_base(gateway: Gateway, tmp_path: Path) -> None:
marker: Final = "watsonx-" + uuid.uuid4().hex
def respond(request: Request) -> Reply:
if request.method == "POST" and request.target == "/identity/token":
assert f"apikey={_DEPLOYMENT_KEY}".encode() in request.body, request.body
return Reply(body=json.dumps({"access_token": _IAM_TOKEN, "expires_in": 3600}).encode())
if request.method == "POST" and request.target.startswith("/ml/v1/text/generation?"):
assert request.headers["authorization"] == f"Bearer {_IAM_TOKEN}", request.headers
assert marker.encode() in request.body, request.body[:400]
return Reply(
body=json.dumps(
{"model_id": "ibm/granite-3-8b-instruct", "results": [{"generated_text": marker}]}
).encode()
)
return Reply(status=404)
with wire_server(respond) as wire:
with owned_proxy(
gateway,
tmp_path,
{"WATSONX_IAM_URL": f"{wire.url}/identity/token"},
config=_config(tmp_path, wire.url),
remove_environment=_WATSONX_ENV,
) as candidate:
response: Final = candidate.request(
"POST",
"/watsonx/ml/v1/text/generation",
{"model_id": "ibm/granite-3-8b-instruct", "input": marker, "project_id": "synthetic-project"},
)
assert response.status_code == 200, response.text
assert response.json()["results"][0]["generated_text"] == marker
forwarded: Final = tuple((request.method, request.target.split("?")[0]) for request in wire.drain())
assert forwarded == (("POST", "/identity/token"), ("POST", "/ml/v1/text/generation")), forwarded

10
tests/integration/streaming/test_decision_model_buffering.py
View file

@ -229,11 +229,17 @@ def _frames(identity: str, endpoint: str) -> tuple[str, ...]:
'object': 'chat.completion.chunk',
'created': 1,
'model': 'scripted-chat',
'choices': [{'index': 0, 'delta': {'content': text}, 'finish_reason': None}],
'choices': [
{
'index': 0,
'delta': {'content': text},
'finish_reason': 'stop' if index == len(pieces) - 1 else None,
}
],
}
)
}"
for text in pieces
for index, text in enumerate(pieces)
),
"data: [DONE]",
)

564
tests/integration/streaming/test_stream_clean_eof_without_finish_reason.py Normal file
View file

@ -0,0 +1,564 @@
import asyncio
import itertools
import json
import uuid
from collections import Counter
from collections.abc import Callable, Iterator, Mapping
from concurrent.futures import ThreadPoolExecutor
from contextlib import ExitStack
from pathlib import Path
from typing import Final, Literal
import anthropic
import httpx
import openai
import pytest
from integration._support.anthropic_sse import CONTENT_BLOCK_START, CONTENT_TAIL, message_start, text_delta
from integration._support.client import Gateway, Scenario, eventually, string_value
from integration._support.database import read_rows, scratch_database
from integration._support.database_relay import database_relay, plaintext_database_url
from integration._support.openai_wire import answering_model_discovery, posted_targets
from integration._support.process import owned_proxy, owned_proxy_process
from integration._support.responses_stream import data_frames, spend_row, sse_events
from integration._support.wire import Reply, Request, Wire, wire_server
from integration.providers._count_tokens_system_lift import (
anthropic_client,
async_anthropic_client,
async_openai_client,
openai_client,
)
from pydantic import JsonValue
Ending = Literal["ending-complete", "ending-clean-eof", "ending-done-without-finish", "ending-done-only"]
Family = Literal["openai", "anthropic", "hosted_vllm", "gemini", "langgraph", "agentcore"]
_ENDINGS: Final[tuple[Ending, ...]] = (
"ending-complete",
"ending-clean-eof",
"ending-done-without-finish",
"ending-done-only",
)
_PARTIAL: Final = ('{"findings":[{"title":', '"unfinished')
_BACKEND: Final = "gpt-5.6"
_CLAUDE: Final = "claude-sonnet-5-5"
_GEMINI: Final = "gemini-3.8-flash"
_AGENT_RUNTIME: Final = "arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/integration-agent"
_PROVIDER_KEY: Final = "synthetic-provider-key"
_CALL_ID_HEADER: Final = "x-litellm-call-id"
_USAGE: Final = {"prompt_tokens": 9, "completion_tokens": 7, "total_tokens": 16}
_TARGETS: Final[Mapping[Family, str]] = {
"openai": "/chat/completions",
"hosted_vllm": "/chat/completions",
"anthropic": "/v1/messages",
"gemini": f"/models/{_GEMINI}:streamGenerateContent",
"langgraph": "/runs/stream",
"agentcore": "/runtimes/",
}
_DEPLOYMENTS: Final[Mapping[Family, str]] = {
"openai": f"openai/{_BACKEND}",
"hosted_vllm": f"hosted_vllm/{_BACKEND}",
"anthropic": f"anthropic/{_CLAUDE}",
"gemini": f"gemini/{_GEMINI}",
"langgraph": "langgraph/agent",
"agentcore": f"bedrock/agentcore/{_AGENT_RUNTIME}",
}
def _data(frame: JsonValue) -> bytes:
return b"data: " + json.dumps(frame).encode() + b"\n\n"
def _chat_frames(identity: str, ending: Ending) -> tuple[bytes, ...]:
chunk: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": _BACKEND}
content: Final = (
_data({**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": _PARTIAL[0]}}]}),
_data({**chunk, "choices": [{"index": 0, "delta": {"content": _PARTIAL[1]}, "finish_reason": None}]}),
)
match ending:
case "ending-complete":
finish: Final = _data({**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]})
usage: Final = _data({**chunk, "choices": [], "usage": _USAGE})
return (*content, finish + usage + b"data: [DONE]\n\n")
case "ending-clean-eof":
return content
case "ending-done-without-finish":
return (*content, b"data: [DONE]\n\n")
case "ending-done-only":
return (b"data: [DONE]\n\n",)
def _anthropic_frames(identity: str, ending: Ending) -> tuple[bytes, ...]:
content: Final = (message_start(identity, _CLAUDE), CONTENT_BLOCK_START, *map(text_delta, _PARTIAL))
return (*content, CONTENT_TAIL) if ending == "ending-complete" else content
def _gemini_frames(ending: Ending) -> tuple[bytes, ...]:
"""Gemini puts finishReason on the last text chunk and can follow it with a usage-only chunk."""
finished: Final = ending == "ending-complete"
content: Final = tuple(
_data(
{
"candidates": [
{
"content": {"parts": [{"text": part}], "role": "model"},
"index": 0,
**({"finishReason": "STOP"} if finished and part == _PARTIAL[-1] else {}),
}
]
}
)
for part in _PARTIAL
)
usage: Final = _data({"usageMetadata": {"promptTokenCount": 9, "candidatesTokenCount": 7, "totalTokenCount": 16}})
return (*content, usage) if finished else content
def _event(name: str, frame: JsonValue) -> bytes:
return f"event: {name}\n".encode() + _data(frame)
def _langgraph_frames() -> tuple[bytes, ...]:
"""A LangGraph server's run stream in its default mode, its run metadata first and then the state after each step"""
question: Final = {"type": "human", "content": "Answer in two steps."}
answers: Final = tuple({"type": "ai", "content": part} for part in _PARTIAL)
return (
_event("metadata", {"run_id": "run-1", "attempt": 1}),
*(_event("values", {"messages": [question, *answers[:step]]}) for step in range(len(answers) + 1)),
)
def _agentcore_frames(ending: Ending) -> tuple[bytes, ...]:
"""A runtime that forwards Bedrock Converse events without the metadata event, so its turn ends on messageStop"""
start: Final = _data({"event": {"messageStart": {"role": "assistant"}}})
content: Final = tuple(
_data({"event": {"contentBlockDelta": {"delta": {"text": part}, "contentBlockIndex": 0}}}) for part in _PARTIAL
)
end: Final = (
_data({"event": {"contentBlockStop": {"contentBlockIndex": 0}}}),
_data({"event": {"messageStop": {"stopReason": "end_turn"}}}),
)
return (start, *content, *end) if ending == "ending-complete" else (start, *content)
def _frames(family: Family, identity: str, ending: Ending) -> tuple[bytes, ...]:
match family:
case "agentcore":
return _agentcore_frames(ending)
case "anthropic":
return _anthropic_frames(identity, ending)
case "gemini":
return _gemini_frames(ending)
case "langgraph":
return _langgraph_frames()
case "openai" | "hosted_vllm":
return _chat_frames(identity, ending)
def _ending(request: Request) -> Ending:
endings: Final = [ending for ending in _ENDINGS if ending.encode() in request.body]
assert len(endings) == 1, request.body
return endings[0]
def _upstream(family: Family, ending: Ending | None = None) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
assert request.method == "POST" and _TARGETS[family] in request.target, request.target
frames: Final = _frames(family, f"upstream-{uuid.uuid4().hex}", ending or _ending(request))
return Reply(content_type="text/event-stream", chunks=frames)
return answering_model_discovery(respond)
def _deploy(scenario: Scenario, family: Family, wire: Wire) -> str:
if family == "agentcore":
return scenario.model(
model=_DEPLOYMENTS[family],
aws_bedrock_runtime_endpoint=wire.url,
aws_access_key_id="AKIAINTEGRATIONAGENT",
aws_secret_access_key="synthetic-agentcore-secret",
aws_region_name="us-west-2",
)
return scenario.model(model=_DEPLOYMENTS[family], api_base=wire.url, api_key=_PROVIDER_KEY)
def _prompt(ending: Ending) -> str:
return f"Return JSON {ending} {uuid.uuid4().hex}"
def _post(gateway: Gateway, path: str, body: Mapping[str, JsonValue]) -> tuple[str, str]:
response: Final = gateway.request("POST", path, {**body, "stream": True, "num_retries": 0})
assert response.status_code == 200, response.text
return response.text, response.headers[_CALL_ID_HEADER]
def _chat(gateway: Gateway, model: str, ending: Ending) -> tuple[str, str]:
return _post(
gateway,
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": _prompt(ending)}],
"stream_options": {"include_usage": True},
},
)
def _choices(text: str) -> Iterator[Mapping[str, JsonValue]]:
for frame in data_frames(text):
yield from (choice for choice in frame.get("choices") or [] if isinstance(choice, dict))
def _finish_reasons(text: str) -> tuple[str, ...]:
return tuple(string_value(choice["finish_reason"]) for choice in _choices(text) if choice.get("finish_reason"))
def _chat_text(text: str) -> str:
deltas: Final = (choice.get("delta") for choice in _choices(text))
return "".join(string_value(delta.get("content") or "") for delta in deltas if isinstance(delta, dict))
def _chat_errors(text: str) -> tuple[Mapping[str, JsonValue], ...]:
return tuple(frame for frame in data_frames(text) if "error" in frame)
def _event_types(text: str) -> tuple[str, ...]:
return tuple(event for event, _ in sse_events(text))
@pytest.mark.parametrize("family", ["openai", "anthropic", "agentcore"])
def test_chat_stream_that_ends_cleanly_without_a_finish_reason_reaches_the_client_as_an_error(
gateway: Gateway, family: Family
) -> None:
with wire_server(_upstream(family)) as wire, gateway.scenario() as scenario:
model: Final = _deploy(scenario, family, wire)
text, call_id = _chat(gateway, model, "ending-clean-eof")
assert _chat_text(text) == "".join(_PARTIAL), text
assert _finish_reasons(text) == (), text
assert len(_chat_errors(text)) == 1, text
assert spend_row(call_id)["status"] == "failure"
def test_chat_stream_that_sends_done_without_a_finish_reason_reaches_the_client_as_an_error(gateway: Gateway) -> None:
with wire_server(_upstream("openai")) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=_DEPLOYMENTS["openai"], api_base=wire.url, api_key=_PROVIDER_KEY)
text, call_id = _chat(gateway, model, "ending-done-without-finish")
assert _chat_text(text) == "".join(_PARTIAL), text
assert _finish_reasons(text) == (), text
assert len(_chat_errors(text)) == 1, text
assert spend_row(call_id)["status"] == "failure"
def test_chat_stream_that_sends_only_done_still_completes_empty(gateway: Gateway) -> None:
"""An upstream can answer with only [DONE]: Gemini's OpenAI-compatible endpoint did on 2026-10-20 when
reasoning used up max_tokens."""
with wire_server(_upstream("openai")) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=_DEPLOYMENTS["openai"], api_base=wire.url, api_key=_PROVIDER_KEY)
text, _ = _chat(gateway, model, "ending-done-only")
assert _chat_text(text) == "", text
assert _finish_reasons(text) == ("stop",), text
assert _chat_errors(text) == (), text
assert text.rstrip().endswith("data: [DONE]"), text
@pytest.mark.parametrize("family", ["openai", "anthropic", "gemini", "langgraph", "agentcore"])
def test_chat_stream_with_a_provider_finish_reason_still_completes(gateway: Gateway, family: Family) -> None:
with wire_server(_upstream(family)) as wire, gateway.scenario() as scenario:
model: Final = _deploy(scenario, family, wire)
text, call_id = _chat(gateway, model, "ending-complete")
assert _chat_text(text) == "".join(_PARTIAL), text
assert _finish_reasons(text) == ("stop",), text
assert _chat_errors(text) == (), text
assert text.rstrip().endswith("data: [DONE]"), text
assert spend_row(call_id)["status"] == "success"
def test_chat_stream_from_a_mock_response_deployment_still_completes(gateway: Gateway) -> None:
"""mock_completion streams never carry a finish reason, so a mocked deployment must keep its synthetic stop."""
with gateway.scenario() as scenario:
model: Final = scenario.model(model=_DEPLOYMENTS["openai"], mock_response="".join(_PARTIAL))
text, call_id = _chat(gateway, model, "ending-clean-eof")
assert _chat_text(text) == "".join(_PARTIAL), text
assert _finish_reasons(text) == ("stop",), text
assert _chat_errors(text) == (), text
assert spend_row(call_id)["status"] == "success"
@pytest.mark.parametrize(("ending", "terminal"), [("ending-complete", True), ("ending-clean-eof", False)])
def test_messages_over_a_chat_bridged_deployment_reports_end_turn_only_for_a_finished_stream(
gateway: Gateway, ending: Ending, terminal: bool
) -> None:
with wire_server(_upstream("hosted_vllm")) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=_DEPLOYMENTS["hosted_vllm"], api_base=wire.url, api_key=_PROVIDER_KEY)
text, call_id = _post(
gateway,
"/v1/messages",
{"model": model, "max_tokens": 64, "messages": [{"role": "user", "content": _prompt(ending)}]},
)
types: Final = _event_types(text)
stop_reasons: Final = tuple(
data["delta"]["stop_reason"]
for event, data in sse_events(text)
if event == "message_delta" and isinstance(data.get("delta"), dict)
)
assert types[0] == "message_start", types
assert ("message_stop" in types) is terminal, types
assert ("error" in types) is not terminal, types
assert stop_reasons == (("end_turn",) if terminal else ()), types
assert spend_row(call_id)["status"] == ("success" if terminal else "failure")
@pytest.mark.parametrize(("ending", "terminal"), [("ending-complete", True), ("ending-clean-eof", False)])
def test_responses_over_a_chat_bridged_anthropic_deployment_completes_only_a_finished_stream(
gateway: Gateway, ending: Ending, terminal: bool
) -> None:
with wire_server(_upstream("anthropic")) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=_DEPLOYMENTS["anthropic"], api_base=wire.url, api_key=_PROVIDER_KEY)
text, call_id = _post(gateway, "/v1/responses", {"model": model, "input": _prompt(ending)})
types: Final = tuple(string_value(frame.get("type") or "") for frame in data_frames(text))
assert types[:1] == ("response.created",), text
assert ("response.completed" in types) is terminal, text
assert (types[-1] in {"error", "response.failed"}) is not terminal, text
assert spend_row(call_id)["status"] == ("success" if terminal else "failure")
def _drain_openai_sdk(gateway: Gateway, model: str, sync_mode: bool) -> None:
messages: Final = [{"role": "user", "content": _prompt("ending-clean-eof")}]
if sync_mode:
for _ in openai_client(gateway).chat.completions.create(model=model, messages=messages, stream=True):
pass
return
async def drain() -> None:
stream: Final = await async_openai_client(gateway).chat.completions.create(
model=model, messages=messages, stream=True
)
async for _ in stream:
pass
asyncio.run(drain())
def _drain_anthropic_sdk(gateway: Gateway, model: str, sync_mode: bool) -> None:
messages: Final = [{"role": "user", "content": _prompt("ending-clean-eof")}]
if sync_mode:
for _ in anthropic_client(gateway).messages.create(model=model, max_tokens=64, messages=messages, stream=True):
pass
return
async def drain() -> None:
stream: Final = await async_anthropic_client(gateway).messages.create(
model=model, max_tokens=64, messages=messages, stream=True
)
async for _ in stream:
pass
asyncio.run(drain())
@pytest.mark.parametrize("sync_mode", [True, False])
def test_openai_sdk_stream_cut_before_a_finish_reason_raises(gateway: Gateway, sync_mode: bool) -> None:
with wire_server(_upstream("openai")) as wire, gateway.scenario() as scenario:
model: Final = _deploy(scenario, "openai", wire)
with pytest.raises(openai.APIError, match="Stream ended before a finish_reason was received"):
_drain_openai_sdk(gateway, model, sync_mode)
@pytest.mark.parametrize("sync_mode", [True, False])
def test_anthropic_sdk_messages_stream_cut_before_a_finish_reason_raises(gateway: Gateway, sync_mode: bool) -> None:
with wire_server(_upstream("hosted_vllm")) as wire, gateway.scenario() as scenario:
model: Final = _deploy(scenario, "hosted_vllm", wire)
with pytest.raises(anthropic.APIError, match="Stream ended before a finish_reason was received"):
_drain_anthropic_sdk(gateway, model, sync_mode)
def test_responses_stream_cut_on_its_deployment_finishes_on_the_request_fallback(gateway: Gateway) -> None:
with (
wire_server(_upstream("anthropic")) as primary,
wire_server(_upstream("anthropic", "ending-complete")) as backup,
gateway.scenario() as scenario,
):
model: Final = _deploy(scenario, "anthropic", primary)
fallback: Final = _deploy(scenario, "anthropic", backup)
text, _ = _post(
gateway, "/v1/responses", {"model": model, "input": _prompt("ending-clean-eof"), "fallbacks": [fallback]}
)
types: Final = tuple(string_value(frame.get("type") or "") for frame in data_frames(text))
assert "response.completed" in types and "response.failed" not in types, text
assert (len(posted_targets(primary)), len(posted_targets(backup))) == (1, 1), text
def test_chat_stream_cut_after_its_content_is_not_replayed_on_the_request_fallback(gateway: Gateway) -> None:
with (
wire_server(_upstream("openai")) as primary,
wire_server(_upstream("openai", "ending-complete")) as backup,
gateway.scenario() as scenario,
):
model: Final = _deploy(scenario, "openai", primary)
fallback: Final = _deploy(scenario, "openai", backup)
text, _ = _post(
gateway,
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": _prompt("ending-clean-eof")}],
"fallbacks": [fallback],
},
)
assert _chat_text(text) == "".join(_PARTIAL), text
assert len(_chat_errors(text)) == 1, text
assert posted_targets(backup) == (), text
Route = Literal["/v1/chat/completions", "/v1/messages", "/v1/responses"]
def _route_body(route: Route, model: str, ending: Ending) -> Mapping[str, JsonValue]:
match route:
case "/v1/chat/completions":
return {"model": model, "messages": [{"role": "user", "content": _prompt(ending)}]}
case "/v1/messages":
return {"model": model, "max_tokens": 64, "messages": [{"role": "user", "content": _prompt(ending)}]}
case "/v1/responses":
return {"model": model, "input": _prompt(ending)}
def _finished(route: Route, text: str) -> bool:
match route:
case "/v1/chat/completions":
return _finish_reasons(text) == ("stop",) and _chat_errors(text) == ()
case "/v1/messages":
return "message_stop" in _event_types(text) and "error" not in _event_types(text)
case "/v1/responses":
return "response.completed" in tuple(string_value(frame.get("type") or "") for frame in data_frames(text))
def _stream_outcome(gateway: Gateway, model: str, route: Route, ending: Ending) -> tuple[bool, str]:
text, call_id = _post(gateway, route, _route_body(route, model, ending))
return _finished(route, text), call_id
_BRIDGED_FAMILY: Final[Mapping[Route, Family]] = {
"/v1/chat/completions": "openai",
"/v1/messages": "hosted_vllm",
"/v1/responses": "anthropic",
}
def test_concurrent_burst_ends_each_stream_by_whether_its_provider_finished_it(gateway: Gateway) -> None:
endings: Final[tuple[Ending, ...]] = ("ending-complete", "ending-clean-eof")
cases: Final = tuple(itertools.product(_BRIDGED_FAMILY, endings, range(5)))
with ExitStack() as wires, gateway.scenario() as scenario:
models: Final = {
route: _deploy(scenario, family, wires.enter_context(wire_server(_upstream(family))))
for route, family in _BRIDGED_FAMILY.items()
}
with ThreadPoolExecutor(max_workers=12) as pool:
outcomes: Final = tuple(
pool.map(lambda case: _stream_outcome(gateway, models[case[0]], case[0], case[1]), cases)
)
assert tuple(finished for finished, _ in outcomes) == tuple(
ending == "ending-complete" for _, ending, _ in cases
), outcomes
assert tuple(spend_row(call_id)["status"] for _, call_id in outcomes) == tuple(
"success" if ending == "ending-complete" else "failure" for _, ending, _ in cases
)
_SPEND_LOG_INSERT: Final = b'INSERT INTO "public"."LiteLLM_SpendLogs"'
def _spend_statuses(database_url: str | None) -> Counter[tuple[str, str]]:
rows: Final = read_rows('SELECT litellm_call_id, status FROM "LiteLLM_SpendLogs"', (), database_url=database_url)
return Counter((string_value(row["litellm_call_id"]), string_value(row["status"])) for row in rows)
def test_spend_log_outage_mid_burst_keeps_each_stream_ending_and_lands_every_row_once(
gateway: Gateway, tmp_path: Path
) -> None:
endings: Final[tuple[Ending, ...]] = ("ending-complete", "ending-clean-eof")
cases: Final = tuple(itertools.product(_BRIDGED_FAMILY, endings))
with (
scratch_database() as database_url,
database_relay(database_url, _SPEND_LOG_INSERT) as (relay, relayed_url),
owned_proxy(gateway, tmp_path, {"DATABASE_URL": plaintext_database_url(relayed_url)}, workers=2) as candidate,
ExitStack() as wires,
candidate.scenario() as scenario,
ThreadPoolExecutor(max_workers=len(cases)) as pool,
):
models: Final = {
route: _deploy(scenario, family, wires.enter_context(wire_server(_upstream(family))))
for route, family in _BRIDGED_FAMILY.items()
}
def burst() -> tuple[tuple[bool, str], ...]:
return tuple(pool.map(lambda case: _stream_outcome(candidate, models[case[0]], case[0], case[1]), cases))
relay.arm()
before: Final = burst()
assert relay.tripped.wait(90), "no spend log write reached the database after the first burst"
down: Final = not relay.reconnected.is_set()
during: Final = burst()
assert relay.reconnected.wait(90), "the proxy never reconnected to the database after the outage"
assert down and relay.refused >= 1, relay.refused
finished: Final = tuple(ending == "ending-complete" for _, ending in cases)
assert tuple(done for done, _ in before + during) == finished * 2, (before, during)
expected: Final = Counter(
(call_id, "success" if ending == "ending-complete" else "failure")
for (_, ending), (_, call_id) in zip(cases * 2, before + during, strict=True)
)
landed: Final = eventually(
lambda: _spend_statuses(database_url),
lambda rows: {call_id for call_id, _ in expected} <= {call_id for call_id, _ in rows},
seconds=120,
)
assert landed == expected
def _attempt(gateway: Gateway, model: str, route: Route, ending: Ending) -> tuple[bool, str] | None:
"""One stream's outcome, or None when the worker serving it was killed before the proxy finished the reply"""
try:
return _stream_outcome(gateway, model, route, ending)
except httpx.TransportError:
return None
def test_worker_killed_mid_burst_leaves_the_proxy_ending_each_stream_by_whether_its_provider_finished_it(
gateway: Gateway, tmp_path: Path
) -> None:
endings: Final[tuple[Ending, ...]] = ("ending-complete", "ending-clean-eof")
cases: Final = tuple(itertools.product(_BRIDGED_FAMILY, endings, range(4)))
probes: Final = tuple(itertools.product(_BRIDGED_FAMILY, endings))
with (
owned_proxy_process(gateway, tmp_path, {}, workers=2) as owned,
ExitStack() as wires,
owned.gateway.scenario() as scenario,
):
candidate: Final = owned.gateway
models: Final = {
route: _deploy(scenario, family, wires.enter_context(wire_server(_upstream(family))))
for route, family in _BRIDGED_FAMILY.items()
}
workers: Final = eventually(owned.worker_pids, lambda pids: len(pids) == 2, seconds=60)
with owned.replacing_worker(min(workers)) as kill, ThreadPoolExecutor(max_workers=len(cases)) as pool:
futures: Final = [
pool.submit(_attempt, candidate, models[route], route, ending) for route, ending, _ in cases
]
kill()
attempts: Final = tuple(future.result() for future in futures)
answered: Final = tuple(
(outcome[0], ending == "ending-complete")
for (_, ending, _), outcome in zip(cases, attempts, strict=True)
if outcome is not None
)
assert all(finished == expected for finished, expected in answered), attempts
after: Final = tuple(_stream_outcome(candidate, models[route], route, ending) for route, ending in probes)
assert tuple(finished for finished, _ in after) == tuple(ending == "ending-complete" for _, ending in probes)
statuses: Final = eventually(
lambda: _spend_statuses(candidate.database_url),
lambda rows: {call_id for _, call_id in after} <= {call_id for call_id, _ in rows},
seconds=120,
)
assert tuple(
statuses[(call_id, "success" if ending == "ending-complete" else "failure")]
for (_, ending), (_, call_id) in zip(probes, after, strict=True)
) == (1,) * len(probes)

12
tests/unit/integrations/test_anthropic_cache_control_hook.py
View file

@ -1116,7 +1116,7 @@ def test_cache_control_hook_counts_tool_call_cache_controls():
assert AnthropicCacheControlHook.count_request_cache_breakpoints([message]) == 3
def test_cache_control_hook_counts_tool_call_marks_except_answered_server_tool_calls():
def test_cache_control_hook_counts_tool_call_marks_except_server_tool_calls_the_client_did_not_answer():
message: Final[AllMessageValues] = {
"role": "assistant",
"content": None,
@ -1170,7 +1170,7 @@ def test_cache_control_hook_counts_tool_call_marks_except_answered_server_tool_c
},
}
assert AnthropicCacheControlHook.count_request_cache_breakpoints([message]) == 2
assert AnthropicCacheControlHook.count_request_cache_breakpoints([message]) == 1
_SEARCH_ANSWER: Final = {"role": "tool", "tool_call_id": "srvtoolu_01Search", "content": "found get_weather"}
@ -1182,6 +1182,12 @@ _SEARCH_CALL: Final = {
"function": {"name": "tool_search_tool_bm25", "arguments": "{}"},
"cache_control": {"type": "ephemeral"},
}
_UNRUN_WEB_SEARCH_CALL: Final = {
"id": "srvtoolu_01Web",
"type": "function",
"function": {"name": "web_search", "arguments": "{}"},
"cache_control": {"type": "ephemeral"},
}
_WEATHER_CALL: Final = {
"id": "toolu_01Weather",
"type": "function",
@ -1238,6 +1244,8 @@ def _replayed_search_call(*tool_calls: dict[str, object]) -> AllMessageValues:
(_SEARCH_CALL,), (_SEARCH_ANSWER, _GO_ON), False, 0, id="client_answer_replaced_by_the_kept_result"
),
pytest.param((_SEARCH_CALL,), (_GO_ON,), True, 0, id="placeholder_replaced_by_the_kept_result"),
pytest.param((_UNRUN_WEB_SEARCH_CALL,), (), False, 0, id="trailing_unrun_search_replayed_without_its_mark"),
pytest.param((_UNRUN_WEB_SEARCH_CALL,), (_GO_ON,), False, 0, id="unrun_search_dropped_before_a_new_message"),
pytest.param(
(_SEARCH_CALL, _WEATHER_CALL),
(_SEARCH_ANSWER,),

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

@ -2946,6 +2946,44 @@ async def test_async_success_handler_runs_async_callbacks_in_the_post_response_p
assert in_post_response_phase() is False, "the phase must not leak into the request task"
@pytest.mark.parametrize("stream", [True, False], ids=["streaming", "non_streaming"])
def test_reused_logging_obj_does_not_report_the_previous_attempts_first_token(logging_obj, stream: bool):
"""A fallback reuses the request's logging object but runs in its own window.
When the first attempt already recorded a first-token time, the fallback must not
report it. Otherwise ``completion_start_time`` lands before the attempt's
``start_time``, the callbacks see a negative TTFT, and latency-based routing caches
a negative sample that empties its buffer window.
"""
logging_obj.stream = stream
logging_obj.model_call_details["litellm_params"] = {}
logging_obj.litellm_params = {}
first_attempt_start: Final = datetime.datetime(2026, 1, 1, 12, 0, 0)
logging_obj.update_completion_start_time(first_attempt_start + datetime.timedelta(seconds=1))
fallback_start: Final = first_attempt_start + datetime.timedelta(seconds=600)
fallback_end: Final = fallback_start + datetime.timedelta(seconds=5)
observed: dict[str, object] = {}
class Recorder(CustomLogger):
def log_success_event(self, kwargs, response_obj, start_time, end_time) -> None:
observed["completion_start_time"] = kwargs.get("completion_start_time")
with patch.object(logging_obj, "get_combined_callback_list", return_value=[Recorder()]):
logging_obj.success_handler(
result=_success_response(),
start_time=fallback_start,
end_time=fallback_end,
)
completion_start_time: Final = observed["completion_start_time"]
assert isinstance(completion_start_time, datetime.datetime)
assert completion_start_time >= fallback_start, (
"the fallback attempt reported a first-token time that predates its own start"
)
def test_is_sync_litellm_request():
assert LitellmLogging.is_sync_litellm_request({}) is True
assert LitellmLogging.is_sync_litellm_request({"acompletion": True}) is False

330
tests/unit/litellm_core_utils/test_streaming_handler.py
View file

@ -32,8 +32,12 @@ from litellm.litellm_core_utils.streaming_handler import (
_ProviderChunkEarlyReturn,
_ProviderChunkParsed,
)
from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicModelResponseIterator
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator as GeminiModelResponseIterator,
)
from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response
from litellm.types.utils import (
ChoiceLogprobs,
@ -572,11 +576,6 @@ _GET_WEATHER_CALL: Final = {
"<think>a</think>b<think>c</think>d",
0,
),
(
(_stream_chunk({"reasoning_content": "Still thinking"}),),
"<think>Still thinking</think>",
0,
),
(
(_stream_chunk({"content": "hi"}), _stream_chunk({}, "stop")),
"hi",
@ -588,7 +587,6 @@ _GET_WEATHER_CALL: Final = {
"reasoning_then_function_call",
"reasoning_cut_by_length",
"reasoning_after_answer",
"stream_ends_mid_reasoning",
"answer_without_reasoning",
],
)
@ -622,6 +620,26 @@ async def test_merged_reasoning_stream_keeps_think_tags_balanced(
)
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_mode", [True, False])
async def test_merged_reasoning_stream_cut_mid_reasoning_fails_instead_of_closing_the_think_block(
logging_obj: Logging, sync_mode: bool
) -> None:
response: Final = CustomStreamWrapper(
completion_stream=ModelResponseListIterator(
model_responses=[_stream_chunk({"reasoning_content": "Still thinking"})]
),
model="deepseek-reasoner",
logging_obj=logging_obj,
)
response.merge_reasoning_content_in_choices = True
with pytest.raises(MidStreamFallbackError) as excinfo:
await _collect(response, sync_mode)
assert excinfo.value.generated_content == "<think>Still thinking"
def test_strip_sse_data_from_chunk():
"""Test the static method that strips 'data: ' prefix from SSE chunks"""
# Test with string inputs
@ -2270,6 +2288,242 @@ async def test_mid_stream_error_keeps_credentials_in_the_api_base_out(routed_api
assert "deployment-secret" not in debug_info
def _incomplete_openai_chat_chunk() -> ModelResponseStream:
return ModelResponseStream(
id="chatcmpl-incomplete-eof",
created=1,
model="gpt-5.6",
object="chat.completion.chunk",
choices=[
StreamingChoices(
index=0,
delta=Delta(content='{"findings":[{"title":"unfinished', role="assistant"),
finish_reason=None,
)
],
)
@pytest.mark.parametrize("custom_llm_provider", ["openai", "azure", "hosted_vllm"])
@pytest.mark.asyncio
async def test_async_openai_compatible_chat_clean_eof_without_finish_reason_raises_midstream_fallback(
logging_obj: Logging, custom_llm_provider: str
) -> None:
"""An OpenAI-compatible stream without a provider finish reason must fail at clean EOF."""
logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider
response: Final = CustomStreamWrapper(
completion_stream=_as_async(iter([_incomplete_openai_chat_chunk()])),
model="gpt-5.6",
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
stream_options={"include_usage": True},
)
first_chunk: Final = await response.__anext__()
with pytest.raises(MidStreamFallbackError) as excinfo:
await response.__anext__()
assert first_chunk.choices[0].delta.content == '{"findings":[{"title":"unfinished'
assert excinfo.value.generated_content == '{"findings":[{"title":"unfinished'
assert excinfo.value.is_pre_first_chunk is False
assert response.received_finish_reason is None
@pytest.mark.parametrize("custom_llm_provider", ["openai", "azure", "hosted_vllm"])
def test_sync_openai_compatible_chat_clean_eof_without_finish_reason_raises_midstream_fallback(
logging_obj: Logging, custom_llm_provider: str
) -> None:
"""The synchronous OpenAI-compatible stream must fail at clean EOF too."""
logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider
response: Final = CustomStreamWrapper(
completion_stream=iter([_incomplete_openai_chat_chunk()]),
model="gpt-5.6",
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
stream_options={"include_usage": True},
)
first_chunk: Final = next(response)
with pytest.raises(MidStreamFallbackError) as excinfo:
next(response)
assert first_chunk.choices[0].delta.content == '{"findings":[{"title":"unfinished'
assert excinfo.value.generated_content == '{"findings":[{"title":"unfinished'
assert excinfo.value.is_pre_first_chunk is False
assert response.received_finish_reason is None
async def _collect(wrapper: CustomStreamWrapper, sync_mode: bool) -> list[ModelResponseStream]:
if sync_mode:
return list(wrapper)
return [chunk async for chunk in wrapper]
@pytest.mark.parametrize("custom_llm_provider", ["openai", "azure", "hosted_vllm"])
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_stream_that_sends_no_chunks_keeps_the_default_stop(
logging_obj: Logging, custom_llm_provider: str, sync_mode: bool
) -> None:
"""Gemini's OpenAI-compatible endpoint answers with only [DONE] when reasoning uses up max_tokens."""
wrapper: Final = CustomStreamWrapper(
completion_stream=iter(()) if sync_mode else _as_async(iter(())),
model="gemini-2.5-flash",
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
stream_options={"include_usage": True},
)
received: Final = await _collect(wrapper, sync_mode)
assert [c.choices[0].finish_reason for c in received if c.choices and c.choices[0].finish_reason] == ["stop"]
_ANTHROPIC_PARTIAL_EVENTS: Final[tuple[dict[str, object], ...]] = (
{
"type": "message_start",
"message": {
"id": "msg_incomplete_eof",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-5",
"content": [],
"stop_reason": None,
"usage": {"input_tokens": 9, "output_tokens": 1},
},
},
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "text_delta", "text": '{"findings":[{"title":"unfinished'},
},
)
_ANTHROPIC_TERMINAL_EVENTS: Final[tuple[dict[str, object], ...]] = (
{"type": "content_block_stop", "index": 0},
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 7}},
{"type": "message_stop"},
)
def _anthropic_sse_lines(events: Sequence[Mapping[str, object]]) -> list[str]:
return [line for event in events for line in (f"event: {event['type']}", f"data: {json.dumps(event)}", "")]
def _anthropic_stream(
events: Sequence[Mapping[str, object]], sync_mode: bool, logging_obj: Logging
) -> CustomStreamWrapper:
lines: Final = _anthropic_sse_lines(events)
return CustomStreamWrapper(
completion_stream=AnthropicModelResponseIterator(
streaming_response=iter(lines) if sync_mode else _as_async(iter(lines)), sync_stream=sync_mode
),
model="claude-sonnet-4-5",
custom_llm_provider="anthropic",
logging_obj=logging_obj,
)
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_anthropic_stream_that_ends_before_message_delta_fails_with_the_eof_message(
sync_mode: bool, logging_obj: Logging
) -> None:
wrapper: Final = _anthropic_stream(_ANTHROPIC_PARTIAL_EVENTS, sync_mode, logging_obj)
with pytest.raises(MidStreamFallbackError) as excinfo:
await _collect(wrapper, sync_mode)
assert excinfo.value.generated_content == '{"findings":[{"title":"unfinished'
assert excinfo.value.is_pre_first_chunk is False
assert isinstance(excinfo.value.original_exception, litellm.APIConnectionError)
assert "AnthropicException - Stream ended before a finish_reason was received" in str(excinfo.value)
assert wrapper.received_finish_reason is None
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_anthropic_stream_with_message_stop_ends_with_its_finish_reason(
sync_mode: bool, logging_obj: Logging
) -> None:
wrapper: Final = _anthropic_stream(
(*_ANTHROPIC_PARTIAL_EVENTS, *_ANTHROPIC_TERMINAL_EVENTS), sync_mode, logging_obj
)
received: Final = await _collect(wrapper, sync_mode)
assert [c.choices[0].finish_reason for c in received if c.choices and c.choices[0].finish_reason] == ["stop"]
assert "".join(c.choices[0].delta.content or "" for c in received if c.choices) == (
'{"findings":[{"title":"unfinished'
)
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_generic_chunk_stream_without_a_finish_reason_keeps_the_default_stop(
sync_mode: bool, logging_obj: Logging
) -> None:
"""Providers that stream GenericStreamingChunks and never mark the last one finished still end with stop,
but the logged response built from ``chunks`` never gains that synthetic stop."""
chunks: Final = [
{"text": text, "tool_use": None, "is_finished": False, "finish_reason": "", "usage": None, "index": 0}
for text in ("partial", " reply")
]
wrapper: Final = CustomStreamWrapper(
completion_stream=ModelResponseListIterator(model_responses=chunks),
model="sagemaker/tgi-endpoint",
custom_llm_provider="sagemaker",
logging_obj=logging_obj,
)
received: Final = await _collect(wrapper, sync_mode)
assert [c.choices[0].finish_reason for c in received if c.choices and c.choices[0].finish_reason] == ["stop"]
assert "".join(c.choices[0].delta.content or "" for c in received if c.choices) == "partial reply"
assert all(not (c.choices and c.choices[0].finish_reason) for c in wrapper.chunks)
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_gemini_stream_with_a_usage_chunk_after_its_finish_reason_still_completes(
sync_mode: bool, logging_obj: Logging
) -> None:
"""Gemini marks every chunk unfinished, so its finishReason only shows up mid-stream before the usage chunk."""
lines: Final = [
f"data: {json.dumps(event)}"
for event in (
{"candidates": [{"content": {"parts": [{"text": "Hello"}], "role": "model"}, "index": 0}]},
{
"candidates": [
{"content": {"parts": [{"text": " there"}], "role": "model"}, "finishReason": "STOP", "index": 0}
]
},
{"usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 2, "totalTokenCount": 5}},
)
]
logging_obj.update_environment_variables(
model="gemini-2.5-flash", user=None, optional_params={}, litellm_params={"litellm_call_id": "gemini-tail"}
)
iterator: Final = GeminiModelResponseIterator(
streaming_response=iter(lines) if sync_mode else _as_async(iter(lines)), sync_stream=sync_mode, logging_obj=logging_obj
)
wrapper: Final = CustomStreamWrapper(
completion_stream=iter(iterator) if sync_mode else iterator,
model="gemini-2.5-flash",
custom_llm_provider="gemini",
logging_obj=logging_obj,
)
received: Final = await _collect(wrapper, sync_mode)
assert "".join(c.choices[0].delta.content or "" for c in received if c.choices) == "Hello there"
assert [c.choices[0].finish_reason for c in received if c.choices and c.choices[0].finish_reason] == ["stop"]
def test_streaming_handler_with_created_time_propagation(
initialized_custom_stream_wrapper: CustomStreamWrapper, logging_obj: Logging
@ -3638,11 +3892,8 @@ def test_usage_chunk_after_finish_reason_updates_hidden_params(logging_obj):
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.parametrize("ends_with_finish_chunk", [True, False])
@pytest.mark.asyncio
async def test_stream_without_provider_usage_falls_back_to_the_token_estimate(
sync_mode: bool, ends_with_finish_chunk: bool
) -> None:
async def test_stream_without_provider_usage_falls_back_to_the_token_estimate(sync_mode: bool) -> None:
model: Final = "gpt-4o"
messages: Final = [{"role": "user", "content": "Write a long sentence about a fox. " * 30}]
text: Final = "The quick brown fox jumps over the lazy dog. " * 20
@ -3659,9 +3910,7 @@ async def test_stream_without_provider_usage_falls_back_to_the_token_estimate(
choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")],
)
wrapper: Final = CustomStreamWrapper(
completion_stream=ModelResponseListIterator(
model_responses=[content_chunk, finish_chunk] if ends_with_finish_chunk else [content_chunk]
),
completion_stream=ModelResponseListIterator(model_responses=[content_chunk, finish_chunk]),
model=model,
custom_llm_provider="openai",
logging_obj=Logging(
@ -3700,6 +3949,55 @@ async def test_stream_without_provider_usage_falls_back_to_the_token_estimate(
assert litellm.completion_cost(completion_response=assembled) == pytest.approx(spend)
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_stream_cut_before_a_finish_chunk_bills_the_token_estimate_on_its_failure(sync_mode: bool) -> None:
model: Final = "gpt-4o"
messages: Final = [{"role": "user", "content": "Write a long sentence about a fox. " * 30}]
text: Final = "The quick brown fox jumps over the lazy dog. " * 20
logging_obj: Final = Logging(
model=model,
messages=messages,
stream=True,
call_type="completion",
start_time=datetime.datetime(2025, 3, 3, 21, 38, 10),
litellm_call_id="cut-no-usage-call",
function_id="cut-no-usage-fn",
)
logging_obj.update_environment_variables(
model=model, user=None, optional_params={}, litellm_params={"litellm_call_id": "cut-no-usage-call"}
)
wrapper: Final = CustomStreamWrapper(
completion_stream=ModelResponseListIterator(
model_responses=[
ModelResponseStream(
id="chatcmpl-cut-no-usage",
created=1741037890,
model=model,
choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content=text))],
)
]
),
model=model,
custom_llm_provider="openai",
logging_obj=logging_obj,
stream_options=None,
)
with pytest.raises(MidStreamFallbackError):
await _collect(wrapper, sync_mode)
expected: Final = (
litellm.token_counter(model=model, messages=messages),
litellm.token_counter(model=model, text=text, count_response_tokens=True),
)
usage: Final = logging_obj.model_call_details["combined_usage_object"]
assert (usage.prompt_tokens, usage.completion_tokens) == expected
rates: Final = litellm.model_cost[model]
spend: Final = expected[0] * rates["input_cost_per_token"] + expected[1] * rates["output_cost_per_token"]
assert logging_obj.model_call_details["response_cost"] == pytest.approx(spend)
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_stream_with_usage_after_the_finish_chunk_bills_only_the_reported_usage(sync_mode: bool) -> None:
@ -8582,10 +8880,8 @@ async def test_no_synthetic_finish_reason_logged_when_provider_sent_none(sync_mo
custom_llm_provider="bedrock",
logging_obj=logging_obj,
)
if sync_mode:
list(response)
else:
[c async for c in response]
with pytest.raises(MidStreamFallbackError):
await _collect(response, sync_mode)
assert response.received_finish_reason is None
assert all(not (c.choices and c.choices[0].finish_reason) for c in response.chunks)

40
tests/unit/llms/a2a/test_common_utils.py
View file

@ -2,6 +2,7 @@
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
import pytest
@ -50,3 +51,42 @@ async def test_agent_without_entra_credentials_gets_no_bearer():
assert header is None
assert resolver.calls == []
@pytest.mark.parametrize(
("api_key", "expected"),
(
("secret", "Bearer secret"),
("Bearer secret", "Bearer Bearer secret"),
("os.environ/SYNTHETIC_KEY", "Bearer os.environ/SYNTHETIC_KEY"),
("", None),
(None, None),
(42, None),
),
)
async def test_plain_agent_key_is_used_verbatim_only_for_nonempty_strings(
api_key: object, expected: str | None
) -> None:
header: Final = await resolve_a2a_hop_auth_header({"api_key": api_key}, None)
assert header == ({"Authorization": expected} if expected is not None else None)
@pytest.mark.parametrize("provider", ("azure_ai", "databricks_agent"))
async def test_completion_bridge_api_key_is_not_used_for_a2a_hop(provider: str) -> None:
assert await resolve_a2a_hop_auth_header({"api_key": "secret"}, provider) is None
async def test_entra_bearer_overrides_agent_key_and_passthrough() -> None:
resolver: Final = _RecordingEntraResolver()
params: Final = {**_SERVICE_PRINCIPAL, "api_key": "secret"}
header: Final = await resolve_a2a_hop_auth_header(params, None, resolver, {"authorization": "Bearer forwarded"})
assert header == {"Authorization": "Bearer minted-entra-token"}
assert resolver.calls == [params]
@pytest.mark.parametrize("name", ("Authorization", "aUtHoRiZaTiOn"))
async def test_passthrough_authorization_takes_precedence_over_plain_agent_key(name: str) -> None:
assert (
await resolve_a2a_hop_auth_header({"api_key": "secret"}, None, agent_extra_headers={name: "Bearer forwarded"})
is None
)

84
tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py
View file

@ -19,8 +19,9 @@ from typing import Final
from unittest.mock import MagicMock, Mock, patch
import litellm
from litellm.exceptions import MidStreamFallbackError
from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
class TestAgentCoreAcceptHeader:
@ -1271,6 +1272,87 @@ def test_agentcore_runtime_response_becomes_assistant_message(
assert response.choices[0].finish_reason == "stop"
_AGENTCORE_RUNTIME_URL: Final = (
"https://bedrock-agentcore.us-west-2.amazonaws.com/runtimes/"
"arn%3Aaws%3Abedrock-agentcore%3Aus-west-2%3A888602223428%3Aruntime%2Fhosted_agent_13sf6-cALnp38iZD/invocations"
)
_AGENTCORE_DELTAS: Final = (
'data: {"event":{"messageStart":{"role":"assistant"}}}\n\n'
'data: {"event":{"contentBlockDelta":{"delta":{"text":"Machine learning lets "},"contentBlockIndex":0}}}\n\n'
'data: {"event":{"contentBlockDelta":{"delta":{"text":"computers learn from data."},"contentBlockIndex":0}}}\n\n'
)
async def _stream_agentcore(body: str, sync_mode: bool) -> list[tuple[str | None, str | None]]:
"""Stream `body` from the runtime as SSE and return each chunk's (content, finish_reason)."""
kwargs: Final = {
"model": "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/"
"hosted_agent_13sf6-cALnp38iZD",
"messages": [{"role": "user", "content": "Explain machine learning in simple terms"}],
"aws_access_key_id": "AKIAAGENTCOREKEY",
"aws_secret_access_key": "agentcore-secret",
"stream": True,
}
if sync_mode:
return [
(chunk.choices[0].delta.content, chunk.choices[0].finish_reason)
for chunk in litellm.completion(**kwargs, client=HTTPHandler())
]
return [
(chunk.choices[0].delta.content, chunk.choices[0].finish_reason)
async for chunk in await litellm.acompletion(
**kwargs, client=AsyncHTTPHandler(transport=httpx.AsyncHTTPTransport())
)
]
@pytest.mark.parametrize(
"ending",
[
pytest.param(
'data: {"event":{"contentBlockStop":{"contentBlockIndex":0}}}\n\n'
'data: {"event":{"messageStop":{"stopReason":"end_turn"}}}\n\n',
id="message_stop",
),
pytest.param(
'data: {"event":{"messageStop":{"stopReason":"end_turn"}}}\n\n'
'data: {"event":{"metadata":{"usage":{"inputTokens":9,"outputTokens":7,"totalTokens":16}}}}\n\n'
'data: {"message":{"role":"assistant",'
'"content":[{"text":"Machine learning lets computers learn from data."}]}}\n\n',
id="strands_full",
),
],
)
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.respx(assert_all_called=True)
async def test_agentcore_stream_that_signals_its_end_completes(
ending: str, sync_mode: bool, respx_mock: respx.Router
) -> None:
respx_mock.post(_AGENTCORE_RUNTIME_URL).mock(
return_value=httpx.Response(
200, content=(_AGENTCORE_DELTAS + ending).encode(), headers={"content-type": "text/event-stream"}
)
)
chunks: Final = await _stream_agentcore(_AGENTCORE_DELTAS + ending, sync_mode)
assert "".join(content for content, _ in chunks if content) == "Machine learning lets computers learn from data."
assert [finish for _, finish in chunks if finish] == ["stop"]
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.respx(assert_all_called=True)
async def test_agentcore_stream_cut_before_it_signals_its_end_raises(sync_mode: bool, respx_mock: respx.Router) -> None:
respx_mock.post(_AGENTCORE_RUNTIME_URL).mock(
return_value=httpx.Response(
200, content=_AGENTCORE_DELTAS.encode(), headers={"content-type": "text/event-stream"}
)
)
with pytest.raises(MidStreamFallbackError, match="Stream ended before a finish_reason was received"):
await _stream_agentcore(_AGENTCORE_DELTAS, sync_mode)
def test_get_complete_url_sends_a_non_string_runtime_endpoint_to_that_endpoint():
vpc_endpoint: Final = "https://vpce-1.bedrock-agentcore.us-west-2.vpce.amazonaws.com"

105
tests/unit/llms/deepgram/test_deepgram_common_utils.py
View file

@ -7,6 +7,7 @@ import pytest
import litellm
from litellm.llms.deepgram.common_utils import (
DeepgramSpeakMeter,
deepgram_listen_addon_pricing_models,
deepgram_listen_audio_seconds,
deepgram_listen_callback_params,
@ -18,7 +19,12 @@ from litellm.llms.deepgram.common_utils import (
deepgram_listen_requested_model,
deepgram_listen_transcript,
deepgram_listen_websocket_target,
deepgram_speak_character_rate,
deepgram_speak_registry_key,
deepgram_speak_requested_model,
deepgram_speak_websocket_target,
)
from litellm.types.passthrough_endpoints.pass_through_endpoints import WebSocketSessionUsage
NOVA_3_URL: Final = "wss://api.deepgram.com/v1/listen?model=nova-3&encoding=linear16&sample_rate=16000"
@ -324,3 +330,102 @@ def test_deepgram_listen_is_priced(
)
def test_deepgram_listen_addon_pricing_models(upstream_url: str, expected: tuple[str, ...]):
assert deepgram_listen_addon_pricing_models(upstream_url) == expected
@pytest.mark.parametrize(
("api_base", "query_string", "expected"),
[
pytest.param(
None,
"encoding=linear16",
"wss://api.deepgram.com/v1/speak?encoding=linear16&model=aura-asteria-en",
id="no model: Deepgram's own default, named so billing and upstream agree",
),
pytest.param(
None,
"model=aura-asteria-en&model=aura-2-thalia-en",
"wss://api.deepgram.com/v1/speak?model=aura-asteria-en",
id="second model dropped: Deepgram honours the last, auth and pricing read the first",
),
pytest.param(
"http://localhost:8080/v1",
"model=aura-2-thalia-en",
"ws://localhost:8080/v1/speak?model=aura-2-thalia-en",
id="server-configured api base",
),
],
)
def test_deepgram_speak_websocket_target(api_base: str | None, query_string: str, expected: str):
assert deepgram_speak_websocket_target(api_base, query_string) == expected
@pytest.mark.parametrize(
"query_string",
[
"",
"model=",
"model=&model=aura-2-thalia-en",
"model=aura-2-thalia-en",
"model=aura-asteria-en&model=aura-2-thalia-en",
"encoding=mulaw",
],
)
def test_deepgram_speak_requested_voice_is_the_only_voice_the_upstream_target_carries(query_string: str):
"""Authorization runs against the requested voice and the upstream URL is built separately, so they must agree or
a key could be authorized for one voice and reach another"""
target: Final = deepgram_speak_websocket_target(None, query_string)
assert parse_qs(urlparse(target).query)["model"] == [deepgram_speak_requested_model(query_string)]
@pytest.mark.usefixtures("local_model_cost_map")
@pytest.mark.parametrize(
("model", "registry_key", "priced"),
[
("aura-asteria-en", "deepgram/aura", True),
("aura-2-thalia-en", "deepgram/aura-2", True),
("aura-2-celeste-es", "deepgram/aura-2", True),
("aura-99-x-en", "deepgram/aura-99", False),
("aura-2", "deepgram/aura-2", True),
("nova-3", "deepgram/nova-3", False),
],
)
def test_deepgram_speak_prices_each_voice_at_its_generation_row_only(model: str, registry_key: str, priced: bool):
"""A voice is never billed at another generation's rate, and a listen model's row is not a speak price"""
upstream_url: Final = f"wss://api.deepgram.com/v1/speak?model={model}"
assert deepgram_speak_registry_key(upstream_url) == registry_key
assert deepgram_speak_character_rate(upstream_url) == (
litellm.model_cost[registry_key]["input_cost_per_character"] if priced else None
)
@pytest.mark.parametrize(
("frame", "characters"),
[
('{"type": "Speak", "text": "Hello world."}', 12),
('{"type":"Speak","text":"¿qué tal?"}', 9),
('{"type": "Flush"}', 0),
('{"type": "Speak"}', 0),
('{"type": "Speak", "text": 12}', 0),
('["Speak", "Hello"]', 0),
("Hello", 0),
("[" * 100_000, 0),
],
)
def test_deepgram_speak_meter_measures_only_speak_text_and_never_raises(frame: str, characters: int):
meter: Final = DeepgramSpeakMeter("wss://api.deepgram.com/v1/speak?model=aura-2-thalia-en", character_rate=1.0)
assert meter.measure(frame) == characters
@pytest.mark.parametrize(
("upstream_url", "model"),
[
("wss://api.deepgram.com/v1/speak?model=aura-2-thalia-en", "aura-2-thalia-en"),
("wss://api.deepgram.com/v1/speak?encoding=linear16&model=aura-asteria-en", "aura-asteria-en"),
],
)
def test_deepgram_speak_meter_prices_the_total_at_the_rate_it_holds(upstream_url: str, model: str) -> None:
assert DeepgramSpeakMeter(upstream_url, character_rate=0.25).price(47) == WebSocketSessionUsage(
model=model, custom_llm_provider="deepgram", response_cost=47 * 0.25
)

29
tests/unit/llms/langgraph/chat/test_langgraph_transformation.py
View file

@ -22,12 +22,14 @@ import asyncio
import importlib
import json
import os
from collections.abc import AsyncIterator
from typing import Final
import httpx
import pytest
import litellm
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@ -268,3 +270,28 @@ async def test_langgraph_streaming_acompletion_without_a_client_posts_through_th
"assistant_id": "agent",
"input": {"messages": [{"role": "human", "content": "What is 25 * 4?"}]},
}
@pytest.mark.asyncio
async def test_langgraph_async_stream_delivers_every_step_of_the_run_before_its_stop() -> None:
async def events() -> AsyncIterator[bytes]:
question: Final = {"type": "human", "content": "What is 25 * 4?"}
answers: Final = ({"type": "ai", "content": "25 * 4"}, {"type": "ai", "content": " is 100"})
yield b'event: metadata\ndata: {"run_id": "run-1", "attempt": 1}\n\n'
for step in range(len(answers) + 1):
yield b"event: values\ndata: " + json.dumps({"messages": [question, *answers[:step]]}).encode() + b"\n\n"
client: Final = AsyncHTTPHandler(
transport=httpx.MockTransport(lambda request: httpx.Response(200, content=events()))
)
response: Final = await litellm.acompletion(
model="langgraph/agent",
messages=[{"role": "user", "content": "What is 25 * 4?"}],
api_base="http://langgraph.test",
stream=True,
client=client,
)
streamed: Final = [(chunk.choices[0].delta.content, chunk.choices[0].finish_reason) async for chunk in response]
assert streamed == [("25 * 4", None), (" is 100", None), (None, "stop")]

20
tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py
View file

@ -29,6 +29,7 @@ from openai.types.responses import (
ResponseOutputRefusal,
ResponseOutputText,
)
from openai.types.responses.response_compaction_item import ResponseCompactionItem
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -3349,6 +3350,25 @@ class TestGetStructuredMessages:
result = handler.get_structured_messages(data)
assert result is None
def test_should_return_none_for_a_moderation_style_list_of_strings(self):
handler = OpenAIResponsesHandler()
data = {"model": "omni-moderation-latest", "input": ["safe text", "more safe text"]}
assert handler.get_structured_messages(data) is None
def test_should_keep_an_image_item_replayed_with_an_sdk_compaction_item(self):
handler = OpenAIResponsesHandler()
image_url = "https://example.com/cat.png"
data = {
"model": "gpt-5.5",
"input": [
ResponseCompactionItem(type="compaction", id="cmp_1", encrypted_content="abc"),
{"role": "user", "content": [{"type": "input_image", "image_url": image_url}]},
],
}
result = handler.get_structured_messages(data)
assert result is not None
assert image_url in json.dumps(result)
class ToolAppendingGuardrail(CustomGuardrail):
"""Guardrail that appends a new function tool, mimicking a guardrail that

54
tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py
View file

@ -18,6 +18,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import respx
from fastapi.testclient import TestClient
from pydantic import BaseModel, ConfigDict
import litellm
@ -2941,3 +2942,56 @@ async def test_failed_upstream_stream_relays_its_body_only_when_that_is_a_json_o
"id": "req-1",
"method": "tasks/resubscribe",
}
def _a2a_app_refusing_with(refusal: Exception) -> TestClient:
from fastapi import FastAPI
from litellm.proxy.agent_endpoints.a2a_endpoints import router
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
def refuse() -> UserAPIKeyAuth:
raise refusal
app: Final = FastAPI()
app.include_router(router)
app.dependency_overrides[user_api_key_auth] = refuse
return TestClient(app)
@pytest.mark.parametrize("path", ["/a2a/agent-1", "/a2a/agent-1/message/send", "/v1/a2a/agent-1/message/send"])
def test_a2a_auth_refusals_answer_as_jsonrpc_errors_with_the_callers_id(path: str):
from litellm.proxy._types import ProxyException
client: Final = _a2a_app_refusing_with(
ProxyException(message="Budget has been exceeded!", type="budget_exceeded", param=None, code=422)
)
response: Final = client.post(path, json={"jsonrpc": "2.0", "id": "req-7", "method": "message/send"})
assert response.status_code == 422
assert response.json() == {
"jsonrpc": "2.0",
"id": "req-7",
"error": {
"code": -32000,
"message": "Budget has been exceeded!",
"data": {"type": "budget_exceeded", "param": None, "code": "422"},
},
}
def test_a2a_http_refusals_answer_as_jsonrpc_errors_with_a_null_id_for_unparseable_bodies():
from fastapi import HTTPException
client: Final = _a2a_app_refusing_with(HTTPException(status_code=403, detail="Agent 'agent-1' is not allowed"))
response: Final = client.post("/a2a/agent-1", content=b"not json")
assert response.status_code == 403
body: Final = response.json()
assert body["jsonrpc"] == "2.0"
assert body["id"] is None
assert body["error"]["code"] == -32000
assert body["error"]["message"] == "Agent 'agent-1' is not allowed"
assert body["error"]["data"]["code"] == "403"

44
tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py
View file

@ -11202,9 +11202,47 @@ def test_route_skips_budget_checks_marks_only_spend_free_routes() -> None:
assert route_skips_budget_checks(route="/v1/chat/completions") is False
def test_request_skips_budget_checks_extends_route_rule_with_zero_cost_models() -> None:
assert request_skips_budget_checks(route="/v1/models", model=None, llm_router=None) is True
assert request_skips_budget_checks(route="/v1/chat/completions", model=None, llm_router=None) is False
@pytest.mark.parametrize(
"route, request_body, model, expected",
[
("/v1/models", None, None, True),
("/v1/models", None, "paid", True),
("/v1/chat/completions", None, None, False),
("/v1/chat/completions", None, "free", True),
("/v1/chat/completions", None, "paid", False),
("/mcp", {"method": "tools/list"}, "paid", True),
("/mcp", {"method": "tools/call"}, "paid", False),
("/mcp", None, "paid", False),
("/mcp/tools/call", {"method": "initialize"}, "paid", False),
("/health/services", None, "paid", False),
],
)
def test_request_skips_budget_checks_for_non_spending_requests(
route: str, request_body: Mapping[str, object] | None, model: str | None, expected: bool
) -> None:
router: Final = litellm.Router(
model_list=[
{
"model_name": "free",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"input_cost_per_token": 0,
"output_cost_per_token": 0,
},
},
{
"model_name": "paid",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"input_cost_per_token": 1,
"output_cost_per_token": 1,
},
},
]
)
assert (
request_skips_budget_checks(route=route, model=model, llm_router=router, request_body=request_body) is expected
)
def _agent_model_ceiling_resolver(

141
tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py
View file

@ -5,7 +5,7 @@ import os
import subprocess
import sys
import time
from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from contextlib import contextmanager, nullcontext
from datetime import datetime, timedelta, timezone
from functools import partial
@ -17,6 +17,7 @@ from typing import Final
from unittest.mock import ANY, AsyncMock, MagicMock, patch
import fastapi
import orjson
import pytest
from fastapi import HTTPException, Request, status
@ -24,6 +25,8 @@ from starlette.datastructures import UploadFile
import litellm
import litellm.proxy.proxy_server
from litellm.proxy import utils as proxy_utils
from litellm.proxy.common_utils import user_api_key_cache as key_cache_utils
from litellm.caching.dual_cache import DualCache
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
from litellm.proxy.common_utils.user_api_key_cache import (
@ -85,6 +88,7 @@ from litellm.proxy.auth.user_api_key_auth import (
user_api_key_auth,
user_api_key_auth_websocket_for_model,
)
from litellm.proxy.pass_through_endpoints import pass_through_endpoints as pass_through
from litellm.proxy.spend_tracking.carried_budget_state import carried_budget_metadata
from tests.unit.proxy.db.fake_prisma_engine import engine_call
from tests._master_key import MASTER_KEY
@ -11173,6 +11177,30 @@ async def test_temp_budget_increase_applied_for_cached_key():
assert cached_after.max_budget == 2.0
BUDGETED_PASS_THROUGH_ROUTE: Final = "/custom-provider/infer"
@pytest.fixture
def registered_budgeted_route() -> Iterator[None]:
pass_through.InitPassThroughEndpointHelpers.add_exact_path_route(
app=fastapi.FastAPI(),
path=BUDGETED_PASS_THROUGH_ROUTE,
target="https://provider.example/infer",
custom_headers=None,
forward_headers=None,
merge_query_params=None,
dependencies=None,
cost_per_request=1.0,
endpoint_id="budget-regression",
auth=True,
methods=["POST"],
)
try:
yield
finally:
pass_through.InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_id="budget-regression")
@pytest.mark.asyncio
@pytest.mark.parametrize(
"team_member_spend, expect_blocked",
@ -11182,7 +11210,11 @@ async def test_temp_budget_increase_applied_for_cached_key():
(2.39, False),
],
)
async def test_cached_key_team_member_budget_blocks_at_exact_cap(team_member_spend, expect_blocked):
@pytest.mark.usefixtures("registered_budgeted_route")
@pytest.mark.parametrize(
"route", ["/v1/messages", "/health/services", "/health/test_connection", BUDGETED_PASS_THROUGH_ROUTE]
)
async def test_cached_key_team_member_budget_blocks_at_exact_cap(team_member_spend, expect_blocked, route):
"""A team member counter sitting exactly at the cap (where a resized reservation
lands it) must be rejected by the cached-key auth path like every other budget check."""
from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTableCachedObj
@ -11227,8 +11259,8 @@ async def test_cached_key_team_member_budget_blocks_at_exact_cap(team_member_spe
)
mock_request = MagicMock()
mock_request.url.path = "/v1/messages"
mock_request.method = "POST"
mock_request.url.path = route
mock_request.method = "POST" if route == "/v1/messages" else "GET"
mock_request.headers = {"authorization": f"Bearer {api_key}"}
mock_request.query_params = {}
mock_request.state = SimpleNamespace()
@ -11276,6 +11308,107 @@ async def test_cached_key_team_member_budget_blocks_at_exact_cap(team_member_spe
assert f"TeamMember={user_id}:{team_id}" in exc_info.value.message
@pytest.mark.asyncio
@pytest.mark.parametrize("key_max_budget, member_max_budget", [(None, 1.0), (1.0, None)])
@pytest.mark.parametrize(
"route, expect_blocked",
[
("/models", False),
("/v1/models", False),
("/utils/model_info", False),
("/model/info", False),
("/key/info", False),
("/user/info", False),
("/v1/chat/completions", True),
("/v1/messages", True),
],
)
async def test_cached_key_budget_preserves_non_spending_access(
key_max_budget: float | None,
member_max_budget: float | None,
route: str,
expect_blocked: bool,
monkeypatch: pytest.MonkeyPatch,
) -> None:
api_key: Final = "sk-budget-route-fixture"
hashed_token: Final = proxy_utils.hash_token(api_key)
team_id: Final = "budget-route-team"
user_id: Final = "budget-route-user"
cache: Final = DualCache()
await cache_key_object(
hashed_token=hashed_token,
user_api_key_obj=UserAPIKeyAuth(
token=hashed_token,
user_id=user_id,
team_id=team_id,
team_member_spend=2.0,
spend=2.0,
max_budget=key_max_budget,
),
user_api_key_cache=cache,
proxy_logging_obj=None,
)
await cache.async_set_cache(key=f"team_id:{team_id}", value=LiteLLM_TeamTableCachedObj(team_id=team_id))
await cache.async_set_cache(
key=user_id,
value=LiteLLM_UserTable(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER),
)
await cache.async_set_cache(
key=key_cache_utils.team_membership_auth_cache_key(team_id=team_id, user_id=user_id),
value=LiteLLM_TeamMembership(
user_id=user_id,
team_id=team_id,
spend=2.0,
litellm_budget_table=LiteLLM_BudgetTable(max_budget=member_max_budget),
),
)
await cache.async_set_cache(key=f"spend:team_member:{user_id}:{team_id}", value=2.0)
await cache.async_set_cache(key=f"spend:key:{hashed_token}", value=2.0)
request: Final = Request(
scope={
"type": "http",
"path": route,
"method": ("POST" if route in ("/v1/chat/completions", "/v1/messages") else "GET"),
"headers": [(b"authorization", f"Bearer {api_key}".encode())],
"query_string": b"",
}
)
proxy_logging: Final = MagicMock()
proxy_logging.budget_alerts = AsyncMock()
proxy_logging.post_call_failure_hook = AsyncMock(return_value=None)
monkeypatch.setattr(
litellm.proxy.proxy_server,
"general_settings",
{"disable_budget_reservation": True},
)
monkeypatch.setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", MagicMock())
monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache)
monkeypatch.setattr(litellm.proxy.proxy_server, "spend_counter_cache", cache)
monkeypatch.setattr(litellm.proxy.proxy_server, "proxy_logging_obj", proxy_logging)
async def authenticate() -> UserAPIKeyAuth:
return await user_api_key_auth_builder(
request=request,
api_key=f"Bearer {api_key}",
azure_api_key_header="",
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
request_data={},
)
if expect_blocked:
with pytest.raises(ProxyException) as exc_info:
await authenticate()
assert exc_info.value.type == ProxyErrorTypes.budget_exceeded
return
result: Final = await authenticate()
assert result.user_id == user_id
assert result.team_id == team_id
@pytest.mark.asyncio
@pytest.mark.parametrize(
"expiry_offset, expect_blocked",

200
tests/unit/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py
View file

@ -1,4 +1,5 @@
"""Deepgram ``/v1/listen`` passthrough WebSocket route: registration, auth, credential injection, target URL."""
"""Deepgram ``/v1/listen`` and ``/v1/speak`` passthrough WebSocket routes: registration, auth, credential injection,
target URL."""
import asyncio
from collections.abc import Mapping
@ -15,21 +16,27 @@ from starlette.websockets import WebSocketDisconnect
import litellm
from litellm.caching.dual_cache import DualCache
from litellm.constants import WEBSOCKET_CLOSE_REASON_MAX_BYTES
from litellm.litellm_core_utils.get_model_cost_map import adopt_model_cost_map
from litellm.llms.deepgram.common_utils import DeepgramSpeakMeter
from litellm.proxy._lazy_features import LAZY_FEATURES
from litellm.proxy._types import LiteLLMRoutes, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import cache_key_object
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
_websocket_relay,
deepgram_listen_websocket_route,
deepgram_speak_websocket_route,
router,
)
from litellm.proxy.utils import hash_token
from litellm.types.passthrough_endpoints.pass_through_endpoints import WebSocketSessionMeter
GET_CREDENTIALS: Final = (
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials"
)
USER_API_KEY_AUTH: Final = "litellm.proxy.auth.user_api_key_auth.user_api_key_auth"
LISTEN_PATHS: Final = ("/deepgram/v1/listen", "/deepgram/listen")
SPEAK_PATHS: Final = ("/deepgram/v1/speak", "/deepgram/speak")
NOVA_2_STREAMING_KEY: Final = "deepgram/streaming/nova-2"
pytestmark: Final = pytest.mark.usefixtures("local_model_cost_map")
@ -62,6 +69,7 @@ class _RelayCall:
forward_headers: bool
endpoint: str
accept_websocket: bool
meter: WebSocketSessionMeter | None = None
class _FakeRelay:
@ -78,6 +86,7 @@ class _FakeRelay:
forward_headers: bool,
endpoint: str,
accept_websocket: bool,
meter: WebSocketSessionMeter | None = None,
) -> None:
self.calls.append(
_RelayCall(
@ -87,6 +96,7 @@ class _FakeRelay:
forward_headers=forward_headers,
endpoint=endpoint,
accept_websocket=accept_websocket,
meter=meter,
)
)
@ -472,3 +482,191 @@ def test_deepgram_listen_echoes_the_browser_subprotocol_that_carries_the_litellm
MappingProxyType({"Authorization": "Token dg-provider-key"})
]
assert [call.forward_headers for call in relay.calls] == [False]
async def _serve_speak(websocket: _FakeWebSocket, user_api_key_dict: UserAPIKeyAuth | None = None) -> _FakeRelay:
relay: Final = _FakeRelay()
await deepgram_speak_websocket_route(
websocket=websocket,
user_api_key_dict=user_api_key_dict or UserAPIKeyAuth(),
relay=relay,
)
return relay
@pytest.mark.asyncio
@pytest.mark.parametrize("path", SPEAK_PATHS)
async def test_deepgram_speak_relays_with_provider_auth_and_meters_the_session(
path: str, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
monkeypatch.setenv("DEEPGRAM_API_KEY", "dg-provider-key")
websocket: Final = _FakeWebSocket(path, "encoding=linear16&sample_rate=24000")
caller: Final = UserAPIKeyAuth(api_key="sk-litellm-virtual", team_id="team-tts")
target: Final = "wss://api.deepgram.com/v1/speak?encoding=linear16&sample_rate=24000&model=aura-asteria-en"
relay: Final = await _serve_speak(websocket, caller)
assert relay.calls == [
_RelayCall(
target=target,
custom_headers=MappingProxyType({"Authorization": "Token dg-provider-key"}),
user_api_key_dict=caller,
forward_headers=False,
endpoint=path,
accept_websocket=False,
meter=DeepgramSpeakMeter(target, litellm.model_cost["deepgram/aura"]["input_cost_per_character"]),
)
]
assert websocket.accepts == [None]
assert websocket.closed is None
@pytest.mark.asyncio
async def test_deepgram_speak_bills_at_the_connect_time_rate_after_a_cost_map_reload_drops_the_row(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
monkeypatch.setenv("DEEPGRAM_API_KEY", "dg-provider-key")
rate: Final = litellm.model_cost["deepgram/aura-2"]["input_cost_per_character"]
relay: Final = await _serve_speak(_FakeWebSocket("/deepgram/v1/speak", "model=aura-2-thalia-en"))
meters: Final = [call.meter for call in relay.calls if isinstance(call.meter, DeepgramSpeakMeter)]
adopt_model_cost_map({key: row for key, row in litellm.model_cost.items() if key != "deepgram/aura-2"})
assert [meter.price(47).response_cost for meter in meters] == [47 * rate]
@pytest.mark.asyncio
async def test_deepgram_speak_honours_server_configured_api_base(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("DEEPGRAM_API_BASE", "http://localhost:8080/v1")
monkeypatch.setenv("DEEPGRAM_API_KEY", "dg-provider-key")
relay: Final = await _serve_speak(_FakeWebSocket("/deepgram/v1/speak", "model=aura-2-thalia-en"))
assert [call.target for call in relay.calls] == ["ws://localhost:8080/v1/speak?model=aura-2-thalia-en"]
@pytest.mark.asyncio
async def test_deepgram_speak_closes_cleanly_when_provider_credentials_missing(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.delenv("DEEPGRAM_API_KEY", raising=False)
websocket: Final = _FakeWebSocket("/deepgram/v1/speak", "model=aura-2-thalia-en")
relay: Final = await _serve_speak(websocket)
assert relay.calls == []
assert websocket.accepts == []
assert websocket.closed == (
1011,
"Required 'DEEPGRAM_API_KEY' in environment to make pass-through calls to Deepgram.",
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("query", "missing_key"),
[
pytest.param("model=aura-99-x-en", "deepgram/aura-99", id="generation without a row"),
pytest.param("model=nova-3", "deepgram/nova-3", id="listen model has no per-character rate"),
],
)
async def test_deepgram_speak_refuses_voices_it_cannot_price(
query: str, missing_key: str, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
monkeypatch.setenv("DEEPGRAM_API_KEY", "dg-provider-key")
websocket: Final = _FakeWebSocket("/deepgram/v1/speak", query)
relay: Final = await _serve_speak(websocket)
assert relay.calls == []
assert websocket.closed == (
1008,
f"No streaming price for '{missing_key}': add it to the model cost map to enable it",
)
@pytest.mark.asyncio
async def test_deepgram_speak_refusal_for_a_long_unpriced_model_fits_a_close_frame(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
monkeypatch.setenv("DEEPGRAM_API_KEY", "dg-provider-key")
websocket: Final = _FakeWebSocket("/deepgram/v1/speak", f"model={'x' * 200}")
relay: Final = await _serve_speak(websocket)
assert relay.calls == []
assert websocket.closed is not None
assert (websocket.closed[0], len(websocket.closed[1].encode())) == (1008, WEBSOCKET_CLOSE_REASON_MAX_BYTES)
@pytest.mark.parametrize("path", SPEAK_PATHS)
@pytest.mark.parametrize(
"query",
[
pytest.param("model=aura-2-thalia-en", id="allowed voice"),
pytest.param("model=aura-2-thalia-en&model=aura-asteria-en", id="second voice stripped"),
],
)
def test_deepgram_speak_relays_only_the_voice_the_key_was_checked_against(
path: str, query: str, monkeypatch: pytest.MonkeyPatch
) -> None:
"""A key allowed only one voice must not reach another by repeating ``model``: the real key auth path must see
the voice the upstream target will carry"""
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
monkeypatch.setenv("DEEPGRAM_API_KEY", "dg-provider-key")
monkeypatch.setattr(litellm, "max_budget", 0.0)
cache: Final = asyncio.run(_cache_restricted_key("sk-only-thalia", ["aura-2-thalia-en"]))
relay: Final = _FakeRelay()
client: Final = TestClient(_app_with_relay(relay))
with patch.multiple( # test-quality-ok: the real key auth path reads these proxy_server globals and has no injection seam
"litellm.proxy.proxy_server",
master_key="sk-master",
prisma_client=MagicMock(),
user_api_key_cache=cache,
llm_model_list=None,
llm_router=None,
):
with client.websocket_connect(f"{path}?{query}", headers={"Authorization": "Bearer sk-only-thalia"}):
pass
assert [(call.target, call.endpoint) for call in relay.calls] == [
("wss://api.deepgram.com/v1/speak?model=aura-2-thalia-en", path)
]
@pytest.mark.parametrize(
"query",
[
pytest.param("model=aura-asteria-en", id="denied voice named"),
pytest.param("", id="voice omitted, default denied"),
],
)
def test_deepgram_speak_refuses_a_voice_the_key_is_not_allowed(query: str, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
monkeypatch.setenv("DEEPGRAM_API_KEY", "dg-provider-key")
monkeypatch.setattr(litellm, "max_budget", 0.0)
cache: Final = asyncio.run(_cache_restricted_key("sk-only-thalia", ["aura-2-thalia-en"]))
relay: Final = _FakeRelay()
client: Final = TestClient(_app_with_relay(relay))
with patch.multiple( # test-quality-ok: the real key auth path reads these proxy_server globals and has no injection seam
"litellm.proxy.proxy_server",
master_key="sk-master",
prisma_client=MagicMock(),
user_api_key_cache=cache,
llm_model_list=None,
llm_router=None,
):
with pytest.raises(WebSocketDisconnect) as disconnect:
with client.websocket_connect(
f"/deepgram/v1/speak?{query}", headers={"Authorization": "Bearer sk-only-thalia"}
):
pass
assert disconnect.value.code == 1008
assert relay.calls == []

113
tests/unit/proxy/pass_through_endpoints/test_watsonx_proxy_route.py
View file

@ -437,3 +437,116 @@ class TestWatsonxProxyRoute:
call_args["target"]
== f"https://us-south.ml.cloud.ibm.com/{endpoint_path}"
)
_WATSONX_ENV = (
"WATSONX_API_KEY",
"WATSONX_APIKEY",
"WX_API_KEY",
"WATSONX_API_BASE",
"WATSONX_URL",
"WX_URL",
"WML_URL",
"WATSONX_TOKEN",
"WATSONX_ZENAPIKEY",
)
def _watsonx_deployment(api_key: str, api_base: str, use_in_pass_through: bool) -> dict:
return {
"model_name": "granite",
"litellm_params": {
"model": "watsonx/ibm/granite-3-8b-instruct",
"api_key": api_key,
"api_base": api_base,
"project_id": "proj",
"use_in_pass_through": use_in_pass_through,
},
}
def _post_request(body: dict) -> MagicMock:
request = MagicMock(spec=Request)
request.method = "POST"
request.query_params = {}
request.headers = {"content-type": "application/json"}
request.json = AsyncMock(return_value=body)
return request
def _iam_response(token: str) -> MagicMock:
response = MagicMock()
response.json.return_value = {"access_token": token, "expires_in": 3600}
return response
async def _route_target_and_iam_apikey(
monkeypatch, llm_router: litellm.Router, env: dict[str, str]
) -> tuple[str, str, str]:
from litellm.llms.watsonx.common_utils import iam_token_cache
from litellm.proxy.pass_through_endpoints.passthrough_endpoint_router import PassthroughEndpointRouter
for name in _WATSONX_ENV:
monkeypatch.delenv(name, raising=False)
for name, value in env.items():
monkeypatch.setenv(name, value)
monkeypatch.setattr(litellm, "credential_list", [])
iam_token_cache.flush_cache()
monkeypatch.setattr(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router",
PassthroughEndpointRouter(llm_router_getter=lambda: llm_router),
)
endpoint_func = AsyncMock(return_value={"results": []})
with (
patch.object(litellm.module_level_client, "post", return_value=_iam_response("iam-token-1")) as iam_post,
patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route",
return_value=endpoint_func,
) as create_route,
):
await watsonx_proxy_route(
endpoint="ml/v1/text/generation",
request=_post_request({"input": "hi", "project_id": "proj"}),
fastapi_response=MagicMock(spec=Response),
user_api_key_dict=MagicMock(spec=UserAPIKeyAuth),
)
route_kwargs = create_route.call_args.kwargs
return (
route_kwargs["target"],
route_kwargs["custom_headers"]["Authorization"],
iam_post.call_args.kwargs["data"]["apikey"],
)
@pytest.mark.asyncio
async def test_watsonx_route_uses_the_use_in_pass_through_deployment_key_and_base(monkeypatch):
llm_router = litellm.Router(
model_list=[_watsonx_deployment("wx-deployment-key", "https://eu-de.ml.cloud.ibm.com", use_in_pass_through=True)]
)
target, authorization, iam_apikey = await _route_target_and_iam_apikey(monkeypatch, llm_router, env={})
assert target == "https://eu-de.ml.cloud.ibm.com/ml/v1/text/generation"
assert authorization == "Bearer iam-token-1"
assert iam_apikey == "wx-deployment-key"
@pytest.mark.asyncio
async def test_watsonx_route_falls_back_to_env_when_no_deployment_is_flagged(monkeypatch):
llm_router = litellm.Router(
model_list=[_watsonx_deployment("wx-deployment-key", "https://eu-de.ml.cloud.ibm.com", use_in_pass_through=False)]
)
target, authorization, iam_apikey = await _route_target_and_iam_apikey(
monkeypatch,
llm_router,
env={
"WX_API_KEY": "wx-env-key",
"WATSONX_API_KEY": "watsonx-env-key",
"WATSONX_API_BASE": "https://us-south.ml.cloud.ibm.com",
},
)
assert target == "https://us-south.ml.cloud.ibm.com/ml/v1/text/generation"
assert authorization == "Bearer iam-token-1"
assert iam_apikey == "wx-env-key"

8
tests/unit/router_strategy/test_router_routing_groups.py
View file

@ -236,7 +236,7 @@ async def test_async_dispatch_uses_group_strategy_for_grouped_model():
wraps=group_selector.async_get_available_deployments,
) as latency_spy,
patch(
"litellm.router.simple_shuffle", wraps=litellm.router.simple_shuffle
"litellm.router_python_core.simple_shuffle", wraps=litellm.router_python_core.simple_shuffle
) as shuffle_spy,
):
await router.async_get_available_deployment(
@ -271,7 +271,7 @@ async def test_async_dispatch_falls_back_to_default_for_ungrouped_models():
wraps=group_selector.async_get_available_deployments,
) as latency_spy,
patch(
"litellm.router.simple_shuffle", wraps=litellm.router.simple_shuffle
"litellm.router_python_core.simple_shuffle", wraps=litellm.router_python_core.simple_shuffle
) as shuffle_spy,
):
await router.async_get_available_deployment(
@ -331,7 +331,7 @@ async def test_async_dispatch_uses_default_selector_when_constructed_with_enum()
wraps=default_selector.async_get_available_deployments,
) as latency_spy,
patch(
"litellm.router.simple_shuffle", wraps=litellm.router.simple_shuffle
"litellm.router_python_core.simple_shuffle", wraps=litellm.router_python_core.simple_shuffle
) as shuffle_spy,
):
await router.async_get_available_deployment(
@ -1376,7 +1376,7 @@ async def test_group_call_dispatches_via_group_selector():
"async_get_available_deployments",
wraps=group_selector.async_get_available_deployments,
) as latency_spy,
patch("litellm.router.simple_shuffle", wraps=litellm.router.simple_shuffle) as shuffle_spy,
patch("litellm.router_python_core.simple_shuffle", wraps=litellm.router_python_core.simple_shuffle) as shuffle_spy,
):
deployment = await router.async_get_available_deployment(model="quality", request_kwargs={})

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

@ -2686,7 +2686,7 @@ async def test_plain_tag_exhaustion_with_universal_default_tag_raises_by_default
router = _quality_high_cost_low_router()
with patch(
"litellm.router.async_get_cooldown_deployments",
"litellm.router_python_core.async_get_cooldown_deployments",
new=AsyncMock(return_value=["quality-high-1", "quality-high-2"]),
):
with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info:
@ -2711,7 +2711,7 @@ async def test_plain_tag_exhaustion_with_universal_default_tag_falls_open_when_a
from unittest.mock import AsyncMock, patch
with patch(
"litellm.router.async_get_cooldown_deployments",
"litellm.router_python_core.async_get_cooldown_deployments",
new=AsyncMock(return_value=["quality-high-1", "quality-high-2"]),
):
response = await router.acompletion(

4
tests/unit/router_utils/test_cooldown_handlers.py
View file

@ -1303,7 +1303,7 @@ class TestDeploymentCallbackOnFailureCooldownTimePrecedence:
},
}
with patch("litellm.router.set_cooldown_deployments") as mock_set_cooldown:
with patch("litellm.router_python_core.set_cooldown_deployments") as mock_set_cooldown:
router.deployment_callback_on_failure(
kwargs=kwargs,
completion_response=None,
@ -1339,7 +1339,7 @@ class TestDeploymentCallbackOnFailureCooldownTimePrecedence:
},
}
with patch("litellm.router.set_cooldown_deployments") as mock_set_cooldown:
with patch("litellm.router_python_core.set_cooldown_deployments") as mock_set_cooldown:
router.deployment_callback_on_failure(
kwargs=kwargs,
completion_response=None,

3
tests/unit/router_utils/test_kubernetes_pod_discovery.py
View file

@ -17,6 +17,7 @@ import respx
from openai import AsyncOpenAI, OpenAI
import litellm
import litellm.router_python_core
from litellm.constants import (
CALLER_SESSION_ID_METADATA_KEY,
KUBERNETES_POD_ROUTING_KEY,
@ -281,7 +282,7 @@ def _patch_router_discovery(
*discoveries: KubernetesPodDiscovery,
) -> None:
discovery_iter: Final = iter(discoveries)
monkeypatch.setattr(litellm.router, "KubernetesPodDiscovery", lambda: next(discovery_iter))
monkeypatch.setattr(litellm.router_python_core, "KubernetesPodDiscovery", lambda: next(discovery_iter))
def _stub_sync_dns(monkeypatch: pytest.MonkeyPatch, *ips: str) -> None:

6
tests/unit/test_conftest_isolation.py
View file

@ -117,12 +117,12 @@ def test_live_router_membership_is_scoped_to_this_test():
}
]
)
assert _CanaryRouterHolder.router in litellm_router_module._live_routers
assert _CanaryRouterHolder.router._core in litellm_router_module._live_routers
def test_live_router_membership_was_rolled_back():
assert _CanaryRouterHolder.router is not None
assert _CanaryRouterHolder.router not in litellm_router_module._live_routers
assert _CanaryRouterHolder.router._core not in litellm_router_module._live_routers
def test_aws_cache_reset_sees_new_handlers_and_replaced_caches(monkeypatch: pytest.MonkeyPatch) -> None:
@ -175,7 +175,7 @@ def registered_collection_models() -> Iterator[None]:
yield
litellm_utils_module._runtime_registered_model_cost.clear()
litellm_utils_module._runtime_registered_model_cost.update(original_registrations)
litellm_router_module._live_routers.discard(collection_router)
litellm_router_module._live_routers.discard(collection_router._core)
litellm_router_module._live_routers.update(original_routers)

0
tests/unit/test_router/core/__init__.py Normal file
View file

26
tests/unit/test_router/core/conftest.py Normal file
View file

@ -0,0 +1,26 @@
from collections.abc import Callable
import pytest
from litellm.router_core import RouterCore
from tests.unit.test_router.core.scripted_host import CORE_BUILDERS, CoreDeclined, ScriptedHost
@pytest.fixture
def host() -> ScriptedHost:
return ScriptedHost()
@pytest.fixture(params=sorted(CORE_BUILDERS))
def make_core(request: pytest.FixtureRequest, host: ScriptedHost) -> Callable[..., RouterCore]:
"""``make_core(**router_arguments)``: the core under test, on the ``host`` fixture. Every test that
uses it runs once per core in ``CORE_BUILDERS`` (``-k python`` or ``-k rust`` picks one)."""
builder = CORE_BUILDERS[request.param]
def build(**arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments
try:
return builder(host, **arguments)
except CoreDeclined as declined:
pytest.skip(f"{request.param} core declines this configuration: {declined}")
return build

487
tests/unit/test_router/core/scripted_host.py Normal file
View file

@ -0,0 +1,487 @@
"""
A host for driving ``PythonCore`` in isolation.
``PythonCore`` reaches outside itself only through its host (``litellm/router_python_host.py``):
inference calls, backoff sleeps, the process settings, callback registration and the process-wide
cache. ``ScriptedHost`` serves all of those from the test. Inference ops answer from a per-model
script and record every call; sleeps are recorded and return at once; the process settings are a
fixed ``RouterDynamicConfig`` the test can replace; callbacks live in the host's own lists, and an
answered call runs the success or failure callbacks the core registered, as litellm's logging does.
Ops that are not inference (``get_model_info``, ``should_retry``, ...) are litellm's own functions.
"""
import dataclasses
import itertools
from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping
from types import MethodType
from typing import Final
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.logging_callback_manager import LoggingCallbackManager
from litellm.router_config import RouterDynamicConfig, build_router_config
from litellm.router_core import RouterCore
from litellm.router_python_core import PythonCore
from litellm.router_python_host import CallbackList, Endpoint, PythonHost
from litellm.types.utils import EmbeddingResponse, ModelResponse
INFERENCE_OPS: Final = frozenset(
{
"acompletion",
"completion",
"atext_completion",
"text_completion",
"aembedding",
"embedding",
"aimage_generation",
"image_generation",
"aspeech",
"atranscription",
"arerank",
"rerank",
"amoderation",
"aresponses",
"responses",
"anthropic_messages",
"agenerate_content",
"aadapter_generate_content",
"aocr",
"asearch",
"avideo_generation",
"acreate_batch",
"acreate_file",
"afile_content",
"afile_delete",
"acreate_fine_tuning_job",
"aget_responses",
"adelete_responses",
"acancel_responses",
"alist_input_items",
}
)
Outcome = object
"""What a scripted call answers: a ``BaseException`` instance is raised, a callable is called with
the request's keyword arguments (and its result used in turn), anything else is returned as is.
``None`` builds a default response for the op."""
@dataclasses.dataclass(frozen=True, slots=True)
class Call:
op: str
model: str
kwargs: Mapping[str, object]
@property
def deployment_id(self) -> str | None:
model_info: Final = self.kwargs.get("model_info")
return str(model_info["id"]) if isinstance(model_info, Mapping) and "id" in model_info else None
@property
def model_group(self) -> str | None:
metadata: Final = self.kwargs.get("metadata") or self.kwargs.get("litellm_metadata")
return metadata.get("model_group") if isinstance(metadata, Mapping) else None
def completion_response(content: str = "ok", total_tokens: int = 10, model: str = "gpt-4o") -> ModelResponse:
return ModelResponse(
model=model,
choices=[{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": content}}],
usage={
"prompt_tokens": total_tokens // 2,
"completion_tokens": total_tokens - total_tokens // 2,
"total_tokens": total_tokens,
},
)
def embedding_response(model: str = "text-embedding-3-small") -> EmbeddingResponse:
return EmbeddingResponse(
model=model,
data=[{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}],
usage={"prompt_tokens": 2, "total_tokens": 2},
)
def _fresh_error(template: BaseException) -> BaseException:
"""A copy of ``template`` with the same attributes (``copy.copy`` rebuilds litellm errors from
their args and drops ``status_code`` and the rest)."""
error: Final = BaseException.__new__(type(template))
error.args = template.args
error.__dict__.update(template.__dict__)
return error
def _default_response(op: str, model: str) -> object:
if op in {"aembedding", "embedding"}:
return embedding_response(model)
return completion_response(model=model)
def rate_limit(model: str = "gpt-4o", message: str = "slow down") -> litellm.RateLimitError:
return litellm.RateLimitError(message=message, llm_provider="openai", model=model)
def bad_request(model: str = "gpt-4o", message: str = "bad input") -> litellm.BadRequestError:
return litellm.BadRequestError(message=message, llm_provider="openai", model=model)
def context_window(model: str = "gpt-4o") -> litellm.ContextWindowExceededError:
return litellm.ContextWindowExceededError(message="too long", llm_provider="openai", model=model)
def content_policy(model: str = "gpt-4o") -> litellm.ContentPolicyViolationError:
return litellm.ContentPolicyViolationError(message="blocked", llm_provider="openai", model=model)
def server_error(model: str = "gpt-4o") -> litellm.InternalServerError:
return litellm.InternalServerError(message="boom", llm_provider="openai", model=model)
def timeout(model: str = "gpt-4o") -> litellm.Timeout:
return litellm.Timeout(message="timed out", llm_provider="openai", model=model)
def auth_error(model: str = "gpt-4o") -> litellm.AuthenticationError:
return litellm.AuthenticationError(message="bad key", llm_provider="openai", model=model)
def dynamic_config(**changes: object) -> RouterDynamicConfig: # kwargs-ok: RouterDynamicConfig fields
"""Process settings as a fresh litellm process has them, with no callbacks, cache or fallbacks."""
base: Final = RouterDynamicConfig(
version=1,
num_retries=None,
max_fallbacks=None,
router_max_fallbacks=5,
allowed_fails=3,
request_timeout=6000.0,
fallbacks=None,
default_fallbacks=None,
context_window_fallbacks=None,
content_policy_fallbacks=None,
drop_params=None,
provider_list=litellm.provider_list,
expose_router_debug_in_errors=False,
redact_user_api_key_info=None,
model_alias_map={},
model_cost=litellm.model_cost,
known_custom_logger_compatible_callbacks=litellm._known_custom_logger_compatible_callbacks,
exception_types=litellm.LITELLM_EXCEPTION_TYPES,
callbacks=[],
success_callback=[],
failure_callback=[],
input_callback=[],
callback_manager=litellm.logging_callback_manager,
cache=None,
)
return dataclasses.replace(base, **changes)
class HostCallbacks:
"""The callback registration the core asks its host for, over the host's own lists."""
def __init__(self, host: "ScriptedHost") -> None:
self._host: Final = host
self._keys: Final = LoggingCallbackManager()
def _add(self, name: CallbackList, callback: object) -> None:
target: Final = self._host.lists[name]
if callback not in target:
target.append(callback)
def add_litellm_callback(self, callback: object) -> None:
self._add_to("callbacks", callback)
def _add_to(self, name: str, callback: object) -> None:
target: Final = self._host.lists.setdefault(name, [])
if callback not in target:
target.append(callback)
def add_litellm_success_callback(self, callback: object) -> None:
self._add("success_callback", callback)
def add_litellm_async_success_callback(self, callback: object) -> None:
self._add("_async_success_callback", callback)
def add_litellm_failure_callback(self, callback: object) -> None:
self._add("failure_callback", callback)
def add_litellm_async_failure_callback(self, callback: object) -> None:
self._add("_async_failure_callback", callback)
def add_litellm_input_callback(self, callback: object) -> None:
self._add("input_callback", callback)
def get_custom_logger_key(self, custom_logger: CustomLogger) -> str:
return self._keys.get_custom_logger_key(custom_logger)
def with_custom_logger_keys(
self, candidates: Iterable[CustomLogger], keys: frozenset[str]
) -> tuple[CustomLogger, ...]:
return self._keys.with_custom_logger_keys(candidates, keys)
def without_custom_logger_keys(self, callbacks: Iterable[object], keys: frozenset[str]) -> Iterator[object]:
return self._keys.without_custom_logger_keys(callbacks, keys) # pyright: ignore[reportArgumentType] # any callback list
def get_all_callbacks(self) -> list[object]: # mutable-ok: matches LoggingCallbackManager
return [callback for callbacks in self._host.lists.values() for callback in callbacks]
def get_custom_loggers_for_type(
self, callback_type: type[CustomLogger]
) -> list[CustomLogger]: # mutable-ok: matches LoggingCallbackManager
found: Final[list[CustomLogger]] = []
for callback in self.get_all_callbacks():
if isinstance(callback, callback_type) and callback not in found:
found.append(callback)
return found
def drop_callbacks(self, callbacks: Iterable[object], method_owners: Iterable[object] = ()) -> frozenset[str]:
dropped_ids: Final = frozenset(map(id, callbacks))
owner_ids: Final = frozenset(map(id, method_owners))
def keep(callback: object) -> bool:
if id(callback) in dropped_ids:
return False
return not (isinstance(callback, MethodType) and id(callback.__self__) in owner_ids)
vacated: Final = frozenset(
self.get_custom_logger_key(callback)
for callback in self._host.lists.get("callbacks", [])
if isinstance(callback, CustomLogger) and id(callback) in dropped_ids
)
for name, callbacks_in_list in self._host.lists.items():
self._host.lists[name] = [callback for callback in callbacks_in_list if keep(callback)]
return vacated
class ScriptedHost(PythonHost):
"""
``script`` maps a deployment's ``litellm_params.model`` (or ``"*"`` for any) to the outcomes its
calls answer in turn; the last outcome repeats. A model with no script answers a default response.
"""
__slots__ = ("_dynamic", "_scripts", "_turns", "cache", "calls", "lists", "raised", "sleeps", "suppressed")
def __init__(
self, script: Mapping[str, Iterable[Outcome]] | None = None, **dynamic: object
) -> None: # kwargs-ok: RouterDynamicConfig fields
self._scripts: dict[str, tuple[Outcome, ...]] = {k: tuple(v) for k, v in (script or {}).items()}
self._turns: dict[str, Iterator[int]] = {}
self._dynamic: RouterDynamicConfig = dynamic_config(**dynamic)
self.calls: Final[list[Call]] = []
self.sleeps: Final[list[float]] = []
# The errors the scripted calls raised, in order: each is a fresh copy of its scripted template.
self.raised: Final[list[BaseException]] = []
self.lists: Final[dict[str, list[object]]] = {
"callbacks": [],
"success_callback": [],
"_async_success_callback": [],
"failure_callback": [],
"_async_failure_callback": [],
"input_callback": [],
}
self.cache: object = None
self.suppressed = False
# --- what the test controls -------------------------------------------------------------
def answer(self, model: str, *outcomes: Outcome) -> None:
"""Script the calls to ``model`` from now on."""
self._scripts[model] = outcomes
self._turns.pop(model, None)
def set_dynamic(self, **changes: object) -> None: # kwargs-ok: RouterDynamicConfig fields
self._dynamic = dataclasses.replace(self._dynamic, **changes, version=self._dynamic.version + 1)
def calls_to(self, op: str | None = None) -> list[Call]: # mutable-ok: a fresh list
return [call for call in self.calls if op is None or call.op == op]
def deployments_called(self, op: str | None = None) -> list[str | None]: # mutable-ok: a fresh list
return [call.deployment_id for call in self.calls_to(op)]
# --- the host operations ----------------------------------------------------------------
def invoke(self, op: str) -> Endpoint:
if op not in INFERENCE_OPS:
return super().invoke(op)
if op.startswith("a") or op in {"anthropic_messages", "_arealtime"}:
async def endpoint(*args: object, **kwargs: object) -> object: # kwargs-ok: the endpoint's request
return await self._answer_async(op, args, kwargs)
return endpoint
def sync_endpoint(*args: object, **kwargs: object) -> object: # kwargs-ok: the endpoint's request
return self._answer_sync(op, args, kwargs)
return sync_endpoint
def sleep(self, seconds: float) -> Awaitable[None]:
self.sleeps.append(seconds)
async def done() -> None:
return None
return done()
def dynamic(self) -> RouterDynamicConfig:
return self._dynamic
@property
def callbacks(self) -> HostCallbacks: # pyright: ignore[reportIncompatibleVariableOverride] # the core uses the same methods
return HostCallbacks(self)
def callback_list(self, name: CallbackList) -> list[object]: # mutable-ok: the host's own list
return self.lists.setdefault(name, [])
def set_callback_list(self, name: CallbackList, value: list[object]) -> None: # mutable-ok: assigned as given
self.lists[name] = value
def global_cache(self) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # any cache object
return self.cache
def set_global_cache(self, cache: object) -> None: # pyright: ignore[reportIncompatibleMethodOverride] # any cache object
self.cache = cache
def suppress_debug_info(self) -> None:
self.suppressed = True
# --- answering ------------------------------------------------------------------------
def _outcome(self, op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> tuple[Call, object]:
model: Final = str(kwargs.get("model") or (args[0] if args else ""))
call: Final = Call(op=op, model=model, kwargs=dict(kwargs))
self.calls.append(call)
key: Final = model if model in self._scripts else "*" if "*" in self._scripts else None
if key is None:
return call, None
outcomes: Final = self._scripts[key]
turn: Final = next(self._turns.setdefault(key, itertools.count()))
outcome: object = outcomes[min(turn, len(outcomes) - 1)] if outcomes else None
while callable(outcome) and not isinstance(outcome, BaseException | type):
outcome = outcome(**kwargs) # rebind-ok: a callable outcome answers with what it returns
if isinstance(outcome, BaseException):
# A fresh error per call: the core stamps attributes (num_retries, ...) on the one it raises.
outcome = _fresh_error(outcome) # rebind-ok: the scripted error is a template
self.raised.append(outcome)
return call, outcome
def _logging_kwargs(
self, call: Call, error: BaseException | None, response: object
) -> dict[str, object]: # mutable-ok: litellm's logging kwargs
kwargs: Final = call.kwargs
usage: Final = getattr(response, "usage", None)
metadata: Final = kwargs.get("metadata") or kwargs.get("litellm_metadata") or {}
model_info: Final = kwargs.get("model_info") or {}
return {
"call_type": call.op,
"model": call.model,
"exception": error,
"litellm_params": {
"metadata": metadata,
"model_info": model_info,
"api_base": kwargs.get("api_base"),
**({"cooldown_time": kwargs["cooldown_time"]} if "cooldown_time" in kwargs else {}),
},
"standard_logging_object": {
"total_tokens": getattr(usage, "total_tokens", 0) or 0,
"model_group": metadata.get("model_group"),
"model_id": model_info.get("id"),
"hidden_params": {"litellm_model_name": call.model},
},
}
def _finish(self, op: str, call: Call, outcome: object) -> tuple[BaseException | None, object]:
if isinstance(outcome, type) and issubclass(outcome, BaseException):
outcome = outcome("scripted failure") # rebind-ok: an exception class raises an instance
if isinstance(outcome, BaseException):
return outcome, None
return None, _default_response(op, call.model) if outcome is None else outcome
async def _answer_async(self, op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> object:
call, outcome = self._outcome(op, args, kwargs)
error, response = self._finish(op, call, outcome)
logged: Final = self._logging_kwargs(call, error, response)
if error is not None:
for callback in list(self.lists.get("failure_callback", ())):
if callable(callback):
callback(logged, None, 0.0, 0.0)
for callback in list(self.lists.get("_async_failure_callback", ())):
if callable(callback):
await callback(logged, None, 0.0, 0.0)
raise error
for callback in list(self.lists.get("_async_success_callback", ())):
if callable(callback):
await callback(logged, response, 0.0, 0.0)
return response
def _answer_sync(self, op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> object:
call, outcome = self._outcome(op, args, kwargs)
error, response = self._finish(op, call, outcome)
logged: Final = self._logging_kwargs(call, error, response)
if error is not None:
for callback in list(self.lists.get("failure_callback", ())):
if callable(callback):
callback(logged, None, 0.0, 0.0)
raise error
for callback in list(self.lists.get("success_callback", ())):
if callable(callback):
callback(logged, response, 0.0, 0.0)
return response
def deployment(
model_name: str,
deployment_id: str,
model: str | None = None,
**params: object, # kwargs-ok: extra litellm_params
) -> dict[str, object]: # mutable-ok: a model_list entry
"""One ``model_list`` entry; ``model`` defaults to ``openai/<deployment_id>`` so each deployment
answers its own script."""
model_info: Final = params.pop("model_info", None)
return {
"model_name": model_name,
"litellm_params": {"model": model or f"openai/{deployment_id}", "api_key": "sk-test", **params},
"model_info": {"id": deployment_id, **(model_info if isinstance(model_info, dict) else {})},
}
def python_core(host: ScriptedHost, **arguments: object) -> PythonCore: # kwargs-ok: the router's constructor arguments
"""A ``PythonCore`` built the way ``Router`` builds one, on ``host``."""
config: Final = build_router_config(arguments)
return PythonCore(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the same arguments Router binds
CoreBuilder = Callable[..., RouterCore]
"""``builder(host, **router_arguments)``: a router core on ``host``."""
CORE_BUILDERS: Final[dict[str, CoreBuilder]] = {"python": python_core}
"""Every core the suite runs on. Another core adds its builder here; a builder that cannot take a
configuration raises ``CoreDeclined``, and the test is skipped for that core."""
class CoreDeclined(Exception):
"""The core does not serve this configuration (the router would stay on the Python core)."""
__all__ = [
"Call",
"HostCallbacks",
"INFERENCE_OPS",
"CORE_BUILDERS",
"CoreBuilder",
"CoreDeclined",
"ScriptedHost",
"auth_error",
"bad_request",
"completion_response",
"content_policy",
"context_window",
"deployment",
"dynamic_config",
"embedding_response",
"python_core",
"rate_limit",
"server_error",
"timeout",
]

445
tests/unit/test_router/core/test_endpoints.py Normal file
View file

@ -0,0 +1,445 @@
"""
What a router core sends to its host for each endpoint, and what it hands back.
Every test builds its core through ``make_core`` on the ``host`` fixture, so it runs once per core.
"""
import datetime
import io
from collections.abc import AsyncIterator, Callable, Coroutine, Mapping
from typing import Final
import pytest
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.router_core import RouterCore
from litellm.types.utils import EmbeddingResponse, ModelResponse, ModelResponseStream
from tests.unit.test_router.core.scripted_host import (
ScriptedHost,
bad_request,
completion_response,
deployment,
embedding_response,
)
MESSAGES: Final = [{"role": "user", "content": "hi"}]
MakeCore = Callable[..., RouterCore]
EndpointCall = Callable[[RouterCore], object]
ENDPOINTS: Final[dict[str, tuple[str, EndpointCall]]] = {
"acompletion": ("acompletion", lambda core: core.acompletion(model="g", messages=MESSAGES)),
"completion": ("completion", lambda core: core.completion(model="g", messages=MESSAGES)),
"aembedding": ("aembedding", lambda core: core.aembedding(model="g", input=["a"])),
"embedding": ("embedding", lambda core: core.embedding(model="g", input=["a"])),
"atext_completion": ("atext_completion", lambda core: core.atext_completion(model="g", prompt="hi")),
"aimage_generation": ("aimage_generation", lambda core: core.aimage_generation(model="g", prompt="a cat")),
"image_generation": ("image_generation", lambda core: core.image_generation(model="g", prompt="a cat")),
"arerank": ("arerank", lambda core: core.arerank(model="g", query="q", documents=["a"])),
"aspeech": ("aspeech", lambda core: core.aspeech(model="g", input="hi", voice="alloy")),
"atranscription": ("atranscription", lambda core: core.atranscription(model="g", file=io.BytesIO(b"x"))),
"aresponses": ("aresponses", lambda core: core.aresponses(model="g", input="hi")),
"aanthropic_messages": (
"anthropic_messages",
lambda core: core.aanthropic_messages(model="g", messages=MESSAGES, max_tokens=5),
),
"call-acompletion": ("acompletion", lambda core: core.call("acompletion", {"model": "g", "messages": MESSAGES})),
"call-completion": ("completion", lambda core: core.call("completion", {"model": "g", "messages": MESSAGES})),
}
SYNC_ENDPOINTS: Final = frozenset({"completion", "embedding", "image_generation"})
LITELLM_METADATA_ENDPOINTS: Final = frozenset({"aresponses", "aanthropic_messages"})
EMBEDDING_ENDPOINTS: Final = frozenset({"aembedding", "embedding"})
async def _invoke(core: RouterCore, endpoint: str) -> object:
result: Final = ENDPOINTS[endpoint][1](core)
if endpoint in SYNC_ENDPOINTS:
return result
return await result # pyright: ignore[reportGeneralTypeIssues] # the async endpoints return awaitables
def _one_group(
make_core: MakeCore, **litellm_params: object
) -> RouterCore: # kwargs-ok: the deployment's litellm_params
return make_core(model_list=[deployment("g", "d1", **litellm_params)])
def _metadata(kwargs: Mapping[str, object], endpoint: str) -> Mapping[str, object]:
bucket: Final = kwargs["litellm_metadata" if endpoint in LITELLM_METADATA_ENDPOINTS else "metadata"]
assert isinstance(bucket, Mapping)
return bucket
@pytest.mark.parametrize("endpoint", sorted(ENDPOINTS))
async def test_each_endpoint_sends_the_deployment_litellm_params_to_the_host(
make_core: MakeCore, host: ScriptedHost, endpoint: str
):
core: Final = _one_group(make_core, api_base="https://d1.example", temperature=0.3, seed=7)
await _invoke(core, endpoint)
sent: Final = host.calls[0].kwargs
assert (sent["model"], sent["api_key"], sent["api_base"], sent["temperature"], sent["seed"]) == (
"openai/d1",
"sk-test",
"https://d1.example",
0.3,
7,
)
@pytest.mark.parametrize("endpoint", sorted(ENDPOINTS))
async def test_each_endpoint_invokes_its_host_op_once_with_routing_metadata(
make_core: MakeCore, host: ScriptedHost, endpoint: str
):
core: Final = _one_group(make_core, api_base="https://d1.example")
await _invoke(core, endpoint)
assert [call.op for call in host.calls] == [ENDPOINTS[endpoint][0]]
call: Final = host.calls[0]
metadata: Final = _metadata(call.kwargs, endpoint)
assert (call.model_group, call.deployment_id) == ("g", "d1")
assert (metadata["model_group"], metadata["deployment"], metadata["api_base"]) == (
"g",
"openai/d1",
"https://d1.example",
)
assert metadata["model_info"]["id"] == "d1" # pyright: ignore[reportIndexIssue] # metadata values are untyped
@pytest.mark.parametrize("endpoint", sorted(ENDPOINTS))
async def test_each_endpoint_returns_the_host_response_object_with_routing_hidden_params(
make_core: MakeCore, host: ScriptedHost, endpoint: str
):
answer: Final[ModelResponse | EmbeddingResponse] = (
embedding_response() if endpoint in EMBEDDING_ENDPOINTS else completion_response("scripted")
)
host.answer("openai/d1", answer)
core: Final = _one_group(make_core)
response: Final = await _invoke(core, endpoint)
assert response is answer
assert response._hidden_params["model_id"] == "d1"
assert response._hidden_params["additional_headers"]["x-litellm-model-group"] == "g"
@pytest.mark.parametrize("endpoint", sorted(SYNC_ENDPOINTS | {"call-completion"}))
async def test_sync_endpoints_return_the_response_not_an_awaitable(
make_core: MakeCore, host: ScriptedHost, endpoint: str
):
answer: Final = embedding_response() if endpoint == "embedding" else completion_response("sync")
host.answer("openai/d1", answer)
core: Final = _one_group(make_core)
if endpoint == "call-completion":
assert await core.call("completion", {"model": "g", "messages": MESSAGES}) is answer
else:
assert ENDPOINTS[endpoint][1](core) is answer
@pytest.mark.parametrize(
"endpoint, request_body",
[("acompletion", {"messages": MESSAGES}), ("completion", {"messages": MESSAGES})],
)
async def test_router_only_include_fallback_errors_is_not_sent_to_the_host(
make_core: MakeCore, host: ScriptedHost, endpoint: str, request_body: dict[str, object]
):
core: Final = _one_group(make_core)
await core.call(endpoint, {"model": "g", "include_fallback_errors": True, **request_body})
assert host.calls[0].op == endpoint
assert "include_fallback_errors" not in host.calls[0].kwargs
@pytest.mark.parametrize("router_only_param", ["rpm", "tpm", "max_parallel_requests"])
async def test_deployment_limit_params_are_forwarded_with_the_rest_of_litellm_params(
make_core: MakeCore, host: ScriptedHost, router_only_param: str
):
core: Final = _one_group(make_core, **{router_only_param: 50})
await core.acompletion(model="g", messages=MESSAGES)
assert host.calls[0].kwargs[router_only_param] == 50
@pytest.mark.parametrize("endpoint", ["acompletion", "aembedding", "completion", "embedding"])
async def test_request_metadata_reaches_the_host_beside_the_routing_metadata(
make_core: MakeCore, host: ScriptedHost, endpoint: str
):
core: Final = _one_group(make_core)
request: Final[dict[str, object]] = {"messages": MESSAGES} if "completion" in endpoint else {"input": ["a"]}
await core.call(endpoint, {"model": "g", "metadata": {"trace": "t-1"}, **request})
metadata: Final = _metadata(host.calls[0].kwargs, endpoint)
assert (metadata["trace"], metadata["model_group"]) == ("t-1", "g")
async def test_aresponses_sends_request_metadata_as_metadata_and_routing_metadata_as_litellm_metadata(
make_core: MakeCore, host: ScriptedHost
):
core: Final = _one_group(make_core)
await core.aresponses(model="g", input="hi", metadata={"trace": "t-1"})
sent: Final = host.calls[0].kwargs
assert sent["metadata"] == {"trace": "t-1"}
assert sent["litellm_metadata"]["model_group"] == "g" # pyright: ignore[reportIndexIssue] # untyped kwargs
@pytest.mark.parametrize(
"router_arguments, deployment_params, expected_timeout",
[
({}, {}, 6000.0),
({"timeout": 42}, {}, 42),
({"timeout": 42}, {"timeout": 11}, 11.0),
({}, {"timeout": 11}, 11.0),
],
)
async def test_acompletion_timeout_comes_from_deployment_then_router_then_process_setting(
make_core: MakeCore,
host: ScriptedHost,
router_arguments: dict[str, object],
deployment_params: dict[str, object],
expected_timeout: float,
):
core: Final = make_core(model_list=[deployment("g", "d1", **deployment_params)], **router_arguments)
await core.acompletion(model="g", messages=MESSAGES)
assert host.calls[0].kwargs["timeout"] == expected_timeout
@pytest.mark.parametrize("endpoint", ["acompletion", "aembedding", "aresponses", "aanthropic_messages"])
async def test_request_timeout_overrides_router_and_deployment_timeouts(
make_core: MakeCore, host: ScriptedHost, endpoint: str
):
core: Final = make_core(model_list=[deployment("g", "d1", timeout=11)], timeout=42)
request: Final[dict[str, object]] = (
{"input": ["a"]}
if endpoint == "aembedding"
else {"input": "hi"}
if endpoint == "aresponses"
else {"messages": MESSAGES, "max_tokens": 5}
)
await core.call(endpoint, {"model": "g", "timeout": 7, **request})
assert host.calls[0].kwargs["timeout"] == 7
async def test_aresponses_timeout_falls_back_to_router_timeout(make_core: MakeCore, host: ScriptedHost):
core: Final = make_core(model_list=[deployment("g", "d1")], timeout=42)
await core.aresponses(model="g", input="hi")
assert host.calls[0].kwargs["timeout"] == 42
@pytest.mark.parametrize(
"request_stream_timeout, expected_stream_timeout",
[(None, 3.0), (2, 2)],
)
async def test_streaming_acompletion_sends_deployment_or_request_stream_timeout(
make_core: MakeCore, host: ScriptedHost, request_stream_timeout: int | None, expected_stream_timeout: float
):
core: Final = _one_group(make_core, timeout=11, stream_timeout=3)
extra: Final = {} if request_stream_timeout is None else {"stream_timeout": request_stream_timeout}
await core.acompletion(model="g", messages=MESSAGES, stream=True, **extra)
sent: Final = host.calls[0].kwargs
assert (sent["stream"], sent["timeout"], sent["stream_timeout"]) == (True, 11.0, expected_stream_timeout)
@pytest.mark.parametrize(
"request_params, expected",
[
({}, ("u1", 0.9)),
({"temperature": 0.1}, ("u1", 0.1)),
({"user": "u2"}, ("u2", 0.9)),
],
)
async def test_default_litellm_params_fill_what_the_request_leaves_unset(
make_core: MakeCore, host: ScriptedHost, request_params: dict[str, object], expected: tuple[str, float]
):
core: Final = make_core(
model_list=[deployment("g", "d1")], default_litellm_params={"user": "u1", "temperature": 0.9}
)
await core.acompletion(model="g", messages=MESSAGES, **request_params)
sent: Final = host.calls[0].kwargs
assert (sent["user"], sent["temperature"]) == expected
@pytest.mark.parametrize(
"request_credentials, expected_api_base",
[
({"api_key": "sk-client"}, "https://d1.example"),
({"api_key": "sk-client", "api_base": "https://client.example"}, "https://client.example"),
],
)
async def test_request_level_api_key_replaces_the_deployment_key_on_the_host_call(
make_core: MakeCore, host: ScriptedHost, request_credentials: dict[str, str], expected_api_base: str
):
core: Final = _one_group(make_core, api_base="https://d1.example")
await core.acompletion(model="g", messages=MESSAGES, **request_credentials)
call: Final = host.calls[0]
assert (call.kwargs["api_key"], call.kwargs["api_base"], call.model_group) == (
"sk-client",
expected_api_base,
"g",
)
async def test_construction_registers_one_success_and_one_failure_callback_per_mode(
make_core: MakeCore, host: ScriptedHost
):
_one_group(make_core)
registered: Final = {name: len(callbacks) for name, callbacks in host.lists.items()}
assert registered == {
"callbacks": 0,
"success_callback": 1,
"_async_success_callback": 1,
"failure_callback": 1,
"_async_failure_callback": 1,
"input_callback": 0,
}
assert all(callable(callbacks[0]) for callbacks in host.lists.values() if callbacks)
async def test_retire_removes_the_callbacks_the_core_registered(make_core: MakeCore, host: ScriptedHost):
core: Final = _one_group(make_core)
core.retire()
assert {name: callbacks for name, callbacks in host.lists.items() if callbacks} == {}
async def test_sync_completion_counts_its_tokens_and_request_in_group_usage(make_core: MakeCore, host: ScriptedHost):
core: Final = _one_group(make_core, tpm=100, rpm=5)
assert await core.health().usage("g") == (None, None)
core.completion(model="g", messages=MESSAGES)
assert await core.health().usage("g") == (10, 1)
assert await core.health().usage("g", kind="remaining") == {
"x-ratelimit-remaining-tokens": 90,
"x-ratelimit-limit-tokens": 100,
"x-ratelimit-remaining-requests": 4,
"x-ratelimit-limit-requests": 5,
}
class DeferredLoggingHost(ScriptedHost):
"""``ScriptedHost`` that, like litellm, runs async success callbacks after the endpoint returns."""
__slots__ = ("pending",)
def __init__(self) -> None:
super().__init__()
self.pending: list[Coroutine[object, object, None]] = []
async def _answer_async(self, op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> object:
callbacks: Final = self.lists["_async_success_callback"]
self.lists["_async_success_callback"] = []
try:
response: Final = await super()._answer_async(op, args, kwargs)
finally:
self.lists["_async_success_callback"] = callbacks
call: Final = self.calls[-1]
logged: Final = self._logging_kwargs(call, None, response)
self.pending.extend(callback(logged, response, 0.0, 0.0) for callback in callbacks if callable(callback))
return response
async def drain(self) -> None:
while self.pending:
await self.pending.pop(0)
class TestUsageWhenTheHostLogsAfterTheResponse:
"""litellm runs async success callbacks after the endpoint has returned; this host does the same."""
@pytest.fixture
def host(self) -> DeferredLoggingHost:
return DeferredLoggingHost()
async def test_one_acompletion_counts_its_tokens_and_request_once(
self, make_core: MakeCore, host: DeferredLoggingHost
):
core: Final = _one_group(make_core, tpm=100, rpm=5)
await core.acompletion(model="g", messages=MESSAGES)
assert await core.health().usage("g") == (10, 1)
await host.drain()
assert await core.health().usage("g") == (10, 1)
async def test_usage_based_routing_registers_its_logger_on_the_host(make_core: MakeCore, host: ScriptedHost):
make_core(model_list=[deployment("g", "d1", rpm=1)], routing_strategy="usage-based-routing-v2")
assert [type(callback).__name__ for callback in host.lists["callbacks"]] == ["LowestTPMLoggingHandler_v2"]
async def test_a_successful_call_uses_up_the_deployment_rpm_for_the_next_selection(
make_core: MakeCore, host: ScriptedHost
):
core: Final = make_core(model_list=[deployment("g", "d1", rpm=1)], routing_strategy="usage-based-routing-v2")
await core.acompletion(model="g", messages=MESSAGES)
with pytest.raises(litellm.RateLimitError, match="No deployments available"):
await core.acompletion(model="g", messages=MESSAGES)
assert host.deployments_called() == ["d1"]
async def test_host_error_reaches_the_caller_as_raised(make_core: MakeCore, host: ScriptedHost):
error: Final = bad_request(message="scripted bad input")
host.answer("openai/d1", error)
core: Final = _one_group(make_core)
with pytest.raises(litellm.BadRequestError, match="scripted bad input"):
await core.acompletion(model="g", messages=MESSAGES)
assert host.deployments_called("acompletion")[0] == "d1"
@pytest.mark.parametrize("endpoint", ["acompletion", "completion"])
async def test_unknown_model_group_raises_without_calling_the_host(
make_core: MakeCore, host: ScriptedHost, endpoint: str
):
core: Final = _one_group(make_core)
with pytest.raises(litellm.BadRequestError, match="model=missing"):
await core.call(endpoint, {"model": "missing", "messages": MESSAGES})
assert host.calls == []
@pytest.mark.parametrize("op", ["no_such_endpoint", "acompletions", ""])
async def test_call_with_an_unknown_op_raises_attribute_error_without_calling_the_host(
make_core: MakeCore, host: ScriptedHost, op: str
):
core: Final = _one_group(make_core)
with pytest.raises(AttributeError):
await core.call(op, {"model": "g"})
assert host.calls == []
async def test_sync_text_completion_invokes_the_sync_host_op_with_deployment_credentials(
make_core: MakeCore, host: ScriptedHost
):
core: Final = _one_group(make_core, api_base="https://d1.example")
core.text_completion(model="g", prompt="hi")
call: Final = host.calls[0]
assert (call.op, call.kwargs["api_base"], call.kwargs["prompt"], call.model_group) == (
"text_completion",
"https://d1.example",
"hi",
"g",
)
def _scripted_stream(*texts: str) -> Callable[..., CustomStreamWrapper]:
def answer(**_request: object) -> CustomStreamWrapper: # kwargs-ok: the host passes the request kwargs
async def chunks() -> AsyncIterator[ModelResponseStream]:
for text in texts:
yield ModelResponseStream(choices=[{"index": 0, "delta": {"role": "assistant", "content": text}}])
logging_obj: Final = LiteLLMLogging(
model="d1",
messages=MESSAGES,
stream=True,
call_type="acompletion",
start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc),
litellm_call_id="stream-test",
function_id="stream-test",
)
logging_obj.update_environment_variables(model="d1", user=None, optional_params={}, litellm_params={})
return CustomStreamWrapper(
completion_stream=chunks(), model="d1", custom_llm_provider="openai", logging_obj=logging_obj
)
return answer
async def test_streaming_acompletion_yields_the_host_stream_chunks_with_routing_hidden_params(
make_core: MakeCore, host: ScriptedHost
):
host.answer("openai/d1", _scripted_stream("he", "llo"))
core: Final = _one_group(make_core)
stream: Final = await core.acompletion(model="g", messages=MESSAGES, stream=True)
texts: Final = [chunk.choices[0].delta.content async for chunk in stream] # pyright: ignore[reportUnknownMemberType] # stream chunks are untyped
assert "".join(text for text in texts if text) == "hello"
assert stream._hidden_params["model_id"] == "d1"
assert stream._hidden_params["additional_headers"]["x-litellm-model-group"] == "g"
assert (host.calls[0].op, host.calls[0].kwargs["stream"]) == ("acompletion", True)

756
tests/unit/test_router/core/test_fallbacks.py Normal file
View file

@ -0,0 +1,756 @@
"""Cross-group fallbacks of a router core: where the chain comes from, how it is walked, and what is raised."""
from collections.abc import Callable, Mapping
import pytest
import litellm
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
from litellm.router_api import ChangeSet
from litellm.router_core import RouterCore
from tests.unit.test_router.core.scripted_host import (
ScriptedHost,
bad_request,
content_policy,
context_window,
deployment,
rate_limit,
server_error,
timeout,
)
MakeCore = Callable[..., RouterCore]
MESSAGES = [{"role": "user", "content": "hi"}]
GROUPS_ABC = [deployment("a", "a"), deployment("b", "b"), deployment("c", "c")]
GROUPS_ABCD = [*GROUPS_ABC, deployment("d", "d")]
CHAIN_ABCD = [{"a": ["b"]}, {"b": ["c"]}, {"c": ["d"]}]
TIERED_A = [deployment("a", "a1", order=1), deployment("a", "a2", order=2), deployment("b", "b")]
def build(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments
return make_core(**{"model_list": GROUPS_ABC, "num_retries": 0, **arguments})
async def complete(core: RouterCore, model: str = "a", **request: object) -> object: # kwargs-ok: request kwargs
return await core.acompletion(model=model, messages=MESSAGES, **request) # pyright: ignore[reportAttributeAccessIssue] # endpoint method
def headers(response: object) -> Mapping[str, object]:
return getattr(response, "_hidden_params")["additional_headers"]
def show_router_debug(host: ScriptedHost) -> None:
host.set_dynamic(expose_router_debug_in_errors=True)
# --- where the fallback chain comes from -------------------------------------------------------
@pytest.mark.asyncio
@pytest.mark.parametrize(
("constructor", "process", "request_kwargs", "expected"),
[
({"fallbacks": [{"a": ["b"]}]}, None, {}, ["a", "b"]),
({}, [{"a": ["c"]}], {}, ["a", "c"]),
({}, None, {"fallbacks": [{"a": ["b"]}]}, ["a", "b"]),
({}, None, {"fallbacks": ["c"]}, ["a", "c"]),
({"fallbacks": [{"a": ["b"]}]}, [{"a": ["c"]}], {}, ["a", "b"]),
({"fallbacks": [{"a": ["b"]}]}, None, {"fallbacks": [{"a": ["c"]}]}, ["a", "c"]),
],
ids=[
"constructor",
"process-setting",
"request-kwarg",
"request-plain-list",
"constructor-over-process-setting",
"request-over-constructor",
],
)
async def test_a_failed_group_falls_back_to_the_chain_from_the_winning_source(
make_core: MakeCore,
host: ScriptedHost,
constructor: dict[str, object],
process: list[object] | None,
request_kwargs: dict[str, object],
expected: list[str],
) -> None:
if process is not None:
host.set_dynamic(fallbacks=process)
core = build(make_core, **constructor)
host.answer("openai/a", rate_limit())
response = await complete(core, **request_kwargs)
assert host.deployments_called("acompletion") == expected
assert headers(response)["x-litellm-model-group"] == expected[-1]
@pytest.mark.asyncio
async def test_an_empty_request_fallback_list_overrides_the_constructor_chain(
make_core: MakeCore, host: ScriptedHost
) -> None:
core = build(make_core, fallbacks=[{"a": ["b"]}])
host.answer("openai/a", rate_limit())
with pytest.raises(litellm.RateLimitError):
await complete(core, fallbacks=[])
assert host.deployments_called("acompletion") == ["a"]
@pytest.mark.asyncio
@pytest.mark.parametrize("field", ["assign", "settings"])
async def test_fallbacks_applied_by_changeset_take_effect(make_core: MakeCore, host: ScriptedHost, field: str) -> None:
core = build(make_core)
core.apply(ChangeSet(**{field: {"fallbacks": [{"a": ["c"]}]}})) # pyright: ignore[reportArgumentType] # assign or settings
host.answer("openai/a", rate_limit())
await complete(core)
assert host.deployments_called("acompletion") == ["a", "c"]
# --- resolving the group's chain ----------------------------------------------------------------
@pytest.mark.asyncio
async def test_the_generic_star_chain_catches_a_group_with_no_entry(make_core: MakeCore, host: ScriptedHost) -> None:
core = build(make_core, fallbacks=[{"b": ["a"]}, {"*": ["c"]}])
host.answer("openai/a", rate_limit())
await complete(core)
assert host.deployments_called("acompletion") == ["a", "c"]
@pytest.mark.asyncio
@pytest.mark.parametrize(
"fallbacks",
[[{"*": ["c"]}, {"a": ["b"]}], [{"a": ["b"]}, {"*": ["c"]}]],
ids=["star-first", "star-last"],
)
async def test_a_group_specific_chain_wins_over_the_generic_star_chain(
make_core: MakeCore, host: ScriptedHost, fallbacks: list[object]
) -> None:
core = build(make_core, fallbacks=fallbacks)
host.answer("openai/a", rate_limit())
await complete(core)
assert host.deployments_called("acompletion") == ["a", "b"]
@pytest.mark.asyncio
@pytest.mark.parametrize(
("group", "fallback_key"),
[("openai/gpt-x", "gpt-x"), ("gpt-4o", "openai/gpt-4o")],
ids=["provider-stripped-key", "provider-prefixed-key"],
)
async def test_a_chain_keyed_with_or_without_the_provider_prefix_matches_the_group(
make_core: MakeCore, host: ScriptedHost, group: str, fallback_key: str
) -> None:
core = build(
make_core, model_list=[deployment(group, "p"), deployment("b", "b")], fallbacks=[{fallback_key: ["b"]}]
)
host.answer("openai/p", rate_limit())
await complete(core, model=group)
assert host.deployments_called("acompletion") == ["p", "b"]
@pytest.mark.asyncio
async def test_a_prefixed_key_does_not_match_a_group_whose_provider_cannot_be_inferred(
make_core: MakeCore, host: ScriptedHost
) -> None:
core = build(
make_core, model_list=[deployment("mygroup", "p"), deployment("b", "b")], fallbacks=[{"openai/mygroup": ["b"]}]
)
host.answer("openai/p", rate_limit())
with pytest.raises(litellm.RateLimitError):
await complete(core, model="mygroup")
assert host.deployments_called("acompletion") == ["p"]
@pytest.mark.asyncio
async def test_a_dict_target_names_its_group_under_model(make_core: MakeCore, host: ScriptedHost) -> None:
core = build(make_core, fallbacks=[{"a": [{"model": "c"}]}])
host.answer("openai/a", rate_limit())
await complete(core)
assert host.deployments_called("acompletion") == ["a", "c"]
# --- walking the chain ------------------------------------------------------------------------
@pytest.mark.asyncio
@pytest.mark.parametrize(
("failing", "expected"),
[(["a"], ["a", "b"]), (["a", "b"], ["a", "b", "c"])],
ids=["first-target-answers", "second-target-answers"],
)
async def test_targets_are_tried_in_listed_order_until_one_answers(
make_core: MakeCore, host: ScriptedHost, failing: list[str], expected: list[str]
) -> None:
core = build(make_core, fallbacks=[{"a": ["b", "c"]}])
for group in failing:
host.answer(f"openai/{group}", rate_limit())
response = await complete(core)
assert host.deployments_called("acompletion") == expected
assert headers(response)["x-litellm-model-group"] == expected[-1]
@pytest.mark.asyncio
async def test_a_failed_fallback_group_follows_its_own_chain(make_core: MakeCore, host: ScriptedHost) -> None:
core = build(make_core, fallbacks=[{"a": ["b"]}, {"b": ["c"]}])
host.answer("openai/a", rate_limit())
host.answer("openai/b", rate_limit())
response = await complete(core)
assert host.deployments_called("acompletion") == ["a", "b", "c"]
assert headers(response)["x-litellm-model-group"] == "c"
@pytest.mark.asyncio
@pytest.mark.parametrize(
("fallbacks", "expected"),
[
([{"a": ["b"]}, {"b": ["a", "c"]}, {"c": ["a", "b"]}], ["a", "b", "c"]),
([{"a": ["a", "b"]}], ["a", "b"]),
([{"a": ["b", "b", "c"]}], ["a", "b", "c"]),
],
ids=["cyclic-graph", "group-lists-itself", "target-listed-twice"],
)
async def test_no_group_is_tried_twice_in_one_request(
make_core: MakeCore, host: ScriptedHost, fallbacks: list[object], expected: list[str]
) -> None:
core = build(make_core, fallbacks=fallbacks)
host.answer("*", rate_limit())
with pytest.raises(litellm.RateLimitError):
await complete(core)
assert host.deployments_called("acompletion") == expected
@pytest.mark.asyncio
@pytest.mark.parametrize(
("constructor", "process", "expected"),
[
({"max_fallbacks": 0}, {}, ["a"]),
({"max_fallbacks": 1}, {}, ["a", "b"]),
({"max_fallbacks": 2}, {}, ["a", "b", "c"]),
({}, {"max_fallbacks": 2}, ["a", "b", "c"]),
({}, {"router_max_fallbacks": 2}, ["a", "b", "c"]),
({}, {}, ["a", "b", "c", "d"]),
({"max_fallbacks": 3}, {"max_fallbacks": 1}, ["a", "b", "c", "d"]),
({}, {"max_fallbacks": 1, "router_max_fallbacks": 3}, ["a", "b"]),
],
ids=[
"constructor-0",
"constructor-1",
"constructor-2",
"process-max_fallbacks-2",
"process-router_max_fallbacks-2",
"default",
"constructor-over-process",
"process-max_fallbacks-over-router_max_fallbacks",
],
)
async def test_max_fallbacks_bounds_the_depth_of_a_chained_walk(
make_core: MakeCore,
host: ScriptedHost,
constructor: dict[str, object],
process: dict[str, object],
expected: list[str],
) -> None:
if process:
host.set_dynamic(**process)
core = build(make_core, model_list=GROUPS_ABCD, fallbacks=CHAIN_ABCD, **constructor)
host.answer("*", rate_limit())
with pytest.raises(litellm.RateLimitError):
await complete(core)
assert host.deployments_called("acompletion") == expected
@pytest.mark.asyncio
async def test_max_fallbacks_applied_by_changeset_bounds_the_walk(make_core: MakeCore, host: ScriptedHost) -> None:
core = build(make_core, model_list=GROUPS_ABCD, fallbacks=CHAIN_ABCD)
core.apply(ChangeSet(assign={"max_fallbacks": 1}))
host.answer("*", rate_limit())
with pytest.raises(litellm.RateLimitError):
await complete(core)
assert host.deployments_called("acompletion") == ["a", "b"]
@pytest.mark.asyncio
async def test_a_fallback_group_with_no_deployments_is_passed_over(make_core: MakeCore, host: ScriptedHost) -> None:
core = build(make_core, fallbacks=[{"a": ["missing", "c"]}])
host.answer("openai/a", rate_limit())
response = await complete(core)
assert host.deployments_called("acompletion") == ["a", "c"]
assert headers(response)["x-litellm-model-group"] == "c"
@pytest.mark.asyncio
async def test_a_chain_of_only_empty_groups_raises_the_primary_error(make_core: MakeCore, host: ScriptedHost) -> None:
core = build(make_core, fallbacks=[{"a": ["missing"]}])
host.answer("openai/a", rate_limit())
with pytest.raises(litellm.RateLimitError):
await complete(core)
assert host.deployments_called("acompletion") == ["a"]
@pytest.mark.asyncio
async def test_a_fallback_group_runs_after_the_primary_group_spends_its_retries(
make_core: MakeCore, host: ScriptedHost
) -> None:
core = build(make_core, fallbacks=[{"a": ["b"]}], num_retries=2)
host.answer("openai/a", rate_limit())
await complete(core)
assert host.deployments_called("acompletion") == ["a", "a", "a", "b"]
# --- error-specific chains ---------------------------------------------------------------------
@pytest.mark.asyncio
@pytest.mark.parametrize(
("error", "field"),
[(context_window, "context_window_fallbacks"), (content_policy, "content_policy_fallbacks")],
ids=["context-window", "content-policy"],
)
@pytest.mark.parametrize("source", ["constructor", "process-setting", "request-kwarg"])
async def test_an_error_specific_chain_wins_over_the_general_chain(
make_core: MakeCore, host: ScriptedHost, error: Callable[[], Exception], field: str, source: str
) -> None:
chain = [{"a": ["c"]}]
if source == "process-setting":
host.set_dynamic(**{field: chain})
core = build(make_core, fallbacks=[{"a": ["b"]}], **({field: chain} if source == "constructor" else {}))
host.answer("openai/a", error())
response = await complete(core, **({field: chain} if source == "request-kwarg" else {}))
assert host.deployments_called("acompletion") == ["a", "c"]
assert headers(response)["x-litellm-model-group"] == "c"
@pytest.mark.asyncio
@pytest.mark.parametrize(
("error", "field", "raised"),
[
(context_window, "context_window_fallbacks", litellm.ContextWindowExceededError),
(content_policy, "content_policy_fallbacks", litellm.ContentPolicyViolationError),
],
ids=["context-window", "content-policy"],
)
async def test_an_error_specific_chain_without_the_group_raises_without_trying_the_general_chain(
make_core: MakeCore, host: ScriptedHost, error: Callable[[], Exception], field: str, raised: type[Exception]
) -> None:
core = build(make_core, fallbacks=[{"a": ["b"]}], **{field: [{"other": ["c"]}]})
host.answer("openai/a", error())
with pytest.raises(raised):
await complete(core)
assert host.deployments_called("acompletion") == ["a"]
@pytest.mark.asyncio
@pytest.mark.parametrize("error", [context_window, content_policy], ids=["context-window", "content-policy"])
async def test_without_an_error_specific_chain_the_general_chain_applies(
make_core: MakeCore, host: ScriptedHost, error: Callable[[], Exception]
) -> None:
core = build(make_core, fallbacks=[{"a": ["b"]}])
host.answer("openai/a", error())
await complete(core)
assert host.deployments_called("acompletion") == ["a", "b"]
@pytest.mark.asyncio
async def test_a_context_window_fallback_that_fails_otherwise_follows_the_general_chain(
make_core: MakeCore, host: ScriptedHost
) -> None:
core = build(make_core, context_window_fallbacks=[{"a": ["b"]}], fallbacks=[{"b": ["c"]}])
host.answer("openai/a", context_window())
host.answer("openai/b", rate_limit())
await complete(core)
assert host.deployments_called("acompletion") == ["a", "b", "c"]
# --- default fallbacks -------------------------------------------------------------------------
@pytest.mark.asyncio
@pytest.mark.parametrize("source", ["constructor", "process-setting"])
async def test_default_fallbacks_catch_any_group(make_core: MakeCore, host: ScriptedHost, source: str) -> None:
if source == "process-setting":
host.set_dynamic(default_fallbacks=["c"])
core = build(make_core, **({"default_fallbacks": ["c"]} if source == "constructor" else {}))
host.answer("openai/a", rate_limit())
await complete(core)
assert host.deployments_called("acompletion") == ["a", "c"]
@pytest.mark.asyncio
async def test_default_fallbacks_join_the_configured_chains_as_the_star_entry(
make_core: MakeCore, host: ScriptedHost
) -> None:
core = build(make_core, fallbacks=[{"b": ["a"]}], default_fallbacks=["c"])
host.answer("openai/a", rate_limit())
await complete(core)
assert core.settings().fallbacks == [{"b": ["a"]}, {"*": ["c"]}]
assert host.deployments_called("acompletion") == ["a", "c"]
def test_default_fallbacks_given_as_a_string_are_refused(make_core: MakeCore) -> None:
with pytest.raises(TypeError, match=r"default_fallbacks must be a list of model names, got 'c'"):
build(make_core, default_fallbacks="c")
# --- order tiers inside one group ----------------------------------------------------------------
@pytest.mark.asyncio
@pytest.mark.parametrize(
("failing", "expected", "model_group"),
[
(["a1"], ["a1", "a2"], "a"),
(["a1", "a2"], ["a1", "a2", "b"], "b"),
],
ids=["next-tier-answers", "external-fallback-after-every-tier"],
)
async def test_order_tiers_are_tried_before_the_external_fallbacks(
make_core: MakeCore, host: ScriptedHost, failing: list[str], expected: list[str], model_group: str
) -> None:
core = build(make_core, model_list=TIERED_A, fallbacks=[{"a": ["b"]}])
for deployment_id in failing:
host.answer(f"openai/{deployment_id}", rate_limit())
response = await complete(core)
assert host.deployments_called("acompletion") == expected
assert headers(response)["x-litellm-model-group"] == model_group
@pytest.mark.asyncio
async def test_order_tiers_are_walked_without_any_configured_fallbacks(make_core: MakeCore, host: ScriptedHost) -> None:
core = build(make_core, model_list=[*TIERED_A, deployment("a", "a3", order=3)])
host.answer("*", rate_limit())
with pytest.raises(litellm.RateLimitError):
await complete(core)
assert host.deployments_called("acompletion") == ["a1", "a2", "a3"]
@pytest.mark.asyncio
async def test_a_context_window_error_skips_the_order_tiers_for_its_own_chain(
make_core: MakeCore, host: ScriptedHost
) -> None:
core = build(make_core, model_list=TIERED_A, context_window_fallbacks=[{"a": ["b"]}])
host.answer("openai/a1", context_window())
await complete(core)
assert host.deployments_called("acompletion") == ["a1", "b"]
# --- when no fallback runs -----------------------------------------------------------------------
@pytest.mark.asyncio
async def test_disable_fallbacks_on_the_request_raises_the_primary_error(
make_core: MakeCore, host: ScriptedHost
) -> None:
core = build(make_core, fallbacks=[{"a": ["b"]}])
host.answer("openai/a", rate_limit())
with pytest.raises(litellm.RateLimitError):
await complete(core, disable_fallbacks=True)
assert host.deployments_called("acompletion") == ["a"]
@pytest.mark.asyncio
@pytest.mark.parametrize(
"error",
[
GuardrailRaisedException(guardrail_name="g", message="no"),
BlockedPiiEntityError(entity_type="EMAIL", guardrail_name="g"),
],
ids=["guardrail-raised", "blocked-pii"],
)
async def test_a_guardrail_intervention_never_falls_back(
make_core: MakeCore, host: ScriptedHost, error: Exception
) -> None:
core = build(make_core, fallbacks=[{"a": ["b"]}])
host.answer("openai/a", error)
with pytest.raises(type(error)) as raised:
await complete(core)
assert raised.value is host.raised[-1]
assert host.deployments_called("acompletion") == ["a"]
def unsupported_params() -> litellm.UnsupportedParamsError:
return litellm.UnsupportedParamsError(message="nope", llm_provider="openai", model="a")
@pytest.mark.asyncio
@pytest.mark.parametrize("error", [bad_request, unsupported_params], ids=["bad-request", "unsupported-params"])
async def test_a_bad_request_does_not_fall_back_when_fallback_on_unsupported_params_is_off(
make_core: MakeCore, host: ScriptedHost, error: Callable[[], Exception]
) -> None:
core = build(make_core, fallbacks=[{"a": ["b"]}], fallback_on_unsupported_params=False)
host.answer("openai/a", error())
with pytest.raises(type(error())):
await complete(core)
assert host.deployments_called("acompletion") == ["a"]
@pytest.mark.asyncio
@pytest.mark.parametrize(
("error", "fallback_on_unsupported_params"),
[
(bad_request, True),
(unsupported_params, True),
(context_window, False),
(content_policy, False),
(rate_limit, False),
(server_error, False),
],
ids=[
"bad-request-flag-on",
"unsupported-params-flag-on",
"context-window-flag-off",
"content-policy-flag-off",
"rate-limit-flag-off",
"server-error-flag-off",
],
)
async def test_errors_outside_the_unsupported_params_rule_fall_back(
make_core: MakeCore, host: ScriptedHost, error: Callable[[], Exception], fallback_on_unsupported_params: bool
) -> None:
core = build(make_core, fallbacks=[{"a": ["b"]}], fallback_on_unsupported_params=fallback_on_unsupported_params)
host.answer("openai/a", error())
await complete(core)
assert host.deployments_called("acompletion") == ["a", "b"]
# --- the error when every group fails --------------------------------------------------------------
@pytest.mark.asyncio
async def test_when_every_fallback_fails_the_primary_groups_error_is_raised(
make_core: MakeCore, host: ScriptedHost
) -> None:
core = build(make_core, fallbacks=[{"a": ["b", "c"]}], max_fallbacks=1)
host.answer("openai/a", rate_limit())
host.answer("openai/b", timeout())
host.answer("openai/c", server_error())
with pytest.raises(litellm.RateLimitError) as raised:
await complete(core)
assert str(raised.value) == "litellm.RateLimitError: slow down"
assert host.deployments_called("acompletion") == ["a", "b", "c"]
@pytest.mark.asyncio
async def test_without_router_debug_the_error_message_is_not_extended(make_core: MakeCore, host: ScriptedHost) -> None:
core = build(make_core, fallbacks=[{"a": ["b"]}])
host.answer("*", rate_limit())
with pytest.raises(litellm.RateLimitError) as raised:
await complete(core)
assert str(raised.value) == "litellm.RateLimitError: slow down"
@pytest.mark.asyncio
@pytest.mark.parametrize(
("arguments", "suffix"),
[
({}, "\n\nLiteLLM: model group 'a' failed with the error above. No fallback was attempted."),
(
{"fallbacks": [{"a": ["b", "c"]}]},
"\n\nLiteLLM: model group 'a' failed with the error above. "
"Fallback to b, c also failed: litellm.RateLimitError: slow down",
),
(
{"fallbacks": [{"a": ["b"]}, {"b": ["c"]}]},
"\n\nLiteLLM: model group 'a' failed with the error above. "
"Fallback to b also failed: litellm.RateLimitError: slow down",
),
(
{"fallbacks": [{"b": ["c"]}]},
"\n\nLiteLLM: model group 'a' failed with the error above and no fallback model group was found for it, "
"so the request was not retried on another model. Fallbacks are configured for: b. "
"Add a fallbacks entry for that model group (Router fallbacks or proxy router_settings.fallbacks) "
"to retry on another model.",
),
],
ids=["no-fallbacks", "fallback-list-failed", "chained-fallback-failed", "no-chain-for-group"],
)
async def test_router_debug_explains_the_fallback_outcome_once(
make_core: MakeCore, host: ScriptedHost, arguments: dict[str, object], suffix: str
) -> None:
core = build(make_core, **arguments)
show_router_debug(host)
host.answer("*", rate_limit())
with pytest.raises(litellm.RateLimitError) as raised:
await complete(core)
assert str(raised.value) == f"litellm.RateLimitError: slow down{suffix}"
@pytest.mark.asyncio
async def test_router_debug_names_the_first_fallback_groups_error(make_core: MakeCore, host: ScriptedHost) -> None:
core = build(make_core, fallbacks=[{"a": ["b"]}])
show_router_debug(host)
host.answer("openai/a", rate_limit())
host.answer("openai/b", server_error())
with pytest.raises(litellm.RateLimitError) as raised:
await complete(core)
assert str(raised.value).endswith(
"LiteLLM: model group 'a' failed with the error above. Fallback to b also failed: litellm.InternalServerError: boom"
)
@pytest.mark.asyncio
async def test_router_debug_reports_a_fallback_group_with_no_deployments(
make_core: MakeCore, host: ScriptedHost
) -> None:
core = build(make_core, fallbacks=[{"a": ["missing"]}])
show_router_debug(host)
host.answer("openai/a", rate_limit())
with pytest.raises(litellm.RateLimitError) as raised:
await complete(core)
assert (
"LiteLLM: model group 'a' failed with the error above. Fallback to missing also failed: "
"litellm.BadRequestError: You passed in model=missing. There are no healthy deployments for this model"
) in str(raised.value)
@pytest.mark.asyncio
async def test_router_debug_on_an_unmatched_context_window_chain_says_no_fallback_was_attempted(
make_core: MakeCore, host: ScriptedHost
) -> None:
core = build(make_core, context_window_fallbacks=[{"other": ["c"]}])
show_router_debug(host)
host.answer("openai/a", context_window())
with pytest.raises(litellm.ContextWindowExceededError) as raised:
await complete(core)
assert str(raised.value).endswith(
"\n\nLiteLLM: model group 'a' failed with the error above. No fallback was attempted."
)
@pytest.mark.asyncio
async def test_router_debug_on_content_policy_without_any_chain_hints_at_content_policy_fallbacks(
make_core: MakeCore, host: ScriptedHost
) -> None:
core = build(make_core)
show_router_debug(host)
host.answer("openai/a", content_policy())
with pytest.raises(litellm.ContentPolicyViolationError) as raised:
await complete(core)
message = str(raised.value)
assert "model=a. content_policy_fallback=None. fallbacks=None." in message
assert message.endswith("LiteLLM: model group 'a' failed with the error above. No fallback was attempted.")
# --- response headers ------------------------------------------------------------------------------
@pytest.mark.asyncio
@pytest.mark.parametrize(
("failing", "attempted", "model_group"),
[([], 0, "a"), (["a"], 1, "b")],
ids=["primary-answers", "one-fallback"],
)
async def test_success_headers_name_the_answering_group_and_the_fallbacks_attempted(
make_core: MakeCore, host: ScriptedHost, failing: list[str], attempted: int, model_group: str
) -> None:
core = build(make_core, fallbacks=[{"a": ["b"]}])
for group in failing:
host.answer(f"openai/{group}", rate_limit())
response = await complete(core)
assert headers(response)["x-litellm-attempted-fallbacks"] == attempted
assert headers(response)["x-litellm-model-group"] == model_group
@pytest.mark.asyncio
async def test_each_attempt_carries_its_own_model_group_to_the_host(make_core: MakeCore, host: ScriptedHost) -> None:
core = build(make_core, fallbacks=[{"a": ["b"]}])
host.answer("openai/a", rate_limit())
await core.call("acompletion", {"model": "a", "messages": MESSAGES})
assert [call.model_group for call in host.calls_to("acompletion")] == ["a", "b"]
assert [call.model for call in host.calls_to("acompletion")] == ["openai/a", "openai/b"]
@pytest.mark.asyncio
async def test_embedding_falls_back_like_completion(make_core: MakeCore, host: ScriptedHost) -> None:
core = build(make_core, model_list=[deployment("e", "e"), deployment("f", "f")], fallbacks=[{"e": ["f"]}])
host.answer("openai/e", rate_limit())
response = await core.aembedding(model="e", input=["x"]) # pyright: ignore[reportAttributeAccessIssue] # endpoint method
assert host.deployments_called("aembedding") == ["e", "f"]
assert headers(response)["x-litellm-attempted-fallbacks"] == 1
def test_sync_completion_falls_back_like_async(make_core: MakeCore, host: ScriptedHost) -> None:
core = build(make_core, fallbacks=[{"a": ["b"]}])
host.answer("openai/a", rate_limit())
response = core.completion(model="a", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # endpoint method
assert host.deployments_called("completion") == ["a", "b"]
assert headers(response)["x-litellm-attempted-fallbacks"] == 1

657
tests/unit/test_router/core/test_registry_changesets.py Normal file
View file

@ -0,0 +1,657 @@
"""
The core's registry and its settings, which are static after construction: ``apply(ChangeSet)``
(and ``aapply``) is the only way to change them, and a direct attribute write is routed into one.
"""
import copy
from collections.abc import Callable, Mapping
import pytest
import litellm
from litellm.router_api import ApplyResult, ChangeSet
from litellm.router_core import RouterCore
from litellm.types.router import Deployment
from tests.unit.test_router.core.scripted_host import (
ScriptedHost,
context_window,
deployment,
server_error,
)
MakeCore = Callable[..., RouterCore]
MESSAGES = [{"role": "user", "content": "hi"}]
def _model_list() -> list[dict[str, object]]: # mutable-ok: a fresh model_list
return [deployment("gpt", "a"), deployment("gpt", "b"), deployment("emb", "e")]
def _fresh(error: Callable[[], BaseException]) -> Callable[..., BaseException]:
"""A scripted outcome raising a new exception per call: the core stamps num_retries on the instance it saw."""
return lambda **_: error()
def _ids(deployments: object) -> list[str]: # mutable-ok: a fresh list
return [entry["model_info"]["id"] for entry in deployments] # pyright: ignore[reportGeneralTypeIssues, reportIndexIssue] # deployment dicts
def _deployment(model_name: str, deployment_id: str, **params: object) -> Deployment: # kwargs-ok: litellm_params
return Deployment(**deployment(model_name, deployment_id, **params)) # pyright: ignore[reportArgumentType] # a model_list entry
# -- registry lookups after construction --------------------------------------------------------
@pytest.mark.parametrize(("deployment_id", "group"), [("b", "gpt"), ("e", "emb")])
def test_registry_finds_each_configured_deployment_by_id(make_core: MakeCore, deployment_id: str, group: str) -> None:
registry = make_core(model_list=_model_list()).registry()
found = registry.deployment(deployment_id)
assert found is not None
assert (found.model_name, found.model_info.id, found.litellm_params.model) == (
group,
deployment_id,
f"openai/{deployment_id}",
)
assert registry.has_deployment(deployment_id)
def test_registry_has_no_deployment_for_an_unknown_id(make_core: MakeCore) -> None:
registry = make_core(model_list=_model_list()).registry()
assert registry.deployment("missing") is None
assert not registry.has_deployment("missing")
@pytest.mark.parametrize(("group", "ids"), [("gpt", ["a", "b"]), ("unknown", [])])
def test_registry_lists_a_groups_deployments_in_configured_order(
make_core: MakeCore, group: str, ids: list[str]
) -> None:
registry = make_core(model_list=_model_list()).registry()
assert _ids(registry.deployments(group)) == ids
assert registry.names(group, kind="ids") == ids
def test_registry_names_every_group_and_deployment_id(make_core: MakeCore) -> None:
registry = make_core(model_list=_model_list()).registry()
assert registry.names(kind="group_set") == {"gpt", "emb"}
assert registry.names(kind="ids") == ["a", "b", "e"]
assert _ids(registry.deployments()) == ["a", "b", "e"]
@pytest.mark.parametrize(
("name", "kind", "resolved"),
[
("fast", "group", "gpt"),
("fast", "alias", "gpt"),
("gpt", "alias", None),
("e", "from_id", "emb"),
("fast", "recognized", True),
("nope", "recognized", False),
],
)
def test_registry_resolves_names_through_groups_aliases_and_ids(
make_core: MakeCore, name: str, kind: str, resolved: object
) -> None:
registry = make_core(model_list=_model_list(), model_group_alias={"fast": "gpt"}).registry()
assert registry.resolve(name, kind=kind) == resolved # pyright: ignore[reportCallIssue, reportArgumentType] # kind is one of the literals
def test_registry_aliases_are_the_configured_model_group_alias(make_core: MakeCore) -> None:
registry = make_core(model_list=_model_list(), model_group_alias={"fast": "gpt"}).registry()
assert registry.aliases() == {"fast": "gpt"}
assert _ids(registry.deployments("fast")) == ["a", "b"]
def test_registry_version_is_unchanged_by_reads(make_core: MakeCore) -> None:
core = make_core(model_list=_model_list())
before = core.registry().version
core.registry().deployments("gpt")
core.registry().names()
core.registry().resolve("gpt")
core.settings()
assert core.registry().version == before
# -- registry().version on every write path -----------------------------------------------------
async def _apply_upsert_new(core: RouterCore) -> None:
core.apply(ChangeSet(upsert=[_deployment("gpt", "c")]))
async def _apply_upsert_changed(core: RouterCore) -> None:
core.apply(ChangeSet(upsert=[_deployment("gpt", "a", api_base="https://changed.example")]))
async def _apply_delete(core: RouterCore) -> None:
core.apply(ChangeSet(delete=["a"]))
async def _aapply_upsert_new(core: RouterCore) -> None:
await core.aapply(ChangeSet(upsert=[_deployment("gpt", "c")])) # pyright: ignore[reportAttributeAccessIssue] # aapply is on every core
async def _add_deployment(core: RouterCore) -> None:
core.add_deployment(_deployment("gpt", "c")) # pyright: ignore[reportAttributeAccessIssue] # public on the core
async def _set_model_list(core: RouterCore) -> None:
core.set_model_list([deployment("gpt", "c")]) # pyright: ignore[reportAttributeAccessIssue] # public on the core
async def _upsert_deployment(core: RouterCore) -> None:
with pytest.deprecated_call():
core.upsert_deployment(_deployment("gpt", "c")) # pyright: ignore[reportAttributeAccessIssue] # deprecated, still public
async def _delete_deployment(core: RouterCore) -> None:
with pytest.deprecated_call():
core.delete_deployment("a") # pyright: ignore[reportAttributeAccessIssue] # deprecated, still public
@pytest.mark.asyncio
@pytest.mark.parametrize(
("write", "ids_after"),
[
(_apply_upsert_new, ["a", "b", "e", "c"]),
(_apply_upsert_changed, ["a", "b", "e"]),
(_apply_delete, ["b", "e"]),
(_aapply_upsert_new, ["a", "b", "e", "c"]),
(_add_deployment, ["a", "b", "e", "c"]),
(_set_model_list, ["c"]),
(_upsert_deployment, ["a", "b", "e", "c"]),
(_delete_deployment, ["b", "e"]),
],
ids=lambda value: getattr(value, "__name__", None),
)
async def test_every_deployment_write_advances_the_registry_version(
make_core: MakeCore, write: Callable[[RouterCore], object], ids_after: list[str]
) -> None:
core = make_core(model_list=_model_list())
before = core.registry().version
await write(core) # pyright: ignore[reportGeneralTypeIssues] # each write is a coroutine function
assert core.registry().version > before
assert sorted(core.registry().names(kind="ids")) == sorted(ids_after)
@pytest.mark.asyncio
async def test_a_changed_upsert_is_what_the_registry_returns_afterwards(make_core: MakeCore) -> None:
core = make_core(model_list=_model_list())
await _apply_upsert_changed(core)
changed = core.registry().deployment("a")
assert changed is not None
assert changed.litellm_params.api_base == "https://changed.example"
@pytest.mark.parametrize(
"changes",
[
ChangeSet(),
ChangeSet(upsert=[_deployment("gpt", "a")]),
ChangeSet(delete=["missing"]),
ChangeSet(assign={"num_retries": 4, "cooldown_time": 9}),
ChangeSet(settings={"num_retries": 4, "allowed_fails": 2}),
],
ids=["empty", "identical_upsert", "unknown_delete", "assign", "settings"],
)
def test_a_changeset_that_changes_no_deployment_keeps_the_registry_version(
make_core: MakeCore, changes: ChangeSet
) -> None:
core = make_core(model_list=_model_list())
before = core.registry().version
result = core.apply(changes)
assert core.registry().version == before
assert result.version == before
def test_adding_a_deployment_whose_id_exists_changes_nothing(make_core: MakeCore) -> None:
core = make_core(model_list=_model_list())
before = core.registry().version
added = core.add_deployment(_deployment("emb", "a")) # pyright: ignore[reportAttributeAccessIssue] # public on the core
assert added is None
assert core.registry().version == before
assert core.registry().resolve("a", kind="from_id") == "gpt"
# -- ApplyResult --------------------------------------------------------------------------------
def test_apply_result_carries_the_new_version_and_the_upserted_deployment(make_core: MakeCore) -> None:
core = make_core(model_list=_model_list())
result = core.apply(ChangeSet(upsert=[_deployment("gpt", "c")]))
assert isinstance(result, ApplyResult)
assert result.version == core.registry().version
assert [None if d is None else d.model_info.id for d in result.upserted] == ["c"]
assert list(result.deleted) == []
def test_apply_result_reports_none_for_an_upsert_that_changed_nothing(make_core: MakeCore) -> None:
core = make_core(model_list=_model_list())
result = core.apply(ChangeSet(upsert=[_deployment("gpt", "a"), _deployment("gpt", "c")]))
assert [None if d is None else d.model_info.id for d in result.upserted] == [None, "c"]
def test_apply_result_lists_each_delete_in_order_with_none_for_an_unknown_id(make_core: MakeCore) -> None:
core = make_core(model_list=_model_list())
result = core.apply(ChangeSet(delete=["e", "missing", "a"]))
assert [None if d is None else d["model_info"]["id"] for d in result.deleted] == ["e", None, "a"] # pyright: ignore[reportIndexIssue] # the core returns the model_list dict
assert core.registry().names(kind="ids") == ["b"]
def test_apply_upserts_before_it_deletes(make_core: MakeCore) -> None:
core = make_core(model_list=_model_list())
result = core.apply(ChangeSet(upsert=[_deployment("gpt", "c")], delete=["c"]))
assert [d is not None for d in result.upserted] == [True]
assert [d is not None for d in result.deleted] == [True]
assert not core.registry().has_deployment("c")
def test_settings_and_deployments_change_together_in_one_apply(make_core: MakeCore) -> None:
core = make_core(model_list=_model_list())
result = core.apply(ChangeSet(assign={"cooldown_time": 17}, upsert=[_deployment("emb", "f")], delete=["b"]))
assert result.version == core.registry().version
assert core.settings().cooldown_time == 17
assert core.registry().names("emb", kind="ids") == ["e", "f"]
assert core.registry().names("gpt", kind="ids") == ["a"]
# -- settings changed by a changeset change later calls -----------------------------------------
def _assign(**values: object) -> ChangeSet: # kwargs-ok: settings by name
return ChangeSet(assign=values)
def _settings(**values: object) -> ChangeSet: # kwargs-ok: settings by name
return ChangeSet(settings=values)
CHANGESET_KINDS = pytest.mark.parametrize("changeset", [_assign, _settings], ids=["assign", "settings"])
@pytest.mark.asyncio
@CHANGESET_KINDS
@pytest.mark.parametrize("num_retries", [0, 3])
async def test_num_retries_from_a_changeset_governs_the_next_failing_call(
make_core: MakeCore, host: ScriptedHost, changeset: Callable[..., ChangeSet], num_retries: int
) -> None:
core = make_core(model_list=[deployment("gpt", "a")], num_retries=2)
host.answer("openai/a", _fresh(server_error))
with pytest.raises(litellm.InternalServerError):
await core.acompletion(model="gpt", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # an endpoint method
host.calls.clear()
core.apply(changeset(num_retries=num_retries))
with pytest.raises(litellm.InternalServerError):
await core.acompletion(model="gpt", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # an endpoint method
assert host.deployments_called("acompletion") == ["a"] * (num_retries + 1)
assert core.settings().as_dict()["num_retries"] == num_retries
@pytest.mark.asyncio
@CHANGESET_KINDS
async def test_fallbacks_from_a_changeset_are_used_by_the_next_call(
make_core: MakeCore, host: ScriptedHost, changeset: Callable[..., ChangeSet]
) -> None:
core = make_core(model_list=[deployment("gpt", "a"), deployment("backup", "f")], num_retries=0)
host.answer("openai/a", _fresh(server_error))
with pytest.raises(litellm.InternalServerError):
await core.acompletion(model="gpt", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # an endpoint method
host.calls.clear()
core.apply(changeset(fallbacks=[{"gpt": ["backup"]}]))
response = await core.acompletion(model="gpt", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # an endpoint method
assert host.deployments_called("acompletion") == ["a", "f"]
assert response._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 1
assert core.settings().fallbacks == [{"gpt": ["backup"]}]
@pytest.mark.asyncio
@CHANGESET_KINDS
async def test_context_window_fallbacks_from_a_changeset_are_used_by_the_next_call(
make_core: MakeCore, host: ScriptedHost, changeset: Callable[..., ChangeSet]
) -> None:
core = make_core(model_list=[deployment("gpt", "a"), deployment("long", "l")], num_retries=0)
host.answer("openai/a", _fresh(context_window))
with pytest.raises(litellm.ContextWindowExceededError):
await core.acompletion(model="gpt", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # an endpoint method
host.calls.clear()
core.apply(changeset(context_window_fallbacks=[{"gpt": ["long"]}]))
response = await core.acompletion(model="gpt", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # an endpoint method
assert host.deployments_called("acompletion") == ["a", "l"]
assert response._hidden_params["additional_headers"]["x-litellm-model-group"] == "long"
assert core.settings().context_window_fallbacks == [{"gpt": ["long"]}]
@pytest.mark.asyncio
@CHANGESET_KINDS
async def test_timeout_from_a_changeset_is_passed_on_the_next_call(
make_core: MakeCore, host: ScriptedHost, changeset: Callable[..., ChangeSet]
) -> None:
core = make_core(model_list=[deployment("gpt", "a")], timeout=30)
await core.acompletion(model="gpt", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # an endpoint method
core.apply(changeset(timeout=12))
await core.acompletion(model="gpt", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # an endpoint method
assert [call.kwargs["timeout"] for call in host.calls_to("acompletion")] == [30, 12]
assert core.settings().as_dict()["timeout"] == 12
@pytest.mark.asyncio
@CHANGESET_KINDS
async def test_model_group_alias_from_a_changeset_routes_the_next_call(
make_core: MakeCore, host: ScriptedHost, changeset: Callable[..., ChangeSet]
) -> None:
core = make_core(model_list=[deployment("gpt", "a"), deployment("emb", "e")], model_group_alias={"fast": "gpt"})
core.apply(changeset(model_group_alias={"fast": "emb"}))
await core.acompletion(model="fast", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # an endpoint method
assert host.deployments_called("acompletion") == ["e"]
assert core.registry().aliases() == {"fast": "emb"}
assert core.registry().resolve("fast") == "emb"
@pytest.mark.asyncio
@CHANGESET_KINDS
async def test_allowed_fails_and_cooldown_time_from_a_changeset_govern_the_next_cooldown(
make_core: MakeCore, changeset: Callable[..., ChangeSet]
) -> None:
core = make_core(model_list=[deployment("gpt", "a"), deployment("gpt", "b")])
assert core.health().record_failure("a", server_error(), 500) is False
core.apply(changeset(allowed_fails=0, cooldown_time=42))
assert core.health().record_failure("b", server_error(), 500) is True
assert [
(i, entry["cooldown_time"]) for i, entry in core.health().cooldowns.get_active_cooldowns(["a", "b"], None)
] == [("b", 42)]
assert (core.settings().allowed_fails, core.settings().cooldown_time) == (0, 42)
@pytest.mark.asyncio
async def test_a_deployment_cooled_down_under_applied_allowed_fails_is_not_picked_again(
make_core: MakeCore, host: ScriptedHost
) -> None:
core = make_core(model_list=[deployment("gpt", "a"), deployment("gpt", "b")], num_retries=0)
core.apply(ChangeSet(assign={"allowed_fails": 0, "cooldown_time": 60}))
host.answer("*", _fresh(server_error))
for _ in range(2):
with pytest.raises(litellm.InternalServerError):
await core.acompletion(model="gpt", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # an endpoint method
assert sorted(host.deployments_called("acompletion")) == ["a", "b"]
@pytest.mark.asyncio
@CHANGESET_KINDS
async def test_tag_filtering_switched_on_by_a_changeset_restricts_the_next_pick(
make_core: MakeCore, host: ScriptedHost, changeset: Callable[..., ChangeSet]
) -> None:
core = make_core(model_list=[deployment("gpt", "a", tags=["paid"]), deployment("gpt", "b", tags=["free"])])
tagged = {"metadata": {"tags": ["paid"]}}
assert sorted(_ids(await core.select("gpt", dict(tagged), kind="healthy"))) == ["a", "b"] # pyright: ignore[reportCallIssue] # select takes kind
core.apply(changeset(enable_tag_filtering=True))
assert _ids(await core.select("gpt", dict(tagged), kind="healthy")) == ["a"] # pyright: ignore[reportCallIssue] # select takes kind
for _ in range(3):
await core.acompletion(model="gpt", messages=MESSAGES, metadata={"tags": ["paid"]}) # pyright: ignore[reportAttributeAccessIssue] # an endpoint method
assert host.deployments_called("acompletion") == ["a", "a", "a"]
assert core.settings().enable_tag_filtering is True
@pytest.mark.asyncio
async def test_routing_strategy_from_settings_is_reported_and_still_routes(
make_core: MakeCore, host: ScriptedHost
) -> None:
core = make_core(model_list=[deployment("gpt", "a")])
strategy = "least-busy"
core.apply(ChangeSet(settings={"routing_strategy": strategy}))
await core.acompletion(model="gpt", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # an endpoint method
assert core.settings().routing_strategy == strategy
assert core.settings().as_dict()["routing_strategy"] == strategy
assert host.deployments_called("acompletion") == ["a"]
@pytest.mark.asyncio
async def test_default_litellm_params_from_a_changeset_are_sent_on_the_next_call(
make_core: MakeCore, host: ScriptedHost
) -> None:
core = make_core(model_list=[deployment("gpt", "a")], default_litellm_params={"temperature": 0.2})
core.apply(ChangeSet(assign={"default_litellm_params": {"temperature": 0.7}}))
await core.acompletion(model="gpt", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # an endpoint method
assert [call.kwargs.get("temperature") for call in host.calls_to("acompletion")] == [0.7]
@pytest.mark.parametrize(
("name", "value"),
[
("cooldown_time", 7.5),
("allowed_fails", 9),
("disable_cooldowns", True),
("content_policy_fallbacks", [{"gpt": ["emb"]}]),
],
)
def test_settings_reflect_an_assigned_field(make_core: MakeCore, name: str, value: object) -> None:
core = make_core(model_list=_model_list())
core.apply(ChangeSet(assign={name: value}))
assert getattr(core.settings(), name) == value
assert core.settings().is_set(name)
# -- what apply refuses -------------------------------------------------------------------------
@pytest.mark.parametrize("name", ["bogus", "model_list", "_core_num_retries"])
def test_assigning_an_unknown_setting_raises_and_applies_nothing(make_core: MakeCore, name: str) -> None:
core = make_core(model_list=_model_list(), num_retries=1)
settings_before = core.settings()
with pytest.raises(ValueError, match=name):
core.apply(ChangeSet(assign={"num_retries": 5, "cooldown_time": 3, name: 1}, upsert=[_deployment("gpt", "c")]))
assert core.settings().as_dict()["num_retries"] == 1
assert core.settings() is settings_before
assert not core.registry().has_deployment("c")
def test_apply_refuses_refresh_model_info_and_applies_nothing(make_core: MakeCore) -> None:
core = make_core(model_list=_model_list(), num_retries=1)
version_before = core.registry().version
with pytest.raises(TypeError, match="aapply"):
core.apply(ChangeSet(refresh_model_info=True, assign={"num_retries": 5}, delete=["a"]))
assert core.settings().as_dict()["num_retries"] == 1
assert core.registry().version == version_before
assert core.registry().has_deployment("a")
# -- the settings() snapshot --------------------------------------------------------------------
def test_settings_snapshot_is_reused_between_reads(make_core: MakeCore) -> None:
core = make_core(model_list=_model_list())
assert core.settings() is core.settings()
@CHANGESET_KINDS
def test_applying_a_setting_replaces_the_snapshot_and_leaves_the_old_one_as_it_was(
make_core: MakeCore, changeset: Callable[..., ChangeSet]
) -> None:
core = make_core(model_list=_model_list(), cooldown_time=5)
before = core.settings()
core.apply(changeset(cooldown_time=11))
after = core.settings()
assert after is not before
assert (before.cooldown_time, after.cooldown_time) == (5, 11)
assert core.settings() is after
@pytest.mark.parametrize(
"changes",
[ChangeSet(upsert=[_deployment("gpt", "c")]), ChangeSet(delete=["a"])],
ids=["upsert", "delete"],
)
def test_a_changeset_without_settings_keeps_the_snapshot(make_core: MakeCore, changes: ChangeSet) -> None:
core = make_core(model_list=_model_list())
before = core.settings()
core.apply(changes)
assert core.settings() is before
# -- a direct attribute write is a changeset ----------------------------------------------------
@pytest.mark.parametrize(
("name", "value"),
[
("num_retries", 4),
("timeout", 9.0),
("retry_after", 3),
("allowed_fails", 1),
("cooldown_time", 11.0),
("fallbacks", [{"gpt": ["emb"]}]),
("context_window_fallbacks", [{"gpt": ["emb"]}]),
("enable_tag_filtering", True),
("model_group_alias", {"fast": "emb"}),
],
)
def test_writing_a_setting_attribute_is_the_same_as_assigning_it_in_a_changeset(
make_core: MakeCore, name: str, value: object
) -> None:
written = make_core(model_list=_model_list())
applied = make_core(model_list=_model_list())
snapshot_before = written.settings()
version_before = written.registry().version
setattr(written, name, copy.deepcopy(value))
applied.apply(ChangeSet(assign={name: copy.deepcopy(value)}))
assert written.settings().as_dict() == applied.settings().as_dict()
assert written.registry().aliases() == applied.registry().aliases()
if written.settings().is_set(name):
assert getattr(written.settings(), name) == getattr(applied.settings(), name) == value
assert written.settings() is not snapshot_before
assert written.registry().version == version_before
# -- constructor arguments ----------------------------------------------------------------------
@pytest.mark.parametrize(
("name", "value"),
[
("model_list", _model_list()),
("fallbacks", [{"gpt": ["emb"]}]),
("model_group_alias", {"fast": "gpt"}),
("model_group_retry_policy", {"gpt": {"RateLimitErrorRetries": 1}}),
("allowed_fails_policy", {"RateLimitErrorAllowedFails": 2}),
],
)
def test_construction_leaves_the_arguments_it_was_given_unchanged(
make_core: MakeCore, name: str, value: object
) -> None:
given = copy.deepcopy(value)
arguments: dict[str, object] = {"model_list": _model_list(), name: given}
make_core(**arguments)
assert given == value
@pytest.mark.asyncio
async def test_mutating_the_model_list_after_construction_changes_nothing(
make_core: MakeCore, host: ScriptedHost
) -> None:
model_list = _model_list()
core = make_core(model_list=model_list)
version_before = core.registry().version
model_list.append(deployment("gpt", "late"))
model_list[2]["model_name"] = "gpt"
model_list[0]["litellm_params"]["model"] = "openai/elsewhere" # pyright: ignore[reportIndexIssue] # a model_list entry
model_list[0]["model_info"]["id"] = "renamed" # pyright: ignore[reportIndexIssue] # a model_list entry
await core.aembedding(model="emb", input=["x"]) # pyright: ignore[reportAttributeAccessIssue] # an endpoint method
assert core.registry().names(kind="ids") == ["a", "b", "e"]
assert core.registry().names(kind="group_set") == {"gpt", "emb"}
assert core.registry().version == version_before
assert [(call.model, call.deployment_id) for call in host.calls_to("aembedding")] == [("openai/e", "e")]
# -- retire --------------------------------------------------------------------------------------
def _registered(host: ScriptedHost) -> Mapping[str, int]:
return {name: len(callbacks) for name, callbacks in host.lists.items() if callbacks}
def test_retire_releases_every_callback_the_core_registered_on_its_host(
make_core: MakeCore, host: ScriptedHost
) -> None:
core = make_core(model_list=_model_list())
assert _registered(host) != {}
core.retire()
assert _registered(host) == {}
def test_retire_leaves_another_cores_callbacks_registered(make_core: MakeCore, host: ScriptedHost) -> None:
kept = make_core(model_list=_model_list())
retired = make_core(model_list=[deployment("other", "o")])
registered_by_both = _registered(host)
retired.retire()
assert _registered(host) != {}
assert sum(_registered(host).values()) < sum(registered_by_both.values())
kept.retire()
assert _registered(host) == {}

832
tests/unit/test_router/core/test_retries_cooldowns.py Normal file
View file

@ -0,0 +1,832 @@
"""
Retries and cooldowns of a router core, driven through its public surface on a ``ScriptedHost``.
Backoff values carry litellm's random jitter (up to ``JITTER`` seconds), so sleeps are asserted
inside ``[base, base + JITTER)`` rather than exactly.
"""
from collections.abc import Callable, Mapping
from typing import Final
import httpx
import pytest
import litellm
from litellm.constants import DEFAULT_COOLDOWN_TIME_SECONDS, INITIAL_RETRY_DELAY, JITTER
from litellm.router_api import ChangeSet
from litellm.router_core import RouterCore
from litellm.types.router import AllowedFailsPolicy, RetryPolicy
from tests.unit.test_router.core.scripted_host import (
Call,
ScriptedHost,
auth_error,
bad_request,
content_policy,
context_window,
deployment,
rate_limit,
server_error,
timeout,
)
MESSAGES: Final = [{"role": "user", "content": "hi"}]
MakeCore = Callable[..., RouterCore]
ErrorBuilder = Callable[[], Exception]
class LitellmParamsLoggingHost(ScriptedHost):
"""``ScriptedHost`` whose failure logging carries the call's ``cooldown_time`` in ``litellm_params``, as
litellm's own logging does (``get_litellm_params``)."""
def _logging_kwargs(
self, call: Call, error: BaseException | None, response: object
) -> dict[str, object]: # mutable-ok: litellm's logging kwargs
logged: Final = super()._logging_kwargs(call, error, response)
litellm_params: Final = logged["litellm_params"]
if isinstance(litellm_params, dict) and "cooldown_time" in call.kwargs:
litellm_params["cooldown_time"] = call.kwargs["cooldown_time"]
return logged
@pytest.fixture
def host() -> ScriptedHost:
return LitellmParamsLoggingHost()
def fresh(builder: ErrorBuilder) -> Callable[..., Exception]:
"""A scripted outcome that raises a new error per call, so stamps from one call never leak into the next."""
return lambda **_: builder()
def rate_limited(retry_after: str | None = None) -> litellm.RateLimitError:
headers: Final = {"retry-after": retry_after} if retry_after is not None else {}
return litellm.RateLimitError(
message="slow down", llm_provider="openai", model="gpt-4o", response=httpx.Response(429, headers=headers)
)
async def attempt(core: RouterCore, model: str = "g", **kwargs: object) -> object: # kwargs-ok: request kwargs
"""The response, or the error the core raised."""
try:
return await core.acompletion(model=model, messages=MESSAGES, **kwargs) # pyright: ignore[reportAttributeAccessIssue] # every core serves acompletion
except Exception as error:
return error
def cooled_down(core: RouterCore, ids: list[str]) -> dict[str, Mapping[str, object]]: # mutable-ok: a fresh dict
return dict(core.health().cooldowns.get_active_cooldowns(ids, None))
def headers_of(response: object) -> Mapping[str, object]:
hidden: Final = getattr(response, "_hidden_params", None) or {}
return hidden.get("additional_headers") or {}
def two_deployments(
**params: object,
) -> list[dict[str, object]]: # kwargs-ok: litellm_params # mutable-ok: a model_list
return [deployment("g", "a", **params), deployment("g", "b", **params)]
# --- how many retries ---------------------------------------------------------------------------
@pytest.mark.parametrize(
("constructor", "request_value", "process", "expected_retries"),
[
(None, None, None, 2),
(3, None, None, 3),
(None, None, 4, 4),
(1, None, 4, 1),
(3, 1, None, 1),
(3, 0, None, 0),
(None, 2, 5, 2),
(0, 2, None, 2),
],
)
@pytest.mark.asyncio
async def test_num_retries_comes_from_request_then_constructor_then_process_then_default(
make_core: MakeCore,
host: ScriptedHost,
constructor: int | None,
request_value: int | None,
process: int | None,
expected_retries: int,
) -> None:
if process is not None:
host.set_dynamic(num_retries=process)
arguments: Final = {"num_retries": constructor} if constructor is not None else {}
core: Final = make_core(model_list=[deployment("g", "a")], **arguments)
host.answer("openai/a", fresh(server_error))
request: Final = {"num_retries": request_value} if request_value is not None else {}
error: Final = await attempt(core, **request)
assert isinstance(error, litellm.InternalServerError)
assert host.deployments_called("acompletion") == ["a"] * (expected_retries + 1)
@pytest.mark.asyncio
async def test_process_num_retries_is_read_once_when_the_core_is_built(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(model_list=[deployment("g", "a")])
host.set_dynamic(num_retries=5)
host.answer("openai/a", fresh(server_error))
await attempt(core)
assert len(host.calls_to("acompletion")) == 3
@pytest.mark.asyncio
async def test_num_retries_applied_through_settings_changes_later_requests(
make_core: MakeCore, host: ScriptedHost
) -> None:
core: Final = make_core(model_list=[deployment("g", "a")], num_retries=3)
core.apply(ChangeSet(settings={"num_retries": 1}))
host.answer("openai/a", fresh(server_error))
await attempt(core)
assert len(host.calls_to("acompletion")) == 2
# --- which errors retry -------------------------------------------------------------------------
@pytest.mark.parametrize(
("builder", "expected_calls"),
[
(bad_request, 1),
(auth_error, 1),
(context_window, 1),
(content_policy, 1),
(rate_limit, 3),
(server_error, 3),
(timeout, 3),
],
)
@pytest.mark.asyncio
async def test_lone_deployment_retries_only_retryable_errors(
make_core: MakeCore, host: ScriptedHost, builder: ErrorBuilder, expected_calls: int
) -> None:
core: Final = make_core(model_list=[deployment("g", "a")], num_retries=2)
host.answer("openai/a", fresh(builder))
error: Final = await attempt(core)
assert type(error) is type(builder())
assert len(host.calls_to("acompletion")) == expected_calls
@pytest.mark.parametrize(
("builder", "expected_calls"),
[(bad_request, 1), (context_window, 1), (content_policy, 1), (server_error, 3), (timeout, 3)],
)
@pytest.mark.asyncio
async def test_two_deployment_group_retries_only_retryable_errors(
make_core: MakeCore, host: ScriptedHost, builder: ErrorBuilder, expected_calls: int
) -> None:
core: Final = make_core(model_list=two_deployments(), num_retries=2)
host.answer("*", fresh(builder))
error: Final = await attempt(core)
assert type(error) is type(builder())
assert len(host.calls_to("acompletion")) == expected_calls
@pytest.mark.asyncio
async def test_auth_error_retries_on_the_other_deployment_of_the_group(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(model_list=two_deployments(), num_retries=2)
host.answer("*", fresh(auth_error))
error: Final = await attempt(core)
assert isinstance(error, litellm.AuthenticationError)
assert sorted(host.deployments_called("acompletion")) == ["a", "b"]
@pytest.mark.asyncio
async def test_rate_limits_stop_retrying_once_every_deployment_is_cooling_down(
make_core: MakeCore, host: ScriptedHost
) -> None:
core: Final = make_core(model_list=two_deployments(), num_retries=5)
host.answer("*", fresh(rate_limit))
error: Final = await attempt(core)
assert isinstance(error, litellm.RateLimitError)
assert sorted(host.deployments_called("acompletion")) == ["a", "b"]
assert set(cooled_down(core, ["a", "b"])) == {"a", "b"}
@pytest.mark.parametrize("builder", [bad_request, auth_error, context_window, content_policy])
@pytest.mark.asyncio
async def test_an_error_raised_without_retrying_carries_no_retry_counts(
make_core: MakeCore, host: ScriptedHost, builder: ErrorBuilder
) -> None:
core: Final = make_core(model_list=[deployment("g", "a")], num_retries=2)
host.answer("openai/a", fresh(builder))
error: Final = await attempt(core)
assert getattr(error, "num_retries", None) is None
assert getattr(error, "max_retries", None) is None
@pytest.mark.parametrize("builder", [rate_limit, server_error, timeout])
@pytest.mark.parametrize("num_retries", [1, 3])
@pytest.mark.asyncio
async def test_exhausted_retries_stamp_num_retries_and_max_retries_on_the_final_error(
make_core: MakeCore, host: ScriptedHost, builder: ErrorBuilder, num_retries: int
) -> None:
core: Final = make_core(model_list=[deployment("g", "a")], num_retries=num_retries)
host.answer("openai/a", fresh(builder))
error: Final = await attempt(core)
assert getattr(error, "num_retries", None) == num_retries
assert getattr(error, "max_retries", None) == num_retries
assert f"LiteLLM Retried: {num_retries} times" in str(error)
@pytest.mark.asyncio
async def test_the_final_error_is_the_latest_one_not_the_first(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(model_list=[deployment("g", "a")], num_retries=2)
first: Final = server_error()
last: Final = timeout()
host.answer("openai/a", first, fresh(server_error), last)
error: Final = await attempt(core)
assert error is host.raised[-1]
assert type(error) is type(last)
# --- retry policies -----------------------------------------------------------------------------
@pytest.mark.parametrize(
("field", "builder", "retries"),
[
("BadRequestErrorRetries", bad_request, 2),
("BadRequestErrorRetries", context_window, 1),
("ContentPolicyViolationErrorRetries", content_policy, 1),
("RateLimitErrorRetries", rate_limit, 0),
("TimeoutErrorRetries", timeout, 1),
("InternalServerErrorRetries", server_error, 4),
("DefaultRetries", server_error, 1),
],
)
@pytest.mark.asyncio
async def test_retry_policy_sets_retries_per_error_type(
make_core: MakeCore, host: ScriptedHost, field: str, builder: ErrorBuilder, retries: int
) -> None:
core: Final = make_core(
model_list=[deployment("g", "a")], num_retries=3, retry_policy=RetryPolicy(**{field: retries})
)
host.answer("openai/a", fresh(builder))
error: Final = await attempt(core)
assert type(error) is type(builder())
assert len(host.calls_to("acompletion")) == retries + 1
@pytest.mark.asyncio
async def test_retry_policy_given_as_a_dict_works_like_a_retry_policy(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(
model_list=[deployment("g", "a")], num_retries=3, retry_policy={"BadRequestErrorRetries": 1}
)
host.answer("openai/a", fresh(bad_request))
await attempt(core)
assert len(host.calls_to("acompletion")) == 2
@pytest.mark.asyncio
async def test_retry_policy_for_one_error_type_leaves_others_to_num_retries(
make_core: MakeCore, host: ScriptedHost
) -> None:
core: Final = make_core(
model_list=[deployment("g", "a")], num_retries=1, retry_policy=RetryPolicy(RateLimitErrorRetries=4)
)
host.answer("openai/a", fresh(timeout))
await attempt(core)
assert len(host.calls_to("acompletion")) == 2
@pytest.mark.asyncio
async def test_model_group_retry_policy_applies_only_to_its_group(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(
model_list=[deployment("g", "a"), deployment("h", "c")],
num_retries=2,
model_group_retry_policy={"g": RetryPolicy(InternalServerErrorRetries=0)},
)
host.answer("*", fresh(server_error))
await attempt(core, model="g")
await attempt(core, model="h")
assert host.deployments_called("acompletion") == ["a", "c", "c", "c"]
@pytest.mark.asyncio
async def test_model_group_retry_policy_wins_over_retry_policy(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(
model_list=[deployment("g", "a")],
retry_policy=RetryPolicy(InternalServerErrorRetries=4),
model_group_retry_policy={"g": RetryPolicy(InternalServerErrorRetries=1)},
)
host.answer("openai/a", fresh(server_error))
await attempt(core)
assert len(host.calls_to("acompletion")) == 2
@pytest.mark.parametrize(("request_value", "expected_calls"), [(4, 2), (0, 1)])
@pytest.mark.asyncio
async def test_retry_policy_wins_over_request_num_retries_unless_the_request_opts_out(
make_core: MakeCore, host: ScriptedHost, request_value: int, expected_calls: int
) -> None:
core: Final = make_core(model_list=[deployment("g", "a")], retry_policy=RetryPolicy(InternalServerErrorRetries=1))
host.answer("openai/a", fresh(server_error))
await attempt(core, num_retries=request_value)
assert len(host.calls_to("acompletion")) == expected_calls
# --- per-deployment num_retries -----------------------------------------------------------------
@pytest.mark.parametrize(
("deployment_value", "request_kwargs", "expected_calls"),
[
(1, {}, 2),
("1", {}, 2),
(0, {}, 1),
(4, {}, 5),
(4, {"num_retries": 0}, 1),
(0, {"num_retries": 2}, 3),
],
)
@pytest.mark.asyncio
async def test_deployment_num_retries_beats_the_router_but_not_the_request(
make_core: MakeCore,
host: ScriptedHost,
deployment_value: object,
request_kwargs: dict[str, object],
expected_calls: int,
) -> None:
core: Final = make_core(model_list=[deployment("g", "a", num_retries=deployment_value)], num_retries=3)
host.answer("openai/a", fresh(server_error))
error: Final = await attempt(core, **request_kwargs)
assert isinstance(error, litellm.InternalServerError)
assert len(host.calls_to("acompletion")) == expected_calls
# --- backoff ------------------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_lone_deployment_backs_off_exponentially(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(model_list=[deployment("g", "a")], num_retries=3)
host.answer("openai/a", fresh(server_error))
await attempt(core)
bases: Final = [INITIAL_RETRY_DELAY * 2**n for n in range(3)]
assert len(host.sleeps) >= 3
assert all(base <= slept < base + JITTER for base, slept in zip(bases, host.sleeps))
@pytest.mark.asyncio
async def test_retry_after_setting_is_the_minimum_backoff(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(model_list=[deployment("g", "a")], num_retries=2, retry_after=3)
host.answer("openai/a", fresh(server_error))
await attempt(core)
assert len(host.sleeps) >= 2
assert all(3 <= slept < 3 + JITTER for slept in host.sleeps[:2])
@pytest.mark.asyncio
async def test_lone_deployment_waits_the_retry_after_header(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(model_list=[deployment("g", "a")], num_retries=2)
host.answer("openai/a", lambda **_: rate_limited(retry_after="7"))
await attempt(core)
assert len(host.sleeps) >= 2
assert all(7 <= slept < 7 + JITTER for slept in host.sleeps[:2])
@pytest.mark.asyncio
async def test_lone_deployment_ignores_a_retry_after_header_over_a_minute(
make_core: MakeCore, host: ScriptedHost
) -> None:
core: Final = make_core(model_list=[deployment("g", "a")], num_retries=2)
host.answer("openai/a", lambda **_: rate_limited(retry_after="120"))
await attempt(core)
assert INITIAL_RETRY_DELAY <= host.sleeps[0] < INITIAL_RETRY_DELAY + JITTER
@pytest.mark.parametrize("failure", [fresh(server_error), fresh(timeout), lambda **_: rate_limited(retry_after="7")])
@pytest.mark.asyncio
async def test_a_group_with_siblings_retries_immediately(
make_core: MakeCore, host: ScriptedHost, failure: Callable[..., Exception]
) -> None:
core: Final = make_core(model_list=two_deployments(), num_retries=2)
host.answer("*", failure, None)
response: Final = await attempt(core)
assert not isinstance(response, Exception)
assert host.sleeps == [0]
@pytest.mark.asyncio
async def test_no_sleep_when_the_error_is_not_retried(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(model_list=[deployment("g", "a")], num_retries=3)
host.answer("openai/a", fresh(bad_request))
await attempt(core)
assert host.sleeps == []
# --- retry headers on success -------------------------------------------------------------------
@pytest.mark.parametrize("failures_first", [1, 2, 3])
@pytest.mark.asyncio
async def test_success_after_retries_reports_attempted_and_max_retries(
make_core: MakeCore, host: ScriptedHost, failures_first: int
) -> None:
core: Final = make_core(model_list=[deployment("g", "a")], num_retries=3)
host.answer("openai/a", *[fresh(server_error)] * failures_first, None)
response: Final = await attempt(core)
headers: Final = headers_of(response)
assert headers["x-litellm-attempted-retries"] == failures_first
assert headers["x-litellm-max-retries"] == 3
assert headers["x-litellm-model-group"] == "g"
@pytest.mark.asyncio
async def test_first_try_success_reports_zero_retries_and_no_max(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(model_list=[deployment("g", "a")], num_retries=3)
response: Final = await attempt(core)
headers: Final = headers_of(response)
assert headers["x-litellm-attempted-retries"] == 0
assert "x-litellm-max-retries" not in headers
@pytest.mark.asyncio
async def test_success_headers_report_the_retry_policy_budget(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(
model_list=[deployment("g", "a")], num_retries=1, retry_policy=RetryPolicy(InternalServerErrorRetries=4)
)
host.answer("openai/a", fresh(server_error), fresh(server_error), None)
response: Final = await attempt(core)
assert headers_of(response)["x-litellm-attempted-retries"] == 2
assert headers_of(response)["x-litellm-max-retries"] == 4
# --- allowed fails and cooldowns ----------------------------------------------------------------
@pytest.mark.parametrize("allowed_fails", [0, 1, 2])
@pytest.mark.asyncio
async def test_a_deployment_cools_down_once_its_failures_pass_allowed_fails(
make_core: MakeCore, host: ScriptedHost, allowed_fails: int
) -> None:
core: Final = make_core(model_list=two_deployments(), num_retries=0, allowed_fails=allowed_fails)
host.answer("*", fresh(server_error))
for _ in range(2 * allowed_fails + 2):
await attempt(core)
called: Final = host.deployments_called("acompletion")
cooled: Final = cooled_down(core, ["a", "b"])
assert cooled
assert {d for d in ("a", "b") if called.count(d) > allowed_fails} == set(cooled)
@pytest.mark.asyncio
async def test_process_allowed_fails_is_the_default(make_core: MakeCore, host: ScriptedHost) -> None:
host.set_dynamic(allowed_fails=0)
core: Final = make_core(model_list=two_deployments(), num_retries=0)
host.answer("*", fresh(server_error))
await attempt(core)
assert core.settings().allowed_fails == 0
assert list(cooled_down(core, ["a", "b"])) == host.deployments_called("acompletion")
@pytest.mark.asyncio
async def test_a_cooled_deployment_is_skipped_by_later_picks(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(model_list=two_deployments(), num_retries=0, allowed_fails=0)
host.answer("*", fresh(server_error))
await attempt(core)
first: Final = host.deployments_called("acompletion")[0]
healthy: Final = await core.select("g", {}, kind="healthy") # pyright: ignore[reportAttributeAccessIssue] # every core serves select
await attempt(core)
assert [d["model_info"]["id"] for d in healthy] == [{"a": "b", "b": "a"}[first]]
assert host.deployments_called("acompletion")[1] != first
@pytest.mark.asyncio
async def test_a_group_with_every_deployment_cooling_down_raises_without_calling(
make_core: MakeCore, host: ScriptedHost
) -> None:
core: Final = make_core(model_list=two_deployments(), num_retries=0, allowed_fails=0)
host.answer("*", fresh(server_error))
await attempt(core)
await attempt(core)
error: Final = await attempt(core)
assert isinstance(error, litellm.RouterRateLimitError)
assert len(host.calls_to("acompletion")) == 2
@pytest.mark.asyncio
async def test_rate_limit_cools_down_a_deployment_with_siblings_at_once(
make_core: MakeCore, host: ScriptedHost
) -> None:
core: Final = make_core(model_list=two_deployments(), num_retries=0)
host.answer("*", fresh(rate_limit))
await attempt(core)
assert list(cooled_down(core, ["a", "b"])) == host.deployments_called("acompletion")
@pytest.mark.parametrize(
"policy",
[AllowedFailsPolicy(RateLimitErrorAllowedFails=2), {"RateLimitErrorAllowedFails": 2}],
ids=["AllowedFailsPolicy", "dict"],
)
@pytest.mark.asyncio
async def test_allowed_fails_policy_lets_rate_limits_through_until_its_count(
make_core: MakeCore, host: ScriptedHost, policy: object
) -> None:
core: Final = make_core(model_list=two_deployments(), num_retries=0, allowed_fails_policy=policy)
host.answer("*", fresh(rate_limit))
await attempt(core)
assert cooled_down(core, ["a", "b"]) == {}
@pytest.mark.asyncio
async def test_allowed_fails_policy_of_zero_cools_down_on_the_first_failure_of_that_type(
make_core: MakeCore, host: ScriptedHost
) -> None:
core: Final = make_core(
model_list=two_deployments(),
num_retries=0,
allowed_fails=5,
allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=0),
)
host.answer("*", fresh(server_error), fresh(timeout))
await attempt(core)
after_server_error: Final = cooled_down(core, ["a", "b"])
await attempt(core)
assert after_server_error == {}
assert list(cooled_down(core, ["a", "b"])) == host.deployments_called("acompletion")[1:]
@pytest.mark.asyncio
async def test_deployment_allowed_fails_policy_in_model_info_cools_down_on_a_bad_request(
make_core: MakeCore, host: ScriptedHost
) -> None:
core: Final = make_core(
model_list=two_deployments(model_info={"allowed_fails_policy": {"BadRequestErrorAllowedFails": 0}}),
num_retries=0,
)
host.answer("*", fresh(bad_request))
await attempt(core)
assert list(cooled_down(core, ["a", "b"])) == host.deployments_called("acompletion")
@pytest.mark.asyncio
async def test_bad_request_does_not_cool_down_without_a_deployment_policy(
make_core: MakeCore, host: ScriptedHost
) -> None:
core: Final = make_core(model_list=two_deployments(), num_retries=0, allowed_fails=0)
host.answer("*", fresh(bad_request))
await attempt(core)
assert cooled_down(core, ["a", "b"]) == {}
# --- cooldown time ------------------------------------------------------------------------------
@pytest.mark.parametrize(
("arguments", "params", "failure", "expected_seconds"),
[
({}, {}, fresh(rate_limit), DEFAULT_COOLDOWN_TIME_SECONDS),
({"cooldown_time": 42}, {}, fresh(rate_limit), 42),
({}, {"cooldown_time": 30}, fresh(rate_limit), 30),
({}, {"model_info": {"cooldown_time": 30}}, fresh(rate_limit), 30),
({}, {}, lambda **_: rate_limited(retry_after="9"), 9),
({"cooldown_time": 42}, {}, lambda **_: rate_limited(retry_after="9"), 9),
({}, {"cooldown_time": 30}, lambda **_: rate_limited(retry_after="9"), 30),
],
ids=[
"default",
"constructor",
"litellm_params",
"model_info",
"retry-after",
"header-over-router",
"deployment-over-header",
],
)
@pytest.mark.asyncio
async def test_cooldown_time_comes_from_deployment_then_retry_after_then_router(
make_core: MakeCore,
host: ScriptedHost,
arguments: dict[str, object],
params: dict[str, object],
failure: Callable[..., Exception],
expected_seconds: float,
) -> None:
core: Final = make_core(model_list=two_deployments(**params), num_retries=0, **arguments)
host.answer("*", failure)
await attempt(core)
(entry,) = cooled_down(core, ["a", "b"]).values()
assert entry["cooldown_time"] == expected_seconds
@pytest.mark.parametrize("change", [{"settings": {"cooldown_time": 20}}, {"assign": {"cooldown_time": 20}}])
@pytest.mark.asyncio
async def test_cooldown_time_applied_at_runtime_is_used_for_later_cooldowns(
make_core: MakeCore, host: ScriptedHost, change: dict[str, Mapping[str, object]]
) -> None:
core: Final = make_core(model_list=two_deployments(), num_retries=0)
core.apply(ChangeSet(**change)) # pyright: ignore[reportArgumentType] # ChangeSet's mapping fields
host.answer("*", fresh(rate_limit))
await attempt(core)
assert core.settings().cooldown_time == 20
(entry,) = cooled_down(core, ["a", "b"]).values()
assert entry["cooldown_time"] == 20
@pytest.mark.asyncio
async def test_a_deployment_cooldown_time_of_zero_skips_the_cooldown(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(model_list=two_deployments(model_info={"cooldown_time": 0}), num_retries=0)
host.answer("*", fresh(rate_limit))
await attempt(core)
assert cooled_down(core, ["a", "b"]) == {}
# --- disable_cooldowns --------------------------------------------------------------------------
@pytest.mark.parametrize("failure", [fresh(rate_limit), fresh(server_error)])
@pytest.mark.asyncio
async def test_disable_cooldowns_keeps_failing_deployments_routable(
make_core: MakeCore, host: ScriptedHost, failure: Callable[..., Exception]
) -> None:
core: Final = make_core(model_list=two_deployments(), num_retries=0, allowed_fails=0, disable_cooldowns=True)
host.answer("*", failure)
for _ in range(4):
await attempt(core)
assert cooled_down(core, ["a", "b"]) == {}
assert len(host.calls_to("acompletion")) == 4
@pytest.mark.asyncio
async def test_disable_cooldowns_assigned_at_runtime_stops_cooldowns(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(model_list=two_deployments(), num_retries=0)
core.apply(ChangeSet(assign={"disable_cooldowns": True}))
host.answer("*", fresh(rate_limit))
await attempt(core)
assert core.settings().disable_cooldowns is True
assert cooled_down(core, ["a", "b"]) == {}
# --- single-deployment groups -------------------------------------------------------------------
@pytest.mark.parametrize("arguments", [{}, {"allowed_fails": 0}], ids=["default", "allowed_fails=0"])
@pytest.mark.parametrize("builder", [rate_limit, server_error, timeout])
@pytest.mark.asyncio
async def test_a_single_deployment_group_is_not_cooled_down(
make_core: MakeCore, host: ScriptedHost, arguments: dict[str, object], builder: ErrorBuilder
) -> None:
core: Final = make_core(model_list=[deployment("g", "a")], num_retries=0, **arguments)
host.answer("openai/a", fresh(builder))
for _ in range(3):
await attempt(core)
assert cooled_down(core, ["a"]) == {}
assert host.deployments_called("acompletion") == ["a", "a", "a"]
@pytest.mark.asyncio
async def test_an_allowed_fails_policy_naming_the_error_cools_down_a_single_deployment(
make_core: MakeCore, host: ScriptedHost
) -> None:
core: Final = make_core(
model_list=[deployment("g", "a")],
num_retries=0,
allowed_fails_policy=AllowedFailsPolicy(InternalServerErrorAllowedFails=0),
)
host.answer("openai/a", fresh(server_error))
await attempt(core)
error: Final = await attempt(core)
assert set(cooled_down(core, ["a"])) == {"a"}
assert isinstance(error, litellm.RouterRateLimitError)
assert host.deployments_called("acompletion") == ["a"]
# --- health() -----------------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_health_cooldown_entry_describes_the_failure(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(model_list=two_deployments(), num_retries=0)
host.answer("*", fresh(rate_limit))
await attempt(core)
(entry,) = cooled_down(core, ["a", "b"]).values()
assert entry["status_code"] == "429"
assert "slow down" in str(entry["exception_received"])
@pytest.mark.parametrize("failures", [1, 3])
@pytest.mark.asyncio
async def test_health_counts_failures_this_minute(make_core: MakeCore, host: ScriptedHost, failures: int) -> None:
core: Final = make_core(model_list=[deployment("g", "a")], num_retries=0)
host.answer("openai/a", fresh(server_error))
for _ in range(failures):
await attempt(core)
assert core.health().failures_this_minute("a") == failures
assert core.health().failures_this_minute("missing") == 0
@pytest.mark.asyncio
async def test_health_counts_every_retried_attempt_as_a_failure(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(model_list=[deployment("g", "a")], num_retries=2)
host.answer("openai/a", fresh(server_error))
await attempt(core)
assert core.health().failures_this_minute("a") == 3
@pytest.mark.asyncio
async def test_health_record_failure_cools_down_past_allowed_fails(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(model_list=two_deployments(), allowed_fails=1)
first: Final = core.health().record_failure("a", server_error(), 500)
second: Final = core.health().record_failure("a", server_error(), 500)
assert (first, second) == (False, True)
assert set(cooled_down(core, ["a", "b"])) == {"a"}
assert core.health().failures_this_minute("a") == 2

624
tests/unit/test_router/core/test_selection.py Normal file
View file

@ -0,0 +1,624 @@
"""
Deployment selection on a router core: which deployment ``select`` picks and which one an endpoint
call reaches, across groups, aliases, wildcard routes, deployment ids, teams, weights, order,
routing strategies and cooldowns.
"""
import re
from collections.abc import Callable, Mapping
from typing import Final
import openai
import pytest
import litellm
from litellm.router_core import RouterCore
from litellm.types.router import RouterRateLimitError
from tests.unit.test_router.core.scripted_host import (
ScriptedHost,
auth_error,
bad_request,
context_window,
deployment,
rate_limit,
server_error,
timeout,
)
MESSAGES: Final = [{"role": "user", "content": "hello there, router"}]
PICKS: Final = 20
MakeCore = Callable[..., RouterCore]
async def _picked_ids(
core: RouterCore, model: str, times: int = PICKS, **select_arguments: object
) -> set[str]: # kwargs-ok: select's keyword arguments
picked: Final[set[str]] = set()
for _ in range(times):
chosen = await core.select(model, **select_arguments) # pyright: ignore[reportCallIssue] # PythonCore's select takes keyword options
assert isinstance(chosen, Mapping)
picked.add(str(chosen["model_info"]["id"]))
return picked
async def _fail_through(core: RouterCore, deployment_id: str, times: int = 1) -> None:
for _ in range(times):
with pytest.raises(openai.APIError):
await core.acompletion(model=deployment_id, messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # endpoint method
# --- picking within a model group ----------------------------------------------------------
@pytest.mark.asyncio
@pytest.mark.parametrize(("model", "expected"), [("g1", "a"), ("g2", "b"), ("g3", "c")])
async def test_select_returns_the_only_deployment_of_the_requested_group(
make_core: MakeCore, model: str, expected: str
) -> None:
core: Final = make_core(model_list=[deployment("g1", "a"), deployment("g2", "b"), deployment("g3", "c")])
chosen: Final = await core.select(model)
assert chosen["model_info"]["id"] == expected # pyright: ignore[reportIndexIssue]
assert chosen["model_name"] == model # pyright: ignore[reportIndexIssue]
assert chosen["litellm_params"]["model"] == f"openai/{expected}" # pyright: ignore[reportIndexIssue]
@pytest.mark.asyncio
async def test_select_only_picks_deployments_of_the_requested_group(make_core: MakeCore) -> None:
core: Final = make_core(
model_list=[deployment("g", "a"), deployment("g", "b"), deployment("other", "c"), deployment("other", "d")]
)
assert await _picked_ids(core, "g") <= {"a", "b"}
assert await _picked_ids(core, "other") <= {"c", "d"}
@pytest.mark.asyncio
async def test_select_does_not_call_the_provider(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(model_list=[deployment("g", "a")])
await core.select("g")
assert host.calls == []
@pytest.mark.asyncio
@pytest.mark.parametrize(
("op", "request_payload"),
[
("acompletion", {"messages": MESSAGES}),
("atext_completion", {"prompt": "hi"}),
("aembedding", {"input": "hi"}),
],
)
async def test_async_endpoint_calls_the_deployment_of_the_requested_group(
make_core: MakeCore, host: ScriptedHost, op: str, request_payload: dict[str, object]
) -> None:
core: Final = make_core(model_list=[deployment("g1", "a"), deployment("g2", "b")])
await core.call(op, {"model": "g2", **request_payload})
assert host.deployments_called(op) == ["b"]
assert host.calls[-1].model == "openai/b"
@pytest.mark.parametrize(
("op", "request_payload"), [("completion", {"messages": MESSAGES}), ("embedding", {"input": "hi"})]
)
def test_sync_endpoint_calls_the_deployment_of_the_requested_group(
make_core: MakeCore, host: ScriptedHost, op: str, request_payload: dict[str, object]
) -> None:
core: Final = make_core(model_list=[deployment("g1", "a"), deployment("g2", "b")])
getattr(core, op)(model="g1", **request_payload)
assert host.deployments_called(op) == ["a"]
@pytest.mark.asyncio
async def test_completion_reports_the_requested_group_in_its_headers(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(model_list=[deployment("g1", "a"), deployment("g2", "b")])
response: Final = await core.acompletion(model="g2", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue]
assert response._hidden_params["additional_headers"]["x-litellm-model-group"] == "g2"
assert host.calls[-1].model_group == "g2"
# --- model_group_alias ---------------------------------------------------------------------
ALIASES: Final = {
"plain": "g",
"visible": {"model": "g", "hidden": False},
"hidden": {"model": "g", "hidden": True},
}
@pytest.mark.asyncio
@pytest.mark.parametrize("alias", sorted(ALIASES))
async def test_select_resolves_an_alias_to_its_target_group(make_core: MakeCore, alias: str) -> None:
core: Final = make_core(
model_list=[deployment("g", "a"), deployment("other", "b")], model_group_alias=dict(ALIASES)
)
chosen: Final = await core.select(alias)
assert chosen["model_info"]["id"] == "a" # pyright: ignore[reportIndexIssue]
assert chosen["model_name"] == "g" # pyright: ignore[reportIndexIssue]
@pytest.mark.asyncio
@pytest.mark.parametrize("alias", sorted(ALIASES))
async def test_completion_through_an_alias_calls_the_target_group(
make_core: MakeCore, host: ScriptedHost, alias: str
) -> None:
core: Final = make_core(
model_list=[deployment("g", "a"), deployment("other", "b")], model_group_alias=dict(ALIASES)
)
await core.acompletion(model=alias, messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue]
assert host.deployments_called("acompletion") == ["a"]
@pytest.mark.parametrize(("alias", "target"), [("plain", "g"), ("visible", "g"), ("hidden", "g"), ("g", None)])
def test_registry_resolves_alias_targets(make_core: MakeCore, alias: str, target: str | None) -> None:
core: Final = make_core(model_list=[deployment("g", "a")], model_group_alias=dict(ALIASES))
assert core.registry().resolve(alias, kind="alias") == target
def test_hidden_alias_is_left_out_of_the_routable_names(make_core: MakeCore) -> None:
core: Final = make_core(model_list=[deployment("g", "a")], model_group_alias=dict(ALIASES))
assert sorted(core.registry().names(kind="groups")) == ["g", "plain", "visible"]
@pytest.mark.asyncio
async def test_alias_to_an_unknown_group_fails_naming_the_target(make_core: MakeCore) -> None:
core: Final = make_core(model_list=[deployment("g", "a")], model_group_alias={"dangling": "missing"})
with pytest.raises(
litellm.BadRequestError, match=re.escape("You passed in model=missing. There are no healthy deployments")
):
await core.select("dangling")
# --- wildcard / pattern groups ---------------------------------------------------------------
WILDCARD_MODEL_LIST: Final = [
deployment("*", "catch_all", model="openai/*"),
deployment("anthropic/*", "anthropic_any", model="anthropic/*"),
deployment("openai/gpt-4o", "exact"),
deployment("openai/gpt-*", "gpt_any", model="openai/gpt-*"),
]
@pytest.mark.asyncio
@pytest.mark.parametrize(
("requested", "expected_id", "expected_group", "expected_model"),
[
("openai/gpt-4o", "exact", "openai/gpt-4o", "openai/exact"),
("openai/gpt-4o-mini", "gpt_any", "openai/gpt-*", "openai/gpt-4o-mini"),
("anthropic/claude-x", "anthropic_any", "anthropic/*", "anthropic/claude-x"),
("gpt-5", "catch_all", "*", "openai/gpt-5"),
("whatever/foo", "catch_all", "*", "openai/whatever/foo"),
],
)
async def test_select_prefers_an_exact_group_then_the_most_specific_pattern(
make_core: MakeCore, requested: str, expected_id: str, expected_group: str, expected_model: str
) -> None:
core: Final = make_core(model_list=list(WILDCARD_MODEL_LIST))
chosen: Final = await core.select(requested)
assert chosen["model_info"]["id"] == expected_id # pyright: ignore[reportIndexIssue]
assert chosen["model_name"] == expected_group # pyright: ignore[reportIndexIssue]
assert chosen["litellm_params"]["model"] == expected_model # pyright: ignore[reportIndexIssue]
@pytest.mark.asyncio
async def test_completion_through_a_wildcard_calls_the_requested_model(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(model_list=[deployment("openai/*", "wild", model="openai/*")])
response: Final = await core.acompletion(model="openai/gpt-x", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue]
assert host.deployments_called("acompletion") == ["wild"]
assert host.calls[-1].model == "openai/gpt-x"
assert response._hidden_params["additional_headers"]["x-litellm-model-group"] == "openai/gpt-x"
@pytest.mark.asyncio
async def test_a_pattern_that_does_not_match_is_an_unknown_model(make_core: MakeCore) -> None:
core: Final = make_core(model_list=[deployment("anthropic/*", "anthropic_any", model="anthropic/*")])
with pytest.raises(litellm.BadRequestError, match="model=openai/gpt-4o"):
await core.select("openai/gpt-4o")
# --- selecting by deployment id or litellm model --------------------------------------------
@pytest.mark.asyncio
@pytest.mark.parametrize("deployment_id", ["a", "b", "c"])
async def test_select_by_deployment_id_returns_that_deployment(make_core: MakeCore, deployment_id: str) -> None:
core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b"), deployment("h", "c")])
assert await _picked_ids(core, deployment_id) == {deployment_id}
@pytest.mark.asyncio
@pytest.mark.parametrize("deployment_id", ["a", "b"])
async def test_completion_by_deployment_id_calls_that_deployment(
make_core: MakeCore, host: ScriptedHost, deployment_id: str
) -> None:
core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")])
for _ in range(5):
await core.acompletion(model=deployment_id, messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue]
assert host.deployments_called("acompletion") == [deployment_id] * 5
@pytest.mark.asyncio
async def test_select_by_litellm_model_name_finds_the_deployment(make_core: MakeCore) -> None:
core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")])
assert await _picked_ids(core, "openai/b") == {"b"}
@pytest.mark.asyncio
async def test_specific_deployment_picks_by_litellm_model_across_groups(make_core: MakeCore) -> None:
core: Final = make_core(
model_list=[
deployment("g1", "c", model="openai/shared"),
deployment("g2", "d", model="openai/shared"),
deployment("g3", "e"),
]
)
assert await _picked_ids(core, "openai/shared", specific_deployment=True) <= {"c", "d"}
assert await _picked_ids(core, "openai/e", specific_deployment=True) == {"e"}
@pytest.mark.asyncio
async def test_completion_with_specific_deployment_calls_the_named_litellm_model(
make_core: MakeCore, host: ScriptedHost
) -> None:
core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")])
for _ in range(5):
await core.acompletion(model="openai/b", messages=MESSAGES, specific_deployment=True) # pyright: ignore[reportAttributeAccessIssue]
assert host.deployments_called("acompletion") == ["b"] * 5
# --- unknown models ---------------------------------------------------------------------------
@pytest.mark.asyncio
@pytest.mark.parametrize(
("op", "request_payload"),
[
("acompletion", {"messages": MESSAGES}),
("aembedding", {"input": "hi"}),
("atext_completion", {"prompt": "hi"}),
("completion", {"messages": MESSAGES}),
("embedding", {"input": "hi"}),
],
)
async def test_unknown_model_raises_bad_request_naming_the_model(
make_core: MakeCore, host: ScriptedHost, op: str, request_payload: dict[str, object]
) -> None:
core: Final = make_core(model_list=[deployment("g", "a")])
with pytest.raises(litellm.BadRequestError) as raised:
await core.call(op, {"model": "nope", **request_payload})
assert (
str(raised.value)
== "litellm.BadRequestError: You passed in model=nope. There are no healthy deployments for this model"
)
assert raised.value.status_code == 400
assert host.calls == []
@pytest.mark.asyncio
async def test_select_of_an_unknown_model_raises_bad_request(make_core: MakeCore) -> None:
core: Final = make_core(model_list=[deployment("g", "a")])
with pytest.raises(litellm.BadRequestError) as raised:
await core.select("nope")
assert raised.value.status_code == 400
assert raised.value.model == "nope"
assert "There are no healthy deployments for this model" in str(raised.value)
# --- team-scoped deployments --------------------------------------------------------------------
TEAM_MODEL_LIST: Final = [
deployment("g", "public"),
deployment("team-x-group", "team_x", model_info={"team_id": "tx", "team_public_model_name": "g"}),
deployment("team-y-group", "team_y", model_info={"team_id": "ty", "team_public_model_name": "g"}),
]
@pytest.mark.asyncio
@pytest.mark.parametrize(
("metadata", "expected"),
[
({}, "public"),
({"user_api_key_team_id": "tx"}, "team_x"),
({"user_api_key_team_id": "ty"}, "team_y"),
({"user_api_key_team_id": "unknown-team"}, "public"),
],
)
async def test_select_serves_a_team_its_own_deployment_under_the_public_name(
make_core: MakeCore, metadata: dict[str, str], expected: str
) -> None:
core: Final = make_core(model_list=list(TEAM_MODEL_LIST))
assert await _picked_ids(core, "g", request_kwargs={"metadata": metadata}) == {expected}
@pytest.mark.asyncio
@pytest.mark.parametrize(("team_id", "expected"), [("tx", "team_x"), ("ty", "team_y"), ("tz", "public")])
async def test_completion_for_a_team_calls_its_deployment_under_the_public_group(
make_core: MakeCore, host: ScriptedHost, team_id: str, expected: str
) -> None:
core: Final = make_core(model_list=list(TEAM_MODEL_LIST))
await core.acompletion(model="g", messages=MESSAGES, metadata={"user_api_key_team_id": team_id}) # pyright: ignore[reportAttributeAccessIssue]
assert host.deployments_called("acompletion") == [expected]
assert host.calls[-1].model_group == "g"
@pytest.mark.asyncio
@pytest.mark.parametrize("metadata", [{}, {"user_api_key_team_id": "other"}])
async def test_team_only_public_name_is_unknown_to_other_callers(make_core: MakeCore, metadata: dict[str, str]) -> None:
core: Final = make_core(
model_list=[deployment("team-x-group", "team_x", model_info={"team_id": "tx", "team_public_model_name": "pub"})]
)
with pytest.raises(litellm.BadRequestError, match="model=pub"):
await core.select("pub", {"metadata": metadata})
assert await _picked_ids(core, "pub", request_kwargs={"metadata": {"user_api_key_team_id": "tx"}}) == {"team_x"}
@pytest.mark.parametrize(("team_id", "names"), [(None, []), ("tx", ["pub"])])
def test_team_public_names_are_routable_only_for_their_team(
make_core: MakeCore, team_id: str | None, names: list[str]
) -> None:
core: Final = make_core(
model_list=[deployment("team-x-group", "team_x", model_info={"team_id": "tx", "team_public_model_name": "pub"})]
)
assert list(core.registry().names(kind="groups", team_id=team_id)) == names
# --- weights, rpm and tpm under simple-shuffle -----------------------------------------------------
@pytest.mark.asyncio
@pytest.mark.parametrize("metric", ["weight", "rpm", "tpm"])
async def test_simple_shuffle_never_picks_a_zero_weighted_deployment_beside_a_positive_one(
make_core: MakeCore, metric: str
) -> None:
core: Final = make_core(
model_list=[
deployment("g", "unset"),
deployment("g", "positive", **{metric: 5}),
deployment("g", "zero", **{metric: 0}),
]
)
assert await _picked_ids(core, "g") == {"positive"}
@pytest.mark.asyncio
async def test_weight_takes_precedence_over_rpm_and_tpm(make_core: MakeCore) -> None:
core: Final = make_core(
model_list=[
deployment("g", "weighted", weight=1, rpm=0, tpm=0),
deployment("g", "by_limits", weight=0, rpm=100, tpm=100),
]
)
assert await _picked_ids(core, "g") == {"weighted"}
@pytest.mark.asyncio
async def test_completion_calls_only_the_weighted_deployment(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(model_list=[deployment("g", "a", weight=0), deployment("g", "b", weight=3)])
for _ in range(10):
await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue]
assert set(host.deployments_called("acompletion")) == {"b"}
@pytest.mark.asyncio
@pytest.mark.parametrize("favoured", ["a", "b"])
async def test_request_router_weights_override_deployment_weights(make_core: MakeCore, favoured: str) -> None:
core: Final = make_core(model_list=[deployment("g", "a", weight=1), deployment("g", "b", weight=1)])
picked: Final = await _picked_ids(core, "g", request_kwargs={"_router_weights": {"g": {favoured: 1.0}}})
assert picked == {favoured}
# --- order --------------------------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_select_picks_only_the_lowest_order(make_core: MakeCore) -> None:
core: Final = make_core(
model_list=[
deployment("g", "second", order=2),
deployment("g", "first_a", order=1),
deployment("g", "first_b", order=1),
deployment("g", "third", order=3),
]
)
assert await _picked_ids(core, "g") <= {"first_a", "first_b"}
@pytest.mark.asyncio
async def test_completion_calls_only_the_lowest_order(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(model_list=[deployment("g", "second", order=2), deployment("g", "first", order=1)])
for _ in range(5):
await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue]
assert host.deployments_called("acompletion") == ["first"] * 5
@pytest.mark.asyncio
async def test_next_order_serves_while_the_lowest_order_cools_down(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(
model_list=[deployment("g", "second", order=2), deployment("g", "first", order=1)],
allowed_fails=0,
num_retries=0,
cooldown_time=60,
)
host.answer("openai/first", rate_limit())
await _fail_through(core, "first")
assert await _picked_ids(core, "g") == {"second"}
# --- routing strategies ----------------------------------------------------------------------------
@pytest.mark.asyncio
@pytest.mark.parametrize("routing_strategy", ["least-busy", "usage-based-routing"])
@pytest.mark.parametrize("ids", [("a", "b"), ("b", "a")])
async def test_strategy_without_traffic_picks_the_first_configured_deployment(
make_core: MakeCore, routing_strategy: str, ids: tuple[str, str]
) -> None:
core: Final = make_core(
model_list=[deployment("g", deployment_id) for deployment_id in ids], routing_strategy=routing_strategy
)
assert await _picked_ids(core, "g", messages=MESSAGES) == {ids[0]}
@pytest.mark.asyncio
@pytest.mark.parametrize("cheap_first", [True, False])
async def test_cost_based_routing_picks_the_cheapest_deployment(make_core: MakeCore, cheap_first: bool) -> None:
cheap: Final = deployment("g", "cheap", input_cost_per_token=0.000001, output_cost_per_token=0.000002)
pricey: Final = deployment("g", "pricey", input_cost_per_token=0.00001, output_cost_per_token=0.00002)
core: Final = make_core(
model_list=[cheap, pricey] if cheap_first else [pricey, cheap], routing_strategy="cost-based-routing"
)
assert await _picked_ids(core, "g", messages=MESSAGES) == {"cheap"}
@pytest.mark.asyncio
@pytest.mark.parametrize("ids", [("tight", "roomy"), ("roomy", "tight")])
async def test_usage_based_v2_skips_a_deployment_whose_tpm_the_request_exceeds(
make_core: MakeCore, ids: tuple[str, str]
) -> None:
limits: Final = {"tight": 1, "roomy": 100_000}
core: Final = make_core(
model_list=[deployment("g", deployment_id, tpm=limits[deployment_id]) for deployment_id in ids],
routing_strategy="usage-based-routing-v2",
)
assert await _picked_ids(core, "g", messages=MESSAGES) == {"roomy"}
@pytest.mark.asyncio
async def test_usage_based_v2_rate_limits_when_the_request_exceeds_every_tpm(make_core: MakeCore) -> None:
core: Final = make_core(
model_list=[deployment("g", "a", tpm=1), deployment("g", "b", tpm=1)],
routing_strategy="usage-based-routing-v2",
)
with pytest.raises(
litellm.RateLimitError, match=re.escape("No deployments available for selected model. Passed model=g")
):
await core.select("g", messages=MESSAGES) # pyright: ignore[reportCallIssue]
# --- cooldowns ----------------------------------------------------------------------------------------
@pytest.mark.asyncio
@pytest.mark.parametrize("error", [rate_limit, auth_error, server_error, timeout], ids=lambda e: e.__name__)
async def test_a_deployment_that_failed_past_allowed_fails_is_not_picked(
make_core: MakeCore, host: ScriptedHost, error: Callable[[], Exception]
) -> None:
core: Final = make_core(
model_list=[deployment("g", "a"), deployment("g", "b")], allowed_fails=0, num_retries=0, cooldown_time=60
)
host.answer("openai/a", error())
await _fail_through(core, "a")
assert await _picked_ids(core, "g") == {"b"}
@pytest.mark.asyncio
async def test_completion_avoids_a_cooling_deployment(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(
model_list=[deployment("g", "a"), deployment("g", "b")], allowed_fails=0, num_retries=0, cooldown_time=60
)
host.answer("openai/a", rate_limit())
await _fail_through(core, "a")
for _ in range(5):
await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue]
assert host.deployments_called("acompletion") == ["a", "b", "b", "b", "b", "b"]
@pytest.mark.asyncio
@pytest.mark.parametrize("error", [bad_request, context_window], ids=lambda e: e.__name__)
async def test_client_errors_do_not_cool_a_deployment_down(
make_core: MakeCore, host: ScriptedHost, error: Callable[[], Exception]
) -> None:
core: Final = make_core(
model_list=[deployment("g", "a"), deployment("g", "b")], allowed_fails=0, num_retries=0, cooldown_time=60
)
host.answer("openai/a", error())
await _fail_through(core, "a", times=3)
assert await core.select("a") is not None
assert "a" in await _picked_ids(core, "g", times=60)
@pytest.mark.asyncio
@pytest.mark.parametrize("allowed_fails", [1, 2, 3])
async def test_a_deployment_stays_pickable_until_it_fails_more_than_allowed_fails(
make_core: MakeCore, host: ScriptedHost, allowed_fails: int
) -> None:
core: Final = make_core(
model_list=[deployment("g", "a"), deployment("g", "b")],
allowed_fails=allowed_fails,
num_retries=0,
cooldown_time=60,
)
host.answer("openai/a", rate_limit())
await _fail_through(core, "a", times=allowed_fails)
assert "a" in await _picked_ids(core, "g", times=60)
await _fail_through(core, "a")
assert await _picked_ids(core, "g") == {"b"}
@pytest.mark.asyncio
async def test_a_single_deployment_group_is_never_cooled_down(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(model_list=[deployment("g", "a")], allowed_fails=0, num_retries=0, cooldown_time=60)
host.answer("openai/a", rate_limit())
await _fail_through(core, "a", times=3)
assert await _picked_ids(core, "g") == {"a"}
@pytest.mark.asyncio
async def test_selecting_a_cooling_deployment_by_id_still_returns_it(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(
model_list=[deployment("g", "a"), deployment("g", "b")], allowed_fails=0, num_retries=0, cooldown_time=60
)
host.answer("openai/a", rate_limit())
await _fail_through(core, "a")
assert await _picked_ids(core, "g") == {"b"}
assert await _picked_ids(core, "a") == {"a"}
async def _all_cooling_core(make_core: MakeCore, host: ScriptedHost) -> RouterCore:
core: Final = make_core(
model_list=[deployment("g", "a"), deployment("g", "b")], allowed_fails=0, num_retries=0, cooldown_time=60
)
host.answer("openai/a", rate_limit())
host.answer("openai/b", rate_limit())
await _fail_through(core, "a")
await _fail_through(core, "b")
return core
ALL_COOLING_MESSAGE: Final = (
"No deployments available for selected model, Try again in 60 seconds. All deployments for selected "
"model are in cooldown. Passed model=g. pre-call-checks=False, cooldown_list=['a', 'b']"
)
@pytest.mark.asyncio
@pytest.mark.parametrize("entry", ["select", "acompletion", "completion"])
async def test_every_deployment_cooling_raises_router_rate_limit_error(
make_core: MakeCore, host: ScriptedHost, entry: str
) -> None:
core: Final = await _all_cooling_core(make_core, host)
calls_before: Final = len(host.calls)
call: Final = core.select("g") if entry == "select" else core.call(entry, {"model": "g", "messages": MESSAGES})
with pytest.raises(RouterRateLimitError) as raised:
await call
assert str(raised.value) == ALL_COOLING_MESSAGE
assert raised.value.model == "g"
assert raised.value.cooldown_time == 60
assert sorted(raised.value.cooldown_list) == ["a", "b"]
assert len(host.calls) == calls_before
@pytest.mark.asyncio
async def test_cooling_down_one_group_leaves_other_groups_pickable(make_core: MakeCore, host: ScriptedHost) -> None:
core: Final = make_core(
model_list=[deployment("g", "a"), deployment("g", "b"), deployment("h", "c"), deployment("h", "d")],
allowed_fails=0,
num_retries=0,
cooldown_time=60,
)
host.answer("openai/a", rate_limit())
host.answer("openai/b", rate_limit())
await _fail_through(core, "a")
await _fail_through(core, "b")
with pytest.raises(RouterRateLimitError):
await core.select("g")
assert await _picked_ids(core, "h") <= {"c", "d"}

28
tests/unit/test_router/test_router.py
View file

@ -2893,11 +2893,11 @@ def test_switch_routing_strategy_installs_lar1_then_restores_the_default_selecto
},
)
assert router.routing_strategy == "lar1"
assert "_core_async_get_available_deployment" in router.__dict__
assert "_core_async_get_available_deployment" in router._core.__dict__
router._switch_routing_strategy("usage-based-routing-v2", {})
assert router.lowesttpm_logger_v2 is not None
assert "async_get_available_deployment" not in router.__dict__
assert "_core_async_get_available_deployment" not in router._core.__dict__
@pytest.mark.parametrize(
@ -22656,6 +22656,9 @@ _SSE_CHUNKS: Final[tuple[bytes, ...]] = tuple(
b'data: {"id":"c","object":"chat.completion.chunk","created":1,"model":"gpt-5.6",'
b'"choices":[{"index":0,"delta":{"content":"x"},"finish_reason":null}]}\n\n'
for _ in range(5)
) + (
b'data: {"id":"c","object":"chat.completion.chunk","created":1,"model":"gpt-5.6",'
b'"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\n',
)
@ -24271,7 +24274,7 @@ def _record_phase_events(monkeypatch: pytest.MonkeyPatch) -> list[tuple[str, dic
def record(name: str, attributes: dict[str, str | int]) -> None:
events.append((name, dict(attributes)))
monkeypatch.setattr(litellm.router, "phase_event", record)
monkeypatch.setattr(litellm.router_python_core, "phase_event", record)
return events
@ -25666,7 +25669,9 @@ def test_get_model_list_from_model_alias_should_not_iterate_for_non_alias_lookup
],
model_group_alias={"alias-1": "gpt-5.5"},
)
router.model_group_alias = NoItemsAliasDict({f"alias-{idx}": "gpt-5.5" for idx in range(200)})
# Assigning model_group_alias copies the value (a changeset detaches it), so the tripwire goes
# straight onto the core, where the lookup reads it.
router._core._core_model_group_alias = NoItemsAliasDict({f"alias-{idx}": "gpt-5.5" for idx in range(200)})
model_alias_list = router.get_model_list_from_model_alias(model_name="gpt-5-mini")
assert model_alias_list == []
@ -25686,7 +25691,9 @@ def test_map_team_model_should_not_iterate_aliases_for_non_alias_team_model_name
],
model_group_alias={"alias-1": "gpt-5.5"},
)
router.model_group_alias = NoItemsAliasDict({f"alias-{idx}": "gpt-5.5" for idx in range(200)})
# Assigning model_group_alias copies the value (a changeset detaches it), so the tripwire goes
# straight onto the core, where the lookup reads it.
router._core._core_model_group_alias = NoItemsAliasDict({f"alias-{idx}": "gpt-5.5" for idx in range(200)})
# map_team_model should return the public name unchanged (not the internal UUID name)
# so the router can find all sibling deployments via team_id filtering
@ -29139,13 +29146,13 @@ def test_reset_custom_routing_strategy():
router = _create_custom_routing_router()
router.set_custom_routing_strategy(SpecialEndpointRoutingStrategy(router))
assert "get_available_deployment" in router.__dict__
assert "_core_async_get_available_deployment" in router.__dict__
assert "get_available_deployment" in router._core.__dict__
assert "_core_async_get_available_deployment" in router._core.__dict__
router._reset_custom_routing_strategy()
assert "get_available_deployment" not in router.__dict__
assert "async_get_available_deployment" not in router.__dict__
assert "get_available_deployment" not in router._core.__dict__
assert "_core_async_get_available_deployment" not in router._core.__dict__
assert router.async_get_available_deployment.__func__ is Router.async_get_available_deployment
router._reset_custom_routing_strategy()
@ -30985,7 +30992,7 @@ async def test_responses_fallback_wrapper_disconnect_billing_keeps_the_primary_s
_SYNC_SECONDS: Final = 10
_ROUTER_FILE: Final = Router.__init__.__code__.co_filename
_ROUTER_FILE: Final = Router._update_deployment_indices_after_removal.__code__.co_filename
_INDEX_LOCK_LINES: Final = frozenset(
number
for number, line in enumerate(linecache.getlines(_ROUTER_FILE), start=1)
@ -31442,6 +31449,7 @@ def test_request_carries_image_sees_chat_parts_and_responses_input_items():
assert litellm.Router._request_carries_image(None, {"input": _IMAGE_RESPONSES_INPUT}) is True
assert litellm.Router._request_carries_image([{"role": "user", "content": "Say hi"}], {}) is False
assert litellm.Router._request_carries_image(None, {"input": "Say hi"}) is False
assert litellm.Router._request_carries_image(None, {"input": ["safe text", "more safe text"]}) is False
@pytest.mark.asyncio

6
tests/unit/test_router/test_router_entry_points.py
View file

@ -359,7 +359,7 @@ def test_settings_snapshot_is_frozen(router):
def test_unset_setting_raises_attribute_error_through_old_name():
r = _router()
try:
del r.__dict__["_core_kubernetes_pod_discovery"]
del r._core.__dict__["_core_kubernetes_pod_discovery"]
with warnings.catch_warnings():
warnings.simplefilter("ignore")
assert getattr(r, "kubernetes_pod_discovery", "missing") == "missing"
@ -372,14 +372,14 @@ def test_unset_setting_raises_attribute_error_through_old_name():
def test_unset_attribute_error_names_the_public_attribute(router):
# Reading (or deleting) an old attribute the router never set raises what reading it raised before.
router.__dict__.pop("_core_slack_alerting_logger", None)
router._core.__dict__.pop("_core_slack_alerting_logger", None)
with warnings.catch_warnings():
warnings.simplefilter("ignore")
with pytest.raises(AttributeError) as read:
_ = router.slack_alerting_logger
with pytest.raises(AttributeError) as deleted:
del router.slack_alerting_logger
del router.__dict__["_core_kubernetes_pod_discovery"]
del router._core.__dict__["_core_kubernetes_pod_discovery"]
with pytest.raises(AttributeError) as setting:
_ = router.kubernetes_pod_discovery
for caught, name in ((read, "slack_alerting_logger"), (setting, "kubernetes_pod_discovery")):

52
tests/unit/test_router/test_router_facade.py Normal file
View file

@ -0,0 +1,52 @@
"""
The ``litellm.Router`` facade's own helpers in litellm/router.py: how it binds the constructor
arguments, which members it leaves to the facade, and the module attributes it forwards to the core.
"""
import pytest
import litellm.router as router_module
from litellm import router_python_core
from litellm.router import Router, _class_descriptor, router_arguments
def test_router_arguments_binds_positional_arguments_to_their_parameter_names():
model_list = [{"model_name": "g", "litellm_params": {"model": "openai/gpt-4o-mini"}}]
arguments = router_arguments(model_list, num_retries=2)
assert dict(arguments) == {"model_list": model_list, "num_retries": 2}
def test_router_arguments_leaves_unset_parameters_out():
assert dict(router_arguments()) == {}
def test_router_arguments_rejects_an_unknown_argument_naming_the_router():
with pytest.raises(TypeError, match=r"^Router\.__init__\(\) got an unexpected keyword argument 'not_a_setting'"):
router_arguments(not_a_setting=1)
def test_router_arguments_are_read_only():
arguments = router_arguments(num_retries=2)
with pytest.raises(TypeError):
arguments["num_retries"] = 3 # pyright: ignore[reportIndexIssue] # proving the mapping is read-only
def test_class_descriptor_is_true_for_a_data_descriptor_on_the_class():
assert _class_descriptor(Router, "_core")
def test_class_descriptor_is_false_for_a_plain_method_or_a_missing_name():
assert not _class_descriptor(Router, "from_config")
assert not _class_descriptor(Router, "no_such_member")
def test_module_getattr_forwards_old_helpers_to_the_core_module():
assert router_module.__getattr__("_stream_pre_content_error") is router_python_core._stream_pre_content_error
def test_module_getattr_raises_attribute_error_for_an_unknown_name():
with pytest.raises(AttributeError):
router_module.__getattr__("no_such_router_helper")

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

@ -1182,7 +1182,7 @@ def test_deployment_callback_respects_cooldown_time(model_list):
},
}
with patch("litellm.router.set_cooldown_deployments") as mock_set:
with patch("litellm.router_python_core.set_cooldown_deployments") as mock_set:
router.deployment_callback_on_failure(
kwargs=kwargs,
completion_response=None,
@ -1554,7 +1554,7 @@ def _rpm_tpm_router(model_id: str) -> Router:
def _pin_router_clock(monkeypatch: pytest.MonkeyPatch, pinned: datetime) -> None:
monkeypatch.setattr("litellm.router.get_utc_datetime", lambda: pinned)
monkeypatch.setattr("litellm.router_python_core.get_utc_datetime", lambda: pinned)
monkeypatch.setattr(
"litellm.router_utils.pre_call_checks.io_token_rate_limit_check.get_utc_datetime", lambda: pinned
)
@ -2398,10 +2398,10 @@ def _routers_seen_in_success_loop_discarding_first(first: Router, second: Router
for callback in litellm.success_callback:
if not isinstance(callback, MethodType):
continue
if callback.__self__ is first:
if callback.__self__ is first._core:
first.discard()
yield first
if callback.__self__ is second:
if callback.__self__ is second._core:
yield second

2
tests/unit/test_router/test_router_provider_endpoints.py
View file

@ -53,7 +53,7 @@ def httpx_transport(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
@pytest.fixture
def router_minute_pinned(monkeypatch: pytest.MonkeyPatch) -> None:
pinned: Final = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc)
monkeypatch.setattr("litellm.router.get_utc_datetime", lambda: pinned)
monkeypatch.setattr("litellm.router_python_core.get_utc_datetime", lambda: pinned)
class _RouterLoggingCapture(CustomLogger):

4
tests/unit/test_router_silent_experiment.py
View file

@ -620,7 +620,7 @@ def test_sync_shadow_gets_kwargs_snapshot_taken_before_primary_mutates_them(reco
router = Router(model_list=_tagged_primary_model_list())
with patch( # test-quality-ok: Router has no thread factory to inject; deferring start is the only deterministic way to expose the race
"litellm.router.threading", SimpleNamespace(Thread=_DeferredThread)
"litellm.router_python_core.threading", SimpleNamespace(Thread=_DeferredThread)
):
response = router.completion(
model="primary-model",
@ -652,7 +652,7 @@ def test_sync_shadow_workers_do_not_share_metadata_with_each_other(recording_log
router = Router(model_list=_streaming_model_list(["shadow-a", "shadow-b"]))
with patch( # test-quality-ok: Router has no thread factory to inject; deferring start is the only deterministic way to expose the race
"litellm.router.threading", SimpleNamespace(Thread=_DeferredThread)
"litellm.router_python_core.threading", SimpleNamespace(Thread=_DeferredThread)
):
router.completion(
model="primary-model",

12
tests/unit/test_router_weighted_failover.py
View file

@ -338,7 +338,7 @@ async def test_maybe_run_weighted_failover_persists_excluded_ids_to_kwargs(monke
async def _stub_run_async_fallback(*args, **kwargs):
return "ok"
monkeypatch.setattr("litellm.router.run_async_fallback", _stub_run_async_fallback)
monkeypatch.setattr("litellm.router_python_core.run_async_fallback", _stub_run_async_fallback)
exc = Exception("fail")
exc.failed_deployment_id = "A"
@ -807,11 +807,11 @@ async def test_maybe_run_weighted_failover_skips_when_remaining_all_in_cooldown(
run_async_fallback_called = True
return "should not reach here"
monkeypatch.setattr("litellm.router.run_async_fallback", _should_not_be_called)
monkeypatch.setattr("litellm.router_python_core.run_async_fallback", _should_not_be_called)
# Patch cooldown so B and C appear in cooldown.
with patch(
"litellm.router.async_get_cooldown_deployments",
"litellm.router_python_core.async_get_cooldown_deployments",
new=AsyncMock(return_value=["B", "C"]),
):
result = await router._maybe_run_weighted_failover(
@ -867,10 +867,10 @@ async def test_maybe_run_weighted_failover_proceeds_when_one_healthy_remains(
run_async_fallback_called = True
return "ok from C"
monkeypatch.setattr("litellm.router.run_async_fallback", _stub_run_async_fallback)
monkeypatch.setattr("litellm.router_python_core.run_async_fallback", _stub_run_async_fallback)
with patch(
"litellm.router.async_get_cooldown_deployments",
"litellm.router_python_core.async_get_cooldown_deployments",
new=AsyncMock(return_value=["B"]),
):
result = await router._maybe_run_weighted_failover(
@ -936,7 +936,7 @@ async def test_failover_falls_through_to_external_fallback_when_remaining_in_coo
# Put B in cooldown so weighted failover can't use it after A fails.
with patch(
"litellm.router.async_get_cooldown_deployments",
"litellm.router_python_core.async_get_cooldown_deployments",
new=AsyncMock(return_value=["B"]),
):
response = await router.acompletion(

55
tests/unit/test_scheduler.py
View file

@ -5,12 +5,20 @@ import asyncio
import importlib
import json
import os
import time
from collections.abc import Sequence
from typing import Final
import pytest
import redis
import redis.asyncio
from fakeredis import FakeRedis, FakeServer
from fakeredis.aioredis import FakeConnection
from redis.asyncio import BlockingConnectionPool
import litellm
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.scheduler import FlowItem, Scheduler, SchedulerCacheKeys
from litellm.utils import _invalidate_model_cost_lowercase_map
@ -207,6 +215,53 @@ class _ExpiringCache:
self.store[key] = value
class _SchedulerClock:
now: float = 1000.0
def __call__(self) -> float:
return self.now
@pytest.mark.asyncio
async def test_waiter_observes_shared_expiry_after_reading_just_before_it(monkeypatch: pytest.MonkeyPatch) -> None:
clock: Final = _SchedulerClock()
server: Final = FakeServer()
class Pool(BlockingConnectionPool):
def __init__(self, **kwargs: object) -> None:
super().__init__(connection_class=FakeConnection, server=server, **kwargs)
monkeypatch.setattr(time, "time", clock)
monkeypatch.setattr(redis, "Redis", FakeRedis)
monkeypatch.setattr(redis.asyncio, "BlockingConnectionPool", Pool)
cache: Final = RedisCache(host="scheduler-expiry")
scheduler: Final = Scheduler(redis_cache=cache)
scheduler.cache.in_memory_cache = InMemoryCache(clock=clock)
waiter: Final = FlowItem(priority=2, request_id="waiter", model_name="expiry")
key: Final = f"{SchedulerCacheKeys.queue.value}:{waiter.model_name}"
try:
await cache.async_set_cache(key, [(1, "dead-replica")])
assert await scheduler.get_queue(waiter.model_name) == [(1, "dead-replica")]
clock.now = 1000.1
await scheduler.add_request(waiter)
clock.now = 1005.05
assert not await scheduler.poll(waiter, [])
clock.now = 1005.2
server.connected = False
assert not await scheduler.poll(waiter, [])
server.connected = True
assert await cache.async_get_cache(key) is None
assert await scheduler.get_queue(waiter.model_name) == []
server.connected = False
assert await scheduler.get_queue(waiter.model_name) == []
server.connected = True
assert await scheduler.poll(waiter, [])
assert await scheduler.get_queue(waiter.model_name) == []
finally:
server.connected = True
await cache.disconnect()
@pytest.mark.asyncio
async def test_wait_for_turn_during_cooldown_survives_the_queue_key_expiring():
cache: Final = _ExpiringCache()