From d69eb7f095e93798954aaf19e2d4aef4fa189687 Mon Sep 17 00:00:00 2001 From: yuneng Date: Mon, 21 Sep 2026 19:47:27 +0000 Subject: [PATCH] fix(proxy): fall back to the credentials table when validating litellm_credential_name Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../model_management_endpoints.py | 27 ++++++++----- .../test_model_management_endpoints.py | 39 ++++++++++++++++++- 2 files changed, 56 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 47e617a20af..2ff77b5008b 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -95,6 +95,7 @@ from litellm.proxy.spend_tracking.ptu_feature_flag import ( is_ptu_cost_attribution_enabled, ) from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.repositories.credentials_repository import CredentialsRepository from litellm.repositories.model_repository import ModelRepository from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ModelTableRepository @@ -333,7 +334,9 @@ def _raise_on_strategy_router_write_violation( ) -def _raise_on_invalid_credential_name(litellm_params: updateLiteLLMParams | None) -> None: +async def _raise_on_invalid_credential_name( + litellm_params: updateLiteLLMParams | None, prisma_client: PrismaClient +) -> None: if litellm_params is None or "litellm_credential_name" not in litellm_params.model_fields_set: return credential_name: Final = litellm_params.litellm_credential_name @@ -346,13 +349,19 @@ def _raise_on_invalid_credential_name(litellm_params: updateLiteLLMParams | None code=status.HTTP_400_BAD_REQUEST, param="litellm_credential_name", ) - if CredentialAccessor.find_credential(credential_name) is None: - raise ProxyException( - message=f"Credential '{credential_name}' not found. Create it via /credentials before attaching it to a model.", - type=ProxyErrorTypes.validation_error.value, - code=status.HTTP_400_BAD_REQUEST, - param="litellm_credential_name", - ) + if CredentialAccessor.find_credential(credential_name) is not None: + return + stored_credential: Final = await CredentialsRepository(WriterPinnedClient(prisma_client.db)).find_by_name( + credential_name + ) + if stored_credential is not None: + return + raise ProxyException( + message=f"Credential '{credential_name}' not found. Create it via /credentials before attaching it to a model.", + type=ProxyErrorTypes.validation_error.value, + code=status.HTTP_400_BAD_REQUEST, + param="litellm_credential_name", + ) AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY: Final = 5_872_301 @@ -1135,7 +1144,7 @@ async def patch_model( existing_litellm_params=db_model.litellm_params, null_detaches=True, ) - _raise_on_invalid_credential_name(patch_data.litellm_params) + await _raise_on_invalid_credential_name(patch_data.litellm_params, prisma_client) ModelManagementAuthChecks.can_user_set_aws_session_tags( litellm_params=patch_data.litellm_params, diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index b6b68d28880..e6b5fb25c3e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -10,6 +10,7 @@ import pytest from fastapi.testclient import TestClient from litellm._uuid import uuid +from litellm.models.credentials import CredentialItem from litellm.proxy._types import ( LiteLLM_ModelTable, @@ -4249,10 +4250,11 @@ class TestPatchModelCredentialName: db_model: Deployment, user_api_key_dict: UserAPIKeyAuth, credential_name: str | None, + db_credential: CredentialItem | None = None, + credentials_repository: MagicMock | None = None, ) -> list[dict[str, object]]: import litellm from litellm.proxy.management_endpoints.model_management_endpoints import patch_model, update_db_model - from litellm.types.utils import CredentialItem monkeypatch.setattr( litellm, @@ -4270,6 +4272,8 @@ class TestPatchModelCredentialName: ), ], ) + credentials_repository = credentials_repository or MagicMock() + credentials_repository.find_by_name = AsyncMock(return_value=db_credential) persisted: Final[list[dict[str, object]]] = [] async def persist_model(**kwargs): @@ -4284,6 +4288,10 @@ class TestPatchModelCredentialName: patch("litellm.proxy.proxy_server.llm_router", MagicMock()), patch("litellm.proxy.proxy_server.store_model_in_db", True), patch("litellm.proxy.proxy_server.premium_user", True), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.CredentialsRepository", + return_value=credentials_repository, + ), patch( "litellm.proxy.management_endpoints.model_management_endpoints.get_db_model", new=AsyncMock(return_value=db_model), @@ -4364,6 +4372,7 @@ class TestPatchModelCredentialName: async def test_patch_model_rejects_unknown_credential_name(self, monkeypatch): from litellm.proxy._types import ProxyException + credentials_repository = MagicMock() db_model: Final = Deployment( model_name="gpt-4", litellm_params=LiteLLM_Params( @@ -4380,10 +4389,38 @@ class TestPatchModelCredentialName: db_model, self._admin_user(), "ghost-credential", + credentials_repository=credentials_repository, ) assert exc_info.value.code == "400" assert "not found" in exc_info.value.message.lower() + credentials_repository.find_by_name.assert_awaited_once_with("ghost-credential") + + @pytest.mark.asyncio + async def test_patch_model_accepts_credential_known_only_in_db(self, monkeypatch): + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1"), + ) + + persisted: Final = await self._patch_model( + monkeypatch, + db_model, + self._admin_user(), + "db-only-credential", + db_credential=CredentialItem( + credential_name="db-only-credential", + credential_info={}, + credential_values={"api_key": "sk-db"}, + ), + ) + params: Final = json.loads(persisted[0]["litellm_params"]) + assert params["litellm_credential_name"] == "db-only-credential" @pytest.mark.asyncio async def test_patch_model_replaces_credential_name_and_preserves_other_params(self, monkeypatch):