diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 967cae27990..a7d7f9ca64f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7781,6 +7781,7 @@ class ProxyConfig: verbose_proxy_logger.debug("_model_list: %s", _model_list) llm_router = litellm.Router( model_list=_model_list, + cache_responses=litellm.cache is not None, router_general_settings=RouterGeneralSettings( async_only_mode=True # only init async clients ), diff --git a/tests/unit/proxy/proxy_server/test_db_router_response_cache.py b/tests/unit/proxy/proxy_server/test_db_router_response_cache.py new file mode 100644 index 00000000000..c2ee97c7bb5 --- /dev/null +++ b/tests/unit/proxy/proxy_server/test_db_router_response_cache.py @@ -0,0 +1,99 @@ +"""Regression coverage for response caching on DB-backed router creation.""" + +import asyncio +import importlib +from unittest.mock import AsyncMock, MagicMock + +import litellm +import pytest + +from litellm.proxy.proxy_server import ProxyConfig + + +def _db_model() -> MagicMock: + record = MagicMock() + record.model_id = "db-id-1" + record.model_name = "fake-model" + record.litellm_params = {"model": "fake-model"} + record.model_info = {"id": "db-id-1"} + record.created_by = "default_user_id" + record.created_at = None + record.updated_at = None + record.updated_by = None + return record + + +@pytest.mark.asyncio +async def test_db_router_rebuild_serves_same_caller_from_cache_and_isolates_callers(monkeypatch): + """A DB router rebuild must use response caching without crossing caller boundaries.""" + proxy_config = ProxyConfig() + fake_model_list = [ + { + "model_name": "fake-model", + "litellm_params": { + "model": "openai/fake-model", + "api_key": "sk-fake", + "mock_response": "cached response", + }, + } + ] + monkeypatch.setattr(litellm, "cache", litellm.Cache()) + provider_calls = 0 + litellm_main = importlib.import_module("litellm.main") + original_mock_completion = litellm_main.mock_completion + + def tracked_mock_completion(*args, **kwargs): + nonlocal provider_calls + provider_calls += 1 + return original_mock_completion(*args, **kwargs) + + monkeypatch.setattr(litellm_main, "mock_completion", tracked_mock_completion) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-test") + monkeypatch.setattr("litellm.proxy.proxy_server.llm_model_list", []) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr(proxy_config, "decrypt_model_list_from_db", lambda new_models: fake_model_list) + monkeypatch.setattr(proxy_config, "_add_router_settings_from_db_config", AsyncMock()) + + await proxy_config._update_llm_router( + new_models=[_db_model()], + proxy_logging_obj=MagicMock(), + ) + router = __import__("litellm.proxy.proxy_server", fromlist=["llm_router"]).llm_router + + request = [{"role": "user", "content": "cache canary"}] + first = await router.acompletion(model="fake-model", messages=request, user="caller-a") + await asyncio.sleep(1) + same_caller = await router.acompletion(model="fake-model", messages=request, user="caller-a") + other_caller = await router.acompletion(model="fake-model", messages=request, user="caller-b") + + assert first.id == same_caller.id + assert first.id != other_caller.id + assert provider_calls == 2 + + +@pytest.mark.asyncio +async def test_db_router_rebuild_keeps_cache_disabled_without_global_cache(monkeypatch): + """The DB router must not enable response caching when the proxy has no cache.""" + proxy_config = ProxyConfig() + fake_model_list = [ + { + "model_name": "fake-model", + "litellm_params": {"model": "openai/fake-model", "api_key": "sk-fake"}, + } + ] + monkeypatch.setattr(litellm, "cache", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-test") + monkeypatch.setattr("litellm.proxy.proxy_server.llm_model_list", []) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr(proxy_config, "decrypt_model_list_from_db", lambda new_models: fake_model_list) + monkeypatch.setattr(proxy_config, "_add_router_settings_from_db_config", AsyncMock()) + + await proxy_config._update_llm_router( + new_models=[_db_model()], + proxy_logging_obj=MagicMock(), + ) + router = __import__("litellm.proxy.proxy_server", fromlist=["llm_router"]).llm_router + + assert router.cache_responses is False