mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
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:
parent
671cca2509
commit
ce87a78459
2 changed files with 107 additions and 69 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue