mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-10 08:28:25 +00:00
fix(proxy): ask each provider catalog once per worker, keep deployed names out of catch-all ids, and match file credentials by name
Listings that start together for one deployment now wait for the first catalog request. A catch-all wildcard no longer lists a name a deployment serves by name. A key wildcard only finds file and batch credentials on deployments named under it, so a bare name keeps the default credentials as before.
This commit is contained in:
parent
aabeffc652
commit
b786c5514d
6 changed files with 153 additions and 3 deletions
|
|
@ -399,7 +399,18 @@ def _deployment_team_id(deployment: Mapping[str, object]) -> str | None:
|
|||
|
||||
|
||||
def _cached_wildcard_routing(llm_router: Router | None) -> Callable[[str, str | None], frozenset[str]] | None:
|
||||
return None if llm_router is None else functools.cache(llm_router.wildcard_routed_models)
|
||||
if llm_router is None:
|
||||
return None
|
||||
deployed_names: Final = frozenset(llm_router.model_names)
|
||||
team_deployed_names: Final = frozenset(llm_router.team_model_to_deployment_indices)
|
||||
wildcard_routed_models: Final = llm_router.wildcard_routed_models
|
||||
|
||||
def routed_models(model: str, team_id: str | None) -> frozenset[str]:
|
||||
if model in deployed_names or (team_id, model) in team_deployed_names:
|
||||
return frozenset()
|
||||
return wildcard_routed_models(model, team_id)
|
||||
|
||||
return functools.cache(routed_models)
|
||||
|
||||
|
||||
def _routes_to_model(
|
||||
|
|
|
|||
|
|
@ -561,7 +561,7 @@ def get_team_provider_credentials(
|
|||
return None
|
||||
|
||||
from litellm.proxy._types import SpecialModelNames
|
||||
from litellm.proxy.auth.auth_checks import model_matches_any_wildcard_pattern_in_list
|
||||
from litellm.proxy.auth.auth_checks import is_model_allowed_by_pattern
|
||||
from litellm.proxy.auth.model_checks import get_complete_model_list, get_key_models
|
||||
|
||||
team_id: Final = user_api_key_dict.team_id
|
||||
|
|
@ -663,7 +663,9 @@ def get_team_provider_credentials(
|
|||
deployments_granted_by_wildcard: Final = (
|
||||
model_name
|
||||
for model_name in proxy_model_list
|
||||
if model_matches_any_wildcard_pattern_in_list(model=model_name, allowed_model_list=granted_models)
|
||||
if any(
|
||||
is_model_allowed_by_pattern(model=model_name, allowed_model_pattern=granted) for granted in granted_models
|
||||
)
|
||||
)
|
||||
return next(
|
||||
(
|
||||
|
|
|
|||
|
|
@ -8432,6 +8432,9 @@ class AvailableModelsCache(InMemoryCache):
|
|||
cache_key: Final = self._get_cache_key(custom_llm_provider, litellm_params)
|
||||
self._store(_failed_listing_key(cache_key), True, ttl=ttl_seconds)
|
||||
|
||||
def listing_lock(self, custom_llm_provider: str, litellm_params: LiteLLM_Params | None) -> threading.RLock:
|
||||
return _provider_listing_lock(self._get_cache_key(custom_llm_provider, litellm_params))
|
||||
|
||||
def _cached(self, key: str) -> Sequence[str] | bool | None:
|
||||
return cast(Sequence[str] | bool | None, self.get_cache(key))
|
||||
|
||||
|
|
@ -8443,6 +8446,19 @@ def _failed_listing_key(cache_key: str) -> str:
|
|||
return f"{cache_key}:failed"
|
||||
|
||||
|
||||
_PROVIDER_LISTING_LOCK_GUARD: Final = threading.Lock()
|
||||
|
||||
|
||||
@lru_cache(maxsize=1024)
|
||||
def _cached_provider_listing_lock(cache_key: str) -> threading.RLock:
|
||||
return threading.RLock()
|
||||
|
||||
|
||||
def _provider_listing_lock(cache_key: str) -> threading.RLock:
|
||||
with _PROVIDER_LISTING_LOCK_GUARD:
|
||||
return _cached_provider_listing_lock(cache_key)
|
||||
|
||||
|
||||
# Global cache instance
|
||||
_model_cache: Final = AvailableModelsCache()
|
||||
|
||||
|
|
@ -8494,6 +8510,15 @@ def _get_valid_models_from_provider_api(
|
|||
provider_config: BaseLLMModelInfo,
|
||||
custom_llm_provider: str,
|
||||
litellm_params: LiteLLM_Params | None = None,
|
||||
) -> list[str]:
|
||||
with _model_cache.listing_lock(custom_llm_provider, litellm_params):
|
||||
return _cached_or_listed_provider_models(provider_config, custom_llm_provider, litellm_params)
|
||||
|
||||
|
||||
def _cached_or_listed_provider_models(
|
||||
provider_config: BaseLLMModelInfo,
|
||||
custom_llm_provider: str,
|
||||
litellm_params: LiteLLM_Params | None,
|
||||
) -> list[str]:
|
||||
try:
|
||||
cached_result: Final = _model_cache.get_cached_model_info(custom_llm_provider, litellm_params)
|
||||
|
|
|
|||
|
|
@ -7,6 +7,8 @@ import httpx
|
|||
import pytest
|
||||
import respx
|
||||
|
||||
from litellm import Router
|
||||
|
||||
|
||||
def test_get_team_models_for_all_models_and_team_only_models():
|
||||
from litellm.proxy.auth.model_checks import get_team_models
|
||||
|
|
@ -1214,6 +1216,73 @@ def test_key_granted_openrouter_wildcard_without_an_openrouter_deployment_lists_
|
|||
assert (listed, catalog.call_count) == ([], 0)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _openai_catalog_behind_a_catch_all(
|
||||
monkeypatch: pytest.MonkeyPatch, concrete_name: str, concrete_model_info: Mapping[str, str]
|
||||
) -> Generator[Router, None, None]:
|
||||
import litellm
|
||||
|
||||
monkeypatch.setattr(litellm, "check_provider_endpoint", True)
|
||||
monkeypatch.delenv("OPENAI_API_BASE", raising=False)
|
||||
monkeypatch.delenv("OPENAI_BASE_URL", raising=False)
|
||||
with respx.mock(assert_all_called=False) as respx_mock:
|
||||
respx_mock.get("https://api.openai.com/v1/models").mock(
|
||||
return_value=httpx.Response(200, json={"data": [{"id": "gpt-4o"}, {"id": "gpt-4o-mini"}]})
|
||||
)
|
||||
yield Router(
|
||||
model_list=[
|
||||
{"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "sk-openai-catch-all"}},
|
||||
{
|
||||
"model_name": concrete_name,
|
||||
"litellm_params": {"model": "openrouter/openai/gpt-4o", "api_key": "sk-or-gpt-4o"},
|
||||
"model_info": dict(concrete_model_info),
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("concrete_name", "concrete_model_info", "team_id"),
|
||||
[
|
||||
("gpt-4o", {}, None),
|
||||
("model_name_team-x_gpt-4o", {"team_id": "team-x", "team_public_model_name": "gpt-4o"}, "team-x"),
|
||||
],
|
||||
)
|
||||
def test_catch_all_wildcard_never_lists_a_name_a_concrete_deployment_serves(
|
||||
monkeypatch: pytest.MonkeyPatch, concrete_name: str, concrete_model_info: Mapping[str, str], team_id: str | None
|
||||
) -> None:
|
||||
from litellm.proxy.auth.model_checks import get_complete_model_list
|
||||
|
||||
with _openai_catalog_behind_a_catch_all(monkeypatch, concrete_name, concrete_model_info) as router:
|
||||
listed: Final = get_complete_model_list(
|
||||
key_models=[],
|
||||
team_models=[],
|
||||
proxy_model_list=router.get_model_names(team_id=team_id),
|
||||
user_model=None,
|
||||
infer_model_from_keys=False,
|
||||
llm_router=router,
|
||||
team_id=team_id,
|
||||
)
|
||||
|
||||
assert sorted(listed) == ["gpt-4o", "gpt-4o-mini"]
|
||||
|
||||
|
||||
def test_catch_all_wildcard_model_info_keeps_only_the_concrete_row_for_a_deployed_name(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from litellm.proxy.auth.model_checks import expand_wildcard_deployments_for_model_info
|
||||
|
||||
with _openai_catalog_behind_a_catch_all(monkeypatch, "gpt-4o", {}) as router:
|
||||
rows: Final = expand_wildcard_deployments_for_model_info(
|
||||
[dict(deployment) for deployment in router.get_model_list() or ()], llm_router=router
|
||||
)
|
||||
|
||||
assert sorted((row["model_name"], row["litellm_params"]["model"]) for row in rows) == [
|
||||
("gpt-4o", "openrouter/openai/gpt-4o"),
|
||||
("gpt-4o-mini", "openai/gpt-4o-mini"),
|
||||
]
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_custom_prefix_that_starts_a_discovered_id_still_prefixes_every_listed_model(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
|
|
|
|||
|
|
@ -940,3 +940,26 @@ def test_a_key_wildcard_finds_the_deployment_it_grants_without_listing_any_provi
|
|||
assert credentials is not None
|
||||
assert (credentials["custom_llm_provider"], credentials["api_key"]) == (custom_llm_provider, api_key)
|
||||
assert catalog_requests.call_count == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("deployed_name", "api_key"), [("gpt-4o-mini", None), ("openai/gpt-4o-mini", "openai-deployment-key")]
|
||||
)
|
||||
def test_a_provider_wildcard_grant_finds_only_deployments_named_under_it(deployed_name: str, api_key: str | None):
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": deployed_name,
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "openai-deployment-key"},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
credentials: Final = get_team_provider_credentials(
|
||||
llm_router=router,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test-key", user_id="test-user", models=["openai/*"]),
|
||||
custom_llm_provider="openai",
|
||||
include_accessible_deployments=True,
|
||||
)
|
||||
|
||||
assert (None if credentials is None else credentials["api_key"]) == api_key
|
||||
|
|
|
|||
|
|
@ -8948,6 +8948,26 @@ def test_get_valid_models_keeps_a_good_listing_when_an_overlapping_listing_fails
|
|||
) == ["openrouter/anthropic/claude-haiku-5-5"]
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_get_valid_models_asks_a_provider_once_when_many_listings_start_together(reset_mock_cache: None) -> None:
|
||||
def slow_catalog(request: httpx.Request) -> httpx.Response:
|
||||
threading.Event().wait(0.2)
|
||||
return httpx.Response(200, json={"data": [{"id": "anthropic/claude-haiku-5-5"}]})
|
||||
|
||||
catalog: Final = respx.get("https://openrouter.ai/api/v1/models/user").mock(side_effect=slow_catalog)
|
||||
deployment: Final = LiteLLM_Params(model="openrouter/*", api_key="sk-or-many-listings")
|
||||
|
||||
def list_models(_: int) -> tuple[str, ...]:
|
||||
return tuple(
|
||||
get_valid_models(check_provider_endpoint=True, custom_llm_provider="openrouter", litellm_params=deployment)
|
||||
)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=8) as pool:
|
||||
listings: Final = frozenset(pool.map(list_models, range(8)))
|
||||
|
||||
assert (listings, catalog.call_count) == (frozenset({("openrouter/anthropic/claude-haiku-5-5",)}), 1)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def isolated_openai_model_sets(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue