ss-tools/backend/src/services/llm_provider.py

# #region llm_provider [C:3] [TYPE Module] [SEMANTICS sqlalchemy, llm, provider, encryption, config]
# @defgroup Services Module group.
# @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]
# @RATIONALE EncryptionManager imported from core/encryption.py (extracted from this module)
#            to decouple core encryption from LLM domain — see [SEC:C-1].
from typing import TYPE_CHECKING

from cryptography.exceptions import InvalidTag
from sqlalchemy.orm import Session

from ..core.encryption import get_encryption_manager
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]
# @ingroup Services
# @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]
# @ingroup Services
# @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


# NOTE: EncryptionManager is now imported from ..core.encryption


# #region LLMProviderService [C:3] [TYPE Class]
# @defgroup Services Module group.
# @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 = get_encryption_manager()

    # 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 ->[EXT:method: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,
                max_images=config.max_images,
                context_window=config.context_window,
                max_output_tokens=config.max_output_tokens,
            )
            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 ->[EXT:method: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
            db_provider.max_images = config.max_images
            db_provider.context_window = config.context_window
            db_provider.max_output_tokens = config.max_output_tokens

            self.db.commit()
            self.db.refresh(db_provider)
            return db_provider

    # endregion update_provider

    # region set_max_images [TYPE Function]
    # @PURPOSE: Updates only the max_images field for a provider.
    # @PRE  provider_id must exist.
    # @POST  Returns updated provider or None if not found.
    # @RELATION DEPENDS_ON ->[LLMProvider]
    def set_max_images(self, provider_id: str, max_images: int | None) -> LLMProvider | None:
        with belief_scope("set_max_images"):
            db_provider = self.get_provider(provider_id)
            if not db_provider:
                return None
            db_provider.max_images = max_images
            self.db.commit()
            self.db.refresh(db_provider)
            return db_provider

    # endregion set_max_images

    # 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 ->[EXT:method: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

    # region get_provider_token_config [TYPE Function]
    # @PURPOSE: Returns provider token limits for batch sizing.
    # @PRE  provider_id must be valid.
    # @POST  Returns dict with model name, context_window, max_output_tokens.
    #        Values from DB take priority; None means "use PROVIDER_DEFAULTS fallback".
    # @RATIONALE Centralised helper — both _batch_proc.py and _batch_sizer.py need
    #            the same resolution logic. Avoids duplicating DB queries and defaults.
    def get_provider_token_config(self, provider_id: str) -> dict:
        provider = self.get_provider(provider_id)
        if not provider:
            return {"model": None, "context_window": None, "max_output_tokens": None}
        return {
            "model": provider.default_model or "gpt-4o-mini",
            "context_window": provider.context_window,
            "max_output_tokens": provider.max_output_tokens,
        }
    # endregion get_provider_token_config


# #endregion LLMProviderService

# #endregion llm_provider