mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
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>
This commit is contained in:
yuneng
parent
177021b2ac
commit
d69eb7f095
2 changed files with 56 additions and 10 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue