scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-27 06:13:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 72

fix(router): unregister strategy selectors from global callbacks on discard (#45795)

Browse source 
* fix(router): unregister strategy selectors from global callbacks on discard

Router.discard only removed its own bound methods from litellm's callback
lists, so the strategy selectors it registered (least-busy, usage-based,
latency-based, cost-based and routing group selectors) stayed behind and
kept receiving events after the router was gone. Reuse the identity based
unregister that update_settings already uses so a discarded router takes
only its own selectors with it

* fix(router): purge discarded selectors from every event callback list

function_setup copies litellm.callbacks into the success and failure
lists on every request, so a router that served traffic still leaked its
selectors there after discard. Route the cleanup through
remove_callback_from_all_lists, which now also covers input_callback, and
fold the selector gathering into discard so the router coverage check
sees no untested helper

* fix(router): keep live routers, in-flight hooks and user_config limits intact on discard

discard() now rebuilds every global callback list without the router's selectors, optional pre-call checks and own usage callbacks instead of removing them in place, so a callback loop already running never skips a hook. It then re-registers what the remaining live routers need, so a same-strategy survivor keeps its deduped handler. The proxy's per-request user_config router is discarded only after its call returns.

* fix(router): bind user_config handlers to the request and tighten discard

When a user_config request carries a logging object, its router is
discarded before the request runs and its handlers ride on that
logging object, so v1 rpm and v2 tpm limits count every request and
concurrent requests count once. Without one, the router is discarded
after the request as before. The router's pre-call loops also read
the request-scoped handlers.

discard() now refills only the handler keys it vacated, guards the
callback lists and the live-router set with a lock, covers routers
built without a model_list, and drops its Slack alerting, adaptive
router hook and budget limiter sync task

* fix(proxy): reuse one user_config router per config and bind its handlers to each request

Every user_config request now goes through a cached router built once per
config. Its strategy handlers leave the global callback lists and are bound
to the request's logging object instead, so limits count each request once
and nothing is left behind after the request.

* fix(router): keep the proxy budget limiter on user_config requests

A user_config router now swaps in only its strategy selectors, so the proxy's
RouterBudgetLimiting stays on the request and keeps counting its spend. Usage
logged after a strategy is retired, such as a request still in flight when its
user_config router is evicted, now schedules one more push to Redis.

* fix(proxy): init the retire fields on the key model budget limiter

The key model budget limiter builds its own state instead of calling the base
init, so with the proxy's Redis attached every spend increment raised on the
missing retire flag and the key's model budget never counted spend.

* fix(router): bind the deployment callbacks to user_config requests

A cached user_config router serves requests through each request's logging
object but bound only its selectors and policies there, so a failing
deployment was never cooled down and its usage counters missed requests.

* fix(router): retry an evicted router's budget writes and keep the adaptive header on user_config requests

A retired budget limiter now keeps retrying its queued Redis increments every
sync interval until Redis accepts them, instead of giving up after one failed
write.

The x-litellm-adaptive-router-model header now comes from one stateless hook in
litellm.callbacks that holds no router, so user_config adaptive routers, whose
per-router hooks are request-scoped, still answer with the header.

* fix(router): lock request callback setup against router cleanup and drain in-flight budget writes on retire

function_setup now edits the global callback lists under the same lock that
router cleanup uses, so a router dropping its callbacks on another thread can
no longer shift the indexes function_setup pops or put a dropped handler back.

A budget limiter retired while a Redis write is still in flight now keeps
retrying that batch if the write fails, instead of leaving the spend queued
with no task left to send it.

* fix(utils): set up loggers after releasing the callback lists lock

function_setup held the callback lists lock while set_callbacks ran its one-time
logger setup, which can block on a Slack package install, so other threads and
Router.discard waited on it. The lock now covers only the list reads and writes.

* test(router): annotate the retire test's locals as Final

* fix(logging): skip replaced global callbacks on the input path too

replace_global_callback already kept a same-key global logger out of
success and failure dispatch. pre_call and post_call still sent it the
input events, so a logger with an input hook could count a request start
it never saw finish. Both now share one filter with the success path

* test(logging): type the input-path regression test

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
Kazuki Hashimoto 2026-10-21 11:39:36 +09:00 • committed by GitHub
parent 8d9b9926d1
commit c8ed5ccca8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
21 changed files with 1656 additions and 201 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

73
litellm/litellm_core_utils/litellm_logging.py
View file

@ -596,6 +596,7 @@ class Logging(LiteLLMLoggingBaseClass):
prometheusLogger, \
slack_app
custom_pricing: bool = False
replaced_global_callback_keys: frozenset[str] = frozenset()
stream_options = None
litellm_request_debug: bool = False
streamed_anthropic_message_id: str | None = None
@ -650,11 +651,11 @@ class Logging(LiteLLMLoggingBaseClass):
litellm_call_id: str,
function_id: str,
litellm_trace_id: str | None = None,
dynamic_input_callbacks: list[str | Callable | CustomLogger] | None = None,
dynamic_input_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None,
dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None,
dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None,
dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = None,
dynamic_async_failure_callbacks: list[str | Callable | CustomLogger] | None = None,
dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None,
dynamic_async_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None,
applied_guardrails: list[str] | None = None,
kwargs: dict | None = None,
log_raw_request_response: bool = False,
@ -728,15 +729,17 @@ class Logging(LiteLLMLoggingBaseClass):
self.raw_request_only = raw_request_only
# Initialize dynamic callbacks
self.dynamic_input_callbacks: list[str | Callable | CustomLogger] | None = dynamic_input_callbacks
self.dynamic_input_callbacks: list[str | Callable[..., object] | CustomLogger] | None = dynamic_input_callbacks
self.dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = (
dynamic_success_callbacks
)
self.dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = (
dynamic_async_success_callbacks
)
self.dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = dynamic_failure_callbacks
self.dynamic_async_failure_callbacks: list[str | Callable | CustomLogger] | None = (
self.dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = (
dynamic_failure_callbacks
)
self.dynamic_async_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = (
dynamic_async_failure_callbacks
)
@ -853,10 +856,32 @@ class Logging(LiteLLMLoggingBaseClass):
self.dynamic_async_failure_callbacks, callback
)
def replace_global_callback(self, callback: CustomLogger) -> None:
self.add_dynamic_callback(callback)
self.replaced_global_callback_keys = self.replaced_global_callback_keys | {
litellm.logging_callback_manager.get_custom_logger_key(callback)
}
def add_dynamic_function_callbacks(
self,
*,
async_success: Callable[..., object],
failure: Callable[..., object],
async_failure: Callable[..., object],
) -> None:
self.dynamic_async_success_callbacks = self._with_dynamic_callback(
self.dynamic_async_success_callbacks, async_success
)
self.dynamic_failure_callbacks = self._with_dynamic_callback(self.dynamic_failure_callbacks, failure)
self.dynamic_async_failure_callbacks = self._with_dynamic_callback(
self.dynamic_async_failure_callbacks, async_failure
)
@staticmethod
def _with_dynamic_callback(
callbacks: Sequence[str | Callable | CustomLogger] | None, callback: CustomLogger
) -> list[str | Callable | CustomLogger]:
callbacks: Sequence[str | Callable[..., object] | CustomLogger] | None,
callback: Callable[..., object] | CustomLogger,
) -> list[str | Callable[..., object] | CustomLogger]:
existing: Final = tuple(callbacks or ())
return [*existing, *(() if callback in existing else (callback,))]
@ -893,9 +918,9 @@ class Logging(LiteLLMLoggingBaseClass):
def _process_dynamic_callback_list(
self,
callback_list: list[str | Callable | CustomLogger] | None,
callback_list: list[str | Callable[..., object] | CustomLogger] | None,
dynamic_callbacks_type: Literal["input", "success", "failure", "async_success", "async_failure"],
) -> list[str | Callable | CustomLogger] | None:
) -> list[str | Callable[..., object] | CustomLogger] | None:
"""
Helper function to initialize CustomLogger compatible callbacks in self.dynamic_* callbacks
@ -907,7 +932,7 @@ class Logging(LiteLLMLoggingBaseClass):
if callback_list is None:
return None
processed_list: Final[list[str | Callable | CustomLogger]] = []
processed_list: Final[list[str | Callable[..., object] | CustomLogger]] = []
for callback in callback_list:
if isinstance(callback, str) and callback in litellm._known_custom_logger_compatible_callbacks:
for callback_instance in self._resolve_dynamic_callback_string(callback):
@ -1568,7 +1593,10 @@ class Logging(LiteLLMLoggingBaseClass):
"api_call_start_time"
]
# Input Integration Logging -> If you want to log the fact that an attempt to call the model was made
callbacks: Final = litellm.input_callback + (self.dynamic_input_callbacks or [])
callbacks: Final = (
*self._without_replaced_global_callbacks(litellm.input_callback),
*(self.dynamic_input_callbacks or ()),
)
for callback in callbacks:
try:
if callback == "supabase" and supabaseClient is not None:
@ -1741,7 +1769,10 @@ class Logging(LiteLLMLoggingBaseClass):
)
# Input Integration Logging -> If you want to log the fact that an attempt to call the model was made
callbacks: Final = litellm.input_callback + (self.dynamic_input_callbacks or [])
callbacks: Final = (
*self._without_replaced_global_callbacks(litellm.input_callback),
*(self.dynamic_input_callbacks or ()),
)
for callback in callbacks:
try:
if callback == "sentry" and add_breadcrumb:
@ -4518,14 +4549,26 @@ class Logging(LiteLLMLoggingBaseClass):
_filtered_failure_callbacks = self._remove_internal_litellm_callbacks(_filtered_failure_callbacks)
return len(_filtered_failure_callbacks) > 0
def _without_replaced_global_callbacks(
self, global_callbacks: Iterable[str | Callable[..., object] | CustomLogger]
) -> tuple[str | Callable[..., object] | CustomLogger, ...]:
if not self.replaced_global_callback_keys:
return tuple(global_callbacks)
return tuple(
litellm.logging_callback_manager.without_custom_logger_keys(
global_callbacks, self.replaced_global_callback_keys
)
)
def get_combined_callback_list(
self,
dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None,
global_callbacks: list[str | Callable[..., object] | CustomLogger],
) -> list[str | Callable[..., object] | CustomLogger]:
kept_global_callbacks: Final = self._without_replaced_global_callbacks(global_callbacks)
if dynamic_success_callbacks is None:
return list(global_callbacks)
return list(dict.fromkeys(dynamic_success_callbacks + global_callbacks))
return list(kept_global_callbacks)
return list(dict.fromkeys((*dynamic_success_callbacks, *kept_global_callbacks)))
def _remove_internal_litellm_callbacks(self, callbacks: list) -> list:
"""

149
litellm/litellm_core_utils/logging_callback_manager.py
View file

