307 lines
12 KiB
Python
307 lines
12 KiB
Python
# #region llm_provider [C:3] [TYPE Module] [SEMANTICS sqlalchemy, llm, provider, encryption, config]
|
|
# @BRIEF Service for managing LLM provider configurations with encrypted API keys.
|
|
# @LAYER: Domain
|
|
# @RELATION DEPENDS_ON -> [LLMProvider]
|
|
# @RELATION DEPENDS_ON -> [EncryptionManager]
|
|
# @RELATION DEPENDS_ON -> [LLMProviderConfig]
|
|
import os
|
|
from typing import TYPE_CHECKING
|
|
|
|
from cryptography.exceptions import InvalidTag
|
|
from cryptography.fernet import Fernet
|
|
from sqlalchemy.orm import Session
|
|
|
|
from ..core.logger import belief_scope, logger
|
|
from ..models.llm import LLMProvider
|
|
|
|
if TYPE_CHECKING:
|
|
from ..plugins.llm_analysis.models import LLMProviderConfig
|
|
|
|
MASKED_API_KEY_PLACEHOLDER = "********"
|
|
|
|
|
|
# #region mask_api_key [C:2] [TYPE Function]
|
|
# @BRIEF Mask an API key for safe display, showing first 4 and last 4 characters.
|
|
# @PRE: api_key is a plaintext string or None.
|
|
# @POST: Returns "****" for very short keys; "{first 2}...{last 2}" for <=8 chars;
|
|
# "{first 4}...{last 4}" for longer keys; "" for None/empty.
|
|
def mask_api_key(api_key: str | None) -> str:
|
|
if not api_key:
|
|
return ""
|
|
if len(api_key) <= 4:
|
|
return "****"
|
|
if len(api_key) <= 8:
|
|
return f"{api_key[:2]}...{api_key[-2:]}"
|
|
return f"{api_key[:4]}...{api_key[-4:]}"
|
|
|
|
|
|
# #endregion mask_api_key
|
|
|
|
|
|
# #region is_masked_or_placeholder [C:2] [TYPE Function]
|
|
# @BRIEF Predicate: True when api_key is None, empty, "********", or contains "...".
|
|
# @PRE: api_key can be None.
|
|
# @POST: Returns True only for non-real-key values.
|
|
def is_masked_or_placeholder(api_key: str | None) -> bool:
|
|
if not api_key:
|
|
return True
|
|
return api_key == MASKED_API_KEY_PLACEHOLDER or "..." in api_key
|
|
|
|
|
|
# #endregion is_masked_or_placeholder
|
|
|
|
|
|
# #region _require_fernet_key [C:5] [TYPE Function]
|
|
# @BRIEF Load and validate the Fernet key used for secret encryption.
|
|
# @PRE: ENCRYPTION_KEY environment variable must be set to a valid Fernet key.
|
|
# @POST: Returns validated key bytes ready for Fernet initialization.
|
|
# @RELATION DEPENDS_ON -> [backend.src.core.logger:Function]
|
|
# @SIDE_EFFECT: Emits belief-state logs for missing or invalid encryption configuration.
|
|
# @DATA_CONTRACT: Input[ENCRYPTION_KEY:str] -> Output[bytes]
|
|
# @INVARIANT: Encryption initialization never falls back to a hardcoded secret.
|
|
def _require_fernet_key() -> bytes:
|
|
with belief_scope("_require_fernet_key"):
|
|
raw_key = os.getenv("ENCRYPTION_KEY", "").strip()
|
|
if not raw_key:
|
|
logger.explore(
|
|
"Missing ENCRYPTION_KEY blocks EncryptionManager initialization"
|
|
)
|
|
raise RuntimeError("ENCRYPTION_KEY must be set")
|
|
|
|
key = raw_key.encode()
|
|
try:
|
|
Fernet(key)
|
|
except Exception as exc:
|
|
logger.explore(
|
|
"Invalid ENCRYPTION_KEY blocks EncryptionManager initialization"
|
|
)
|
|
raise RuntimeError("ENCRYPTION_KEY must be a valid Fernet key") from exc
|
|
|
|
logger.reflect("Validated ENCRYPTION_KEY for EncryptionManager initialization")
|
|
return key
|
|
|
|
|
|
# #endregion _require_fernet_key
|
|
|
|
|
|
# #region EncryptionManager [C:5] [TYPE Class]
|
|
# @BRIEF Handles encryption and decryption of sensitive data like API keys.
|
|
# @RELATION CALLS -> [_require_fernet_key]
|
|
# @PRE: ENCRYPTION_KEY is configured with a valid Fernet key before instantiation.
|
|
# @POST: Manager exposes reversible encrypt/decrypt operations for persisted secrets.
|
|
# @SIDE_EFFECT: Initializes Fernet cryptography state from process environment.
|
|
# @DATA_CONTRACT: Input[str] -> Output[str]
|
|
# @INVARIANT: Uses only a validated secret key from environment.
|
|
#
|
|
# @TEST_CONTRACT: EncryptionManagerModel ->
|
|
# {
|
|
# required_fields: {},
|
|
# invariants: [
|
|
# "encrypted data can be decrypted back to the original string"
|
|
# ]
|
|
# }
|
|
# @TEST_FIXTURE: basic_encryption_cycle -> {"data": "my_secret_key"}
|
|
# @TEST_EDGE: decrypt_invalid_data -> raises Exception
|
|
# @TEST_EDGE: empty_string_encryption -> {"data": ""}
|
|
# @TEST_INVARIANT: symmetric_encryption -> verifies: [basic_encryption_cycle, empty_string_encryption]
|
|
class EncryptionManager:
|
|
# region EncryptionManager_init [TYPE Function]
|
|
# @PURPOSE: Initialize the encryption manager with a Fernet key.
|
|
# @PRE: ENCRYPTION_KEY env var must be set to a valid Fernet key.
|
|
# @POST: Fernet instance ready for encryption/decryption.
|
|
def __init__(self):
|
|
self.key = _require_fernet_key()
|
|
self.fernet = Fernet(self.key)
|
|
|
|
# endregion EncryptionManager_init
|
|
|
|
# region encrypt [TYPE Function]
|
|
# @PURPOSE: Encrypt a plaintext string.
|
|
# @PRE: data must be a non-empty string.
|
|
# @POST: Returns encrypted string.
|
|
def encrypt(self, data: str) -> str:
|
|
with belief_scope("encrypt"):
|
|
return self.fernet.encrypt(data.encode()).decode()
|
|
|
|
# endregion encrypt
|
|
|
|
# region decrypt [TYPE Function]
|
|
# @PURPOSE: Decrypt an encrypted string.
|
|
# @PRE: encrypted_data must be a valid Fernet-encrypted string.
|
|
# @POST: Returns original plaintext string.
|
|
def decrypt(self, encrypted_data: str) -> str:
|
|
with belief_scope("decrypt"):
|
|
return self.fernet.decrypt(encrypted_data.encode()).decode()
|
|
|
|
# endregion decrypt
|
|
|
|
|
|
# #endregion EncryptionManager
|
|
|
|
|
|
# #region LLMProviderService [C:3] [TYPE Class]
|
|
# @BRIEF Service to manage LLM provider lifecycle.
|
|
# @RELATION DEPENDS_ON -> [LLMProvider]
|
|
# @RELATION DEPENDS_ON -> [EncryptionManager]
|
|
# @RELATION DEPENDS_ON -> [LLMProviderConfig]
|
|
class LLMProviderService:
|
|
# region LLMProviderService_init [TYPE Function]
|
|
# @PURPOSE: Initialize the service with database session.
|
|
# @PRE: db must be a valid SQLAlchemy Session.
|
|
# @POST: Service ready for provider operations.
|
|
# @RELATION: [DEPENDS_ON] ->[EncryptionManager]
|
|
def __init__(self, db: Session):
|
|
self.db = db
|
|
self.encryption = EncryptionManager()
|
|
|
|
# endregion LLMProviderService_init
|
|
|
|
# region get_all_providers [TYPE Function]
|
|
# @PURPOSE: Returns all configured LLM providers.
|
|
# @PRE: Database connection must be active.
|
|
# @POST: Returns list of all LLMProvider records.
|
|
# @RELATION: [DEPENDS_ON] ->[LLMProvider]
|
|
def get_all_providers(self) -> list[LLMProvider]:
|
|
with belief_scope("get_all_providers"):
|
|
return self.db.query(LLMProvider).all()
|
|
|
|
# endregion get_all_providers
|
|
|
|
# region get_provider [TYPE Function]
|
|
# @PURPOSE: Returns a single LLM provider by ID.
|
|
# @PRE: provider_id must be a valid string.
|
|
# @POST: Returns LLMProvider or None if not found.
|
|
# @RELATION: [DEPENDS_ON] ->[LLMProvider]
|
|
def get_provider(self, provider_id: str) -> LLMProvider | None:
|
|
with belief_scope("get_provider"):
|
|
return (
|
|
self.db.query(LLMProvider).filter(LLMProvider.id == provider_id).first()
|
|
)
|
|
|
|
# endregion get_provider
|
|
|
|
# region create_provider [TYPE Function]
|
|
# @PURPOSE: Creates a new LLM provider with encrypted API key.
|
|
# @PRE: config must contain valid provider configuration.
|
|
# @POST: New provider created and persisted to database.
|
|
# @RELATION: [DEPENDS_ON] ->[LLMProviderConfig]
|
|
# @RELATION: [DEPENDS_ON] ->[LLMProvider]
|
|
# @RELATION: [CALLS] ->[encrypt]
|
|
def create_provider(self, config: "LLMProviderConfig") -> LLMProvider:
|
|
with belief_scope("create_provider"):
|
|
encrypted_key = self.encryption.encrypt(config.api_key)
|
|
db_provider = LLMProvider(
|
|
provider_type=config.provider_type.value,
|
|
name=config.name,
|
|
base_url=config.base_url,
|
|
api_key=encrypted_key,
|
|
default_model=config.default_model,
|
|
is_active=config.is_active,
|
|
is_multimodal=config.is_multimodal,
|
|
)
|
|
self.db.add(db_provider)
|
|
self.db.commit()
|
|
self.db.refresh(db_provider)
|
|
return db_provider
|
|
|
|
# endregion create_provider
|
|
|
|
# region update_provider [TYPE Function]
|
|
# @PURPOSE: Updates an existing LLM provider.
|
|
# @PRE: provider_id must exist, config must be valid.
|
|
# @POST: Provider updated and persisted to database.
|
|
# @RELATION: [DEPENDS_ON] ->[LLMProviderConfig]
|
|
# @RELATION: [DEPENDS_ON] ->[LLMProvider]
|
|
# @RELATION: [CALLS] ->[encrypt]
|
|
def update_provider(
|
|
self, provider_id: str, config: "LLMProviderConfig"
|
|
) -> LLMProvider | None:
|
|
with belief_scope("update_provider"):
|
|
db_provider = self.get_provider(provider_id)
|
|
if not db_provider:
|
|
return None
|
|
|
|
db_provider.provider_type = config.provider_type.value
|
|
db_provider.name = config.name
|
|
db_provider.base_url = config.base_url
|
|
# Ignore masked/placeholder values; they are display-only and must not overwrite secrets.
|
|
if not is_masked_or_placeholder(config.api_key):
|
|
db_provider.api_key = self.encryption.encrypt(config.api_key)
|
|
db_provider.default_model = config.default_model
|
|
db_provider.is_active = config.is_active
|
|
db_provider.is_multimodal = config.is_multimodal
|
|
|
|
self.db.commit()
|
|
self.db.refresh(db_provider)
|
|
return db_provider
|
|
|
|
# endregion update_provider
|
|
|
|
# region delete_provider [TYPE Function]
|
|
# @PURPOSE: Deletes an LLM provider.
|
|
# @PRE: provider_id must exist.
|
|
# @POST: Provider removed from database.
|
|
# @RELATION: [DEPENDS_ON] ->[LLMProvider]
|
|
def delete_provider(self, provider_id: str) -> bool:
|
|
with belief_scope("delete_provider"):
|
|
db_provider = self.get_provider(provider_id)
|
|
if not db_provider:
|
|
return False
|
|
self.db.delete(db_provider)
|
|
self.db.commit()
|
|
return True
|
|
|
|
# endregion delete_provider
|
|
|
|
# region get_decrypted_api_key [TYPE Function]
|
|
# @PURPOSE: Returns the decrypted API key for a provider.
|
|
# @PRE: provider_id must exist with valid encrypted key.
|
|
# @POST: Returns decrypted API key or None on failure.
|
|
# @RELATION: [DEPENDS_ON] ->[LLMProvider]
|
|
# @RELATION: [CALLS] ->[decrypt]
|
|
def get_decrypted_api_key(self, provider_id: str) -> str | None:
|
|
with belief_scope("get_decrypted_api_key"):
|
|
db_provider = self.get_provider(provider_id)
|
|
if not db_provider:
|
|
logger.warning(
|
|
f"[get_decrypted_api_key] Provider {provider_id} not found in database"
|
|
)
|
|
return None
|
|
|
|
logger.info(f"[get_decrypted_api_key] Provider found: {db_provider.id}")
|
|
logger.info(
|
|
f"[get_decrypted_api_key] Encrypted API key length: {len(db_provider.api_key) if db_provider.api_key else 0}"
|
|
)
|
|
|
|
try:
|
|
decrypted_key = self.encryption.decrypt(db_provider.api_key)
|
|
logger.info(
|
|
f"[get_decrypted_api_key] Decryption successful, key length: {len(decrypted_key) if decrypted_key else 0}"
|
|
)
|
|
return decrypted_key
|
|
except InvalidTag as e:
|
|
logger.error(
|
|
f"[get_decrypted_api_key] Integrity check failed (InvalidTag): {e!s}. "
|
|
"The encrypted API key may be corrupted or the ENCRYPTION_KEY has changed."
|
|
)
|
|
return None
|
|
except ValueError as e:
|
|
logger.error(
|
|
f"[get_decrypted_api_key] Decryption format error (ValueError): {e!s}. "
|
|
"The encrypted data may not be valid Fernet ciphertext."
|
|
)
|
|
return None
|
|
except Exception as e:
|
|
logger.error(
|
|
f"[get_decrypted_api_key] Decryption failed with unexpected error "
|
|
f"({type(e).__name__}): {e!s}"
|
|
)
|
|
return None
|
|
|
|
# endregion get_decrypted_api_key
|
|
|
|
|
|
# #endregion LLMProviderService
|
|
|
|
# #endregion llm_provider
|