diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index be76c2ac5fb..9f73eab09c5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8855,6 +8855,99 @@ def _enrich_model_info_with_litellm_data( return model +async def _get_db_only_models( + all_models: List[Dict[str, Any]], + prisma_client: Any, + proxy_config: Any, + search: Optional[str] = None, + sort_by: Optional[str] = None, + page: int = 1, + size: int = 50, +) -> Tuple[List[Dict[str, Any]], int]: + """ + Fetch models from the database that are NOT already loaded in the router. + + Some DB models may fail to sync into the router (e.g. due to decryption errors + or invalid params with ignore_invalid_deployments=True). This function ensures + those models still appear in the model list. + + Args: + all_models: Models already in the router + prisma_client: Prisma client for database queries + proxy_config: Proxy config for decrypting models + search: Optional search term (case-insensitive) + sort_by: Optional sort field + page: Current page number + size: Page size + + Returns: + Tuple of (db_only_models, db_only_count) + """ + # Collect IDs of DB models already in the router + db_model_ids_in_router: set = set() + for m in all_models: + model_info = m.get("model_info", {}) + if model_info.get("db_model", False): + model_id = model_info.get("id") + if model_id: + db_model_ids_in_router.add(model_id) + + # Build where condition - exclude models already in the router + db_where_condition: Dict[str, Any] = {} + if db_model_ids_in_router: + db_where_condition["model_id"] = { + "not": {"in": list(db_model_ids_in_router)} + } + + # Add search filter if provided + if search and search.strip(): + db_where_condition["model_name"] = { + "contains": search.lower().strip(), + "mode": "insensitive", + } + + try: + db_only_count = await prisma_client.db.litellm_proxymodeltable.count( + where=db_where_condition + ) + + if db_only_count == 0: + return [], 0 + + # Determine how many to fetch + router_models_count = len(all_models) + models_needed_for_page = size * page + + if sort_by: + # Need all for correct sorting + take = db_only_count + else: + # Only fetch what's needed for the current page + take = max(0, models_needed_for_page - router_models_count) + take = min(take, db_only_count) + + if take <= 0: + return [], db_only_count + + db_models_raw = await prisma_client.db.litellm_proxymodeltable.find_many( + where=db_where_condition, + take=take, + ) + + db_models: List[Dict[str, Any]] = [] + for db_model in db_models_raw: + decrypted = proxy_config.decrypt_model_list_from_db([db_model]) + if decrypted: + db_models.extend(decrypted) + + return db_models, db_only_count + except Exception as e: + verbose_proxy_logger.exception( + f"Error querying database for db-only models: {str(e)}" + ) + return [], 0 + + async def _apply_search_filter_to_models( all_models: List[Dict[str, Any]], search: str, @@ -8865,7 +8958,8 @@ async def _apply_search_filter_to_models( sort_by: Optional[str] = None, ) -> Tuple[List[Dict[str, Any]], Optional[int]]: """ - Apply search filter to models, querying database for additional matching models. + Apply search filter to models, querying database for additional matching models + that may not be loaded in the router. Args: all_models: List of models to filter @@ -8877,122 +8971,43 @@ async def _apply_search_filter_to_models( sort_by: Optional sort field - if provided, fetch all matching models instead of paginating at DB level Returns: - Tuple of (filtered_models, total_count). total_count is None if not searching. + Tuple of (filtered_models, total_count). total_count is None if no DB-only models exist and not searching. """ - if not search or not search.strip(): + is_searching = bool(search and search.strip()) + + if is_searching: + search_lower = search.lower().strip() + # Filter router models by search term + filtered_router_models = [ + m for m in all_models if search_lower in m.get("model_name", "").lower() + ] + else: + filtered_router_models = all_models + + # Query database for models not in the router + db_only_models: List[Dict[str, Any]] = [] + db_only_count = 0 + if prisma_client is not None: + db_only_models, db_only_count = await _get_db_only_models( + all_models=all_models, + prisma_client=prisma_client, + proxy_config=proxy_config, + search=search if is_searching else None, + sort_by=sort_by, + page=page, + size=size, + ) + + total_count = len(filtered_router_models) + db_only_count + + # If not searching and no DB-only models, return None for total_count + # to preserve the original pagination behavior + if not is_searching and db_only_count == 0: return all_models, None - search_lower = search.lower().strip() - - # Filter models in router by search term - filtered_router_models = [ - m for m in all_models if search_lower in m.get("model_name", "").lower() - ] - - # Separate filtered models into config vs db models, and track db model IDs - filtered_config_models = [] - db_model_ids_in_router = set() - - for m in filtered_router_models: - model_info = m.get("model_info", {}) - is_db_model = model_info.get("db_model", False) - model_id = model_info.get("id") - - if is_db_model and model_id: - db_model_ids_in_router.add(model_id) - else: - filtered_config_models.append(m) - - config_models_count = len(filtered_config_models) - db_models_in_router_count = len(db_model_ids_in_router) - router_models_count = config_models_count + db_models_in_router_count - - # Query database for additional models with search term - db_models = [] - db_models_total_count = 0 - models_needed_for_page = size * page - - # Only query database if prisma_client is available - if prisma_client is not None: - try: - # Build where condition for database query - db_where_condition: Dict[str, Any] = { - "model_name": { - "contains": search_lower, - "mode": "insensitive", - } - } - # Exclude models already in router if we have any - if db_model_ids_in_router: - db_where_condition["model_id"] = { - "not": {"in": list(db_model_ids_in_router)} - } - - # Get total count of matching database models - db_models_total_count = ( - await prisma_client.db.litellm_proxymodeltable.count( - where=db_where_condition - ) - ) - - # Calculate total count for search results - search_total_count = router_models_count + db_models_total_count - - # If sorting is requested, we need to fetch ALL matching models to sort correctly - # Otherwise, we can optimize by only fetching what's needed for the current page - if sort_by: - # Fetch all matching database models for sorting - if db_models_total_count > 0: - db_models_raw = ( - await prisma_client.db.litellm_proxymodeltable.find_many( - where=db_where_condition, - take=db_models_total_count, # Fetch all matching models - ) - ) - - # Convert database models to router format - for db_model in db_models_raw: - decrypted_models = proxy_config.decrypt_model_list_from_db( - [db_model] - ) - if decrypted_models: - db_models.extend(decrypted_models) - else: - # Fetch database models if we need more for the current page - if router_models_count < models_needed_for_page: - models_to_fetch = min( - models_needed_for_page - router_models_count, - db_models_total_count, - ) - - if models_to_fetch > 0: - db_models_raw = ( - await prisma_client.db.litellm_proxymodeltable.find_many( - where=db_where_condition, - take=models_to_fetch, - ) - ) - - # Convert database models to router format - for db_model in db_models_raw: - decrypted_models = proxy_config.decrypt_model_list_from_db( - [db_model] - ) - if decrypted_models: - db_models.extend(decrypted_models) - except Exception as e: - verbose_proxy_logger.exception( - f"Error querying database models with search: {str(e)}" - ) - # If error, use router models count as fallback - search_total_count = router_models_count - else: - # If no prisma_client, only use router models - search_total_count = router_models_count - - # Combine all models - filtered_models = filtered_router_models + db_models - return filtered_models, search_total_count + # Combine router models with DB-only models + filtered_models = filtered_router_models + db_only_models + return filtered_models, total_count def _normalize_datetime_for_sorting(dt: Any) -> Optional[datetime]: diff --git a/tests/test_litellm/proxy/test_model_list_db_only_models.py b/tests/test_litellm/proxy/test_model_list_db_only_models.py new file mode 100644 index 00000000000..8842a633cbb --- /dev/null +++ b/tests/test_litellm/proxy/test_model_list_db_only_models.py @@ -0,0 +1,260 @@ +""" +Test that models stored in the database but not loaded into the router +still appear in the /v2/model/info endpoint response. + +This addresses a bug where models added via the UI would be saved to the DB +but fail to sync into the router (e.g. due to decryption errors or invalid +params with ignore_invalid_deployments=True), making them invisible in the +model list. +""" + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +from litellm.proxy.proxy_server import ( + _apply_search_filter_to_models, + _get_db_only_models, +) + + +def _make_router_model(model_name: str, model_id: str, db_model: bool = True): + """Helper to create a router-format model dict.""" + return { + "model_name": model_name, + "litellm_params": {"model": model_name}, + "model_info": { + "id": model_id, + "db_model": db_model, + }, + } + + +def _make_db_record(model_name: str, model_id: str): + """Helper to create a mock DB record.""" + record = MagicMock() + record.model_id = model_id + record.model_name = model_name + record.litellm_params = {"model": model_name} + record.model_info = {"id": model_id} + record.created_by = "default_user_id" + return record + + +class TestGetDbOnlyModels: + """Test _get_db_only_models fetches models not in router.""" + + @pytest.mark.asyncio + async def test_no_db_only_models(self): + """When all DB models are in the router, return empty list.""" + router_models = [ + _make_router_model("gpt-4", "id-1", db_model=True), + _make_router_model("gpt-5", "id-2", db_model=True), + ] + prisma_client = MagicMock() + prisma_client.db.litellm_proxymodeltable.count = AsyncMock(return_value=0) + proxy_config = MagicMock() + + db_models, count = await _get_db_only_models( + all_models=router_models, + prisma_client=prisma_client, + proxy_config=proxy_config, + ) + + assert db_models == [] + assert count == 0 + # Should have excluded router model IDs in the query + call_args = prisma_client.db.litellm_proxymodeltable.count.call_args + where = call_args.kwargs.get("where", {}) + assert "model_id" in where + assert set(where["model_id"]["not"]["in"]) == {"id-1", "id-2"} + + @pytest.mark.asyncio + async def test_db_only_models_returned(self): + """Models in DB but not in router should be fetched and returned.""" + router_models = [ + _make_router_model("gpt-4", "id-1", db_model=True), + ] + db_record = _make_db_record("gpt-5.1", "id-missing") + + prisma_client = MagicMock() + prisma_client.db.litellm_proxymodeltable.count = AsyncMock(return_value=1) + prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock( + return_value=[db_record] + ) + + decrypted_model = _make_router_model("gpt-5.1", "id-missing", db_model=True) + proxy_config = MagicMock() + proxy_config.decrypt_model_list_from_db = MagicMock( + return_value=[decrypted_model] + ) + + db_models, count = await _get_db_only_models( + all_models=router_models, + prisma_client=prisma_client, + proxy_config=proxy_config, + ) + + assert count == 1 + assert len(db_models) == 1 + assert db_models[0]["model_name"] == "gpt-5.1" + + @pytest.mark.asyncio + async def test_search_filter_applied(self): + """When search is provided, DB query should filter by model name.""" + router_models = [] + prisma_client = MagicMock() + prisma_client.db.litellm_proxymodeltable.count = AsyncMock(return_value=0) + proxy_config = MagicMock() + + await _get_db_only_models( + all_models=router_models, + prisma_client=prisma_client, + proxy_config=proxy_config, + search="gpt-5", + ) + + call_args = prisma_client.db.litellm_proxymodeltable.count.call_args + where = call_args.kwargs.get("where", {}) + assert "model_name" in where + assert where["model_name"]["contains"] == "gpt-5" + + @pytest.mark.asyncio + async def test_config_models_not_excluded(self): + """Config models (db_model=False) should not be added to the exclusion set.""" + router_models = [ + _make_router_model("gpt-4", "config-id-1", db_model=False), + _make_router_model("gpt-5", "db-id-1", db_model=True), + ] + prisma_client = MagicMock() + prisma_client.db.litellm_proxymodeltable.count = AsyncMock(return_value=0) + proxy_config = MagicMock() + + await _get_db_only_models( + all_models=router_models, + prisma_client=prisma_client, + proxy_config=proxy_config, + ) + + call_args = prisma_client.db.litellm_proxymodeltable.count.call_args + where = call_args.kwargs.get("where", {}) + # Only db_model=True IDs should be excluded + assert "db-id-1" in where["model_id"]["not"]["in"] + assert "config-id-1" not in where["model_id"]["not"]["in"] + + +class TestApplySearchFilterIncludesDbOnlyModels: + """Test that _apply_search_filter_to_models includes DB-only models.""" + + @pytest.mark.asyncio + async def test_no_search_includes_db_only_models(self): + """Even without search, DB-only models should be included in results.""" + router_models = [ + _make_router_model("gpt-4", "id-1", db_model=True), + ] + db_record = _make_db_record("gpt-5.1", "id-missing") + + prisma_client = MagicMock() + prisma_client.db.litellm_proxymodeltable.count = AsyncMock(return_value=1) + prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock( + return_value=[db_record] + ) + + decrypted_model = _make_router_model("gpt-5.1", "id-missing", db_model=True) + proxy_config = MagicMock() + proxy_config.decrypt_model_list_from_db = MagicMock( + return_value=[decrypted_model] + ) + + result, total_count = await _apply_search_filter_to_models( + all_models=router_models, + search="", + page=1, + size=50, + prisma_client=prisma_client, + proxy_config=proxy_config, + ) + + # Should include both router model and DB-only model + model_names = [m["model_name"] for m in result] + assert "gpt-4" in model_names + assert "gpt-5.1" in model_names + assert total_count == 2 + + @pytest.mark.asyncio + async def test_no_search_no_db_only_models_returns_none_count(self): + """When no search and no DB-only models, total_count should be None (original behavior).""" + router_models = [ + _make_router_model("gpt-4", "id-1", db_model=True), + ] + prisma_client = MagicMock() + prisma_client.db.litellm_proxymodeltable.count = AsyncMock(return_value=0) + proxy_config = MagicMock() + + result, total_count = await _apply_search_filter_to_models( + all_models=router_models, + search="", + page=1, + size=50, + prisma_client=prisma_client, + proxy_config=proxy_config, + ) + + assert result == router_models + assert total_count is None # Preserve original behavior + + @pytest.mark.asyncio + async def test_search_with_db_only_models(self): + """Search should filter router models AND include matching DB-only models.""" + router_models = [ + _make_router_model("gpt-4", "id-1", db_model=True), + _make_router_model("gpt-5", "id-2", db_model=True), + ] + db_record = _make_db_record("gpt-5.1", "id-missing") + + prisma_client = MagicMock() + prisma_client.db.litellm_proxymodeltable.count = AsyncMock(return_value=1) + prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock( + return_value=[db_record] + ) + + decrypted_model = _make_router_model("gpt-5.1", "id-missing", db_model=True) + proxy_config = MagicMock() + proxy_config.decrypt_model_list_from_db = MagicMock( + return_value=[decrypted_model] + ) + + result, total_count = await _apply_search_filter_to_models( + all_models=router_models, + search="gpt-5", + page=1, + size=50, + prisma_client=prisma_client, + proxy_config=proxy_config, + ) + + model_names = [m["model_name"] for m in result] + # gpt-4 should be filtered out by search + assert "gpt-4" not in model_names + # gpt-5 (from router) and gpt-5.1 (from DB) should be included + assert "gpt-5" in model_names + assert "gpt-5.1" in model_names + assert total_count == 2 + + @pytest.mark.asyncio + async def test_no_prisma_client_returns_router_only(self): + """Without prisma_client, should return only router models.""" + router_models = [ + _make_router_model("gpt-4", "id-1", db_model=True), + ] + + result, total_count = await _apply_search_filter_to_models( + all_models=router_models, + search="", + page=1, + size=50, + prisma_client=None, + proxy_config=MagicMock(), + ) + + assert result == router_models + assert total_count is None