From b44ac6c682ff21859e36116026f7ebbd14caf2bc Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 24 Jan 2026 09:08:29 -0800 Subject: [PATCH] Fixing ruff check --- litellm/proxy/proxy_server.py | 271 +++++++++++------- tests/test_litellm/proxy/test_proxy_server.py | 214 ++++++++++++++ 2 files changed, 378 insertions(+), 107 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f11b2d67f9e..fbb073a3b2c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7706,6 +7706,154 @@ def _enrich_model_info_with_litellm_data( return model +async def _apply_search_filter_to_models( + all_models: List[Dict[str, Any]], + search: str, + page: int, + size: int, + prisma_client: Optional[Any], + proxy_config: Any, +) -> Tuple[List[Dict[str, Any]], Optional[int]]: + """ + Apply search filter to models, querying database for additional matching models. + + Args: + all_models: List of models to filter + search: Search term (case-insensitive) + page: Current page number + size: Page size + prisma_client: Prisma client for database queries + proxy_config: Proxy config for decrypting models + + Returns: + Tuple of (filtered_models, total_count). total_count is None if not searching. + """ + if not search or not search.strip(): + 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 + + 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 + + # 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 + + # Combine all models + filtered_models = filtered_router_models + db_models + return filtered_models, search_total_count + + +def _paginate_models_response( + all_models: List[Dict[str, Any]], + page: int, + size: int, + total_count: Optional[int], + search: Optional[str], +) -> Dict[str, Any]: + """ + Paginate models and return response dictionary. + + Args: + all_models: List of all models + page: Current page number + size: Page size + total_count: Total count (if None, uses len(all_models)) + search: Search term (for logging) + + Returns: + Paginated response dictionary + """ + if total_count is None: + total_count = len(all_models) + + skip = (page - 1) * size + total_pages = -(-total_count // size) if total_count > 0 else 0 + paginated_models = all_models[skip : skip + size] + + verbose_proxy_logger.debug( + f"Pagination: skip={skip}, take={size}, total_count={total_count}, total_pages={total_pages}, search={search}" + ) + + return { + "data": paginated_models, + "total_count": total_count, + "current_page": page, + "total_pages": total_pages, + "size": size, + } + + @router.get( "/v2/model/info", description="v2 - returns models available to the user based on their API key permissions. Shows model info from config.yaml (except api key and api base). Filter to just user-added models with ?user_models_only=true", @@ -7763,94 +7911,15 @@ async def model_info_v2( if model is not None: all_models = [m for m in all_models if m["model_name"] == model] - # Track total count for search (will be calculated if searching) - search_total_count = None - # Apply search filter if provided - if search is not None and search.strip(): - search_lower = search.lower().strip() - - # First, filter ALL models in router by search term (both config and db models) - 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 (not already in router) - # We need enough models to fill the current page (size * page total models) - db_models = [] - db_models_total_count = 0 - models_needed_for_page = size * page # Total models needed up to current page - - 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 (excluding those already in router) - 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 - - # 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 - - # Combine all models: config models first, then db models from router, then db models from database - # filtered_router_models already contains both config and db models from router, so we can use it directly - all_models = filtered_router_models + db_models - # else: No search - models are already in all_models from llm_router.model_list + all_models, search_total_count = await _apply_search_filter_to_models( + all_models=all_models, + search=search or "", + page=page, + size=size, + prisma_client=prisma_client, + proxy_config=proxy_config, + ) if user_models_only: all_models = await non_admin_all_models( @@ -7867,7 +7936,8 @@ async def model_info_v2( llm_router=llm_router, all_models=all_models, ) - # fill in model info based on config.yaml and litellm model_prices_and_context_window.json + + # Fill in model info based on config.yaml and litellm model_prices_and_context_window.json for i, _model in enumerate(all_models): all_models[i] = _enrich_model_info_with_litellm_data( model=_model, debug=debug if debug is not None else False, llm_router=llm_router @@ -7875,26 +7945,13 @@ async def model_info_v2( verbose_proxy_logger.debug("all_models: %s", all_models) - # Use search_total_count if searching, otherwise use len(all_models) - total_count = search_total_count if search_total_count is not None else len(all_models) - - skip = (page - 1) * size - - total_pages = -(-total_count // size) if total_count > 0 else 0 - - paginated_models = all_models[skip : skip + size] - - verbose_proxy_logger.debug( - f"Pagination: skip={skip}, take={size}, total_count={total_count}, total_pages={total_pages}, search={search}" + return _paginate_models_response( + all_models=all_models, + page=page, + size=size, + total_count=search_total_count, + search=search, ) - - return { - "data": paginated_models, - "total_count": total_count, - "current_page": page, - "total_pages": total_pages, - "size": size, - } @router.get( diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index e2908fe4c3c..34dc0d843ad 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -3720,6 +3720,220 @@ async def test_model_info_v2_search_db_models(monkeypatch): app.dependency_overrides = original_overrides +@pytest.mark.asyncio +async def test_apply_search_filter_to_models(monkeypatch): + """ + Test the _apply_search_filter_to_models helper function. + Tests search filtering logic for config models, db models, and database queries. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy.proxy_server import _apply_search_filter_to_models, proxy_config + + # Create mock models with mix of config and db models + mock_models = [ + { + "model_name": "gpt-4-turbo", + "model_info": {"id": "gpt-4-turbo"}, # Config model + }, + { + "model_name": "db-gpt-3.5", + "model_info": {"id": "db-model-1", "db_model": True}, # DB model in router + }, + { + "model_name": "claude-3-opus", + "model_info": {"id": "claude-3-opus"}, # Config model + }, + ] + + # Mock prisma_client + mock_prisma_client = MagicMock() + mock_db_table = MagicMock() + mock_prisma_client.db.litellm_proxymodeltable = mock_db_table + + # Mock database models + mock_db_model_1 = MagicMock( + model_id="db-model-2", + model_name="db-gemini-pro", + litellm_params='{"model": "gemini-pro"}', + model_info='{"id": "db-model-2", "db_model": true}', + ) + + # Mock proxy_config.decrypt_model_list_from_db + mock_decrypt = MagicMock(return_value=[{"model_name": "db-gemini-pro", "model_info": {"id": "db-model-2", "db_model": True}}]) + + monkeypatch.setattr(proxy_config, "decrypt_model_list_from_db", mock_decrypt) + + # Test Case 1: No search term - should return all models unchanged + result_models, total_count = await _apply_search_filter_to_models( + all_models=mock_models.copy(), + search="", + page=1, + size=50, + prisma_client=mock_prisma_client, + proxy_config=proxy_config, + ) + assert result_models == mock_models + assert total_count is None + + # Test Case 2: Search for "gpt" - should filter router models and query DB + mock_db_table.count = AsyncMock(return_value=0) + mock_db_table.find_many = AsyncMock(return_value=[]) + + result_models, total_count = await _apply_search_filter_to_models( + all_models=mock_models.copy(), + search="gpt", + page=1, + size=50, + prisma_client=mock_prisma_client, + proxy_config=proxy_config, + ) + assert len(result_models) == 2 + model_names = [m["model_name"] for m in result_models] + assert "gpt-4-turbo" in model_names + assert "db-gpt-3.5" in model_names + assert "claude-3-opus" not in model_names + assert total_count == 2 # Only router models match + + # Test Case 3: Search with DB models matching + mock_db_table.count = AsyncMock(return_value=1) + mock_db_table.find_many = AsyncMock(return_value=[mock_db_model_1]) + + result_models, total_count = await _apply_search_filter_to_models( + all_models=mock_models.copy(), + search="gemini", + page=1, + size=50, + prisma_client=mock_prisma_client, + proxy_config=proxy_config, + ) + assert total_count == 1 # Router models (0) + DB models (1) + assert len(result_models) == 1 + assert result_models[0]["model_name"] == "db-gemini-pro" + + # Test Case 4: Case-insensitive search + # Reset mocks - no DB models should match "GPT" + mock_db_table.count = AsyncMock(return_value=0) + mock_db_table.find_many = AsyncMock(return_value=[]) + + result_models, total_count = await _apply_search_filter_to_models( + all_models=mock_models.copy(), + search="GPT", + page=1, + size=50, + prisma_client=mock_prisma_client, + proxy_config=proxy_config, + ) + assert len(result_models) == 2 + model_names = [m["model_name"] for m in result_models] + assert "gpt-4-turbo" in model_names + assert "db-gpt-3.5" in model_names + + # Test Case 5: Database query error - should fallback to router models count + mock_db_table.count = AsyncMock(side_effect=Exception("DB error")) + mock_db_table.find_many = AsyncMock(return_value=[]) + + result_models, total_count = await _apply_search_filter_to_models( + all_models=mock_models.copy(), + search="gpt", + page=1, + size=50, + prisma_client=mock_prisma_client, + proxy_config=proxy_config, + ) + # Should still return filtered router models + assert len(result_models) == 2 + assert total_count == 2 # Fallback to router models count + + +def test_paginate_models_response(): + """ + Test the _paginate_models_response helper function. + Tests pagination calculation and response formatting. + """ + from litellm.proxy.proxy_server import _paginate_models_response + + # Create mock models + mock_models = [ + {"model_name": f"model-{i}", "model_info": {"id": f"model-{i}"}} + for i in range(25) + ] + + # Test Case 1: Basic pagination - first page + result = _paginate_models_response( + all_models=mock_models, + page=1, + size=10, + total_count=None, + search=None, + ) + assert result["total_count"] == 25 + assert result["current_page"] == 1 + assert result["total_pages"] == 3 # ceil(25/10) = 3 + assert result["size"] == 10 + assert len(result["data"]) == 10 + assert result["data"][0]["model_name"] == "model-0" + + # Test Case 2: Second page + result = _paginate_models_response( + all_models=mock_models, + page=2, + size=10, + total_count=None, + search=None, + ) + assert result["current_page"] == 2 + assert len(result["data"]) == 10 + assert result["data"][0]["model_name"] == "model-10" + + # Test Case 3: Last page (partial) + result = _paginate_models_response( + all_models=mock_models, + page=3, + size=10, + total_count=None, + search=None, + ) + assert result["current_page"] == 3 + assert len(result["data"]) == 5 # Only 5 models left + assert result["data"][0]["model_name"] == "model-20" + + # Test Case 4: With explicit total_count (for search scenarios) + result = _paginate_models_response( + all_models=mock_models[:10], # Only 10 models in list + page=1, + size=10, + total_count=50, # But total_count says 50 + search="test", + ) + assert result["total_count"] == 50 + assert result["total_pages"] == 5 # ceil(50/10) = 5 + assert len(result["data"]) == 10 + + # Test Case 5: Empty models list + result = _paginate_models_response( + all_models=[], + page=1, + size=10, + total_count=0, + search=None, + ) + assert result["total_count"] == 0 + assert result["total_pages"] == 0 + assert len(result["data"]) == 0 + + # Test Case 6: Page beyond available data + result = _paginate_models_response( + all_models=mock_models[:10], + page=5, + size=10, + total_count=10, + search=None, + ) + assert result["current_page"] == 5 + assert len(result["data"]) == 0 # No data for page 5 + + def test_enrich_model_info_with_litellm_data(): """ Test the _enrich_model_info_with_litellm_data helper function.