Spaces:
Sleeping
Sleeping
File size: 1,284 Bytes
7db0ae4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
"""
This is a file for the Google KMS integration
Relevant issue: https://github.com/BerriAI/litellm/issues/1235
Requires:
* `os.environ["GOOGLE_APPLICATION_CREDENTIALS"], os.environ["GOOGLE_KMS_RESOURCE_NAME"]`
* `pip install google-cloud-kms`
"""
import litellm, os
from typing import Optional
from litellm.proxy._types import KeyManagementSystem
def validate_environment():
if "GOOGLE_APPLICATION_CREDENTIALS" not in os.environ:
raise ValueError(
"Missing required environment variable - GOOGLE_APPLICATION_CREDENTIALS"
)
if "GOOGLE_KMS_RESOURCE_NAME" not in os.environ:
raise ValueError(
"Missing required environment variable - GOOGLE_KMS_RESOURCE_NAME"
)
def load_google_kms(use_google_kms: Optional[bool]):
if use_google_kms is None or use_google_kms == False:
return
try:
from google.cloud import kms_v1 # type: ignore
validate_environment()
# Create the KMS client
client = kms_v1.KeyManagementServiceClient()
litellm.secret_manager_client = client
litellm._key_management_system = KeyManagementSystem.GOOGLE_KMS
litellm._google_kms_resource_name = os.getenv("GOOGLE_KMS_RESOURCE_NAME")
except Exception as e:
raise e
|