Fixing ruff check

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(

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.