Support master key rotations (#9041)

* feat(key_management_endpoints.py): adding support for rotating master key

* feat(key_management_endpoints.py): support decryption-re-encryption of models in db, when master key rotated

* fix(user_api_key_auth.py): raise valid token is None error earlier

enables easier debugging with api key hash in error message

* feat(key_management_endpoints.py): rotate any env vars

* fix(key_management_endpoints.py): uncomment check

* fix: fix linting error
This commit is contained in:
Krish Dholakia 2025-03-06 23:13:30 -08:00 committed by GitHub
parent fcc57318f8
commit da13ec2b64
8 changed files with 214 additions and 33 deletions

View file

@ -21,6 +21,7 @@ from litellm._logging import verbose_proxy_logger
from litellm.constants import LITELLM_PROXY_ADMIN_NAME
from litellm.proxy._types import (
CommonProxyErrors,
LiteLLM_ProxyModelTable,
LitellmUserRoles,
PrismaCompatibleUpdateDBModel,
ProxyErrorTypes,
@ -227,12 +228,16 @@ async def _add_model_to_db(
model_params: Deployment,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient,
):
new_encryption_key: Optional[str] = None,
should_create_model_in_db: bool = True,
) -> Optional[LiteLLM_ProxyModelTable]:
# encrypt litellm params #
_litellm_params_dict = model_params.litellm_params.dict(exclude_none=True)
_orignal_litellm_model_name = model_params.litellm_params.model
for k, v in _litellm_params_dict.items():
encrypted_value = encrypt_value_helper(value=v)
encrypted_value = encrypt_value_helper(
value=v, new_encryption_key=new_encryption_key
)
model_params.litellm_params[k] = encrypted_value
_data: dict = {
"model_id": model_params.model_info.id,
@ -246,9 +251,12 @@ async def _add_model_to_db(
}
if model_params.model_info.id is not None:
_data["model_id"] = model_params.model_info.id
model_response = await prisma_client.db.litellm_proxymodeltable.create(
data=_data # type: ignore
)
if should_create_model_in_db:
model_response = await prisma_client.db.litellm_proxymodeltable.create(
data=_data # type: ignore
)
else:
model_response = LiteLLM_ProxyModelTable(**_data)
return model_response