scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 129

Merge pull request #36263 from BerriAI/litellm_replica_registry_read_through

Browse source 
fix(proxy): read through to the DB on registry misses so just-created models, guardrails, and agents resolve on sibling replicas
This commit is contained in:
Mateo Wang 2026-08-19 10:46:35 -07:00 • committed by GitHub
commit 133e72c8fd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
18 changed files with 1514 additions and 195 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -168,14 +168,13 @@ def _jsonrpc_error(
)
def _get_agent(agent_id: str):
async def _get_agent(agent_id: str) -> "AgentResponse | None":
"""Look up an agent by ID or name. Returns None if not found."""
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
from litellm.proxy.common_utils.registry_read_through import (
get_agent_with_read_through,
)
agent = global_agent_registry.get_agent_by_id(agent_id=agent_id)
if agent is None:
agent = global_agent_registry.get_agent_by_name(agent_name=agent_id)
return agent
return await get_agent_with_read_through(agent_id)
def _enforce_inbound_trace_id(agent: "AgentResponse", request: Request) -> None:
@ -559,7 +558,7 @@ async def get_agent_card(
)
try:
agent: Final = _get_agent(agent_id)
agent: Final = await _get_agent(agent_id)
if agent is None:
raise HTTPException(status_code=404, detail=f"Agent '{agent_id}' not found")
@ -673,7 +672,7 @@ async def invoke_agent_a2a(
params.pop(key)
# Find the agent
agent: Final = _get_agent(agent_id)
agent: Final = await _get_agent(agent_id)
if agent is None:
return _jsonrpc_error(request_id, -32000, f"Agent '{agent_id}' not found", 404)

10
litellm/proxy/agent_endpoints/a2a_routing.py
View file

@ -25,10 +25,12 @@ async def route_a2a_agent_request(
Returns None if not an A2A request (allows normal routing to continue).
"""
# Import here to avoid circular imports
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
AgentRequestHandler,
)
from litellm.proxy.common_utils.registry_read_through import (
get_agent_with_read_through,
)
from litellm.proxy.route_llm_request import (
ROUTE_ENDPOINT_MAPPING,
ProxyModelNotFoundError,
@ -44,11 +46,11 @@ async def route_a2a_agent_request(
agent_name: Final = model_name[4:]
# Look up agent in registry
agent: Final = global_agent_registry.get_agent_by_name(agent_name)
agent: Final = await get_agent_with_read_through(agent_name)
if agent is None:
verbose_proxy_logger.error("[A2A] Agent '%s' not found in registry", agent_name)
route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type)
raise ProxyModelNotFoundError(route=route_name, model_name=model_name)
raise ProxyModelNotFoundError(route=route_name, model_name=model_name, retryable_with_model_read_through=False)
# Verify the caller is permitted to use this agent (admins bypass the check)
is_admin: Final = user_api_key_dict is not None and (
@ -70,7 +72,7 @@ async def route_a2a_agent_request(
if not agent.agent_card_params or "url" not in agent.agent_card_params:
verbose_proxy_logger.error("[A2A] Agent '%s' has no URL configured", agent_name)
route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type)
raise ProxyModelNotFoundError(route=route_name, model_name=model_name)
raise ProxyModelNotFoundError(route=route_name, model_name=model_name, retryable_with_model_read_through=False)
# Inject API base and route to litellm
data["api_base"] = agent.agent_card_params["url"]

1
litellm/proxy/agent_endpoints/agent_registry.py
View file

@ -600,3 +600,4 @@ class AgentRegistry:
global_agent_registry: Final = AgentRegistry()
AGENT_RECONCILE_LOCK: Final = asyncio.Lock()

224
litellm/proxy/common_utils/registry_read_through.py Normal file
View file

@ -0,0 +1,224 @@
"""Read-through recovery for in-memory registries in multi-replica deployments.
A management write (POST /model/new, /guardrails, /v1/agents) lands on one
replica and reaches Postgres, but sibling replicas only refresh their in-memory
registries on the periodic config reload, so a request using the new object
immediately can land on a sibling that has never heard of it and fail 400/404.
On a registry miss, callers here fetch the missing row from the DB and load it
into the local registry before giving up. A short negative-result TTL per key
plus a global resync budget per window bound the DB load from lookups of
genuinely unknown names.
"""
import asyncio
import time
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING, Final
from litellm._logging import verbose_proxy_logger
from litellm.caching.in_memory_cache import InMemoryCache
if TYPE_CHECKING:
from prisma.types import (
LiteLLM_AgentsTableInclude,
LiteLLM_AgentsTableWhereUniqueInput,
LiteLLM_GuardrailsTableWhereInput,
LiteLLM_ProxyModelTableWhereInput,
)
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.types.agents import AgentResponse
READ_THROUGH_MISS_TTL_SECONDS: Final = 2.0
READ_THROUGH_RESYNC_WINDOW_SECONDS: Final = 5.0
READ_THROUGH_MAX_RESYNCS_PER_WINDOW: Final = 20
class RegistryReadThrough:
__slots__ = (
"_lock",
"_max_resyncs_per_window",
"_miss_ttl_seconds",
"_recent_misses",
"_resync",
"_resync_window_seconds",
"_window_resyncs",
"_window_started_at",
)
def __init__(
self,
resync: Callable[[str], Awaitable[bool]],
miss_ttl_seconds: float = READ_THROUGH_MISS_TTL_SECONDS,
max_resyncs_per_window: int = READ_THROUGH_MAX_RESYNCS_PER_WINDOW,
resync_window_seconds: float = READ_THROUGH_RESYNC_WINDOW_SECONDS,
) -> None:
self._resync = resync
self._miss_ttl_seconds = miss_ttl_seconds
self._max_resyncs_per_window = max_resyncs_per_window
self._resync_window_seconds = resync_window_seconds
self._lock = asyncio.Lock()
self._recent_misses = InMemoryCache(max_size_in_memory=1000)
self._window_started_at = float("-inf")
self._window_resyncs = 0
def _consume_resync_budget(self) -> bool:
now: Final = time.monotonic()
if now - self._window_started_at >= self._resync_window_seconds:
self._window_started_at = now
self._window_resyncs = 0
if self._window_resyncs >= self._max_resyncs_per_window:
return False
self._window_resyncs += 1
return True
async def attempt(self, key: str) -> bool:
if self._recent_misses.get_cache(key) is not None:
return False
async with self._lock:
if self._recent_misses.get_cache(key) is not None:
return False
if not self._consume_resync_budget():
verbose_proxy_logger.warning(
"registry read-through for %r skipped: resync budget of %s per %ss exhausted",
key,
self._max_resyncs_per_window,
self._resync_window_seconds,
)
return False
try:
found: Final = await self._resync(key)
except Exception as e: # noqa: BLE001 # a failed read-through must surface the original miss error, not a 500
verbose_proxy_logger.warning("registry read-through for %r failed: %s", key, e)
return False
if not found:
self._recent_misses.set_cache(key, True, ttl=self._miss_ttl_seconds)
return found
def _db_backed_registries_enabled(object_type: str) -> bool:
from litellm.proxy import proxy_server
if proxy_server.prisma_client is None or proxy_server.store_model_in_db is not True:
return False
return proxy_server.should_load_db_object(object_type=object_type)
async def _resync_model_deployments(model_name: str) -> bool:
from litellm.proxy import proxy_server
from litellm.repositories.model_repository import ModelRepository
if not _db_backed_registries_enabled("models"):
return False
prisma_client: Final = proxy_server.prisma_client
assert prisma_client is not None
table: Final = ModelRepository(prisma_client).table
name_filter: Final[LiteLLM_ProxyModelTableWhereInput] = {"model_name": model_name}
id_filter: Final[LiteLLM_ProxyModelTableWhereInput] = {"model_id": model_name}
rows: Final = await table.find_many(where=name_filter) or await table.find_many(where=id_filter)
if not rows:
return False
router: Final = proxy_server.llm_router
if router is None:
await proxy_server.proxy_config.add_deployment(
prisma_client=prisma_client, proxy_logging_obj=proxy_server.proxy_logging_obj
)
return proxy_server.llm_router is not None
async with proxy_server.MODEL_RECONCILE_LOCK:
proxy_server.proxy_config._add_deployment(db_models=rows)
proxy_server.llm_model_list = router.get_model_list()
return True
async def _resync_guardrails(guardrail_name: str) -> bool:
from litellm.proxy import proxy_server
from litellm.proxy.guardrails.guardrail_registry import (
GUARDRAIL_RECONCILE_LOCK,
IN_MEMORY_GUARDRAIL_HANDLER,
)
from litellm.repositories.table_repositories import GuardrailsRepository
from litellm.types.guardrails import Guardrail
if not _db_backed_registries_enabled("guardrails"):
return False
prisma_client: Final = proxy_server.prisma_client
assert prisma_client is not None
active_row_filter: Final[LiteLLM_GuardrailsTableWhereInput] = {
"guardrail_name": guardrail_name,
"status": "active",
}
row: Final = await GuardrailsRepository(prisma_client).table.find_first(where=active_row_filter)
if row is None:
return False
async with GUARDRAIL_RECONCILE_LOCK:
IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(guardrail=Guardrail(**dict(row)))
return _initialized_guardrail(guardrail_name) is not None
async def _resync_agents(agent_id_or_name: str) -> bool:
from litellm.proxy import proxy_server
from litellm.proxy.agent_endpoints.agent_registry import (
AGENT_RECONCILE_LOCK,
agents_table,
global_agent_registry,
)
from litellm.types.agents import AgentResponse
if not _db_backed_registries_enabled("agents"):
return False
if _agent_from_registry(agent_id_or_name) is not None:
return True
prisma_client: Final = proxy_server.prisma_client
assert prisma_client is not None
table: Final = agents_table(prisma_client)
id_filter: Final[LiteLLM_AgentsTableWhereUniqueInput] = {"agent_id": agent_id_or_name}
name_filter: Final[LiteLLM_AgentsTableWhereUniqueInput] = {"agent_name": agent_id_or_name}
include_permission: Final[LiteLLM_AgentsTableInclude] = {"object_permission": True}
async with AGENT_RECONCILE_LOCK:
if _agent_from_registry(agent_id_or_name) is not None:
return True
row: Final = await table.find_unique(where=id_filter, include=include_permission) or await table.find_unique(
where=name_filter, include=include_permission
)
if row is None:
return False
global_agent_registry.register_agent(agent_config=AgentResponse.model_validate(row.model_dump()))
return True
model_registry_read_through: Final = RegistryReadThrough(resync=_resync_model_deployments)
guardrail_registry_read_through: Final = RegistryReadThrough(resync=_resync_guardrails)
agent_registry_read_through: Final = RegistryReadThrough(resync=_resync_agents)
def _agent_from_registry(agent_id_or_name: str) -> "AgentResponse | None":
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
by_id: Final = global_agent_registry.get_agent_by_id(agent_id=agent_id_or_name)
if by_id is not None:
return by_id
return global_agent_registry.get_agent_by_name(agent_name=agent_id_or_name)
async def get_agent_with_read_through(agent_id_or_name: str) -> "AgentResponse | None":
agent: Final = _agent_from_registry(agent_id_or_name)
if agent is not None:
return agent
if not await agent_registry_read_through.attempt(agent_id_or_name):
return None
return _agent_from_registry(agent_id_or_name)
def _initialized_guardrail(guardrail_name: str) -> "CustomGuardrail | None":
from litellm.proxy.guardrails import guardrail_endpoints
return guardrail_endpoints.GUARDRAIL_REGISTRY.get_initialized_guardrail_callback(guardrail_name=guardrail_name)
async def get_initialized_guardrail_with_read_through(guardrail_name: str) -> "CustomGuardrail | None":
active: Final = _initialized_guardrail(guardrail_name)
if active is not None:
return active
if not await guardrail_registry_read_through.attempt(guardrail_name):
return None
return _initialized_guardrail(guardrail_name)

6
litellm/proxy/guardrails/guardrail_endpoints.py
View file

@ -2305,8 +2305,12 @@ async def apply_guardrail(
litellm_logging_obj = None
start_time: Final = datetime.now(timezone.utc)
from litellm.proxy.common_utils.registry_read_through import (
get_initialized_guardrail_with_read_through,
)
try:
active_guardrail: Final[CustomGuardrail | None] = GUARDRAIL_REGISTRY.get_initialized_guardrail_callback(
active_guardrail: Final[CustomGuardrail | None] = await get_initialized_guardrail_with_read_through(
guardrail_name=request.guardrail_name
)
if active_guardrail is None:

3
litellm/proxy/guardrails/guardrail_registry.py
View file

@ -1,5 +1,6 @@
# litellm/proxy/guardrails/guardrail_registry.py
import asyncio
import importlib
import os
from collections.abc import Callable, Iterator, Mapping
@ -813,4 +814,6 @@ class InMemoryGuardrailHandler:
# In Memory Guardrail Handler for LiteLLM Proxy
########################################################
IN_MEMORY_GUARDRAIL_HANDLER: Final = InMemoryGuardrailHandler()
GUARDRAIL_RECONCILE_LOCK: Final = asyncio.Lock()
########################################################

57
litellm/proxy/management_endpoints/model_access_group_management_endpoints.py
View file

@ -57,20 +57,37 @@ def _model_table(prisma_client: PrismaClient) -> _ModelTableClient:
return ModelRepository(prisma_client).table
def validate_models_exist(model_names: list[str], llm_router: "Router | None") -> tuple[bool, list[str]]:
def validate_models_exist(model_names: Sequence[str], llm_router: "Router | None") -> tuple[bool, Sequence[str]]:
"""
Validate that all requested model names exist in the router.
Checks only exact model name matches.
Returns:
Tuple[bool, List[str]]: (all_valid, missing_models)
(all_valid, missing_models)
"""
if llm_router is None:
return False, model_names
router_model_names: Final = set(llm_router.get_model_names())
missing: Final = [m for m in model_names if m not in router_model_names]
return (len(missing) == 0, missing)
router_model_names: Final = frozenset(llm_router.get_model_names())
missing: Final = tuple(m for m in model_names if m not in router_model_names)
return (not missing, missing)
async def _missing_models_after_read_through(
model_names: Sequence[str], llm_router: "Router | None"
) -> tuple[str, ...]:
from litellm.proxy import proxy_server
from litellm.proxy.common_utils.registry_read_through import (
model_registry_read_through,
)
_, missing = validate_models_exist(model_names=model_names, llm_router=llm_router)
if not missing:
return ()
for name in missing:
await model_registry_read_through.attempt(name)
_, still_missing = validate_models_exist(model_names=model_names, llm_router=proxy_server.llm_router)
return tuple(still_missing)
def add_access_group_to_deployment(model_info: dict[str, Any], access_group: str) -> tuple[dict[str, Any], bool]:
@ -101,13 +118,21 @@ def _raise_http_if_reload_degraded_serving(
before: frozenset[str],
written_models: Sequence[tuple[str, object]],
access_group: str,
still_desired: frozenset[str] | None,
live_after: frozenset[str] | None,
) -> None:
"""Same verdict as the model-write endpoints, expressed through this file's
HTTPException error convention, with the metadata-only obligation: these writes
change group membership, not the models themselves, so a row that was already not
serving before the reload is never blamed here; only a model this reload stopped
serving is reported."""
missing, collateral = reload_serving_verdict(before=before, written_models=written_models, written_must_serve=False)
missing, collateral = reload_serving_verdict(
before=before,
written_models=written_models,
written_must_serve=False,
still_desired=still_desired,
live_after=live_after,
)
gone: Final = tuple(dict.fromkeys((*missing, *collateral)))
if not gone:
return
@ -390,12 +415,12 @@ async def create_model_group(
# Validate model_names exist in router (only if using model_names path)
if not use_model_ids and has_model_names:
assert data.model_names is not None
all_valid, missing_models = validate_models_exist(
missing_models: Final = await _missing_models_after_read_through(
model_names=data.model_names,
llm_router=llm_router,
)
if not all_valid:
if missing_models:
raise HTTPException(
status_code=400,
detail={"error": f"Model(s) not found: {', '.join(missing_models)}"},
@ -439,11 +464,13 @@ async def create_model_group(
live_before_reload: Final = live_model_ids_snapshot()
await clear_cache()
reload_outcome: Final = await clear_cache()
_raise_http_if_reload_degraded_serving(
before=live_before_reload,
written_models=updated_pairs,
access_group=data.access_group,
still_desired=reload_outcome.still_desired,
live_after=reload_outcome.live_after,
)
verbose_proxy_logger.info(
@ -654,12 +681,12 @@ async def update_access_group(
# Validation: Check if all new models exist (only if using model_names path)
if not use_model_ids and has_model_names:
assert data.model_names is not None
all_valid, missing_models = validate_models_exist(
missing_models: Final = await _missing_models_after_read_through(
model_names=data.model_names,
llm_router=llm_router,
)
if not all_valid:
if missing_models:
raise HTTPException(
status_code=400,
detail={"error": f"Model(s) not found: {', '.join(missing_models)}"},
@ -699,11 +726,13 @@ async def update_access_group(
# Clear cache and reload models to pick up the access group changes
live_before_reload: Final = live_model_ids_snapshot()
await clear_cache()
reload_outcome: Final = await clear_cache()
_raise_http_if_reload_degraded_serving(
before=live_before_reload,
written_models=list({**dict(stripped_pairs), **dict(updated_pairs)}.items()),
access_group=access_group,
still_desired=reload_outcome.still_desired,
live_after=reload_outcome.live_after,
)
verbose_proxy_logger.info(
@ -801,11 +830,13 @@ async def delete_access_group(
# Clear cache and reload models to pick up the access group changes
live_before_reload: Final = live_model_ids_snapshot()
await clear_cache()
reload_outcome: Final = await clear_cache()
_raise_http_if_reload_degraded_serving(
before=live_before_reload,
written_models=removed_pairs,
access_group=access_group,
still_desired=reload_outcome.still_desired,
live_after=reload_outcome.live_after,
)
verbose_proxy_logger.info(

116
litellm/proxy/proxy_server.py
View file

@ -4120,6 +4120,31 @@ def _swap_in_model_cost_map(new_model_cost_map: dict) -> int:
return fetched_model_count
def should_load_db_object(object_type: str | SupportedDBObjectType) -> bool:
"""
Check if an object type should be loaded from the database based on general_settings.supported_db_objects.
Args:
object_type: Type of object to check (e.g., SupportedDBObjectType.MODELS, "models", etc.)
Returns:
True if the object should be loaded, False otherwise
"""
supported_db_objects: Final = general_settings.get("supported_db_objects", None)
if supported_db_objects is None:
return True
if not isinstance(supported_db_objects, list):
verbose_proxy_logger.warning(
"supported_db_objects is not a list, got %s. Loading all objects.", type(supported_db_objects)
)
return True
object_type_str: Final = str(object_type)
return any(str(obj) == object_type_str for obj in supported_db_objects)
class ProxyConfig:
"""
Abstraction class on top of config loading/updating logic. Gives us one place to control all config updating logic.
@ -6522,36 +6547,7 @@ class ProxyConfig:
return config
def _should_load_db_object(self, object_type: str | SupportedDBObjectType) -> bool:
"""
Check if an object type should be loaded from the database based on general_settings.supported_db_objects.
Args:
object_type: Type of object to check (e.g., SupportedDBObjectType.MODELS, "models", etc.)
Returns:
True if the object should be loaded, False otherwise
"""
global general_settings
# Get the supported_db_objects configuration
supported_db_objects: Final = general_settings.get("supported_db_objects", None)
# If supported_db_objects is not set, load all objects (default behavior)
if supported_db_objects is None:
return True
# If supported_db_objects is set, only load specified objects
if not isinstance(supported_db_objects, list):
verbose_proxy_logger.warning(
"supported_db_objects is not a list, got %s. Loading all objects.", type(supported_db_objects)
)
return True
# Convert object_type to string for comparison (handles both str and enum)
object_type_str: Final = str(object_type)
# Check if the object type is in the list (supports both str and enum values)
return any(str(obj) == object_type_str for obj in supported_db_objects)
return should_load_db_object(object_type=object_type)
async def _get_models_from_db(self, prisma_client: PrismaClient) -> list | None:
"""
@ -7094,38 +7090,40 @@ class ProxyConfig:
async def _init_guardrails_in_db(self, prisma_client: PrismaClient):
from litellm.proxy.guardrails.guardrail_registry import (
GUARDRAIL_RECONCILE_LOCK,
IN_MEMORY_GUARDRAIL_HANDLER,
Guardrail,
GuardrailRegistry,
)
try:
guardrails_in_db: Final[list[Guardrail]] = await GuardrailRegistry.get_all_guardrails_from_db(
prisma_client=prisma_client
)
verbose_proxy_logger.debug("guardrails from the DB %s", str(guardrails_in_db))
db_guardrail_ids: Final[set] = set()
for guardrail in guardrails_in_db:
guardrail_id = guardrail.get("guardrail_id")
if guardrail_id:
db_guardrail_ids.add(guardrail_id)
try:
IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(
guardrail=cast(Guardrail, guardrail),
)
except Exception as e: # noqa: BLE001 # one unloadable row must not stop the remaining guardrails
verbose_proxy_logger.error(
"litellm.proxy.proxy_server.py::ProxyConfig:_init_guardrails_in_db - "
"skipping guardrail '%s' (ID: %s): %s: %s",
guardrail.get("guardrail_name"),
guardrail_id,
type(e).__name__,
e,
)
async with GUARDRAIL_RECONCILE_LOCK:
guardrails_in_db: Final[list[Guardrail]] = await GuardrailRegistry.get_all_guardrails_from_db(
prisma_client=prisma_client
)
verbose_proxy_logger.debug("guardrails from the DB %s", str(guardrails_in_db))
db_guardrail_ids: Final[set] = set()
for guardrail in guardrails_in_db:
guardrail_id = guardrail.get("guardrail_id")
if guardrail_id:
db_guardrail_ids.add(guardrail_id)
try:
IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(
guardrail=cast(Guardrail, guardrail),
)
except Exception as e: # noqa: BLE001 # one unloadable row must not stop the remaining guardrails
verbose_proxy_logger.error(
"litellm.proxy.proxy_server.py::ProxyConfig:_init_guardrails_in_db - "
"skipping guardrail '%s' (ID: %s): %s: %s",
guardrail.get("guardrail_name"),
guardrail_id,
type(e).__name__,
e,
)
# Drop in-memory DB-backed entries whose row was deleted on another
# pod. Config-loaded entries are never touched.
IN_MEMORY_GUARDRAIL_HANDLER.reconcile_db_guardrails(db_guardrail_ids=db_guardrail_ids)
# Drop in-memory DB-backed entries whose row was deleted on another
# pod. Config-loaded entries are never touched.
IN_MEMORY_GUARDRAIL_HANDLER.reconcile_db_guardrails(db_guardrail_ids=db_guardrail_ids)
except Exception as e:
verbose_proxy_logger.exception("litellm.proxy.proxy_server.py::ProxyConfig:_init_guardrails_in_db - %s", e)
@ -7278,13 +7276,17 @@ class ProxyConfig:
)
async def _init_agents_in_db(self, prisma_client: PrismaClient):
from litellm.proxy.agent_endpoints.agent_registry import (
AGENT_RECONCILE_LOCK,
)
from litellm.proxy.agent_endpoints.agent_registry import (
global_agent_registry as AGENT_REGISTRY,
)
try:
db_agents: Final = await AGENT_REGISTRY.get_all_agents_from_db(prisma_client=prisma_client)
AGENT_REGISTRY.load_agents_from_db_and_config(db_agents=db_agents)
async with AGENT_RECONCILE_LOCK:
db_agents: Final = await AGENT_REGISTRY.get_all_agents_from_db(prisma_client=prisma_client)
AGENT_REGISTRY.load_agents_from_db_and_config(db_agents=db_agents)
except Exception as e:
verbose_proxy_logger.exception("litellm.proxy.proxy_server.py::ProxyConfig:_init_agents_in_db - %s", e)

235
litellm/proxy/route_llm_request.py
View file

@ -146,7 +146,8 @@ ROUTE_ENDPOINT_MAPPING: Final = {
class ProxyModelNotFoundError(HTTPException):
def __init__(self, route: str, model_name: str):
def __init__(self, route: str, model_name: str, retryable_with_model_read_through: bool = True):
self.retryable_with_model_read_through: Final = retryable_with_model_read_through
detail: Final = {
"error": f"{route}: Invalid model name passed in model={model_name}. Call `/v1/models` to view available models for your key."
}
@ -320,112 +321,150 @@ async def add_shared_session_to_data(data: dict) -> None:
pass
RouteType = Literal[
"acompletion",
"atext_completion",
"aembedding",
"aimage_generation",
"aspeech",
"atranscription",
"amoderation",
"arerank",
"aresponses",
"aget_responses",
"adelete_responses",
"acancel_responses",
"acompact_responses",
"acreate_response_reply",
"alist_input_items",
"_arealtime", # private function for realtime API
"acreate_realtime_client_secret",
"arealtime_calls",
"acreate_realtime_transcription_session",
"_aresponses_websocket", # private function for responses WebSocket mode
"aimage_edit",
"agenerate_content",
"agenerate_content_stream",
"allm_passthrough_route",
"acreate_batch",
"aretrieve_batch",
"alist_batches",
"afile_content",
"afile_retrieve",
"acreate_fine_tuning_job",
"acancel_fine_tuning_job",
"alist_fine_tuning_jobs",
"aretrieve_fine_tuning_job",
"avector_store_search",
"avector_store_create",
"avector_store_retrieve",
"avector_store_list",
"avector_store_update",
"avector_store_delete",
"avector_store_file_create",
"avector_store_file_list",
"avector_store_file_retrieve",
"avector_store_file_content",
"avector_store_file_update",
"avector_store_file_delete",
"aocr",
"asearch",
"avideo_generation",
"avideo_list",
"avideo_status",
"avideo_content",
"avideo_remix",
"avideo_create_character",
"avideo_get_character",
"avideo_edit",
"avideo_extension",
"acreate_container",
"alist_containers",
"aretrieve_container",
"adelete_container",
"aupload_container_file",
"alist_container_files",
"aretrieve_container_file",
"adelete_container_file",
"aretrieve_container_file_content",
"acreate_skill",
"alist_skills",
"aget_skill",
"adelete_skill",
"aingest",
"anthropic_messages",
"acreate_interaction",
"aget_interaction",
"adelete_interaction",
"acancel_interaction",
"acreate_agent",
"alist_agents",
"aget_agent",
"adelete_agent",
"alist_agent_versions",
"asend_message",
"call_mcp_tool",
"acancel_batch",
"afile_delete",
"acreate_eval",
"alist_evals",
"aget_eval",
"aupdate_eval",
"adelete_eval",
"acancel_eval",
"acreate_run",
"alist_runs",
"aget_run",
"acancel_run",
"adelete_run",
]
async def route_request(
data: dict,
llm_router: LitellmRouter | None,
user_model: str | None,
route_type: Literal[
"acompletion",
"atext_completion",
"aembedding",
"aimage_generation",
"aspeech",
"atranscription",
"amoderation",
"arerank",
"aresponses",
"aget_responses",
"adelete_responses",
"acancel_responses",
"acompact_responses",
"acreate_response_reply",
"alist_input_items",
"_arealtime", # private function for realtime API
"acreate_realtime_client_secret",
"arealtime_calls",
"acreate_realtime_transcription_session",
"_aresponses_websocket", # private function for responses WebSocket mode
"aimage_edit",
"agenerate_content",
"agenerate_content_stream",
"allm_passthrough_route",
"acreate_batch",
"aretrieve_batch",
"alist_batches",
"afile_content",
"afile_retrieve",
"acreate_fine_tuning_job",
"acancel_fine_tuning_job",
"alist_fine_tuning_jobs",
"aretrieve_fine_tuning_job",
"avector_store_search",
"avector_store_create",
"avector_store_retrieve",
"avector_store_list",
"avector_store_update",
"avector_store_delete",
"avector_store_file_create",
"avector_store_file_list",
"avector_store_file_retrieve",
"avector_store_file_content",
"avector_store_file_update",
"avector_store_file_delete",
"aocr",
"asearch",
"avideo_generation",
"avideo_list",
"avideo_status",
"avideo_content",
"avideo_remix",
"avideo_create_character",
"avideo_get_character",
"avideo_edit",
"avideo_extension",
"acreate_container",
"alist_containers",
"aretrieve_container",
"adelete_container",
"aupload_container_file",
"alist_container_files",
"aretrieve_container_file",
"adelete_container_file",
"aretrieve_container_file_content",
"acreate_skill",
"alist_skills",
"aget_skill",
"adelete_skill",
"aingest",
"anthropic_messages",
"acreate_interaction",
"aget_interaction",
"adelete_interaction",
"acancel_interaction",
"acreate_agent",
"alist_agents",
"aget_agent",
"adelete_agent",
"alist_agent_versions",
"asend_message",
"call_mcp_tool",
"acancel_batch",
"afile_delete",
"acreate_eval",
"alist_evals",
"aget_eval",
"aupdate_eval",
"adelete_eval",
"acancel_eval",
"acreate_run",
"alist_runs",
"aget_run",
"acancel_run",
"adelete_run",
],
route_type: RouteType,
user_api_key_dict: UserAPIKeyAuth | None = None,
):
"""
Common helper to route the request
"""
try:
return await _route_request_single_attempt(
data=data,
llm_router=llm_router,
user_model=user_model,
route_type=route_type,
user_api_key_dict=user_api_key_dict,
)
except ProxyModelNotFoundError as e:
requested_model: Final = data.get("model", "")
if not e.retryable_with_model_read_through or not isinstance(requested_model, str) or not requested_model:
raise
from litellm.proxy import proxy_server
from litellm.proxy.common_utils.registry_read_through import (
model_registry_read_through,
)
if not await model_registry_read_through.attempt(requested_model):
raise
return await _route_request_single_attempt(
data=data,
llm_router=proxy_server.llm_router,
user_model=user_model,
route_type=route_type,
user_api_key_dict=user_api_key_dict,
)
async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited provider coroutines; the inferred union keeps route_request's callers typed
data: dict, # mutable-ok: request body is the proxy-wide mutable dict contract shared with route_request
llm_router: LitellmRouter | None,
user_model: str | None,
route_type: RouteType,
user_api_key_dict: UserAPIKeyAuth | None = None,
):
raise_if_required_body_param_missing(route_type=route_type, data=data)
await add_shared_session_to_data(data)

27
litellm/router.py
View file

@ -8573,11 +8573,9 @@ class Router:
Returns:
- The added/updated deployment
"""
_deployment_model_id: Final = deployment.model_info.id or ""
_deployment_on_router: Final[Deployment | None] = self.get_deployment(model_id=_deployment_model_id)
try:
# check if deployment already exists
_deployment_model_id: Final = deployment.model_info.id or ""
_deployment_on_router: Final[Deployment | None] = self.get_deployment(model_id=_deployment_model_id)
if _deployment_on_router is not None:
# deployment with this model_id exists on the router
if (
@ -8628,10 +8626,31 @@ class Router:
deployment.model_info.id,
e,
)
self._restore_deployment_after_failed_upsert(
previous_deployment=_deployment_on_router, model_id=_deployment_model_id
)
return None
else:
raise e
def _restore_deployment_after_failed_upsert(self, previous_deployment: Deployment | None, model_id: str) -> None:
if previous_deployment is None or self.has_model_id(model_id):
return
try:
self.add_deployment(deployment=previous_deployment)
verbose_router_logger.info(
"Restored deployment %s (id=%s); it keeps serving its previous configuration.",
previous_deployment.model_name,
model_id,
)
except Exception as restore_error: # noqa: BLE001 # best-effort restore: a second failure must not abort the reload
verbose_router_logger.warning(
"Could not restore previously served deployment %s (id=%s) after the failed upsert: %s",
previous_deployment.model_name,
model_id,
restore_error,
)
@staticmethod
def _backend_cost_map_keys(model: str, custom_llm_provider: str | None) -> tuple[str, ...]:
"""The ``litellm.model_cost`` keys a deployment's shared backend info is registered under."""

2
ruff-strict-budget.json
View file

@ -12,7 +12,7 @@
"limit": 2017
},
"ANN202": {
"limit": 855
"limit": 854
},
"ANN204": {
"limit": 711

476
tests/test_litellm/proxy/common_utils/test_registry_read_through.py Normal file
View file

@ -0,0 +1,476 @@
import asyncio
from typing import Final
import pytest
from litellm.proxy.common_utils.registry_read_through import RegistryReadThrough
class ResyncSpy:
def __init__(self, found: bool = True, error: Exception | None = None) -> None:
self.found = found
self.error = error
self.calls: list[str] = []
async def __call__(self, key: str) -> bool:
self.calls.append(key)
if self.error is not None:
raise self.error
return self.found
@pytest.mark.asyncio
async def test_attempt_returns_true_when_resync_finds_object():
spy: Final = ResyncSpy(found=True)
read_through: Final = RegistryReadThrough(resync=spy)
assert await read_through.attempt("new-model") is True
assert spy.calls == ["new-model"]
@pytest.mark.asyncio
async def test_attempt_found_key_is_not_negative_cached():
spy: Final = ResyncSpy(found=True)
read_through: Final = RegistryReadThrough(resync=spy)
assert await read_through.attempt("new-model") is True
assert await read_through.attempt("new-model") is True
assert spy.calls == ["new-model", "new-model"]
@pytest.mark.asyncio
async def test_missing_key_is_negative_cached_within_ttl():
spy: Final = ResyncSpy(found=False)
read_through: Final = RegistryReadThrough(resync=spy, miss_ttl_seconds=60.0)
assert await read_through.attempt("ghost-model") is False
assert await read_through.attempt("ghost-model") is False
assert spy.calls == ["ghost-model"]
@pytest.mark.asyncio
async def test_negative_cache_expires_and_resync_runs_again():
spy: Final = ResyncSpy(found=False)
read_through: Final = RegistryReadThrough(resync=spy, miss_ttl_seconds=0.05)
assert await read_through.attempt("ghost-model") is False
await asyncio.sleep(0.1)
assert await read_through.attempt("ghost-model") is False
assert spy.calls == ["ghost-model", "ghost-model"]
@pytest.mark.asyncio
async def test_resync_exception_returns_false_without_negative_caching():
spy: Final = ResyncSpy(error=RuntimeError("db down"))
read_through: Final = RegistryReadThrough(resync=spy)
assert await read_through.attempt("new-model") is False
assert await read_through.attempt("new-model") is False
assert spy.calls == ["new-model", "new-model"]
@pytest.mark.asyncio
async def test_concurrent_attempts_for_missing_key_resync_once():
class SlowResyncSpy(ResyncSpy):
async def __call__(self, key: str) -> bool:
await asyncio.sleep(0.05)
return await super().__call__(key)
spy: Final = SlowResyncSpy(found=False)
read_through: Final = RegistryReadThrough(resync=spy, miss_ttl_seconds=60.0)
results: Final = await asyncio.gather(*(read_through.attempt("ghost-model") for _ in range(5)))
assert results == [False] * 5
assert spy.calls == ["ghost-model"]
@pytest.mark.asyncio
async def test_distinct_keys_do_not_share_negative_cache():
spy: Final = ResyncSpy(found=False)
read_through: Final = RegistryReadThrough(resync=spy, miss_ttl_seconds=60.0)
assert await read_through.attempt("ghost-a") is False
assert await read_through.attempt("ghost-b") is False
assert spy.calls == ["ghost-a", "ghost-b"]
@pytest.mark.asyncio
async def test_resync_budget_exhausted_blocks_resync_without_negative_caching():
spy: Final = ResyncSpy(found=False)
read_through: Final = RegistryReadThrough(
resync=spy, miss_ttl_seconds=60.0, max_resyncs_per_window=2, resync_window_seconds=60.0
)
assert await read_through.attempt("ghost-a") is False
assert await read_through.attempt("ghost-b") is False
assert await read_through.attempt("ghost-c") is False
assert spy.calls == ["ghost-a", "ghost-b"]
assert read_through._recent_misses.get_cache("ghost-c") is None
@pytest.mark.asyncio
async def test_resync_budget_replenishes_after_window():
spy: Final = ResyncSpy(found=True)
read_through: Final = RegistryReadThrough(resync=spy, max_resyncs_per_window=1, resync_window_seconds=0.05)
assert await read_through.attempt("model-a") is True
assert await read_through.attempt("model-b") is False
await asyncio.sleep(0.1)
assert await read_through.attempt("model-b") is True
assert spy.calls == ["model-a", "model-b"]
class FakeAgentRow:
def __init__(self, agent_id: str, agent_name: str) -> None:
self.agent_id = agent_id
self.agent_name = agent_name
self.object_permission = None
self.spend = 0.0
def model_dump(self):
return {
"agent_id": self.agent_id,
"agent_name": self.agent_name,
"agent_card_params": {"name": self.agent_name, "url": "http://db-agent"},
"litellm_params": {},
"object_permission": None,
"spend": self.spend,
}
@pytest.fixture
def clean_agent_registry():
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
original_agents: Final = list(global_agent_registry.agent_list)
original_config_agents: Final = getattr(global_agent_registry, "config_agents", ())
global_agent_registry.agent_list = []
global_agent_registry.config_agents = ()
try:
yield global_agent_registry
finally:
global_agent_registry.agent_list = original_agents
global_agent_registry.config_agents = original_config_agents
@pytest.mark.asyncio
async def test_get_agent_with_read_through_recovers_agent_created_on_sibling_replica(
clean_agent_registry, monkeypatch
):
from unittest.mock import AsyncMock, MagicMock
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through
agent_id: Final = "read-through-db-agent-id"
prisma_client: Final = MagicMock()
prisma_client.db.litellm_agentstable.find_unique = AsyncMock(
return_value=FakeAgentRow(agent_id, "read-through-db-agent")
)
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
assert clean_agent_registry.get_agent_by_id(agent_id=agent_id) is None
agent: Final = await get_agent_with_read_through(agent_id)
assert agent is not None
assert agent.agent_id == agent_id
prisma_client.db.litellm_agentstable.find_unique.assert_awaited_once_with(
where={"agent_id": agent_id},
include={"object_permission": True},
)
@pytest.mark.asyncio
async def test_get_agent_with_read_through_recovers_agent_by_name(clean_agent_registry, monkeypatch):
from unittest.mock import AsyncMock, MagicMock
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through
agent_name: Final = "read-through-db-agent-by-name"
prisma_client: Final = MagicMock()
prisma_client.db.litellm_agentstable.find_unique = AsyncMock(
side_effect=[None, FakeAgentRow("read-through-name-lookup-id", agent_name)]
)
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
agent: Final = await get_agent_with_read_through(agent_name)
assert agent is not None
assert agent.agent_name == agent_name
prisma_client.db.litellm_agentstable.find_unique.assert_awaited_with(
where={"agent_name": agent_name},
include={"object_permission": True},
)
@pytest.mark.asyncio
async def test_get_agent_with_read_through_returns_none_for_unknown_agent(clean_agent_registry, monkeypatch):
from unittest.mock import AsyncMock, MagicMock
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through
prisma_client: Final = MagicMock()
prisma_client.db.litellm_agentstable.find_unique = AsyncMock(return_value=None)
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
assert await get_agent_with_read_through("agent-nobody-created") is None
assert prisma_client.db.litellm_agentstable.find_unique.await_count == 2
@pytest.mark.asyncio
async def test_resync_agents_already_registered_skips_db(clean_agent_registry, monkeypatch):
from unittest.mock import AsyncMock, MagicMock
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy.common_utils.registry_read_through import _resync_agents
agent_id: Final = "read-through-dedup-agent-id"
prisma_client: Final = MagicMock()
prisma_client.db.litellm_agentstable.find_unique = AsyncMock(
return_value=FakeAgentRow(agent_id, "read-through-dedup-agent")
)
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
assert await _resync_agents(agent_id) is True
assert await _resync_agents(agent_id) is True
assert prisma_client.db.litellm_agentstable.find_unique.await_count == 1
assert len(clean_agent_registry.agent_list) == 1
class FakeGuardrailRow:
def __init__(self, guardrail_id: str, guardrail_name: str) -> None:
self.guardrail_id = guardrail_id
self.guardrail_name = guardrail_name
def __iter__(self):
return iter(
{
"guardrail_id": self.guardrail_id,
"guardrail_name": self.guardrail_name,
"litellm_params": {
"guardrail": "litellm_content_filter",
"mode": "pre_call",
"default_on": True,
"blocked_words": [{"keyword": "secret", "action": "BLOCK"}],
},
"guardrail_info": {},
"status": "active",
}.items()
)
@pytest.mark.asyncio
async def test_get_guardrail_with_read_through_recovers_guardrail_created_on_sibling_replica(monkeypatch):
from unittest.mock import AsyncMock, MagicMock
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy.common_utils.registry_read_through import (
get_initialized_guardrail_with_read_through,
)
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
guardrail_id: Final = "read-through-db-guardrail-id"
guardrail_name: Final = "read-through-db-guardrail"
prisma_client: Final = MagicMock()
prisma_client.db.litellm_guardrailstable.find_first = AsyncMock(
return_value=FakeGuardrailRow(guardrail_id, guardrail_name)
)
prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(
side_effect=AssertionError("full-table guardrail scan on read-through miss")
)
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
try:
guardrail: Final = await get_initialized_guardrail_with_read_through(guardrail_name=guardrail_name)
assert guardrail is not None
assert guardrail.guardrail_name == guardrail_name
prisma_client.db.litellm_guardrailstable.find_first.assert_awaited_once_with(
where={"guardrail_name": guardrail_name, "status": "active"}
)
finally:
IN_MEMORY_GUARDRAIL_HANDLER.delete_in_memory_guardrail(guardrail_id)
@pytest.mark.asyncio
async def test_get_guardrail_with_read_through_returns_none_for_unknown_guardrail(monkeypatch):
from unittest.mock import AsyncMock, MagicMock
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy.common_utils.registry_read_through import (
get_initialized_guardrail_with_read_through,
)
prisma_client: Final = MagicMock()
prisma_client.db.litellm_guardrailstable.find_first = AsyncMock(return_value=None)
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
assert await get_initialized_guardrail_with_read_through(guardrail_name="guardrail-nobody-created") is None
@pytest.mark.asyncio
async def test_resync_guardrails_never_loads_non_active_rows(monkeypatch):
from unittest.mock import AsyncMock, MagicMock
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy.common_utils.registry_read_through import _resync_guardrails
pending_name: Final = "pending-review-guardrail"
prisma_client: Final = MagicMock()
prisma_client.db.litellm_guardrailstable.find_first = AsyncMock(return_value=None)
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
assert await _resync_guardrails(pending_name) is False
prisma_client.db.litellm_guardrailstable.find_first.assert_awaited_once_with(
where={"guardrail_name": pending_name, "status": "active"}
)
@pytest.mark.asyncio
async def test_resync_guardrails_syncs_under_guardrail_reconcile_lock(monkeypatch):
from unittest.mock import AsyncMock, MagicMock
import litellm.proxy.common_utils.registry_read_through as read_through_module
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy.common_utils.registry_read_through import _resync_guardrails
from litellm.proxy.guardrails.guardrail_registry import (
GUARDRAIL_RECONCILE_LOCK,
IN_MEMORY_GUARDRAIL_HANDLER,
)
guardrail_name: Final = "lock-scope-guardrail"
prisma_client: Final = MagicMock()
prisma_client.db.litellm_guardrailstable.find_first = AsyncMock(
return_value=FakeGuardrailRow("lock-scope-guardrail-id", guardrail_name)
)
lock_states: list[bool] = []
def record_sync(guardrail) -> None:
lock_states.append(GUARDRAIL_RECONCILE_LOCK.locked())
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
monkeypatch.setattr(IN_MEMORY_GUARDRAIL_HANDLER, "sync_guardrail_from_db", record_sync)
monkeypatch.setattr(read_through_module, "_initialized_guardrail", lambda guardrail_name: MagicMock())
assert await _resync_guardrails(guardrail_name) is True
assert lock_states == [True]
assert not GUARDRAIL_RECONCILE_LOCK.locked()
@pytest.mark.asyncio
async def test_resync_model_deployments_mutates_router_under_model_reconcile_lock(monkeypatch):
from unittest.mock import AsyncMock, MagicMock
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy.common_utils.registry_read_through import _resync_model_deployments
prisma_client: Final = MagicMock()
prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[MagicMock()])
router: Final = MagicMock()
router.get_model_list.return_value = []
lock_states: list[bool] = []
def record_add_deployment(db_models) -> None:
lock_states.append(proxy_server.MODEL_RECONCILE_LOCK.locked())
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(proxy_server, "llm_model_list", None)
monkeypatch.setattr(proxy_server.proxy_config, "_add_deployment", record_add_deployment)
assert await _resync_model_deployments("lock-scope-model") is True
assert lock_states == [True]
assert not proxy_server.MODEL_RECONCILE_LOCK.locked()
@pytest.mark.asyncio
async def test_resync_model_deployments_respects_supported_db_objects(monkeypatch):
from unittest.mock import AsyncMock, MagicMock
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy.common_utils.registry_read_through import _resync_model_deployments
prisma_client: Final = MagicMock()
prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(
side_effect=AssertionError("db hit for an object type this replica does not load")
)
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["guardrails"]})
assert await _resync_model_deployments("gated-out-model") is False
@pytest.mark.asyncio
async def test_resync_guardrails_respects_supported_db_objects(monkeypatch):
from unittest.mock import AsyncMock, MagicMock
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy.common_utils.registry_read_through import _resync_guardrails
prisma_client: Final = MagicMock()
prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock(
side_effect=AssertionError("db hit for an object type this replica does not load")
)
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["models"]})
assert await _resync_guardrails("gated-out-guardrail") is False
@pytest.mark.asyncio
async def test_resync_agents_respects_supported_db_objects(clean_agent_registry, monkeypatch):
from unittest.mock import AsyncMock, MagicMock
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy.common_utils.registry_read_through import _resync_agents
prisma_client: Final = MagicMock()
prisma_client.db.litellm_agentstable.find_unique = AsyncMock(
side_effect=AssertionError("db hit for an object type this replica does not load")
)
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["models"]})
assert await _resync_agents("gated-out-agent") is False
@pytest.mark.asyncio
async def test_resync_agents_waits_for_agent_reload_and_skips_duplicate_registration(clean_agent_registry, monkeypatch):
from unittest.mock import AsyncMock, MagicMock
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy.agent_endpoints.agent_registry import AGENT_RECONCILE_LOCK
from litellm.proxy.common_utils.registry_read_through import _resync_agents
from litellm.types.agents import AgentResponse
agent_id: Final = "reload-race-agent-id"
prisma_client: Final = MagicMock()
prisma_client.db.litellm_agentstable.find_unique = AsyncMock(
side_effect=AssertionError("db hit while the agent reload held the reconcile lock")
)
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
async with AGENT_RECONCILE_LOCK:
resync_task: Final = asyncio.ensure_future(_resync_agents(agent_id))
await asyncio.sleep(0.05)
assert not resync_task.done()
clean_agent_registry.register_agent(
agent_config=AgentResponse.model_validate(FakeAgentRow(agent_id, "reload-race-agent").model_dump())
)
assert await resync_task is True
assert len(clean_agent_registry.agent_list) == 1

155
tests/test_litellm/proxy/management_endpoints/test_access_group_management.py
View file

@ -13,6 +13,9 @@ sys.path.insert(
) # Adds the parent directory to the system path
from litellm import Router
from litellm.proxy.management_endpoints.model_management_endpoints import (
ReconcileOutcome,
)
@pytest.mark.asyncio
@ -121,7 +124,7 @@ async def test_create_access_group_with_model_ids_tags_only_specific_deployments
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch(
"litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache",
new_callable=AsyncMock,
new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
),
):
response = await create_model_group(
@ -186,7 +189,7 @@ async def test_create_access_group_with_model_names_tags_all_deployments():
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch(
"litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache",
new_callable=AsyncMock,
new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
),
):
response = await create_model_group(
@ -236,7 +239,7 @@ async def test_create_access_group_model_ids_takes_priority_over_model_names():
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch(
"litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache",
new_callable=AsyncMock,
new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
),
):
response = await create_model_group(
@ -313,7 +316,7 @@ async def test_create_access_group_invalid_model_id_returns_400():
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch(
"litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache",
new_callable=AsyncMock,
new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
),
):
with pytest.raises(HTTPException) as exc_info:
@ -352,7 +355,7 @@ async def test_create_access_group_surfaces_dropped_models():
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch(
"litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache",
new=AsyncMock(return_value=None),
new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
),
):
with pytest.raises(HTTPException) as exc_info:
@ -365,6 +368,50 @@ async def test_create_access_group_surfaces_dropped_models():
assert "deploy-A" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_create_access_group_trusts_reload_snapshot_over_post_lock_fresh_read():
"""A concurrent reconcile sampled after the lock is released must not make this
write's reload look like it dropped the tagged model: the verdict has to judge from
the ReconcileOutcome the reload captured under the lock, not a fresh router read."""
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
create_model_group,
)
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
NewModelGroupRequest,
)
deploy_a = MagicMock(model_id="deploy-A", model_name="gpt-4o", model_info={})
mock_prisma = MagicMock()
mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=deploy_a)
mock_prisma.db.litellm_proxymodeltable.update = AsyncMock()
concurrently_wiped_router = MagicMock()
concurrently_wiped_router.get_model_ids.side_effect = [["deploy-A"], []]
with (
patch("litellm.proxy.proxy_server.llm_router", concurrently_wiped_router),
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch(
"litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache",
new=AsyncMock(
return_value=ReconcileOutcome(
still_desired=frozenset({"deploy-A"}), live_after=frozenset({"deploy-A"})
)
),
),
):
response = await create_model_group(
data=NewModelGroupRequest(access_group="production-models", model_ids=["deploy-A"]),
user_api_key_dict=UserAPIKeyAuth(user_id="test_admin", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert response.models_updated == 1
assert concurrently_wiped_router.get_model_ids.call_count == 1
@pytest.mark.asyncio
async def test_tag_deployment_parses_string_model_info_and_refuses_corrupt():
"""The model_info column can arrive as its JSON string; tagging must parse it rather
@ -420,7 +467,7 @@ async def test_delete_access_group_ignores_models_that_were_already_dead():
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch(
"litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache",
new=AsyncMock(return_value=None),
new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
),
):
response = await delete_access_group(
@ -430,3 +477,99 @@ async def test_delete_access_group_ignores_models_that_were_already_dead():
assert response.models_updated == 1
mock_prisma.db.litellm_proxymodeltable.update.assert_awaited_once()
@pytest.mark.asyncio
async def test_create_access_group_read_through_recovers_model_created_on_sibling_replica():
"""Regression: an access group referencing a model that another replica just wrote
to the DB must be created instead of 400ing until the periodic config reload."""
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
create_model_group,
)
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
NewModelGroupRequest,
)
from types import SimpleNamespace
model_name = "e2e-ag-sibling-replica-model"
db_row = SimpleNamespace(
model_id=f"{model_name}-id",
model_name=model_name,
litellm_params={"model": "openai/gpt-4o", "api_key": "fake", "mock_response": "hi"},
model_info={},
blocked=False,
)
mock_router = Router(
model_list=[
{
"model_name": "some-other-model",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"},
}
]
)
mock_prisma = MagicMock()
mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(side_effect=[[db_row], [], [db_row]])
mock_prisma.db.litellm_proxymodeltable.update = AsyncMock()
with (
patch("litellm.proxy.proxy_server.llm_router", mock_router),
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch("litellm.proxy.proxy_server.store_model_in_db", True),
patch(
"litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache",
new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
),
):
response = await create_model_group(
data=NewModelGroupRequest(access_group="replica-lag-group", model_names=[model_name]),
user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert response.models_updated == 1
assert response.model_names == [model_name]
assert mock_prisma.db.litellm_proxymodeltable.find_many.await_args_list[0].kwargs["where"] == {
"model_name": model_name
}
@pytest.mark.asyncio
async def test_create_access_group_model_missing_everywhere_still_400s():
from fastapi import HTTPException
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
create_model_group,
)
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
NewModelGroupRequest,
)
model_name = "e2e-ag-model-nobody-created"
mock_router = Router(
model_list=[
{
"model_name": "some-other-model",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"},
}
]
)
mock_prisma = MagicMock()
mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
with (
patch("litellm.proxy.proxy_server.llm_router", mock_router),
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch("litellm.proxy.proxy_server.store_model_in_db", True),
):
with pytest.raises(HTTPException) as exc_info:
await create_model_group(
data=NewModelGroupRequest(access_group="ghost-group", model_names=[model_name]),
user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert exc_info.value.status_code == 400
assert model_name in str(exc_info.value.detail)

54
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -11100,6 +11100,60 @@ async def test_moderations_reraises_proxy_exception_unwrapped():
mock_logging.post_call_failure_hook.assert_awaited_once()
@pytest.mark.asyncio
async def test_init_agents_in_db_rebuilds_registry_under_agent_reconcile_lock(monkeypatch):
from litellm.proxy.agent_endpoints.agent_registry import (
AGENT_RECONCILE_LOCK,
global_agent_registry,
)
from litellm.proxy.proxy_server import ProxyConfig
lock_states: list[bool] = []
async def fake_get_all_agents_from_db(prisma_client) -> list:
lock_states.append(AGENT_RECONCILE_LOCK.locked())
return []
def fake_load_agents_from_db_and_config(db_agents) -> None:
lock_states.append(AGENT_RECONCILE_LOCK.locked())
monkeypatch.setattr(global_agent_registry, "get_all_agents_from_db", fake_get_all_agents_from_db)
monkeypatch.setattr(global_agent_registry, "load_agents_from_db_and_config", fake_load_agents_from_db_and_config)
await ProxyConfig()._init_agents_in_db(prisma_client=MagicMock())
assert lock_states == [True, True]
assert not AGENT_RECONCILE_LOCK.locked()
@pytest.mark.asyncio
async def test_init_guardrails_in_db_snapshots_and_reconciles_under_guardrail_reconcile_lock(monkeypatch):
from litellm.proxy.guardrails.guardrail_registry import (
GUARDRAIL_RECONCILE_LOCK,
IN_MEMORY_GUARDRAIL_HANDLER,
GuardrailRegistry,
)
from litellm.proxy.proxy_server import ProxyConfig
lock_states: list[bool] = []
async def fake_get_all_guardrails_from_db(prisma_client) -> list:
lock_states.append(GUARDRAIL_RECONCILE_LOCK.locked())
return []
def fake_reconcile_db_guardrails(db_guardrail_ids) -> list:
lock_states.append(GUARDRAIL_RECONCILE_LOCK.locked())
return []
monkeypatch.setattr(GuardrailRegistry, "get_all_guardrails_from_db", fake_get_all_guardrails_from_db)
monkeypatch.setattr(IN_MEMORY_GUARDRAIL_HANDLER, "reconcile_db_guardrails", fake_reconcile_db_guardrails)
await ProxyConfig()._init_guardrails_in_db(prisma_client=MagicMock())
assert lock_states == [True, True]
assert not GUARDRAIL_RECONCILE_LOCK.locked()
class TestEmbeddingsFailureHookRequestData:
@pytest.mark.asyncio
async def test_failure_hook_gets_post_setup_data_with_logging_obj(self):

77
tests/test_litellm/proxy/test_route_a2a_models.py
View file

@ -50,6 +50,7 @@ async def test_route_a2a_model_bypasses_router():
)
mock_registry = Mock()
mock_registry.get_agent_by_id = Mock(return_value=None)
mock_registry.get_agent_by_name = Mock(return_value=mock_agent)
# Mock litellm.acompletion to verify it's called
@ -106,3 +107,79 @@ async def test_route_non_a2a_model_raises_error_if_not_in_router():
user_model=None,
route_type="acompletion",
)
class _DbAgentRow:
def __init__(self, agent_id: str, agent_name: str) -> None:
self.agent_id = agent_id
self.agent_name = agent_name
self.object_permission = None
self.spend = 0.0
def model_dump(self):
return {
"agent_id": self.agent_id,
"agent_name": self.agent_name,
"agent_card_params": {"name": self.agent_name, "url": "http://sibling-db-agent.example.com"},
"litellm_params": {},
"object_permission": None,
"spend": self.spend,
}
def _router_without_models():
mock_router = Mock()
mock_router.model_names = []
mock_router.deployment_names = []
mock_router.has_model_id = Mock(return_value=False)
mock_router.model_group_alias = None
mock_router.router_general_settings = Mock(pass_through_all_models=False)
mock_router.default_deployment = None
mock_router.pattern_router = Mock(patterns=[])
mock_router.map_team_model = Mock(return_value=None)
mock_router.is_recognized_model = Mock(return_value=False)
mock_router.team_public_model_names = []
return mock_router
@pytest.mark.asyncio
async def test_route_a2a_model_read_through_recovers_agent_created_on_sibling_replica(monkeypatch):
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
agent_name = "a2a-sibling-replica-agent"
prisma_client = Mock()
prisma_client.db.litellm_agentstable.find_unique = AsyncMock(
side_effect=[None, _DbAgentRow("a2a-sibling-replica-agent-id", agent_name)]
)
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
original_agents = list(global_agent_registry.agent_list)
original_config_agents = getattr(global_agent_registry, "config_agents", ())
global_agent_registry.agent_list = []
global_agent_registry.config_agents = ()
data = {
"model": f"a2a/{agent_name}",
"messages": [{"role": "user", "content": "Hello"}],
}
mock_acompletion = AsyncMock(return_value={"id": "read-through-response"})
try:
with patch("litellm.acompletion", mock_acompletion):
await route_request(
data=data,
llm_router=_router_without_models(),
user_model=None,
route_type="acompletion",
)
finally:
global_agent_registry.agent_list = original_agents
global_agent_registry.config_agents = original_config_agents
mock_acompletion.assert_called_once()
call_kwargs = mock_acompletion.call_args.kwargs
assert call_kwargs["model"] == f"a2a/{agent_name}"
assert call_kwargs["api_base"] == "http://sibling-db-agent.example.com"
prisma_client.db.litellm_agentstable.find_unique.assert_awaited()

156
tests/test_litellm/proxy/test_route_llm_request.py
View file

@ -1119,6 +1119,126 @@ async def test_route_request_rejects_chat_completion_without_messages():
llm_router.acompletion.assert_not_called()
class FakeProxyModelTable:
def __init__(self, rows):
self.rows = rows
self.find_many_wheres = []
async def find_many(self, where=None, **kwargs):
self.find_many_wheres.append(where)
return list(self.rows)
def _fake_prisma_client_with_models(rows):
from types import SimpleNamespace
table = FakeProxyModelTable(rows)
return SimpleNamespace(db=SimpleNamespace(litellm_proxymodeltable=table)), table
def _db_model_row(model_name: str, mock_response: str):
from types import SimpleNamespace
return SimpleNamespace(
model_id=f"{model_name}-id",
model_name=model_name,
litellm_params={"model": "openai/gpt-4o", "api_key": "fake", "mock_response": mock_response},
model_info={},
blocked=False,
)
@pytest.mark.asyncio
async def test_route_request_read_through_recovers_model_created_on_sibling_replica(monkeypatch):
"""Regression: a model written to the DB by another replica must be served on
first request instead of 400ing until the periodic config reload."""
import litellm
import litellm.proxy.proxy_server as proxy_server
model_name = "e2e-sibling-replica-model"
router = litellm.Router(
model_list=[
{
"model_name": "some-other-model",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"},
}
]
)
fake_prisma, table = _fake_prisma_client_with_models([_db_model_row(model_name, "hello-from-db")])
monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma)
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
monkeypatch.setattr(proxy_server, "llm_router", router)
llm_call = await route_request(
data={"model": model_name, "messages": [{"role": "user", "content": "hi"}]},
llm_router=router,
user_model=None,
route_type="acompletion",
)
response = await llm_call
assert response.choices[0].message.content == "hello-from-db"
assert len(table.find_many_wheres) == 1
assert table.find_many_wheres[0] == {"model_name": model_name}
@pytest.mark.asyncio
async def test_route_request_unknown_model_raises_and_hits_db_once_within_ttl(monkeypatch):
import litellm
import litellm.proxy.proxy_server as proxy_server
model_name = "e2e-model-nobody-created"
router = litellm.Router(
model_list=[
{
"model_name": "some-other-model",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"},
}
]
)
fake_prisma, table = _fake_prisma_client_with_models([])
monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma)
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
monkeypatch.setattr(proxy_server, "llm_router", router)
data = {"model": model_name, "messages": [{"role": "user", "content": "hi"}]}
with pytest.raises(ProxyModelNotFoundError):
await route_request(data=data, llm_router=router, user_model=None, route_type="acompletion")
with pytest.raises(ProxyModelNotFoundError):
await route_request(data=data, llm_router=router, user_model=None, route_type="acompletion")
assert table.find_many_wheres == [{"model_name": model_name}, {"model_id": model_name}]
@pytest.mark.asyncio
async def test_route_request_read_through_disabled_without_store_model_in_db(monkeypatch):
import litellm
import litellm.proxy.proxy_server as proxy_server
model_name = "e2e-config-only-proxy-model"
router = litellm.Router(
model_list=[
{
"model_name": "some-other-model",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"},
}
]
)
fake_prisma, table = _fake_prisma_client_with_models([_db_model_row(model_name, "should-not-load")])
monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma)
monkeypatch.setattr(proxy_server, "store_model_in_db", False)
monkeypatch.setattr(proxy_server, "llm_router", router)
with pytest.raises(ProxyModelNotFoundError):
await route_request(
data={"model": model_name, "messages": [{"role": "user", "content": "hi"}]},
llm_router=router,
user_model=None,
route_type="acompletion",
)
assert table.find_many_wheres == []
@pytest.mark.asyncio
async def test_route_request_routing_group_name_passes_model_gate():
from unittest.mock import AsyncMock, patch
@ -1141,3 +1261,39 @@ async def test_route_request_routing_group_name_passes_model_gate():
assert response == "group_response"
spy.assert_called_once_with(**data)
@pytest.mark.asyncio
async def test_route_request_a2a_agent_miss_does_not_consume_model_read_through(monkeypatch):
from types import SimpleNamespace
from unittest.mock import AsyncMock
import litellm
import litellm.proxy.proxy_server as proxy_server
model_name = "a2a/agent-nobody-created"
router = litellm.Router(
model_list=[
{
"model_name": "some-other-model",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"},
}
]
)
fake_prisma, model_table = _fake_prisma_client_with_models([])
agents_find_unique = AsyncMock(return_value=None)
fake_prisma.db.litellm_agentstable = SimpleNamespace(find_unique=agents_find_unique)
monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma)
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
monkeypatch.setattr(proxy_server, "llm_router", router)
with pytest.raises(ProxyModelNotFoundError):
await route_request(
data={"model": model_name, "messages": [{"role": "user", "content": "hi"}]},
llm_router=router,
user_model=None,
route_type="acompletion",
)
assert agents_find_unique.await_count == 2
assert model_table.find_many_wheres == []

89
tests/test_litellm/test_router.py
View file

@ -7643,6 +7643,95 @@ def test_pre_call_checks_keeps_deployment_when_provider_is_unresolvable(monkeypa
assert len(result) == 1
class TestUpsertDeploymentRollback:
"""
Regression tests: `upsert_deployment` pops the previous deployment before
re-adding the edited one. When the re-add raises under
`ignore_invalid_deployments=True`, the pop must be rolled back so this pod
keeps serving the previous configuration instead of silently dropping a live
deployment (the "Error upserting deployment" drop behind the access-group
reload 500 in the 2-replica e2e suite).
"""
def test_failed_upsert_keeps_previous_deployment_serving(self):
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
router = litellm.Router(
model_list=[
{
"model_name": "prod-model",
"litellm_params": {"model": "gpt-4o", "api_key": "sk-old"},
"model_info": {"id": "prod-1", "db_model": True},
}
],
ignore_invalid_deployments=True,
)
result = router.upsert_deployment(
deployment=Deployment(
model_name="prod-model",
litellm_params=LiteLLM_Params(model="auto_router/broken"),
model_info=ModelInfo(id="prod-1", db_model=True),
)
)
assert result is None
restored = router.get_deployment(model_id="prod-1")
assert restored is not None
assert restored.litellm_params.model == "gpt-4o"
assert [model["model_name"] for model in router.model_list] == ["prod-model"]
def test_failed_fresh_add_returns_none_without_restore(self):
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
router = litellm.Router(model_list=[], ignore_invalid_deployments=True)
result = router.upsert_deployment(
deployment=Deployment(
model_name="fresh-router",
litellm_params=LiteLLM_Params(model="auto_router/broken"),
model_info=ModelInfo(id="fresh-1", db_model=True),
)
)
assert result is None
assert router.get_deployment(model_id="fresh-1") is None
assert router.model_list == []
def test_restore_re_adds_popped_deployment(self):
router = litellm.Router(
model_list=[
{
"model_name": "prod-model",
"litellm_params": {"model": "gpt-4o", "api_key": "sk-old"},
"model_info": {"id": "prod-1", "db_model": True},
}
],
ignore_invalid_deployments=True,
)
previous = router.get_deployment(model_id="prod-1")
router.delete_deployment(id="prod-1")
assert router.has_model_id("prod-1") is False
router._restore_deployment_after_failed_upsert(
previous_deployment=previous, model_id="prod-1"
)
restored = router.get_deployment(model_id="prod-1")
assert restored is not None
assert restored.litellm_params.model == "gpt-4o"
router._restore_deployment_after_failed_upsert(
previous_deployment=previous, model_id="prod-1"
)
assert len(router.model_list) == 1
router._restore_deployment_after_failed_upsert(
previous_deployment=None, model_id="prod-1"
)
assert len(router.model_list) == 1
class TestConsumedRequestTagsStamp:
"""Issue #36621: when a request's tags select a tagged pre-routing strategy, those
tags are consumed by the selection; the hook must stamp the rewritten model group so

6
type-discipline-budget.json
View file

@ -1,9 +1,9 @@
{
"LIT001": {
"limit": 22894
"limit": 22892
},
"LIT002": {
"limit": 26888
"limit": 26886
},
"LIT003": {
"limit": 269
@ -27,7 +27,7 @@
"limit": 0
},
"LIT010": {
"limit": 16700
"limit": 16699
},
"LIT011": {
"limit": 5590