fix(proxy): case-insensitive DB-side search for BYOK models

`_apply_search_filter_to_models` used Prisma's JSON path
`string_contains` to match the BYOK `team_public_model_name` field, but
that operator is case-sensitive in Postgres (no `mode: insensitive`
flag like column-level string filters have). So a search for "claude"
missed a stored "Claude Sonnet" via the DB branch even though the
router-side path matched it case-insensitively.

Widen the JSON branch to "row has a team_public_model_name set" and
filter case-insensitively in Python so DB-only BYOK rows match the
same terms users see in the UI. This also drops the now-unused
DB-level page-size optimization and `sort_by` knob — the in-Python
filter is the source of truth for `db_models_total_count` now.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
shivam 2026-05-16 13:32:57 -07:00
parent 671cca2509
commit ce87a78459
No known key found for this signature in database
2 changed files with 107 additions and 69 deletions

View file

@ -10892,11 +10892,8 @@ def _enrich_model_info_with_litellm_data(
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,
sort_by: Optional[str] = None,
) -> Tuple[List[Dict[str, Any]], Optional[int]]:
"""
Apply search filter to models, querying database for additional matching models.
@ -10904,11 +10901,8 @@ async def _apply_search_filter_to_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
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.
@ -10954,13 +10948,19 @@ async def _apply_search_filter_to_models(
# 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:
# Match either the persisted model_name or, for team BYOK rows,
# the user-facing team_public_model_name stored inside model_info.
# Prisma's JSON path `string_contains` is case-sensitive in
# Postgres (it doesn't accept the `mode: insensitive` flag the
# way column-level string filters do), so the BYOK branch
# below can't match mixed-case stored names like
# "Claude Sonnet" against a lowercased search term. Widen the
# JSON branch to "row has a team_public_model_name set"
# (`string_contains: ""` matches any string at the path) and
# filter case-insensitively in Python below so behavior
# matches the router-side path in `_model_matches_search`.
db_where_condition: Dict[str, Any] = {
"OR": [
{
@ -10972,7 +10972,7 @@ async def _apply_search_filter_to_models(
{
"model_info": {
"path": ["team_public_model_name"],
"string_contains": search_lower,
"string_contains": "",
}
},
]
@ -10983,58 +10983,36 @@ async def _apply_search_filter_to_models(
"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
)
# Fetch all candidates and filter in Python. We can't trust a
# DB-level count because the BYOK branch is over-broad — it
# returns every row with a team_public_model_name regardless
# of whether it matches the search term.
db_models_raw = await prisma_client.db.litellm_proxymodeltable.find_many(
where=db_where_condition,
)
def _db_row_matches_search(db_model: Any) -> bool:
if search_lower in (db_model.model_name or "").lower():
return True
info = (
db_model.model_info if isinstance(db_model.model_info, dict) else {}
)
return (
search_lower in (info.get("team_public_model_name") or "").lower()
)
matching_db_rows = [m for m in db_models_raw if _db_row_matches_search(m)]
db_models_total_count = len(matching_db_rows)
# 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)
# Decrypt matching rows. Done after the in-Python filter so we
# don't decrypt BYOK rows we're going to throw away.
for db_model in matching_db_rows:
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)}"
@ -11511,11 +11489,8 @@ async def model_info_v2(
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,
sort_by=sortBy,
)
if user_models_only:

View file

@ -1199,11 +1199,8 @@ async def test_apply_search_filter_matches_team_public_model_name():
filtered, _ = await _apply_search_filter_to_models(
all_models=[byok_model, unrelated_model],
search="claude",
page=1,
size=50,
prisma_client=None,
proxy_config=MagicMock(),
sort_by=None,
)
filtered_ids = {m["model_info"]["id"] for m in filtered}
assert "byok-id-1" in filtered_ids
@ -1213,11 +1210,8 @@ async def test_apply_search_filter_matches_team_public_model_name():
filtered, _ = await _apply_search_filter_to_models(
all_models=[byok_model, unrelated_model],
search="model_name_team-abc-123",
page=1,
size=50,
prisma_client=None,
proxy_config=MagicMock(),
sort_by=None,
)
assert [m["model_info"]["id"] for m in filtered] == ["byok-id-1"]
@ -1225,15 +1219,84 @@ async def test_apply_search_filter_matches_team_public_model_name():
filtered, _ = await _apply_search_filter_to_models(
all_models=[byok_model, unrelated_model],
search="gemini",
page=1,
size=50,
prisma_client=None,
proxy_config=MagicMock(),
sort_by=None,
)
assert filtered == []
@pytest.mark.asyncio
async def test_apply_search_filter_matches_db_byok_case_insensitively():
"""
Regression test: BYOK rows that only exist in the DB (not in the
in-memory router) must still match search case-insensitively against
their stored `team_public_model_name`. Prisma's JSON path
`string_contains` is case-sensitive in Postgres, so a search like
"claude" must still match a stored value of "Claude Sonnet".
"""
from litellm.proxy.proxy_server import _apply_search_filter_to_models
# Stored team_public_model_name uses mixed case; lowercased search
# would never match it via Prisma's case-sensitive JSON filter.
byok_db_row = MagicMock()
byok_db_row.model_id = "byok-db-only"
byok_db_row.model_name = "model_name_team-xyz_internal"
byok_db_row.model_info = {
"id": "byok-db-only",
"team_id": "team-xyz",
"team_public_model_name": "Claude Sonnet 4.6",
"db_model": True,
}
# Decoy row with team_public_model_name set but not matching the
# search — verifies the Python filter prunes the over-broad DB query.
decoy_db_row = MagicMock()
decoy_db_row.model_id = "decoy-db-only"
decoy_db_row.model_name = "model_name_team-xyz_decoy"
decoy_db_row.model_info = {
"id": "decoy-db-only",
"team_id": "team-xyz",
"team_public_model_name": "Some Gemini Variant",
"db_model": True,
}
prisma_client = MagicMock()
prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(
return_value=[byok_db_row, decoy_db_row]
)
proxy_config = MagicMock()
# decrypt_model_list_from_db echoes back a router-shaped dict; mock
# it so we can identify which DB row(s) survived the Python filter.
def _fake_decrypt(rows):
return [
{
"model_name": r.model_name,
"model_info": r.model_info,
"litellm_params": {"model": "claude-sonnet"},
}
for r in rows
]
proxy_config.decrypt_model_list_from_db = _fake_decrypt
filtered, total_count = await _apply_search_filter_to_models(
all_models=[],
search="claude",
prisma_client=prisma_client,
proxy_config=proxy_config,
)
filtered_ids = {m["model_info"]["id"] for m in filtered}
assert (
"byok-db-only" in filtered_ids
), "mixed-case team_public_model_name must match lowercased search"
assert (
"decoy-db-only" not in filtered_ids
), "non-matching BYOK row fetched by over-broad query must be filtered out"
assert total_count == 1
@pytest.mark.asyncio
async def test_filter_models_by_team_id_excludes_viewer_direct_access():
"""