chore: initial public snapshot for github upload
This commit is contained in:
Your Name
@@ -0,0 +1,252 @@
|
||||
"""
|
||||
Custom Secret Manager Integration
|
||||
|
||||
This module provides a base class for implementing custom secret managers in LiteLLM.
|
||||
|
||||
Usage:
|
||||
from litellm.integrations.custom_secret_manager import CustomSecretManager
|
||||
|
||||
class MySecretManager(CustomSecretManager):
|
||||
def __init__(self):
|
||||
super().__init__(secret_manager_name="my_secret_manager")
|
||||
|
||||
async def async_read_secret(
|
||||
self,
|
||||
secret_name: str,
|
||||
optional_params=None,
|
||||
timeout=None,
|
||||
):
|
||||
# Your implementation here
|
||||
return await self._fetch_secret_from_service(secret_name)
|
||||
|
||||
def sync_read_secret(
|
||||
self,
|
||||
secret_name: str,
|
||||
optional_params=None,
|
||||
timeout=None,
|
||||
):
|
||||
# Your implementation here
|
||||
return self._fetch_secret_from_service_sync(secret_name)
|
||||
|
||||
# Set your custom secret manager
|
||||
import litellm
|
||||
from litellm.types.secret_managers.main import KeyManagementSystem
|
||||
|
||||
litellm.secret_manager_client = MySecretManager()
|
||||
litellm._key_management_system = KeyManagementSystem.CUSTOM
|
||||
"""
|
||||
|
||||
from abc import abstractmethod
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.secret_managers.base_secret_manager import BaseSecretManager
|
||||
|
||||
|
||||
class CustomSecretManager(BaseSecretManager):
|
||||
"""
|
||||
Base class for implementing custom secret managers.
|
||||
|
||||
This class provides a standard interface for implementing custom secret management
|
||||
integrations in LiteLLM. Users can extend this class to integrate their own secret
|
||||
management systems.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from litellm.integrations.custom_secret_manager import CustomSecretManager
|
||||
|
||||
class MyVaultSecretManager(CustomSecretManager):
|
||||
def __init__(self, vault_url: str, token: str):
|
||||
super().__init__(secret_manager_name="my_vault")
|
||||
self.vault_url = vault_url
|
||||
self.token = token
|
||||
|
||||
async def async_read_secret(self, secret_name: str, optional_params=None, timeout=None):
|
||||
# Implementation for reading secrets from your vault
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(
|
||||
f"{self.vault_url}/v1/secret/{secret_name}",
|
||||
headers={"X-Vault-Token": self.token},
|
||||
timeout=timeout
|
||||
)
|
||||
return response.json()["data"]["value"]
|
||||
|
||||
def sync_read_secret(self, secret_name: str, optional_params=None, timeout=None):
|
||||
# Sync implementation
|
||||
with httpx.Client() as client:
|
||||
response = client.get(
|
||||
f"{self.vault_url}/v1/secret/{secret_name}",
|
||||
headers={"X-Vault-Token": self.token},
|
||||
timeout=timeout
|
||||
)
|
||||
return response.json()["data"]["value"]
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
secret_manager_name: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Initialize the CustomSecretManager.
|
||||
|
||||
Args:
|
||||
secret_manager_name: A descriptive name for your secret manager.
|
||||
This is used for logging and debugging purposes.
|
||||
**kwargs: Additional keyword arguments to pass to your secret manager.
|
||||
"""
|
||||
super().__init__()
|
||||
self.secret_manager_name = secret_manager_name or "custom_secret_manager"
|
||||
verbose_logger.info("Initialized custom secret manager")
|
||||
|
||||
@abstractmethod
|
||||
async def async_read_secret(
|
||||
self,
|
||||
secret_name: str,
|
||||
optional_params: Optional[dict] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Asynchronously read a secret from your custom secret manager.
|
||||
|
||||
Args:
|
||||
secret_name: Name/path of the secret to read
|
||||
optional_params: Additional parameters specific to your secret manager
|
||||
timeout: Request timeout
|
||||
|
||||
Returns:
|
||||
The secret value if found, None otherwise
|
||||
|
||||
Raises:
|
||||
Exception: If there's an error reading the secret
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def sync_read_secret(
|
||||
self,
|
||||
secret_name: str,
|
||||
optional_params: Optional[dict] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Synchronously read a secret from your custom secret manager.
|
||||
|
||||
Args:
|
||||
secret_name: Name/path of the secret to read
|
||||
optional_params: Additional parameters specific to your secret manager
|
||||
timeout: Request timeout
|
||||
|
||||
Returns:
|
||||
The secret value if found, None otherwise
|
||||
|
||||
Raises:
|
||||
Exception: If there's an error reading the secret
|
||||
"""
|
||||
pass
|
||||
|
||||
async def async_write_secret(
|
||||
self,
|
||||
secret_name: str,
|
||||
secret_value: str,
|
||||
description: Optional[str] = None,
|
||||
optional_params: Optional[dict] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
tags: Optional[Union[dict, list]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Asynchronously write a secret to your custom secret manager.
|
||||
|
||||
This is optional to implement. If your secret manager supports writing secrets,
|
||||
you can override this method.
|
||||
|
||||
Args:
|
||||
secret_name: Name/path of the secret to write
|
||||
secret_value: Value to store
|
||||
description: Description of the secret
|
||||
optional_params: Additional parameters specific to your secret manager
|
||||
timeout: Request timeout
|
||||
tags: Optional tags to apply to the secret
|
||||
|
||||
Returns:
|
||||
Response from the secret manager containing write operation details
|
||||
|
||||
Raises:
|
||||
NotImplementedError: If write operations are not supported
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
f"Write operations are not implemented for {self.secret_manager_name}. "
|
||||
"Override async_write_secret() to add write support."
|
||||
)
|
||||
|
||||
async def async_delete_secret(
|
||||
self,
|
||||
secret_name: str,
|
||||
recovery_window_in_days: Optional[int] = 7,
|
||||
optional_params: Optional[dict] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Asynchronously delete a secret from your custom secret manager.
|
||||
|
||||
This is optional to implement. If your secret manager supports deleting secrets,
|
||||
you can override this method.
|
||||
|
||||
Args:
|
||||
secret_name: Name of the secret to delete
|
||||
recovery_window_in_days: Number of days before permanent deletion (if supported)
|
||||
optional_params: Additional parameters specific to your secret manager
|
||||
timeout: Request timeout
|
||||
|
||||
Returns:
|
||||
Response from the secret manager containing deletion details
|
||||
|
||||
Raises:
|
||||
NotImplementedError: If delete operations are not supported
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
f"Delete operations are not implemented for {self.secret_manager_name}. "
|
||||
"Override async_delete_secret() to add delete support."
|
||||
)
|
||||
|
||||
def validate_environment(self) -> bool:
|
||||
"""
|
||||
Validate that all required environment variables and configuration are present.
|
||||
|
||||
Override this method to validate your secret manager's configuration.
|
||||
|
||||
Returns:
|
||||
True if the environment is valid
|
||||
|
||||
Raises:
|
||||
ValueError: If required configuration is missing
|
||||
"""
|
||||
verbose_logger.debug(
|
||||
"No environment validation configured for custom secret manager"
|
||||
)
|
||||
return True
|
||||
|
||||
async def async_health_check(
|
||||
self, timeout: Optional[Union[float, httpx.Timeout]] = None
|
||||
) -> bool:
|
||||
"""
|
||||
Perform a health check on your secret manager.
|
||||
|
||||
This is optional to implement. Override this method to add health check support.
|
||||
|
||||
Args:
|
||||
timeout: Request timeout
|
||||
|
||||
Returns:
|
||||
True if the secret manager is healthy, False otherwise
|
||||
"""
|
||||
verbose_logger.debug(
|
||||
f"Health check not implemented for {self.secret_manager_name}"
|
||||
)
|
||||
return True
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<{self.__class__.__name__}(name={self.secret_manager_name})>"
|
||||
Reference in New Issue
Block a user