diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 95fcd988a0..5a939b3e97 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -219,6 +219,9 @@ async def patch_model( where={"model_id": model_id}, data=update_data, ) + + # Clear cache and reload models + await clear_cache() return updated_model @@ -879,3 +882,26 @@ def _deduplicate_litellm_router_models(models: List[Dict]) -> List[Dict]: unique_models.append(model) seen_ids.add(model_id) return unique_models + +async def clear_cache(): + """ + Clear router caches and reload models. + """ + from litellm.proxy.proxy_server import ( + proxy_config, + llm_router, + prisma_client, + proxy_logging_obj, + verbose_proxy_logger, + ) + try: + llm_router.model_list.clear() + + await proxy_config.add_deployment( + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj + ) + except Exception as e: + verbose_proxy_logger.exception( + f"Failed to clear cache and reload models. Due to error - {str(e)}" + ) diff --git a/tests/litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/litellm/proxy/management_endpoints/test_model_management_endpoints.py index 3bed3fc761..ebe99892cc 100644 --- a/tests/litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -2,6 +2,7 @@ import json import os import sys from typing import Optional +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi.testclient import TestClient @@ -17,6 +18,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelManagementAuthChecks, + clear_cache, ) from litellm.proxy.utils import PrismaClient from litellm.types.router import Deployment, LiteLLM_Params, updateDeployment @@ -46,6 +48,33 @@ class MockPrismaClient: return self +class MockLLMRouter: + def __init__(self): + self.model_list = ["model1", "model2"] + self.model_names = {"model1": True, "model2": True} + self.cleared = False + + def get_deployment(self, model_id): + return {"model_id": model_id} if model_id in self.model_list else None + + def delete_deployment(self, id): + if id in self.model_list: + self.model_list.remove(id) + self.model_names.pop(id, None) + + +class MockProxyConfig: + def __init__(self, success=True): + self.success = success + self.deployment_called = False + + async def add_deployment(self, prisma_client, proxy_logging_obj): + self.deployment_called = True + if not self.success: + raise Exception("Failed to add deployment") + return True + + class TestModelManagementAuthChecks: def setup_method(self): """Setup test cases""" @@ -199,3 +228,39 @@ class TestModelManagementAuthChecks: premium_user=True, ) assert "403" in str(exc_info.value) + + +class TestClearCache: + """ + Tests for the clear_cache function in model_management_endpoints.py + """ + + @pytest.mark.asyncio + async def test_clear_cache_success(self): + """ + Test that clear_cache successfully clears router model caches and reloads models. + """ + mock_router = MagicMock() + mock_router.model_list = ["openai/gpt-4o", "openai/gpt-4o-mini"] + + mock_config = MagicMock() + mock_config.add_deployment = AsyncMock(return_value=True) + + mock_prisma = MagicMock() + mock_logging = MagicMock() + + with patch("litellm.proxy.proxy_server.llm_router", mock_router), \ + patch("litellm.proxy.proxy_server.proxy_config", mock_config), \ + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), \ + patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_logging), \ + patch("litellm.proxy.proxy_server.verbose_proxy_logger"): + + await clear_cache() + + + assert len(mock_router.model_list) == 0 + + mock_config.add_deployment.assert_called_once_with( + prisma_client=mock_prisma, + proxy_logging_obj=mock_logging + )