
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: -LAN- <laipz8200@outlook.com>
53 lines
1.5 KiB
Python
53 lines
1.5 KiB
Python
import json
|
|
from enum import Enum
|
|
from json import JSONDecodeError
|
|
from typing import Optional
|
|
|
|
from extensions.ext_redis import redis_client
|
|
|
|
|
|
class ProviderCredentialsCacheType(Enum):
|
|
PROVIDER = "provider"
|
|
MODEL = "provider_model"
|
|
LOAD_BALANCING_MODEL = "load_balancing_provider_model"
|
|
|
|
|
|
class ProviderCredentialsCache:
|
|
def __init__(self, tenant_id: str, identity_id: str, cache_type: ProviderCredentialsCacheType):
|
|
self.cache_key = f"{cache_type.value}_credentials:tenant_id:{tenant_id}:id:{identity_id}"
|
|
|
|
def get(self) -> Optional[dict]:
|
|
"""
|
|
Get cached model provider credentials.
|
|
|
|
:return:
|
|
"""
|
|
cached_provider_credentials = redis_client.get(self.cache_key)
|
|
if cached_provider_credentials:
|
|
try:
|
|
cached_provider_credentials = cached_provider_credentials.decode("utf-8")
|
|
cached_provider_credentials = json.loads(cached_provider_credentials)
|
|
except JSONDecodeError:
|
|
return None
|
|
|
|
return dict(cached_provider_credentials)
|
|
else:
|
|
return None
|
|
|
|
def set(self, credentials: dict):
|
|
"""
|
|
Cache model provider credentials.
|
|
|
|
:param credentials: provider credentials
|
|
:return:
|
|
"""
|
|
redis_client.setex(self.cache_key, 86400, json.dumps(credentials))
|
|
|
|
def delete(self):
|
|
"""
|
|
Delete cached model provider credentials.
|
|
|
|
:return:
|
|
"""
|
|
redis_client.delete(self.cache_key)
|