@ -1,4 +1,6 @@
from collections.abc import Callable
import threading
from collections.abc import Callable, Iterable, Iterator
from types import MethodType
from typing import TYPE_CHECKING, Final
import litellm
@ -15,6 +17,7 @@ else:
_custom_logger_compatible_callbacks_literal = str
_generic_api_logger_cache: Final[dict[str, GenericAPILogger]] = {}
_callback_lists_lock: Final = threading.RLock()
class LoggingCallbackManager:
@ -43,16 +46,18 @@ class LoggingCallbackManager:
Add a input callback to litellm.input_callback.
Auto-routes async callbacks to litellm._async_input_callback.
"""
if not isinstance(callback, str) and self._is_async_callable(callback):
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_input_callback)
else:
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.input_callback)
with _callback_lists_lock:
if not isinstance(callback, str) and self._is_async_callable(callback):
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_input_callback)
else:
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.input_callback)
def add_litellm_service_callback(self, callback: CustomLogger | str | Callable[..., object]):
"""
Add a service callback to litellm.service_callback
"""
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.service_callback)
with _callback_lists_lock:
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.service_callback)
def add_litellm_callback(self, callback: CustomLogger | str | Callable[..., object]):
"""
@ -60,10 +65,11 @@ class LoggingCallbackManager:
Ensures no duplicates are added.
"""
self._safe_add_callback_to_list(
callback=callback,
parent_list=litellm.callbacks,
)
with _callback_lists_lock:
self._safe_add_callback_to_list(
callback=callback,
parent_list=litellm.callbacks,
)
def add_litellm_success_callback(self, callback: CustomLogger | str | Callable[..., object]):
"""
@ -71,37 +77,41 @@ class LoggingCallbackManager:
Auto-routes async callbacks to litellm._async_success_callback.
Special-cases 'dynamodb' and 'openmeter' as async callbacks.
"""
if (
isinstance(callback, str)
and callback in ("dynamodb", "openmeter")
or not isinstance(callback, str)
and self._is_async_callable(callback)
):
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback)
else:
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.success_callback)
with _callback_lists_lock:
if (
isinstance(callback, str)
and callback in ("dynamodb", "openmeter")
or not isinstance(callback, str)
and self._is_async_callable(callback)
):
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback)
else:
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.success_callback)
def add_litellm_failure_callback(self, callback: CustomLogger | str | Callable[..., object]):
"""
Add a failure callback to `litellm.failure_callback`.
Auto-routes async callbacks to litellm._async_failure_callback.
"""
if not isinstance(callback, str) and self._is_async_callable(callback):
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback)
else:
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.failure_callback)
with _callback_lists_lock:
if not isinstance(callback, str) and self._is_async_callable(callback):
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback)
else:
self._safe_add_callback_to_list(callback=callback, parent_list=litellm.failure_callback)
def add_litellm_async_success_callback(self, callback: CustomLogger | Callable[..., object] | str):
"""
Add a success callback to litellm._async_success_callback
"""
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback)
with _callback_lists_lock:
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback)
def add_litellm_async_failure_callback(self, callback: CustomLogger | Callable[..., object] | str):
"""
Add a failure callback to litellm._async_failure_callback
"""
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback)
with _callback_lists_lock:
self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback)
def remove_callback_from_list_by_object(self, callback_list, obj, require_self=True):
"""
@ -110,13 +120,14 @@ class LoggingCallbackManager:
if not isinstance(callback_list, list): # Not list -> do nothing
return
if require_self:
remove_list = [c for c in callback_list if hasattr(c, "__self__") and c.__self__ == obj]
else:
remove_list = [c for c in callback_list if c == obj]
with _callback_lists_lock:
if require_self:
remove_list = [c for c in callback_list if hasattr(c, "__self__") and c.__self__ == obj]
else:
remove_list = [c for c in callback_list if c == obj]
for c in remove_list:
callback_list.remove(c)
for c in remove_list:
callback_list.remove(c)
def remove_callbacks_by_type(self, callback_list, callback_type):
"""
@ -134,10 +145,11 @@ class LoggingCallbackManager:
if not isinstance(callback_list, list):
return
remove_list: Final = [c for c in callback_list if isinstance(c, callback_type)]
with _callback_lists_lock:
remove_list: Final = [c for c in callback_list if isinstance(c, callback_type)]
for c in remove_list:
callback_list.remove(c)
for c in remove_list:
callback_list.remove(c)
def _add_string_callback_to_list(
self, callback: str, parent_list: list[CustomLogger | Callable[..., object] | str]
@ -336,6 +348,55 @@ class LoggingCallbackManager:
return "-".join(key_parts)
get_custom_logger_key = _get_custom_logger_key
callback_lists_lock: Final = _callback_lists_lock
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)
with _callback_lists_lock:
vacated_keys: Final = frozenset(
self._get_custom_logger_key(callback)
for callback in litellm.callbacks
if isinstance(callback, CustomLogger) and id(callback) in dropped_ids
)
litellm.callbacks = [callback for callback in litellm.callbacks if keep(callback)]
litellm.input_callback = [callback for callback in litellm.input_callback if keep(callback)]
litellm.success_callback = [callback for callback in litellm.success_callback if keep(callback)]
litellm.failure_callback = [callback for callback in litellm.failure_callback if keep(callback)]
litellm.service_callback = [callback for callback in litellm.service_callback if keep(callback)]
litellm._async_success_callback = [
callback for callback in litellm._async_success_callback if keep(callback)
]
litellm._async_failure_callback = [
callback for callback in litellm._async_failure_callback if keep(callback)
]
return vacated_keys
def with_custom_logger_keys(
self, candidates: Iterable[CustomLogger], keys: frozenset[str]
) -> tuple[CustomLogger, ...]:
if not keys:
return ()
class_names: Final = frozenset(key.split("-", 1)[0] for key in keys)
same_class: Final = (candidate for candidate in candidates if type(candidate).__name__ in class_names)
return tuple(candidate for candidate in same_class if self._get_custom_logger_key(candidate) in keys)
def without_custom_logger_keys(
self, callbacks: Iterable[str | Callable[..., object] | CustomLogger], keys: frozenset[str]
) -> Iterator[str | Callable[..., object] | CustomLogger]:
return (
callback
for callback in callbacks
if not (isinstance(callback, CustomLogger) and self._get_custom_logger_key(callback) in keys)
)
def _reset_all_callbacks(self):
"""
Reset all callbacks to an empty list
@ -363,19 +424,21 @@ class LoggingCallbackManager:
_get_all_callbacks = get_all_callbacks
def remove_callback_from_all_lists(self, obj, require_self=False) -> None:
def remove_callback_from_all_lists(self, obj: object, require_self: bool = False) -> None:
"""
Remove a callback object from every callback list it may have been
promoted into, so a re-initialized callback leaves no stale instance behind.
"""
for callback_list in (
litellm.callbacks,
litellm.success_callback,
litellm.failure_callback,
litellm._async_success_callback,
litellm._async_failure_callback,
):
self.remove_callback_from_list_by_object(callback_list, obj, require_self=require_self)
with _callback_lists_lock:
for callback_list in (
litellm.callbacks,
litellm.input_callback,
litellm.success_callback,
litellm.failure_callback,
litellm._async_success_callback,
litellm._async_failure_callback,
):
self.remove_callback_from_list_by_object(callback_list, obj, require_self=require_self)
def get_active_additional_logging_utils_from_custom_logger(
self,

2
litellm/proxy/hooks/model_max_budget_limiter.py
View file

@ -374,6 +374,8 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
self._redis_increment_queue_lock = asyncio.Lock()
self._redis_increment_flush_lock = asyncio.Lock()
self._detached_increment_operations = None
self._retired = False
self._flush_after_retire = None
self.deployment_budget_config = None
@with_service_target("model_budgets")

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

@ -880,8 +880,8 @@ def _build_passthrough_failure_request_payload(
@dataclass(frozen=True, slots=True)
class _TeamCallbackWiring:
success_callbacks: "list[str | Callable | CustomLogger] | None" = None # mutable-ok: Logging.__init__ arg
failure_callbacks: "list[str | Callable | CustomLogger] | None" = None # mutable-ok: Logging.__init__ arg
success_callbacks: "list[str | Callable[..., object] | CustomLogger] | None" = None # mutable-ok: Logging init arg
failure_callbacks: "list[str | Callable[..., object] | CustomLogger] | None" = None # mutable-ok: Logging init arg
logging_kwargs: dict[str, str | dict[str, str]] | None = None # mutable-ok: Logging.__init__ arg

59
litellm/proxy/route_llm_request.py
View file

@ -1,8 +1,13 @@
import asyncio
from collections.abc import Mapping
import hashlib
import json
from collections import OrderedDict
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal
from uuid import uuid4
import httpx
from fastapi import HTTPException, status
@ -10,6 +15,7 @@ from pydantic import BeforeValidator, TypeAdapter, ValidationError
import litellm
from litellm.litellm_core_utils.core_helpers import normalize_drop_params
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.router_utils.common_utils import is_proxy_admin_request
@ -54,7 +60,34 @@ _ROUTER_ONLY_SETTINGS: Final = frozenset(
)
def _route_user_config_request(data: dict, route_type: str):
class UserConfigRouters:
def __init__(self, max_size: int) -> None:
self._max_size: Final = max_size
self._routers: Final[OrderedDict[str, LitellmRouter]] = OrderedDict() # mutable-ok: LRU evicts in place
def get(self, config: Mapping[str, object], build: Callable[[], LitellmRouter]) -> LitellmRouter:
key: Final = hashlib.sha256(json.dumps(config, sort_keys=True, default=str).encode()).hexdigest()
cached: Final = self._routers.get(key)
if cached is not None:
self._routers.move_to_end(key)
return cached
router: Final = build()
router.serve_through_request_logging()
self._routers[key] = router
if len(self._routers) > self._max_size:
self._routers.popitem(last=False)[1].discard()
return router
def clear(self) -> None:
for router in self._routers.values():
router.discard()
self._routers.clear()
user_config_routers: Final = UserConfigRouters(max_size=32)
async def _route_user_config_request(data: dict, route_type: str):
"""Route a request using the user-provided router config."""
router_config: Final = data.pop("user_config")
@ -63,10 +96,24 @@ def _route_user_config_request(data: dict, route_type: str):
valid_args: Final = litellm.Router.get_valid_args()
filtered_config: Final = {k: v for k, v in router_config.items() if k in valid_args}
user_router: Final = litellm.Router(**filtered_config)
ret_val: Final = getattr(user_router, f"{route_type}")(**data)
user_router.discard()
return ret_val
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.bind_to_request_logging(logging_obj)
return await getattr(cached_router, f"{route_type}")(**{**data, "litellm_logging_obj": logging_obj})
def _request_logging_obj(data: Mapping[str, object], route_type: str) -> LiteLLMLoggingObj:
arrived_with: Final = data.get("litellm_logging_obj")
if isinstance(arrived_with, LiteLLMLoggingObj):
return arrived_with
logging_obj, _ = litellm.utils.function_setup(
original_function=route_type,
rules_obj=litellm.utils.Rules(),
start_time=datetime.now(), # noqa: DTZ005 # Logging subtracts it from naive end times
is_async_call=True,
**{"litellm_call_id": str(uuid4()), **data},
)
return logging_obj
def _is_a2a_agent_model(model_name: object) -> bool:

178
litellm/router.py
View file

@ -38,6 +38,7 @@ from contextvars import ContextVar
from dataclasses import dataclass
from datetime import datetime, timezone
from functools import lru_cache, partial
from itertools import chain
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Generic, Literal, NoReturn, Optional, TypeAlias, TypeVar, Union, cast
@ -553,6 +554,7 @@ _PreRoutingStrategyT = TypeVar("_PreRoutingStrategyT")
_CallbackT = TypeVar("_CallbackT")
_DeploymentT = TypeVar("_DeploymentT", bound=Mapping[str, object])
_ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset(
{"model", "api_base", "api_key", "api_version", "responses_api_version"}
)
@ -1047,6 +1049,8 @@ class RoutingArgs(enum.Enum):
# as the per-request one built from a caller-supplied user_config, drops out on its
# own rather than leaving entries behind that nothing can withdraw.
_live_routers: Final["weakref.WeakSet[Router]"] = weakref.WeakSet()
_routers_holding_global_callbacks: Final["weakref.WeakSet[Router]"] = weakref.WeakSet()
_routers_holding_global_callbacks_lock: Final = threading.Lock()
def _replay_live_router_model_cost() -> None:
@ -1579,22 +1583,18 @@ class Router:
self._init_routing_groups(self._routing_groups_input)
self._override_selectors: dict[str, RouterStrategySelector | None] = {}
self._override_selectors_lock = threading.Lock()
self._request_scoped_selectors: tuple[CustomLogger, ...] = ()
self._request_scoped_policies: tuple[CustomLogger, ...] = ()
self.access_groups = None
with _routers_holding_global_callbacks_lock:
_routers_holding_global_callbacks.add(self)
## USAGE TRACKING ##
if isinstance(litellm._async_success_callback, list):
litellm.logging_callback_manager.add_litellm_async_success_callback(self.deployment_callback_on_success)
else:
litellm.logging_callback_manager.add_litellm_async_success_callback(self.deployment_callback_on_success)
litellm.logging_callback_manager.add_litellm_async_success_callback(self.deployment_callback_on_success)
if isinstance(litellm.success_callback, list):
litellm.logging_callback_manager.add_litellm_success_callback(self.sync_deployment_callback_on_success)
else:
litellm.success_callback = [self.sync_deployment_callback_on_success]
if isinstance(litellm._async_failure_callback, list):
litellm.logging_callback_manager.add_litellm_async_failure_callback(
self.async_deployment_callback_on_failure
)
else:
litellm._async_failure_callback = [self.async_deployment_callback_on_failure]
litellm.logging_callback_manager.add_litellm_async_failure_callback(self.async_deployment_callback_on_failure)
## COOLDOWNS ##
if isinstance(litellm.failure_callback, list):
litellm.logging_callback_manager.add_litellm_failure_callback(self.deployment_callback_on_failure)
@ -1690,20 +1690,116 @@ class Router:
# Stop contributing to cost-map rebuilds straight away rather than waiting
# for this router to be collected.
_live_routers.discard(self)
litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm._async_success_callback, self)
litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.success_callback, self)
litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm._async_failure_callback, self)
litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.failure_callback, self)
litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.input_callback, self)
litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.service_callback, self)
litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, self)
self.retire_routing_state()
self._release_global_callbacks()
# Remove ForwardClientSideHeadersByModelGroup if it exists
if self.optional_callbacks is not None:
for callback in self.optional_callbacks:
litellm.logging_callback_manager.remove_callback_from_list_by_object(
litellm.callbacks, callback, require_self=False
)
def _release_global_callbacks(self) -> None:
with _routers_holding_global_callbacks_lock:
_routers_holding_global_callbacks.discard(self)
slack_alerting_loggers: Final = self._slack_alerting_loggers()
adaptive_router_hooks: Final = self._adaptive_router_hooks()
vacated_keys: Final = litellm.logging_callback_manager.drop_callbacks(
(
*self._callback_selectors(),
*self._override_selectors.values(),
*(self.optional_callbacks or ()),
*slack_alerting_loggers,
*adaptive_router_hooks,
),
method_owners=(self, *slack_alerting_loggers),
)
with _routers_holding_global_callbacks_lock:
survivors: Final = tuple(_routers_holding_global_callbacks)
candidates: Final = chain.from_iterable(router._global_callback_candidates() for router in survivors)
for candidate in litellm.logging_callback_manager.with_custom_logger_keys(candidates, vacated_keys):
self._register_router_selector(candidate)
adaptive_keys: Final = frozenset(
map(litellm.logging_callback_manager.get_custom_logger_key, adaptive_router_hooks)
)
adaptive_survivor: Final = next((router for router in survivors if router.adaptive_routers), None)
if adaptive_survivor is not None and vacated_keys & adaptive_keys:
adaptive_survivor._sync_adaptive_router_hooks()
def serve_through_request_logging(self) -> None:
self._request_scoped_selectors = (*self._callback_selectors(), *self._request_adaptive_router_hooks())
self._request_scoped_policies = (*self._optional_loggers(), *self._slack_alerting_loggers())
self._release_global_callbacks()
def bind_to_request_logging(self, logging_obj: LiteLLMLogging) -> None:
logging_obj.add_dynamic_function_callbacks(
async_success=self.deployment_callback_on_success,
failure=self.deployment_callback_on_failure,
async_failure=self.async_deployment_callback_on_failure,
)
for selector in self._request_scoped_selectors:
logging_obj.replace_global_callback(selector)
for policy in self._request_policies_not_set_globally():
logging_obj.add_dynamic_callback(policy)
def _request_policies_not_set_globally(self) -> tuple[CustomLogger, ...]:
global_keys: Final = frozenset(
map(
litellm.logging_callback_manager.get_custom_logger_key,
(callback for callback in litellm.callbacks if isinstance(callback, CustomLogger)),
)
)
return tuple(
policy
for policy in self._request_scoped_policies
if litellm.logging_callback_manager.get_custom_logger_key(policy) not in global_keys
)
def _request_adaptive_router_hooks(self) -> tuple[CustomLogger, ...]:
from litellm.router_strategy.adaptive_router.hooks import (
AdaptiveRouterPostCallHook,
)
return tuple(
AdaptiveRouterPostCallHook(adaptive_router=tagged.strategy)
for tagged in chain.from_iterable(self.adaptive_routers.values())
)
def _routing_callbacks(self) -> Iterable[object]:
if not (self._request_scoped_selectors or self._request_scoped_policies):
return litellm.callbacks
replaced_keys: Final = frozenset(
map(litellm.logging_callback_manager.get_custom_logger_key, self._request_scoped_selectors)
)
return chain(
litellm.logging_callback_manager.without_custom_logger_keys(litellm.callbacks, replaced_keys),
self._request_scoped_selectors,
self._request_policies_not_set_globally(),
)
def retire_routing_state(self) -> None:
for strategy in (*self._global_callback_candidates(), *self._override_selectors.values()):
if isinstance(strategy, BaseRoutingStrategy | RouterBudgetLimiting):
strategy.retire()
def _slack_alerting_loggers(self) -> tuple[CustomLogger, ...]:
slack_alerting_logger: Final = getattr(self, "slack_alerting_logger", None)
return (slack_alerting_logger,) if isinstance(slack_alerting_logger, CustomLogger) else ()
def _adaptive_router_hooks(self) -> tuple[CustomLogger, ...]:
from litellm.router_strategy.adaptive_router.hooks import (
AdaptiveRouterPostCallHook,
)
own_strategy_ids: Final = frozenset(
id(tagged.strategy) for tagged in chain.from_iterable(self.adaptive_routers.values())
)
hooks: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook)
return tuple(
hook
for hook in hooks
if isinstance(hook, AdaptiveRouterPostCallHook) and id(hook.adaptive_router) in own_strategy_ids
)
def _optional_loggers(self) -> tuple[CustomLogger, ...]:
return tuple(callback for callback in self.optional_callbacks or () if isinstance(callback, CustomLogger))
def _global_callback_candidates(self) -> tuple[CustomLogger, ...]:
return (*self._callback_selectors(), *self._optional_loggers(), *self._slack_alerting_loggers())
@staticmethod
def _create_redis_cache(
@ -1809,7 +1905,7 @@ class Router:
return selector
@staticmethod
def _register_router_selector(selector: RouterStrategySelector) -> None:
def _register_router_selector(selector: CustomLogger) -> None:
if isinstance(selector, LeastBusyLoggingHandler):
if isinstance(litellm.input_callback, list):
litellm.logging_callback_manager.add_litellm_input_callback(selector)
@ -1818,23 +1914,28 @@ class Router:
if isinstance(litellm.callbacks, list):
litellm.logging_callback_manager.add_litellm_callback(selector)
def _callback_selectors(self) -> tuple[RouterStrategySelector, ...]:
return tuple(
selector
for selector in (
*(getattr(self, attr, None) for attr in self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.values()),
*chain.from_iterable(selectors.values() for selectors in self._group_selectors.values()),
)
if selector is not None
)
def _unregister_router_selectors(self, selectors: Sequence[object]) -> None:
"""
Drop router-owned strategy selectors from litellm's global callback
lists by identity. Used before re-init (`routing_strategy_init` /
`_init_routing_groups`) so repeated `update_settings` calls don't
accumulate dead selectors that keep receiving callback events.
lists. Used before re-init (`routing_strategy_init` /
`_init_routing_groups`) and by `discard` so dead selectors stop
receiving callback events. `function_setup` copies `litellm.callbacks`
into the per-event lists on every request, so those are purged too.
"""
for selector in selectors:
if isinstance(selector, BaseRoutingStrategy):
selector.retire()
selector_ids: Final = {id(s) for s in selectors if s is not None}
if not selector_ids:
return
if isinstance(litellm.callbacks, list):
litellm.callbacks = [c for c in litellm.callbacks if id(c) not in selector_ids]
if isinstance(litellm.input_callback, list):
litellm.input_callback = [c for c in litellm.input_callback if id(c) not in selector_ids]
litellm.logging_callback_manager.drop_callbacks(selectors)
def _apply_updated_routing_strategy_args(self) -> None:
"""
@ -9804,7 +9905,7 @@ class Router:
Raises:
- Rate Limit Exception - If the deployment is over it's tpm/rpm limits
"""
for _callback in litellm.callbacks:
for _callback in self._routing_callbacks():
if isinstance(_callback, CustomLogger):
_callback.pre_call_check(deployment)
@ -9825,7 +9926,7 @@ class Router:
Raises:
- Rate Limit Exception - If the deployment is over it's tpm/rpm limits
"""
for _callback in litellm.callbacks:
for _callback in self._routing_callbacks():
if isinstance(_callback, CustomLogger):
try:
await _callback.async_pre_call_check(deployment, parent_otel_span)
@ -9962,7 +10063,7 @@ class Router:
- Rate Limit Exception - If the deployment is over it's tpm/rpm limits
"""
returned_healthy_deployments = healthy_deployments
for _callback in litellm.callbacks:
for _callback in self._routing_callbacks():
if isinstance(_callback, CustomLogger):
try:
returned_healthy_deployments = await _callback.async_filter_deployments(
@ -10617,6 +10718,7 @@ class Router:
changes, otherwise a released router keeps recording turns through its hook."""
from litellm.router_strategy.adaptive_router.hooks import (
AdaptiveRouterPostCallHook,
AdaptiveRouterResponseHeaderHook,
)
for callback in litellm.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook):
@ -10626,6 +10728,8 @@ class Router:
litellm.logging_callback_manager.add_litellm_callback(
AdaptiveRouterPostCallHook(adaptive_router=tagged.strategy)
)
if self.adaptive_routers:
litellm.logging_callback_manager.add_litellm_callback(AdaptiveRouterResponseHeaderHook())
def init_adaptive_router_deployment(self, deployment: Deployment) -> None:
"""

7
litellm/router_strategy/adaptive_router/__init__.py
View file

@ -1,6 +1,9 @@
"""Adaptive router strategy. See README.md for design overview."""
from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter
from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook
from litellm.router_strategy.adaptive_router.hooks import (
AdaptiveRouterPostCallHook,
AdaptiveRouterResponseHeaderHook,
)
__all__ = ["AdaptiveRouter", "AdaptiveRouterPostCallHook"]
__all__ = ["AdaptiveRouter", "AdaptiveRouterPostCallHook", "AdaptiveRouterResponseHeaderHook"]

14
litellm/router_strategy/adaptive_router/hooks.py
View file

@ -170,12 +170,7 @@ def _assistant_content_and_tool_calls(response_obj: Any) -> tuple[object, Sequen
return content, tool_calls
class AdaptiveRouterPostCallHook(CustomLogger):
"""One hook instance per AdaptiveRouter. Registered into litellm.callbacks."""
def __init__(self, adaptive_router: AdaptiveRouter) -> None:
self.adaptive_router = adaptive_router
class AdaptiveRouterResponseHeaderHook(CustomLogger):
async def async_post_call_response_headers_hook(
self,
data: Mapping[str, object],
@ -200,6 +195,13 @@ class AdaptiveRouterPostCallHook(CustomLogger):
return None
return {ADAPTIVE_ROUTER_RESPONSE_HEADER: chosen}
class AdaptiveRouterPostCallHook(CustomLogger):
"""One hook instance per AdaptiveRouter. Registered into litellm.callbacks."""
def __init__(self, adaptive_router: AdaptiveRouter) -> None:
self.adaptive_router = adaptive_router
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
await self._record(kwargs, response_obj, response_status=200)

20
litellm/router_strategy/base_routing_strategy.py
View file

@ -24,6 +24,8 @@ class BaseRoutingStrategy(ABC):
self.dual_cache = dual_cache
self.redis_increment_operation_queue: list[RedisPipelineIncrementOperation] = []
self._sync_task: asyncio.Task[None] | None = None
self._retired = False
self._flush_after_retire: asyncio.Task[None] | None = None
if should_batch_redis_writes:
self.setup_sync_task(default_sync_interval)
@ -47,13 +49,25 @@ class BaseRoutingStrategy(ABC):
def retire(self) -> None:
self.cancel_sync_task()
if not self.redis_increment_operation_queue:
self._retired = True
self._schedule_flush_after_retire()
def _schedule_flush_after_retire(self) -> None:
if self._flush_after_retire is not None and not self._flush_after_retire.done():
return
if not self.redis_increment_operation_queue or self.dual_cache.redis_cache is None:
return
try:
loop: Final = asyncio.get_running_loop()
except RuntimeError:
return
loop.create_task(self._push_in_memory_increments_to_redis())
self._flush_after_retire = loop.create_task(self._drain_increments_after_retire())
async def _drain_increments_after_retire(self) -> None:
if not self.redis_increment_operation_queue:
return
await self._push_in_memory_increments_to_redis()
await self._drain_increments_after_retire()
async def cleanup(self):
"""Cleanup method to be called when shutting down"""
@ -98,6 +112,8 @@ class BaseRoutingStrategy(ABC):
self.redis_increment_operation_queue.append(increment_op)
self.add_to_in_memory_keys_to_update(key=key)
if self._retired:
self._schedule_flush_after_retire()
return result
@with_service_target("router_usage")

30
litellm/router_strategy/budget_limiter.py
View file

@ -128,7 +128,9 @@ class RouterBudgetLimiting(CustomLogger):
self._redis_increment_queue_lock = asyncio.Lock()
self._redis_increment_flush_lock = asyncio.Lock()
self._detached_increment_operations: tuple[RedisPipelineIncrementOperation, ...] | None = None
asyncio.create_task(self.periodic_sync_in_memory_spend_with_redis())
self._sync_task: asyncio.Task[None] = asyncio.create_task(self.periodic_sync_in_memory_spend_with_redis())
self._retired = False
self._flush_after_retire: asyncio.Task[None] | None = None
self.provider_budget_config: GenericBudgetConfigType | None = provider_budget_config
self.deployment_budget_config: GenericBudgetConfigType | None = None
self.tag_budget_config: GenericBudgetConfigType | None = None
@ -140,6 +142,30 @@ class RouterBudgetLimiting(CustomLogger):
if isinstance(litellm.callbacks, list):
litellm.logging_callback_manager.add_litellm_callback(self)
def retire(self) -> None:
self._sync_task.cancel()
self._retired = True
self._schedule_flush_after_retire()
def _schedule_flush_after_retire(self) -> None:
if self._flush_after_retire is not None and not self._flush_after_retire.done():
return
if not self._has_unwritten_increments():
return
try:
loop: Final = asyncio.get_running_loop()
except RuntimeError:
return
self._flush_after_retire = loop.create_task(self._drain_increments_after_retire())
def _has_unwritten_increments(self) -> bool:
return bool(self.redis_increment_operation_queue) or self._detached_increment_operations is not None
async def _drain_increments_after_retire(self) -> None:
while self._has_unwritten_increments():
if not await self._push_in_memory_increments_to_redis():
await asyncio.sleep(DEFAULT_REDIS_SYNC_INTERVAL)
@with_service_target("router_budgets")
async def async_filter_deployments(
self,
@ -427,6 +453,8 @@ class RouterBudgetLimiting(CustomLogger):
if self.dual_cache.redis_cache is None:
return
self.redis_increment_operation_queue.append(increment_op)
if self._retired:
self._schedule_flush_after_retire()
def _get_redis_increment_queue_lock(self) -> asyncio.Lock:
return self._redis_increment_queue_lock

155
litellm/utils.py
View file

@ -1126,83 +1126,90 @@ def function_setup(
dynamic_callbacks: Final[list[str | Callable | CustomLogger] | None] = kwargs.pop("callbacks", None)
all_callbacks: Final = get_dynamic_callbacks(dynamic_callbacks=dynamic_callbacks)
if len(all_callbacks) > 0:
for callback in all_callbacks:
# check if callback is a string - e.g. "lago", "openmeter"
if isinstance(callback, str):
callback = litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class(
callback,
internal_usage_cache=None,
llm_router=None,
)
if callback is None or any(
type(cb) is type(callback) for cb in litellm._async_success_callback
): # don't double add a callback
continue
if callback not in litellm.input_callback:
litellm.input_callback.append(callback)
if callback not in litellm.success_callback:
litellm.logging_callback_manager.add_litellm_success_callback(callback)
if callback not in litellm.failure_callback:
litellm.logging_callback_manager.add_litellm_failure_callback(callback)
if callback not in litellm._async_success_callback:
litellm.logging_callback_manager.add_litellm_async_success_callback(callback)
if callback not in litellm._async_failure_callback:
litellm.logging_callback_manager.add_litellm_async_failure_callback(callback)
print_verbose(f"Initialized litellm callbacks, Async Success Callbacks: {litellm._async_success_callback}")
with litellm.logging_callback_manager.callback_lists_lock:
if len(all_callbacks) > 0:
for callback in all_callbacks:
# check if callback is a string - e.g. "lago", "openmeter"
if isinstance(callback, str):
callback = litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class(
callback,
internal_usage_cache=None,
llm_router=None,
)
if callback is None or any(
type(cb) is type(callback) for cb in litellm._async_success_callback
): # don't double add a callback
continue
if callback not in litellm.input_callback:
litellm.input_callback.append(callback)
if callback not in litellm.success_callback:
litellm.logging_callback_manager.add_litellm_success_callback(callback)
if callback not in litellm.failure_callback:
litellm.logging_callback_manager.add_litellm_failure_callback(callback)
if callback not in litellm._async_success_callback:
litellm.logging_callback_manager.add_litellm_async_success_callback(callback)
if callback not in litellm._async_failure_callback:
litellm.logging_callback_manager.add_litellm_async_failure_callback(callback)
print_verbose(
f"Initialized litellm callbacks, Async Success Callbacks: {litellm._async_success_callback}"
)
if (
len(litellm.input_callback) > 0 or len(litellm.success_callback) > 0 or len(litellm.failure_callback) > 0
) and len(callback_list) == 0:
callback_list = list(set(litellm.input_callback + litellm.success_callback + litellm.failure_callback))
callbacks_to_set_up: Final = (
list(set(litellm.input_callback + litellm.success_callback + litellm.failure_callback))
if len(callback_list) == 0
else []
)
if callbacks_to_set_up:
callback_list = callbacks_to_set_up
## ASYNC CALLBACKS - safety net for callbacks added via direct append
if len(litellm.input_callback) > 0:
removed_async_items = []
for index, callback in enumerate(litellm.input_callback):
if coroutine_checker.is_async_callable(callback):
litellm._async_input_callback.append(callback)
removed_async_items.append(index)
# Pop the async items from input_callback in reverse order to avoid index issues
for index in reversed(removed_async_items):
litellm.input_callback.pop(index)
if len(litellm.success_callback) > 0:
removed_async_items = []
for index, callback in enumerate(litellm.success_callback):
if coroutine_checker.is_async_callable(callback):
litellm.logging_callback_manager.add_litellm_async_success_callback(callback)
removed_async_items.append(index)
elif callback == "dynamodb" or callback == "openmeter":
# dynamo is an async callback, it's used for the proxy and needs to be async
# we only support async dynamo db logging for acompletion/aembedding since that's used on proxy
litellm.logging_callback_manager.add_litellm_async_success_callback(callback)
removed_async_items.append(index)
elif callback in litellm._known_custom_logger_compatible_callbacks and isinstance(callback, str):
add_custom_logger_callback_to_specific_event(callback, "success")
# Pop the async items from success_callback in reverse order to avoid index issues
for index in reversed(removed_async_items):
litellm.success_callback.pop(index)
if len(litellm.failure_callback) > 0:
removed_async_items = []
for index, callback in enumerate(litellm.failure_callback):
if coroutine_checker.is_async_callable(callback):
litellm.logging_callback_manager.add_litellm_async_failure_callback(callback)
removed_async_items.append(index)
elif callback in litellm._known_custom_logger_compatible_callbacks and isinstance(callback, str):
add_custom_logger_callback_to_specific_event(callback, "failure")
# Pop the async items from failure_callback in reverse order to avoid index issues
for index in reversed(removed_async_items):
litellm.failure_callback.pop(index)
if callbacks_to_set_up:
get_set_callbacks: Final[Callable[[], Callable[..., None]]] = litellm_utils.get_set_callbacks
get_set_callbacks()(callback_list=callback_list, function_id=function_id)
## ASYNC CALLBACKS - safety net for callbacks added via direct append
if len(litellm.input_callback) > 0:
removed_async_items = []
for index, callback in enumerate(litellm.input_callback):
if coroutine_checker.is_async_callable(callback):
litellm._async_input_callback.append(callback)
removed_async_items.append(index)
# Pop the async items from input_callback in reverse order to avoid index issues
for index in reversed(removed_async_items):
litellm.input_callback.pop(index)
if len(litellm.success_callback) > 0:
removed_async_items = []
for index, callback in enumerate(litellm.success_callback):
if coroutine_checker.is_async_callable(callback):
litellm.logging_callback_manager.add_litellm_async_success_callback(callback)
removed_async_items.append(index)
elif callback == "dynamodb" or callback == "openmeter":
# dynamo is an async callback, it's used for the proxy and needs to be async
# we only support async dynamo db logging for acompletion/aembedding since that's used on proxy
litellm.logging_callback_manager.add_litellm_async_success_callback(callback)
removed_async_items.append(index)
elif callback in litellm._known_custom_logger_compatible_callbacks and isinstance(callback, str):
add_custom_logger_callback_to_specific_event(callback, "success")
# Pop the async items from success_callback in reverse order to avoid index issues
for index in reversed(removed_async_items):
litellm.success_callback.pop(index)
if len(litellm.failure_callback) > 0:
removed_async_items = []
for index, callback in enumerate(litellm.failure_callback):
if coroutine_checker.is_async_callable(callback):
litellm.logging_callback_manager.add_litellm_async_failure_callback(callback)
removed_async_items.append(index)
elif callback in litellm._known_custom_logger_compatible_callbacks and isinstance(callback, str):
add_custom_logger_callback_to_specific_event(callback, "failure")
# Pop the async items from failure_callback in reverse order to avoid index issues
for index in reversed(removed_async_items):
litellm.failure_callback.pop(index)
get_set_callbacks()(callback_list=callbacks_to_set_up, function_id=function_id)
### DYNAMIC CALLBACKS ###
dynamic_success_callbacks: list[str | Callable | CustomLogger] | None = None
dynamic_async_success_callbacks: list[str | Callable | CustomLogger] | None = None
dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = None
dynamic_async_failure_callbacks: Final[list[str | Callable | CustomLogger] | None] = None
dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None
dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None
dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None
dynamic_async_failure_callbacks: Final[list[str | Callable[..., object] | CustomLogger] | None] = None
if kwargs.get("success_callback", None) is not None and isinstance(kwargs["success_callback"], list):
removed_async_items = []
for index, callback in enumerate(kwargs["success_callback"]):

12
tests/code_coverage_tests/router_code_coverage.py
View file

@ -146,6 +146,18 @@ ignored_function_names = [
"async_admit_deployment_turn", # Bound as the WebSocket turn_admission partial, tested through it in test_router_helper_utils.py
"_deployment_as_configured_now", # Read by async_admit_deployment_turn, tested through turn_admission in test_router_helper_utils.py
"_admit_turn_through_routing_strategy", # Run by async_admit_deployment_turn, tested through turn_admission in test_router_helper_utils.py
"_callback_selectors", # Run by discard, tested through the discard tests in test_router_helper_utils.py
"bind_to_request_logging", # Run by the proxy user_config route, tested through the user_config tests in test_route_llm_request.py
"serve_through_request_logging", # Run by the proxy user_config route, tested through the user_config tests in test_route_llm_request.py
"_request_adaptive_router_hooks", # Run by serve_through_request_logging, tested through the user_config adaptive test
"_release_global_callbacks", # Run by discard, tested through the discard tests in test_router_helper_utils.py
"_routing_callbacks", # Read by the pre-call checks, tested through the user_config limit tests in test_route_llm_request.py
"_request_policies_not_set_globally", # Read by bind_to_request_logging, tested through the user_config budget test in test_route_llm_request.py
"_optional_loggers", # Read by discard and serve_through_request_logging, tested through the budget limiter discard test
"retire_routing_state", # Run by discard, tested through the budget limiter discard test
"_slack_alerting_loggers", # Run by discard, tested through the Slack alerting discard test in test_router_helper_utils.py
"_adaptive_router_hooks", # Run by discard, tested through the adaptive router discard tests in test_router_helper_utils.py
"_global_callback_candidates", # Run by discard, tested through the survivor discard tests in test_router_helper_utils.py
"without_deployments_that_cannot_batch", # Tested through the Azure and OpenAI mixed-group acreate_batch tests in test_router.py
"batch_create_routing", # Tested through the mixed-group acreate_batch tests in test_router.py
"is_batch_create_routing", # Tested through the chat-after-batch-create test in test_router.py

4
tests/unit/conftest.py
View file

@ -252,6 +252,7 @@ def isolate_host_environment(isolated_aws_config_files: tuple[Path, Path]) -> It
@pytest.fixture(autouse=True)
def isolate_litellm_globals() -> Iterator[None]:
original_callbacks: Final = {name: list(getattr(litellm, name) or []) for name in CALLBACK_LISTS}
original_callback_holders: Final = frozenset(litellm_router_module._routers_holding_global_callbacks)
original_reset: Final = {name: getattr(litellm, name) for name in RESET_TO_NONE_GLOBALS}
original_restored: Final = {name: getattr(litellm, name) for name in RESTORED_GLOBALS if hasattr(litellm, name)}
original_clients: Final = {name: litellm.__dict__[name] for name in MODULE_LEVEL_CLIENTS if name in litellm.__dict__}
@ -263,6 +264,7 @@ def isolate_litellm_globals() -> Iterator[None]:
_flush_client_caches()
for name in CALLBACK_LISTS:
setattr(litellm, name, [])
litellm_router_module._routers_holding_global_callbacks.clear()
for name in RESET_TO_NONE_GLOBALS:
setattr(litellm, name, None)
for name in MODULE_LEVEL_CLIENTS:
@ -277,6 +279,8 @@ def isolate_litellm_globals() -> Iterator[None]:
litellm.__dict__.update(original_clients)
for name, value in (original_callbacks | original_reset | original_restored).items():
setattr(litellm, name, value)
litellm_router_module._routers_holding_global_callbacks.clear()
litellm_router_module._routers_holding_global_callbacks.update(original_callback_holders)
for logger, (level, disabled, propagate, handlers, filters) in original_loggers.items():
logger.setLevel(level)
logger.disabled = disabled

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

@ -402,6 +402,32 @@ def test_get_combined_callback_list_preserves_insertion_order(logging_obj):
) == ["prometheus", "langfuse", "datadog", "otel", "s3", "gcs_bucket", "arize", "logfire"]
class _InputEventRecorder(CustomLogger):
def __init__(self, label: str, events: list[str]) -> None:
self._label = label
self.events = events
def log_pre_api_call(self, model: object, messages: object, kwargs: object) -> None:
self.events.append(f"{self._label} pre")
def log_post_api_call(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None:
self.events.append(f"{self._label} post")
def test_input_events_reach_the_replacement_and_not_the_global_callback_it_replaced(
logging_obj: LitellmLogging, monkeypatch: pytest.MonkeyPatch
) -> None:
events: Final[list[str]] = []
monkeypatch.setattr(litellm, "input_callback", [_InputEventRecorder("global", events)])
logging_obj.replace_global_callback(_InputEventRecorder("replacement", events))
logging_obj.update_environment_variables(litellm_params={}, optional_params={})
logging_obj.pre_call(input="", api_key="")
logging_obj.post_call(original_response="ok", input="", api_key="")
assert events == ["replacement pre", "replacement post"]
def test_get_masked_api_base(logging_obj):
api_base = "https://api.openai.com/v1"
masked_api_base = logging_obj._get_masked_api_base(api_base)

2
tests/unit/litellm_core_utils/test_logging_callback_manager.py
View file

@ -145,6 +145,7 @@ def test_remove_callback_from_all_lists():
obj = TestLogger()
manager.add_litellm_callback(obj)
manager.add_litellm_input_callback(obj)
manager.add_litellm_success_callback(obj)
manager.add_litellm_failure_callback(obj)
manager.add_litellm_async_success_callback(obj)
@ -153,6 +154,7 @@ def test_remove_callback_from_all_lists():
manager.remove_callback_from_all_lists(obj)
assert obj not in litellm.callbacks
assert obj not in litellm.input_callback
assert obj not in litellm.success_callback
assert obj not in litellm.failure_callback
assert obj not in litellm._async_success_callback

595
tests/unit/proxy/test_route_llm_request.py
View file

@ -1,19 +1,42 @@
import asyncio
import contextlib
import gc
import threading
import weakref
import httpx
import json
import litellm
import pytest
import redis
import respx
from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping
from datetime import datetime
from types import SimpleNamespace
from uuid import uuid4
from typing import Final
from unittest.mock import MagicMock
from fakeredis import TcpFakeServer
from fastapi import HTTPException
from litellm.proxy.route_llm_request import JSON_OBJECT_ADAPTER, ProxyModelNotFoundError, route_request
from litellm.types.router import RouterGeneralSettings
from litellm.caching.caching import DualCache
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.route_llm_request import (
JSON_OBJECT_ADAPTER,
ProxyModelNotFoundError,
UserConfigRouters,
route_request,
user_config_routers,
)
from litellm.proxy.utils import ProxyLogging
from litellm.router_strategy.adaptive_router.config import (
ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY,
ADAPTIVE_ROUTER_RESPONSE_HEADER,
)
from litellm.types.router import RouterErrors, RouterGeneralSettings, RouterRateLimitError
@pytest.mark.parametrize(
@ -2434,3 +2457,571 @@ async def test_team_key_search_of_registered_bedrock_store_reaches_bedrock_with_
posted_url
)
assert "Credential=AKIAREGISTRYKEY00000/" in signed_auth, signed_auth
@pytest.fixture(autouse=True)
def discard_cached_user_config_routers() -> Iterator[None]:
yield
user_config_routers.clear()
@pytest.fixture
def fake_redis_port() -> Iterator[int]:
server: Final = TcpFakeServer(("127.0.0.1", 0), server_type="redis")
thread: Final = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
yield server.server_address[1]
server.shutdown()
server.server_close()
async def _wait_out_a_rate_limit_minute_rollover() -> None:
seconds_left: Final = 60 - datetime.now().second
if seconds_left < 10:
await asyncio.sleep(seconds_left)
def _proxy_logging_obj(model: str = "rate-limited") -> LiteLLMLoggingObj:
logging_obj, _ = litellm.utils.function_setup(
original_function="acompletion",
rules_obj=litellm.utils.Rules(),
start_time=datetime.now(),
model=model,
messages=[{"role": "user", "content": "hi"}],
litellm_call_id=str(uuid4()),
)
assert isinstance(logging_obj, LiteLLMLoggingObj)
return logging_obj
def _mock_deployment(model_name: str, **deployment_limits: int) -> dict[str, object]:
return {
"model_name": model_name,
**deployment_limits,
"litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test", "mock_response": "ok"},
}
async def _user_config_request_outcome(
redis_port: int | None,
routing_settings: Mapping[str, object],
deployment_limits: Mapping[str, int],
logging_obj: LiteLLMLoggingObj | None,
) -> str:
data: Final = {
"model": "rate-limited",
"messages": [{"role": "user", "content": "hi"}],
"num_retries": 0,
**({"litellm_logging_obj": logging_obj} if logging_obj is not None else {}),
"user_config": {
**routing_settings,
**({"redis_host": "127.0.0.1", "redis_port": redis_port} if redis_port is not None else {}),
"model_list": [_mock_deployment("rate-limited", **deployment_limits)],
},
}
try:
await (await route_request(data, None, None, "acompletion"))
except (litellm.RateLimitError, RouterRateLimitError):
return "429"
finally:
await asyncio.sleep(0.3)
return "200"
async def _tpm_for_three_requests() -> int:
messages: Final = [{"role": "user", "content": "hi"}]
probe: Final = await litellm.acompletion(model="openai/gpt-5.5", messages=messages, mock_response="ok")
assert isinstance(probe, litellm.ModelResponse)
probe_usage: Final = getattr(probe, "usage", None)
assert isinstance(probe_usage, litellm.Usage)
return 2 * probe_usage.total_tokens + litellm.token_counter(messages=messages)
async def _request_cost() -> float:
probe: Final = await litellm.acompletion(
model="openai/gpt-5.5", messages=[{"role": "user", "content": "hi"}], mock_response="ok"
)
return litellm.completion_cost(completion_response=probe)
async def _proxy_router_outcomes(
redis_port: int | None,
routing_settings: Mapping[str, object],
deployment_limits: Mapping[str, int],
requests: int,
) -> list[str]:
router: Final = litellm.Router(
**routing_settings,
num_retries=0,
model_list=[_mock_deployment("proxy-rate-limited", **deployment_limits)],
**({"redis_host": "127.0.0.1", "redis_port": redis_port} if redis_port is not None else {}),
)
async def outcome() -> str:
try:
await router.acompletion(model="proxy-rate-limited", messages=[{"role": "user", "content": "hi"}])
except (litellm.RateLimitError, RouterRateLimitError):
return "429"
finally:
await asyncio.sleep(0.3)
return "200"
try:
return [await outcome() for _ in range(requests)]
finally:
router.discard()
@pytest.mark.asyncio
async def test_a_user_config_adaptive_request_answers_with_the_model_it_picked() -> None:
data: Final = {
"model": "smart-cheap-router",
"messages": [{"role": "user", "content": "Write a one line python function that adds two numbers"}],
"metadata": {},
"user_config": {
"model_list": [
{
"model_name": "smart-cheap-router",
"litellm_params": {
"model": "auto_router/adaptive_router",
"adaptive_router_config": {"available_models": ["fast", "smart"]},
},
},
{
**_mock_deployment("fast"),
"model_info": {"adaptive_router_preferences": {"quality_tier": 2, "strengths": []}},
},
{
**_mock_deployment("smart"),
"model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": ["code_generation"]}},
},
]
},
}
await (await route_request(data, None, None, "acompletion"))
headers: Final = await ProxyLogging(user_api_key_cache=DualCache()).post_call_response_headers_hook(
data=data, user_api_key_dict=UserAPIKeyAuth(), response=None
)
assert headers == {ADAPTIVE_ROUTER_RESPONSE_HEADER: data["metadata"][ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY]}
def _global_router_callbacks() -> tuple[list[object], list[object]]:
return list(litellm.callbacks), list(litellm.input_callback)
@pytest.mark.parametrize(
"routing_settings, deployment_limits",
[
({"routing_strategy": "usage-based-routing-v2"}, {"rpm": 2}),
({"routing_strategy": "usage-based-routing-v2"}, {"tpm": 30}),
({"routing_strategy": "usage-based-routing"}, {"rpm": 2}),
({"optional_pre_call_checks": ["enforce_model_rate_limits"]}, {"rpm": 2}),
],
)
@pytest.mark.parametrize("arrives_with_logging_obj", [True, False])
@pytest.mark.parametrize("with_redis", [True, False])
@pytest.mark.asyncio
async def test_user_config_router_enforces_its_limit_across_proxy_requests(
fake_redis_port: int,
routing_settings: Mapping[str, object],
deployment_limits: Mapping[str, int],
arrives_with_logging_obj: bool,
with_redis: bool,
) -> None:
await _wait_out_a_rate_limit_minute_rollover()
callbacks_before: Final = _global_router_callbacks()
outcomes: Final = [
await _user_config_request_outcome(
fake_redis_port if with_redis else None,
routing_settings,
deployment_limits,
_proxy_logging_obj() if arrives_with_logging_obj else None,
)
for _ in range(6)
]
assert "429" in outcomes, outcomes
assert _global_router_callbacks() == callbacks_before
assert outcomes == await _proxy_router_outcomes(
fake_redis_port if with_redis else None, routing_settings, deployment_limits, len(outcomes)
)
@pytest.mark.parametrize("arrives_with_logging_obj", [True, False])
@pytest.mark.asyncio
async def test_concurrent_user_config_requests_each_count_once_against_the_tpm_limit(
fake_redis_port: int, arrives_with_logging_obj: bool
) -> None:
tpm_for_three_requests: Final = await _tpm_for_three_requests()
def request() -> Coroutine[None, None, str]:
return _user_config_request_outcome(
fake_redis_port,
{"routing_strategy": "usage-based-routing-v2"},
{"tpm": tpm_for_three_requests},
_proxy_logging_obj() if arrives_with_logging_obj else None,
)
await _wait_out_a_rate_limit_minute_rollover()
concurrent_outcomes: Final = await asyncio.gather(request(), request())
later_outcomes: Final = [await request() for _ in range(2)]
assert [*concurrent_outcomes, *later_outcomes] == ["200", "200", "200", "429"]
@pytest.mark.asyncio
async def test_overlapping_user_config_requests_share_one_rpm_limit(fake_redis_port: int) -> None:
def request() -> Coroutine[None, None, str]:
return _user_config_request_outcome(
fake_redis_port, {"routing_strategy": "usage-based-routing-v2"}, {"rpm": 2}, _proxy_logging_obj()
)
await _wait_out_a_rate_limit_minute_rollover()
outcomes: Final = await asyncio.gather(*(request() for _ in range(4)))
assert outcomes.count("200") <= 2, outcomes
assert "429" in outcomes, outcomes
@pytest.mark.parametrize(
"routing_strategy, limited_by",
[("usage-based-routing-v2", "rpm"), ("usage-based-routing-v2", "tpm"), ("usage-based-routing", "rpm")],
)
@pytest.mark.parametrize("proxy_router_shares_redis", [True, False])
@pytest.mark.asyncio
async def test_a_proxy_router_on_the_same_strategy_leaves_the_user_config_limit_unchanged(
fake_redis_port: int, routing_strategy: str, limited_by: str, proxy_router_shares_redis: bool
) -> None:
deployment_limits: Final = {"rpm": 3} if limited_by == "rpm" else {"tpm": await _tpm_for_three_requests()}
def request() -> Coroutine[None, None, str]:
return _user_config_request_outcome(
fake_redis_port, {"routing_strategy": routing_strategy}, deployment_limits, _proxy_logging_obj()
)
async def outcomes() -> list[str]:
first_instance: Final = [await request() for _ in range(2)]
user_config_routers.clear()
second_instance: Final = [await request() for _ in range(3)]
return [*first_instance, *second_instance]
await _wait_out_a_rate_limit_minute_rollover()
alone: Final = await outcomes()
user_config_routers.clear()
await asyncio.sleep(0.3)
redis.Redis(host="127.0.0.1", port=fake_redis_port).flushall()
proxy_router: Final = litellm.Router(
routing_strategy=routing_strategy,
model_list=[_mock_deployment("proxy-model")],
redis_host="127.0.0.1" if proxy_router_shares_redis else None,
redis_port=fake_redis_port if proxy_router_shares_redis else None,
)
try:
beside_proxy_router: Final = await outcomes()
finally:
proxy_router.discard()
assert "429" in alone, alone
assert beside_proxy_router == alone
@pytest.mark.parametrize("user_config_settings", [{}, {"optional_pre_call_checks": ["router_budget_limiting"]}])
@pytest.mark.asyncio
async def test_user_config_requests_stay_under_the_proxy_provider_budget(
user_config_settings: Mapping[str, object],
) -> None:
proxy_router: Final = litellm.Router(
model_list=[_mock_deployment("proxy-model")],
provider_budget_config={"openai": {"budget_limit": await _request_cost() / 2, "time_period": "1d"}},
)
try:
await proxy_router.acompletion(model="proxy-model", messages=[{"role": "user", "content": "hi"}])
await asyncio.sleep(0.3)
with pytest.raises(ValueError, match=RouterErrors.no_deployments_with_provider_budget_routing.value):
await _user_config_request_outcome(None, user_config_settings, {}, _proxy_logging_obj())
finally:
proxy_router.discard()
@pytest.mark.asyncio
async def test_a_user_config_provider_budget_limits_its_own_requests() -> None:
budget_limit: Final = await _request_cost() / 2
def request() -> Coroutine[None, None, str]:
budget: Final = {"openai": {"budget_limit": budget_limit, "time_period": "1d"}}
return _user_config_request_outcome(None, {"provider_budget_config": budget}, {}, _proxy_logging_obj())
first: Final = await request()
with pytest.raises(ValueError, match=RouterErrors.no_deployments_with_provider_budget_routing.value):
await request()
assert first == "200"
@pytest.mark.asyncio
async def test_a_request_in_flight_when_its_user_config_router_is_dropped_still_counts_against_its_budget(
fake_redis_port: int,
) -> None:
budget_limit: Final = 1.5 * await _request_cost()
async def request() -> None:
await (
await route_request(
{
"model": "budgeted",
"messages": [{"role": "user", "content": "hi"}],
"litellm_logging_obj": _proxy_logging_obj("budgeted"),
"user_config": {
"redis_host": "127.0.0.1",
"redis_port": fake_redis_port,
"provider_budget_config": {"openai": {"budget_limit": budget_limit, "time_period": "1d"}},
"model_list": [
{
"model_name": "budgeted",
"litellm_params": {
"model": "openai/gpt-5.5",
"api_key": "sk-test",
"mock_response": "ok",
"mock_delay": 0.5,
},
}
],
},
},
None,
None,
"acompletion",
)
)
await request()
in_flight: Final = asyncio.create_task(request())
await asyncio.sleep(0.2)
user_config_routers.clear()
await in_flight
await asyncio.sleep(0.3)
with pytest.raises(ValueError, match=RouterErrors.no_deployments_with_provider_budget_routing.value):
await request()
async def _user_config_request(model: str, user_config: Mapping[str, object], arrives_with_logging_obj: bool) -> object:
return await (
await route_request(
{
"model": model,
"messages": list(_HELLO_MESSAGES),
"num_retries": 0,
**({"litellm_logging_obj": _proxy_logging_obj(model)} if arrives_with_logging_obj else {}),
"user_config": user_config,
},
None,
None,
"acompletion",
)
)
def _ordered_pair(model_name: str, first_base: str, second_base: str) -> list[dict[str, object]]:
return [
{
"model_name": model_name,
"litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test", "api_base": api_base, "order": order},
}
for order, api_base in ((1, first_base), (2, second_base))
]
@pytest.mark.parametrize("arrives_with_logging_obj", [True, False])
@pytest.mark.asyncio
async def test_a_rate_limited_user_config_deployment_cools_down_for_the_next_requests(
monkeypatch: pytest.MonkeyPatch, arrives_with_logging_obj: bool
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
rate_limited_base: Final = "https://rate-limited.example/v1"
healthy_base: Final = "https://healthy.example/v1"
user_config: Final = {"model_list": _ordered_pair("flaky", rate_limited_base, healthy_base)}
with respx.mock(assert_all_called=False) as respx_mock:
rate_limited: Final = respx_mock.post(f"{rate_limited_base}/chat/completions").mock(
return_value=httpx.Response(429, json={"error": {"message": "slow down"}})
)
respx_mock.post(f"{healthy_base}/chat/completions").mock(return_value=_chat_completion_response("gpt-5.5"))
for _ in range(3):
await _user_config_request("flaky", user_config, arrives_with_logging_obj)
await asyncio.sleep(0.3)
assert rate_limited.call_count == 1
async def _remaining_usage_headers_after_429_then_200s(
send: Callable[[], Awaitable[object]], api_base: str
) -> list[tuple[object, object] | str]:
async def remaining_usage() -> tuple[object, object] | str:
try:
response: Final = await send()
except litellm.RateLimitError:
return "429"
finally:
await asyncio.sleep(0.3)
assert isinstance(response, litellm.ModelResponse)
headers: Final = response._hidden_params["additional_headers"]
return headers.get("x-ratelimit-remaining-requests"), headers.get("x-ratelimit-remaining-tokens")
with respx.mock(assert_all_called=False) as respx_mock:
respx_mock.post(f"{api_base}/chat/completions").mock(
side_effect=[
httpx.Response(429, json={"error": {"message": "slow down"}}),
_chat_completion_response("gpt-5.5"),
_chat_completion_response("gpt-5.5"),
]
)
return [await remaining_usage() for _ in range(3)]
@pytest.mark.parametrize("arrives_with_logging_obj", [True, False])
@pytest.mark.asyncio
async def test_user_config_remaining_usage_headers_match_a_proxy_router_with_the_same_deployment(
monkeypatch: pytest.MonkeyPatch, arrives_with_logging_obj: bool
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
await _wait_out_a_rate_limit_minute_rollover()
proxy_base: Final = "https://proxy-limited.example/v1"
user_config_base: Final = "https://user-config-limited.example/v1"
def deployment(model_name: str, api_base: str) -> dict[str, object]:
return {
"model_name": model_name,
"litellm_params": {
"model": "openai/gpt-5.5",
"api_key": "sk-test",
"api_base": api_base,
"rpm": 10,
"tpm": 1000,
},
}
proxy_router: Final = litellm.Router(model_list=[deployment("proxy-limited", proxy_base)])
try:
proxy_router_headers: Final = await _remaining_usage_headers_after_429_then_200s(
lambda: proxy_router.acompletion(model="proxy-limited", messages=list(_HELLO_MESSAGES), num_retries=0),
proxy_base,
)
finally:
proxy_router.discard()
user_config: Final = {"model_list": [deployment("user-config-limited", user_config_base)]}
user_config_headers: Final = await _remaining_usage_headers_after_429_then_200s(
lambda: _user_config_request("user-config-limited", user_config, arrives_with_logging_obj), user_config_base
)
assert user_config_headers == proxy_router_headers
async def _first_deployment_calls_over_successes_then_failures(
send: Callable[[], Awaitable[object]], first_base: str, second_base: str
) -> int:
successes: Final = [_chat_completion_response("gpt-5.5") for _ in range(6)]
failures: Final = [httpx.Response(500, json={"error": {"message": "boom"}}) for _ in range(6)]
with respx.mock(assert_all_called=False) as respx_mock:
first: Final = respx_mock.post(f"{first_base}/chat/completions").mock(side_effect=[*successes, *failures])
respx_mock.post(f"{second_base}/chat/completions").mock(return_value=_chat_completion_response("gpt-5.5"))
for _ in range(len(successes) + len(failures)):
await send()
await asyncio.sleep(0.3)
return first.call_count
@pytest.mark.parametrize("arrives_with_logging_obj", [True, False])
@pytest.mark.asyncio
async def test_user_config_successes_keep_a_deployment_out_of_cooldown_like_a_proxy_router(
monkeypatch: pytest.MonkeyPatch, arrives_with_logging_obj: bool
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
await _wait_out_a_rate_limit_minute_rollover()
proxy_router: Final = litellm.Router(
model_list=_ordered_pair("proxy-pair", "https://proxy-first.example/v1", "https://proxy-second.example/v1")
)
try:
proxy_router_calls: Final = await _first_deployment_calls_over_successes_then_failures(
lambda: proxy_router.acompletion(model="proxy-pair", messages=list(_HELLO_MESSAGES), num_retries=0),
"https://proxy-first.example/v1",
"https://proxy-second.example/v1",
)
finally:
proxy_router.discard()
user_config: Final = {
"model_list": _ordered_pair("user-pair", "https://user-first.example/v1", "https://user-second.example/v1")
}
user_config_calls: Final = await _first_deployment_calls_over_successes_then_failures(
lambda: _user_config_request("user-pair", user_config, arrives_with_logging_obj),
"https://user-first.example/v1",
"https://user-second.example/v1",
)
assert user_config_calls == proxy_router_calls
@pytest.mark.asyncio
async def test_user_config_adaptive_router_learns_from_proxy_requests() -> None:
user_config: Final = {
"model_list": [
_mock_deployment("smart"),
_mock_deployment("fast"),
{
"model_name": "adaptive",
"litellm_params": {
"model": "auto_router/adaptive_router",
"adaptive_router_config": {"available_models": ["smart", "fast"]},
},
},
]
}
await (
await route_request(
{
"model": "adaptive",
"messages": [{"role": "user", "content": "write a sorting function"}],
"metadata": {"session_id": "session-1"},
"litellm_logging_obj": _proxy_logging_obj("adaptive"),
"user_config": user_config,
},
None,
None,
"acompletion",
)
)
await asyncio.sleep(0.3)
router: Final = user_config_routers.get(user_config, lambda: pytest.fail("the user_config router was not reused"))
snapshot: Final = await router.adaptive_routers["adaptive"][0].strategy.get_state_snapshot()
assert snapshot["feedback_contexts_live"] == 1, snapshot
@pytest.mark.asyncio
async def test_user_config_routers_discard_the_least_recently_used_router_past_capacity() -> None:
routers: Final = UserConfigRouters(max_size=2)
def config(model_name: str) -> dict[str, object]:
return {"routing_strategy": "usage-based-routing-v2", "model_list": [_mock_deployment(model_name)]}
def build(model_name: str) -> litellm.Router:
return litellm.Router(**config(model_name))
def must_be_cached() -> litellm.Router:
pytest.fail("a cached router was rebuilt")
first_handler: Final = weakref.ref(routers.get(config("first"), lambda: build("first")).lowesttpm_logger_v2)
second_handler: Final = weakref.ref(routers.get(config("second"), lambda: build("second")).lowesttpm_logger_v2)
routers.get(config("first"), must_be_cached)
routers.get(config("third"), lambda: build("third"))
await asyncio.sleep(0.2)
gc.collect()
assert routers.get(config("first"), must_be_cached) is not None
assert first_handler() is not None
assert second_handler() is None
routers.clear()

13
tests/unit/router_strategy/adaptive_router/test_hooks.py
View file

@ -10,6 +10,7 @@ from litellm.router_strategy.adaptive_router.config import (
)
from litellm.router_strategy.adaptive_router.hooks import (
AdaptiveRouterPostCallHook,
AdaptiveRouterResponseHeaderHook,
_recent_tool_results,
_resolve_session_key,
)
@ -337,8 +338,7 @@ async def test_post_call_response_headers_hook_returns_chosen_model_header():
"""The header hook returns the `x-litellm-adaptive-router-model` header
so proxy header construction picks it up (works for both streaming and
non-streaming; `async_post_call_success_hook` is too late for streaming)."""
hook = _make_hook()
headers = await hook.async_post_call_response_headers_hook(
headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook(
data={"metadata": {"adaptive_router_chosen_model": "smart"}},
user_api_key_dict=MagicMock(),
response=MagicMock(),
@ -348,8 +348,7 @@ async def test_post_call_response_headers_hook_returns_chosen_model_header():
@pytest.mark.asyncio
async def test_post_call_response_headers_hook_noop_when_metadata_missing_key():
hook = _make_hook()
headers = await hook.async_post_call_response_headers_hook(
headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook(
data={"metadata": {"litellm_session_id": "sess-A"}},
user_api_key_dict=MagicMock(),
response=MagicMock(),
@ -359,8 +358,7 @@ async def test_post_call_response_headers_hook_noop_when_metadata_missing_key():
@pytest.mark.asyncio
async def test_post_call_response_headers_hook_noop_when_no_metadata():
hook = _make_hook()
headers = await hook.async_post_call_response_headers_hook(
headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook(
data={},
user_api_key_dict=MagicMock(),
response=MagicMock(),
@ -370,8 +368,7 @@ async def test_post_call_response_headers_hook_noop_when_no_metadata():
@pytest.mark.asyncio
async def test_post_call_response_headers_hook_noop_when_metadata_not_dict():
hook = _make_hook()
headers = await hook.async_post_call_response_headers_hook(
headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook(
data={"metadata": "not-a-dict"},
user_api_key_dict=MagicMock(),
response=MagicMock(),

28
tests/unit/router_strategy/test_base_routing_strategy.py
View file

@ -1,15 +1,20 @@
import json
import logging
from typing import Any, Dict, List, Optional, Set, Union
from typing import Any, Dict, Final, List, Optional, Set, Union
import pytest
import asyncio
import threading
from unittest.mock import MagicMock, patch
from fakeredis import TcpFakeServer
from litellm.caching.caching import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, RedisPipelineIncrementOperation
from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy
@ -162,3 +167,24 @@ async def test_push_refused_by_the_open_circuit_breaker_is_not_logged_as_an_erro
assert caplog.records == []
assert base_strategy.redis_increment_operation_queue == []
@pytest.mark.asyncio
async def test_an_increment_counted_after_retire_still_reaches_redis():
server: Final = TcpFakeServer(("127.0.0.1", 0), server_type="redis")
threading.Thread(target=server.serve_forever, daemon=True).start()
redis_cache: Final = RedisCache(host="127.0.0.1", port=server.server_address[1])
strategy: Final = BaseRoutingStrategy(
dual_cache=DualCache(redis_cache=redis_cache, in_memory_cache=InMemoryCache()),
should_batch_redis_writes=True,
default_sync_interval=60,
)
try:
strategy.retire()
await strategy._increment_value_in_current_window("counted-after-retire", 3, 60)
await asyncio.sleep(0.2)
assert float(await redis_cache.async_get_cache("counted-after-retire")) == 3
finally:
server.shutdown()
server.server_close()

57
tests/unit/router_strategy/test_budget_limiter_hotpath.py
View file

@ -1,4 +1,5 @@
import asyncio
import contextlib
import gc
import logging
from types import SimpleNamespace
@ -492,6 +493,8 @@ def _new_router_budget_limiter(
budget_limiter._redis_increment_queue_lock = queue_lock if queue_lock is not None else asyncio.Lock()
budget_limiter._redis_increment_flush_lock = asyncio.Lock()
budget_limiter._detached_increment_operations = None
budget_limiter._retired = False
budget_limiter._flush_after_retire = None
return budget_limiter
@ -784,3 +787,57 @@ async def test_cancelled_flush_does_not_requeue_an_applied_batch(cancellations:
await limiter._push_in_memory_increments_to_redis()
assert redis_cache.values[_SPEND_KEY] == 10.0
assert limiter.redis_increment_operation_queue == []
@pytest.mark.asyncio
async def test_spend_queued_when_a_limiter_is_retired_reaches_redis_after_a_failed_write() -> None:
pipeline_started: Final = asyncio.Event()
pipeline_completed: Final = asyncio.Event()
redis_cache: Final = _MockRedisCache(
initial_values={},
pipeline_started=pipeline_started,
pipeline_completed=pipeline_completed,
should_fail_pipeline=True,
)
budget_limiter: Final = _new_router_budget_limiter(
redis_cache=redis_cache, redis_increment_operation_queue=[_increment(10.0)]
)
budget_limiter._sync_task = asyncio.create_task(asyncio.sleep(3600))
budget_limiter.retire()
await asyncio.wait_for(pipeline_started.wait(), timeout=1)
assert redis_cache.values == {}
redis_cache.should_fail_pipeline = False
await asyncio.wait_for(pipeline_completed.wait(), timeout=5)
assert redis_cache.values == {_SPEND_KEY: 10.0}
assert budget_limiter.redis_increment_operation_queue == []
@pytest.mark.asyncio
async def test_spend_in_flight_when_a_limiter_is_retired_reaches_redis_after_that_write_fails() -> None:
pipeline_started: Final = asyncio.Event()
allow_pipeline_to_complete: Final = asyncio.Event()
pipeline_completed: Final = asyncio.Event()
redis_cache: Final = _MockRedisCache(
initial_values={},
pipeline_started=pipeline_started,
allow_pipeline_to_complete=allow_pipeline_to_complete,
should_fail_pipeline=True,
pipeline_completed=pipeline_completed,
)
budget_limiter: Final = _new_router_budget_limiter(
redis_cache=redis_cache, redis_increment_operation_queue=[_increment(10.0)]
)
budget_limiter._sync_task = asyncio.create_task(budget_limiter.periodic_sync_in_memory_spend_with_redis())
await asyncio.wait_for(pipeline_started.wait(), timeout=1)
budget_limiter.retire()
allow_pipeline_to_complete.set()
with contextlib.suppress(asyncio.CancelledError):
await budget_limiter._sync_task
redis_cache.should_fail_pipeline = False
await asyncio.wait_for(pipeline_completed.wait(), timeout=5)
assert redis_cache.values == {_SPEND_KEY: 10.0}
assert budget_limiter.redis_increment_operation_queue == []

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

@ -1,8 +1,12 @@
import asyncio
import functools
import gc
import json
import os
import weakref
from datetime import datetime, timezone
from collections.abc import Mapping
from collections.abc import Iterator, Mapping
from types import MethodType
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
@ -12,8 +16,12 @@ from litellm import Router
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo
from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook
from litellm.types.litellm_params import RoutingStrategyName
from litellm.types.router import AlertingConfig, Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo
from litellm.types.utils import (
CallTypes,
ModelResponse,
@ -1852,6 +1860,315 @@ def test_discard(model_list):
assert len(litellm.service_callback) == 0
_CALLBACK_LIST_NAMES: Final = (
"callbacks",
"input_callback",
"success_callback",
"failure_callback",
"_async_success_callback",
"_async_failure_callback",
)
def _callback_ids(list_name: str) -> list[int]:
return [id(callback) for callback in getattr(litellm, list_name)]
@pytest.mark.parametrize(
"routing_strategy, selector_attr",
[
("least-busy", "leastbusy_logger"),
("usage-based-routing", "lowesttpm_logger"),
("usage-based-routing-v2", "lowesttpm_logger_v2"),
("latency-based-routing", "lowestlatency_logger"),
("cost-based-routing", "lowestcost_logger"),
],
)
@pytest.mark.asyncio
async def test_discard_releases_strategy_selector_across_router_lifecycles(
routing_strategy: RoutingStrategyName, selector_attr: str
):
def retire() -> weakref.ReferenceType[object]:
router: Final = Router(model_list=[], routing_strategy=routing_strategy)
selector: Final = getattr(router, selector_attr)
assert any(callback is selector for callback in litellm.callbacks)
router.discard()
return weakref.ref(selector)
selector_refs: Final = tuple(retire() for _ in range(3))
await asyncio.sleep(0)
gc.collect()
assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES)
assert all(selector_ref() is None for selector_ref in selector_refs)
@pytest.mark.parametrize(
"routing_strategy, selector_attr",
[
("least-busy", "leastbusy_logger"),
("usage-based-routing", "lowesttpm_logger"),
("usage-based-routing-v2", "lowesttpm_logger_v2"),
("latency-based-routing", "lowestlatency_logger"),
("cost-based-routing", "lowestcost_logger"),
],
)
@pytest.mark.asyncio
async def test_discard_purges_selector_from_event_lists_after_serving_a_request(
routing_strategy: RoutingStrategyName, selector_attr: str
):
router: Final = Router(
model_list=[{"model_name": "gpt-5.5", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}}],
routing_strategy=routing_strategy,
)
selector: Final = getattr(router, selector_attr)
await router.acompletion(model="gpt-5.5", messages=[{"role": "user", "content": "hi"}], mock_response="ok")
assert any(callback is selector for callback in litellm.success_callback)
router.discard()
assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES)
def test_discard_keeps_callbacks_owned_by_other_live_routers():
survivor: Final = Router(model_list=[], routing_strategy="latency-based-routing")
unrelated: Final = CustomLogger()
litellm.logging_callback_manager.add_litellm_callback(unrelated)
before_ids: Final = {list_name: _callback_ids(list_name) for list_name in _CALLBACK_LIST_NAMES}
retired: Final = Router(model_list=[], routing_strategy="least-busy")
assert any(callback is retired.leastbusy_logger for callback in litellm.callbacks)
retired.discard()
assert {list_name: _callback_ids(list_name) for list_name in before_ids} == before_ids
assert any(callback is survivor.lowestlatency_logger for callback in litellm.callbacks)
survivor.discard()
@pytest.mark.parametrize(
"routing_strategy, selector_attr",
[
("usage-based-routing", "lowesttpm_logger"),
("usage-based-routing-v2", "lowesttpm_logger_v2"),
("latency-based-routing", "lowestlatency_logger"),
("cost-based-routing", "lowestcost_logger"),
],
)
@pytest.mark.asyncio
async def test_discard_leaves_same_strategy_survivor_receiving_events(
routing_strategy: RoutingStrategyName, selector_attr: str
):
model_list: Final[list[DeploymentTypedDict]] = [
{"model_name": "gpt-5.5", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}}
]
retired: Final = Router(model_list=model_list, routing_strategy=routing_strategy)
survivor: Final = Router(model_list=model_list, routing_strategy=routing_strategy)
retired.discard()
await survivor.acompletion(model="gpt-5.5", messages=[{"role": "user", "content": "hi"}], mock_response="ok")
assert any(callback is getattr(survivor, selector_attr) for callback in litellm.success_callback)
survivor.discard()
def test_discard_drops_routing_group_selectors():
router: Final = Router(
model_list=[
{"model_name": "grouped-model", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}}
],
routing_groups=[
{"group_name": "fast", "models": ["grouped-model"], "routing_strategy": "latency-based-routing"}
],
)
_, group_selector = router._get_routing_context("grouped-model")
assert any(callback is group_selector for callback in litellm.callbacks)
router.discard()
assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES)
_GROUPED_MODEL_LIST: Final[list[DeploymentTypedDict]] = [
{"model_name": "grouped-model", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}}
]
class _LabelledHook(CustomLogger):
def __init__(self, label: str) -> None:
super().__init__()
self.label = label
def _visit_hooks_reloading_routing_groups_mid_loop(
router: Router, hooks: tuple[_LabelledHook, ...]
) -> Iterator[CustomLogger]:
for callback in litellm.callbacks:
if callback is hooks[0]:
router.update_settings(
routing_groups=[{"group_name": "fast", "models": ["grouped-model"], "routing_strategy": "least-busy"}]
)
if callback in hooks:
yield callback
def test_routing_group_reload_keeps_every_hook_of_an_in_flight_callback_loop():
router: Final = Router(
model_list=_GROUPED_MODEL_LIST,
routing_groups=[{"group_name": "fast", "models": ["grouped-model"], "routing_strategy": "least-busy"}],
)
hooks: Final = tuple(_LabelledHook(label) for label in ("guardrail_a", "guardrail_b", "guardrail_c"))
for hook in hooks:
litellm.logging_callback_manager.add_litellm_callback(hook)
visited: Final = tuple(_visit_hooks_reloading_routing_groups_mid_loop(router, hooks))
assert visited == hooks
router.discard()
def _routers_seen_in_success_loop_discarding_first(first: Router, second: Router) -> Iterator[Router]:
for callback in litellm.success_callback:
if not isinstance(callback, MethodType):
continue
if callback.__self__ is first:
first.discard()
yield first
if callback.__self__ is second:
yield second
def test_discard_mid_success_loop_keeps_the_next_router_callback():
first: Final = Router(model_list=_GROUPED_MODEL_LIST)
second: Final = Router(model_list=_GROUPED_MODEL_LIST)
seen: Final = tuple(_routers_seen_in_success_loop_discarding_first(first, second))
assert seen == (first, second)
second.discard()
@pytest.mark.asyncio
async def test_discard_purges_optional_pre_call_checks_from_event_lists_after_serving_a_request():
router: Final = Router(
model_list=_GROUPED_MODEL_LIST,
optional_pre_call_checks=["enforce_model_rate_limits", "prompt_caching", "deployment_affinity"],
)
await router.acompletion(model="grouped-model", messages=[{"role": "user", "content": "hi"}], mock_response="ok")
router.discard()
assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES)
def test_discard_registers_same_config_survivor_optional_pre_call_checks():
retired: Final = Router(model_list=_GROUPED_MODEL_LIST, optional_pre_call_checks=["enforce_model_rate_limits"])
survivor: Final = Router(model_list=_GROUPED_MODEL_LIST, optional_pre_call_checks=["enforce_model_rate_limits"])
retired.discard()
assert survivor.optional_callbacks
assert all(_is_registered(check) for check in survivor.optional_callbacks)
survivor.discard()
def _is_registered(callback: object) -> bool:
return any(registered is callback for registered in litellm.callbacks)
@pytest.mark.asyncio
async def test_discard_releases_per_request_override_selector():
def build_override_then_discard() -> weakref.ReferenceType[object]:
router: Final = Router(model_list=_GROUPED_MODEL_LIST)
selector: Final = router._get_override_strategy_selector("usage-based-routing-v2")
assert selector is not None
router.discard()
return weakref.ref(selector)
selector_ref: Final = build_override_then_discard()
await asyncio.sleep(0)
gc.collect()
assert selector_ref() is None
def test_discard_registers_the_selector_of_a_same_strategy_survivor_built_without_a_model_list():
retired: Final = Router(model_list=[], routing_strategy="latency-based-routing")
survivor: Final = Router(routing_strategy="latency-based-routing")
assert not _is_registered(survivor.lowestlatency_logger)
retired.discard()
assert _is_registered(survivor.lowestlatency_logger)
survivor.discard()
def test_discard_drops_its_slack_alerting_callbacks():
router: Final = Router(model_list=_GROUPED_MODEL_LIST, alerting_config=AlertingConfig(webhook_url="test"))
assert any(isinstance(callback, SlackAlerting) for callback in litellm.callbacks)
router.discard()
assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES)
def _adaptive_router_model_list(name: str) -> list[DeploymentTypedDict]:
return [
{"model_name": "gpt-5.5", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}},
{
"model_name": name,
"litellm_params": {
"model": "auto_router/adaptive_router",
"adaptive_router_config": {"available_models": ["gpt-5.5"]},
},
},
]
def _registered_adaptive_router_strategies() -> tuple[object, ...]:
return tuple(
callback.adaptive_router for callback in litellm.callbacks if isinstance(callback, AdaptiveRouterPostCallHook)
)
def test_discard_drops_its_adaptive_router_hook():
router: Final = Router(model_list=_adaptive_router_model_list("adaptive"), ignore_invalid_deployments=True)
assert _registered_adaptive_router_strategies()
router.discard()
assert _registered_adaptive_router_strategies() == ()
def test_discard_registers_the_adaptive_router_hook_of_a_survivor():
survivor: Final = Router(model_list=_adaptive_router_model_list("adaptive"), ignore_invalid_deployments=True)
retired: Final = Router(model_list=_adaptive_router_model_list("adaptive"), ignore_invalid_deployments=True)
retired.discard()
survivor_strategies: Final = tuple(tagged.strategy for tagged in survivor.adaptive_routers["adaptive"])
assert _registered_adaptive_router_strategies() == survivor_strategies
survivor.discard()
@pytest.mark.asyncio
async def test_discard_releases_its_router_budget_limiter():
def build_budget_limiter_then_discard() -> weakref.ReferenceType[object]:
router: Final = Router(
model_list=_GROUPED_MODEL_LIST,
provider_budget_config={"openai": {"budget_limit": 10.0, "time_period": "1d"}},
)
budget_limiter: Final = router.get_router_deployment_budget_limiter()
assert budget_limiter is not None
router.discard()
return weakref.ref(budget_limiter)
budget_limiter_ref: Final = build_budget_limiter_then_discard()
await asyncio.sleep(0)
gc.collect()
assert budget_limiter_ref() is None
def test_initialize_assistants_endpoint(model_list):
"""Test if the 'initialize_assistants_endpoint' function is working correctly"""
router = Router(model_list=model_list)

108
tests/unit/test_utils.py
View file

@ -3891,6 +3891,114 @@ class TestCallbackAsyncSyncSeparation:
assert my_sync_cb not in litellm._async_input_callback
class _DropsACallbackFromAnotherThreadWhenComparedToDynamodb:
def __init__(self, dropped: Callable[..., object]) -> None:
self.dropped = dropped
self.drop_thread: threading.Thread | None = None
def __eq__(self, other: object) -> bool:
if isinstance(other, str) and other == "dynamodb" and self.drop_thread is None:
self.drop_thread = threading.Thread(
target=litellm.logging_callback_manager.drop_callbacks, args=((self.dropped,),)
)
self.drop_thread.start()
self.drop_thread.join(timeout=0.2)
return self is other
__hash__ = object.__hash__
def test_function_setup_moves_async_callbacks_while_another_thread_drops_a_callback(
monkeypatch: pytest.MonkeyPatch,
) -> None:
def dropped_callback(*args: object, **kwargs: object) -> None:
return None
async def async_callback(*args: object, **kwargs: object) -> None:
return None
drops_on_compare: Final = _DropsACallbackFromAnotherThreadWhenComparedToDynamodb(dropped_callback)
monkeypatch.setattr(litellm.utils, "callback_list", ["initialized"])
for list_name in (
"callbacks",
"input_callback",
"failure_callback",
"service_callback",
"_async_input_callback",
"_async_success_callback",
"_async_failure_callback",
):
monkeypatch.setattr(litellm, list_name, [])
monkeypatch.setattr(litellm, "success_callback", [dropped_callback, drops_on_compare, async_callback])
litellm.utils.function_setup(
"completion",
litellm.utils.Rules(),
datetime.now(),
model="gpt-4o",
messages=[{"role": "user", "content": "hi"}],
litellm_call_id="callback-drop-race",
)
assert drops_on_compare.drop_thread is not None
drops_on_compare.drop_thread.join()
assert litellm.success_callback == [drops_on_compare]
assert litellm._async_success_callback == [async_callback]
class _DropsACallbackFromAnotherThreadWhenComparedToSentry:
def __init__(self, dropped: Callable[..., object]) -> None:
self.dropped = dropped
self.drop_thread: threading.Thread | None = None
self.drop_finished_during_logger_setup = False
def __eq__(self, other: object) -> bool:
if isinstance(other, str) and other == "sentry" and self.drop_thread is None:
self.drop_thread = threading.Thread(
target=litellm.logging_callback_manager.drop_callbacks, args=((self.dropped,),)
)
self.drop_thread.start()
self.drop_thread.join(timeout=5)
self.drop_finished_during_logger_setup = not self.drop_thread.is_alive()
return self is other
__hash__ = object.__hash__
def test_function_setup_lets_other_threads_drop_callbacks_while_it_sets_up_loggers(
monkeypatch: pytest.MonkeyPatch,
) -> None:
def dropped_callback(*args: object, **kwargs: object) -> None:
return None
drops_on_compare: Final = _DropsACallbackFromAnotherThreadWhenComparedToSentry(dropped_callback)
monkeypatch.setattr(litellm.utils, "callback_list", [])
for list_name in (
"callbacks",
"input_callback",
"failure_callback",
"service_callback",
"_async_input_callback",
"_async_success_callback",
"_async_failure_callback",
):
monkeypatch.setattr(litellm, list_name, [])
monkeypatch.setattr(litellm, "success_callback", [drops_on_compare])
litellm.utils.function_setup(
"completion",
litellm.utils.Rules(),
datetime.now(),
model="gpt-4o",
messages=[{"role": "user", "content": "hi"}],
litellm_call_id="logger-setup-outside-lock",
)
assert drops_on_compare.drop_thread is not None
drops_on_compare.drop_thread.join()
assert drops_on_compare.drop_finished_during_logger_setup
class TestMetadataNoneHandling:
"""
Test that metadata=None in kwargs doesn't cause TypeError.