use helper to check when _should_read_secret_from_secret_manager

This commit is contained in:
Ishaan Jaff 2024-11-13 14:31:35 -08:00
parent 33dc97df93
commit 3a072cd75a
3 changed files with 27 additions and 2 deletions

View file

@ -304,7 +304,7 @@ secret_manager_client: Optional[Any] = (
)
_google_kms_resource_name: Optional[str] = None
_key_management_system: Optional[KeyManagementSystem] = None
_key_management_settings: Optional[KeyManagementSettings] = None
_key_management_settings: KeyManagementSettings = KeyManagementSettings()
#### PII MASKING ####
output_parse_pii: bool = False
#############################################

View file

@ -1134,6 +1134,11 @@ class KeyManagementSettings(LiteLLMBase):
If True, virtual keys created by litellm will be stored in the secret manager
"""
access_mode: Literal["read_only", "write_only", "read_and_write"] = "read_only"
"""
Access mode for the secret manager, when write_only will only use for writing secrets
"""
class TeamDefaultSettings(LiteLLMBase):
team_id: str

View file

@ -198,7 +198,10 @@ def get_secret( # noqa: PLR0915
raise ValueError("Unsupported OIDC provider")
try:
if litellm.secret_manager_client is not None:
if (
_should_read_secret_from_secret_manager()
and litellm.secret_manager_client is not None
):
try:
client = litellm.secret_manager_client
key_manager = "local"
@ -321,3 +324,20 @@ def get_secret( # noqa: PLR0915
return default_value
else:
raise e
def _should_read_secret_from_secret_manager() -> bool:
"""
Returns True if the secret manager should be used to read the secret, False otherwise
- If the secret manager client is not set, return False
- If the `_key_management_settings` access mode is "read_only" or "read_and_write", return True
- Otherwise, return False
"""
if litellm.secret_manager_client is not None:
if (
litellm._key_management_settings.access_mode == "read_only"
or litellm._key_management_settings.access_mode == "read_and_write"
):
return True
return False